You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by GitBox <gi...@apache.org> on 2018/09/10 23:50:34 UTC
[GitHub] r39132 closed pull request #2140: WIP - Waiting endpoint
r39132 closed pull request #2140: WIP - Waiting endpoint
URL: https://github.com/apache/incubator-airflow/pull/2140
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/airflow/models.py b/airflow/models.py
index 37f88230a0..7eeaae6ee5 100755
--- a/airflow/models.py
+++ b/airflow/models.py
@@ -1019,7 +1019,7 @@ def key(self):
"""
return self.dag_id, self.task_id, self.execution_date
- def set_state(self, state, session):
+ def set_state(self, state, session):
self.state = state
self.start_date = datetime.now()
self.end_date = datetime.now()
@@ -1393,13 +1393,17 @@ def signal_handler(signum, frame):
Stats.incr('operator_successes_{}'.format(
self.task.__class__.__name__), 1, 1)
self.state = State.SUCCESS
+ if (task.on_success_set_state_to
+ and task.on_success_set_state_to in State.hold()):
+ self.state = task.on_success_set_state_to
+
except AirflowSkipException:
self.state = State.SKIPPED
except (Exception, KeyboardInterrupt) as e:
self.handle_failure(e, test_mode, context)
raise
- # Recording SUCCESS
+ # Recording SUCCESS or on_success_set_state_to
self.end_date = datetime.now()
self.set_duration()
if not test_mode:
@@ -1883,6 +1887,12 @@ class derived from this one results in the creation of a task object,
:type resources: dict
:param run_as_user: unix username to impersonate while running the task
:type run_as_user: str
+ :param on_success_set_state_to: sets the exit state of the operator to this
+ state. Can be used to allow an external event to trigger continuation
+ :type on_success_set_state_to: str
+ :param on_standby_timeout: the amount of time the scheduler waits on the
+ external trigger. Will be converted to seconds.
+ :type on_standby_timeout: timedelta
"""
# For derived classes to define which fields will get jinjaified
@@ -1925,6 +1935,8 @@ def __init__(
trigger_rule=TriggerRule.ALL_SUCCESS,
resources=None,
run_as_user=None,
+ on_success_set_state_to=None,
+ on_standby_timeout=timedelta(3600),
*args,
**kwargs):
@@ -1976,6 +1988,11 @@ def __init__(
self.execution_timeout = execution_timeout
self.on_failure_callback = on_failure_callback
self.on_success_callback = on_success_callback
+ self.on_success_set_state_to = on_success_set_state_to
+ if isinstance(on_standby_timeout, timedelta):
+ self.on_standby_timeout = on_standby_timeout
+ else:
+ self.on_standby_timeout = timedelta(seconds=retry_delay)
self.on_retry_callback = on_retry_callback
if isinstance(retry_delay, timedelta):
self.retry_delay = retry_delay
@@ -2019,6 +2036,7 @@ def __init__(
'on_failure_callback',
'on_success_callback',
'on_retry_callback',
+ 'on_success_set_state',
}
def __eq__(self, other):
diff --git a/airflow/ti_deps/deps/onstandby_dep.py b/airflow/ti_deps/deps/onstandby_dep.py
new file mode 100644
index 0000000000..ad8c579d65
--- /dev/null
+++ b/airflow/ti_deps/deps/onstandby_dep.py
@@ -0,0 +1,57 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from datetime import datetime
+
+from airflow.ti_deps.deps.base_ti_dep import BaseTIDep
+from airflow.utils.db import provide_session
+from airflow.utils.state import State
+
+
+class OnStandByDep(BaseTIDep):
+ """
+ Determines if a task's upstream tasks are in a state that allows a given task instance
+ to run.
+ """
+ NAME = "On Standby"
+ IGNOREABLE = True
+ IS_TASK_DEP = True
+
+ @provide_session
+ def _get_dep_statuses(self, ti, session, dep_context):
+ if dep_context.ignore_ti_state:
+ yield self._passing_status(
+ reason="Context specified that state should be ignored.")
+ return
+
+ if ti.state is not State.ON_STANDBY:
+ yield self._passing_status(
+ reason="The TaskInstance is not ON_STANDBY"
+ )
+ return
+
+ cur_date = datetime.now()
+ timeout = ti.end_date + ti.on_standby_timeout
+ if cur_date < timeout:
+ yield self._failing_status(
+ reason="Task in on stand by and will expire on {}".format(timeout.isoformat())
+ )
+ return
+
+ # if we get here the external trigger has not fired
+ ti.set_state(state=State.FAILED, session=session)
+
+ yield self._failing_status(
+ reason="External trigger time window expired"
+ )
+ return
diff --git a/airflow/utils/state.py b/airflow/utils/state.py
index 4a1dfb620a..90cee04642 100644
--- a/airflow/utils/state.py
+++ b/airflow/utils/state.py
@@ -40,6 +40,7 @@ class State(object):
UP_FOR_RETRY = "up_for_retry"
UPSTREAM_FAILED = "upstream_failed"
SKIPPED = "skipped"
+ ON_STANDBY = 'on_standby'
task_states = (
SUCCESS,
@@ -48,6 +49,7 @@ class State(object):
UPSTREAM_FAILED,
UP_FOR_RETRY,
QUEUED,
+ ON_STANDBY,
)
dag_states = (
@@ -67,6 +69,7 @@ class State(object):
SKIPPED: 'pink',
REMOVED: 'lightgrey',
SCHEDULED: 'white',
+ ON_STANDBY: 'purple'
}
@classmethod
@@ -109,5 +112,15 @@ def unfinished(cls):
cls.SCHEDULED,
cls.QUEUED,
cls.RUNNING,
- cls.UP_FOR_RETRY
+ cls.UP_FOR_RETRY,
+ cls.ON_STANDBY
+ ]
+
+ @classmethod
+ def hold(cls):
+ """
+ List of states indicating a task is waiting for external action
+ """
+ return [
+ cls.ON_STANDBY,
]
diff --git a/airflow/www/api/experimental/endpoints.py b/airflow/www/api/experimental/endpoints.py
index 56b9d79bc5..0768bb311b 100644
--- a/airflow/www/api/experimental/endpoints.py
+++ b/airflow/www/api/experimental/endpoints.py
@@ -18,6 +18,8 @@
from airflow.api.common.experimental import trigger_dag as trigger
from airflow.exceptions import AirflowException
from airflow.www.app import csrf
+from airflow.models import DagRun
+from airflow.settings import Session
from flask import (
g, Markup, Blueprint, redirect, jsonify, abort, request, current_app, send_file
@@ -110,3 +112,73 @@ def task_info(dag_id, task_id):
task = dag.get_task(task_id)
fields = {k: str(v) for k, v in vars(task).items() if not k.startswith('_')}
return jsonify(fields)
+
+@api_experimental.route('/dags/<string:dag_id>/dag_runs/<string:execution_date>/tasks/<string:task_id>', methods=['PUT'])
+@requires_authentication
+def task_callback(dag_id, execution_date, task_id):
+ """
+
+ :param dag_id:
+ :param execution_date:
+ :param task_id:
+ :return:
+ """
+ from airflow.www.views import dagbag
+
+ data = request.get_json(force=True)
+ if not 'state' in data:
+ response = jsonify(error="No state defined")
+ response.status_code = 400
+ return response
+
+ state = data['state']
+ if state is not 'SUCCESS' or state is not 'FAILED':
+ response = jsonify(error="State {} is not a valid state".format(state))
+ response.status_code = 400
+ return response
+
+ if dag_id not in dagbag.dags:
+ response = jsonify(error='Dag {} not found'.format(dag_id))
+ response.status_code = 404
+ return response
+
+ dag = dagbag.dags[dag_id]
+ if not dag.has_task(task_id):
+ response = (jsonify(error='Task {} not found in dag {}'
+ .format(task_id, dag_id)))
+ response.status_code = 404
+ return response
+
+ # Convert string datetime into actual datetime
+ try:
+ execution_date = datetime.strptime(execution_date,
+ '%Y-%m-%dT%H:%M:%S')
+ except ValueError:
+ error_message = (
+ 'Given execution date, {}, could not be identified '
+ 'as a date. Example date format: 2015-11-16T14:34:15'
+ .format(execution_date))
+ _log.info(error_message)
+ response = jsonify({'error': error_message})
+ response.status_code = 400
+
+ return response
+
+ dr = DagRun.find(dag_id=dag_id, execution_date=execution_date)
+ if not dr:
+ response = jsonify(error="DagRun {} not found for dag {} on execution date {}",
+ dag_id=dag_id, execution_date=execution_date)
+ response.status_code = 404
+ return response
+
+ ti = dr.get_task_instance(task_id=task_id)
+ if not ti:
+ response = jsonify(error="TaskInstance {} not found in DagRun for dag {} and execution date {}",
+ task_id=task_id, dag_id=dag_id, execution_date=execution_date)
+ response.status_code = 404
+ return response
+
+ session = Session()
+ ti.refresh_from_db(lock_for_update=True, session=session)
+ ti.state = state
+ session.commit()
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services