Repository: incubator-airflow
Updated Branches:
  refs/heads/master fe7881656 -> 6613676d7


[AIRFLOW-886] Pass result to post_execute() hook

The post_execute() hook should receive
the Operator result in addition to the
execution context.


Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/4da3611c
Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/4da3611c
Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/4da3611c

Branch: refs/heads/master
Commit: 4da3611c46e5d5fec5ac988d5c15d2edf2be5768
Parents: 21d775a
Author: Jeremiah Lowin <[email protected]>
Authored: Sat Feb 18 12:04:22 2017 -0500
Committer: Jeremiah Lowin <[email protected]>
Committed: Sat Feb 18 18:38:58 2017 -0500

----------------------------------------------------------------------
 UPDATING.md       |  9 +++++++++
 airflow/models.py | 27 +++++++++++++++++++++------
 tests/models.py   | 26 ++++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 6 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/4da3611c/UPDATING.md
----------------------------------------------------------------------
diff --git a/UPDATING.md b/UPDATING.md
index ba708cd..6fd7afb 100644
--- a/UPDATING.md
+++ b/UPDATING.md
@@ -11,6 +11,15 @@ assists people when migrating to a new version.
 
 A new DaskExecutor allows Airflow tasks to be run in Dask Distributed clusters.
 
+### Deprecated Features
+These features are marked for deprecation. They may still work (and raise a 
`DeprecationWarning`), but are no longer
+supported and will be removed entirely in Airflow 2.0
+
+- `post_execute()` hooks now take two arguments, `context` and `result`
+  (AIRFLOW-886)
+
+  Previously, post_execute() only took one argument, `context`.
+
 ## Airflow 1.8
 
 ### Database

http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/4da3611c/airflow/models.py
----------------------------------------------------------------------
diff --git a/airflow/models.py b/airflow/models.py
index a6c100e..6ed6115 100755
--- a/airflow/models.py
+++ b/airflow/models.py
@@ -1374,7 +1374,22 @@ class TaskInstance(Base):
                 if result is not None:
                     self.xcom_push(key=XCOM_RETURN_KEY, value=result)
 
-                task_copy.post_execute(context=context)
+                # TODO remove deprecated behavior in Airflow 2.0
+                try:
+                    task_copy.post_execute(context=context, result=result)
+                except TypeError as e:
+                    if 'unexpected keyword argument' in str(e):
+                        warnings.warn(
+                            'BaseOperator.post_execute() now takes two '
+                            'arguments, `context` and `result`, but "{}" only '
+                            'expected one. This behavior is deprecated and '
+                            'will be removed in a future version of '
+                            'Airflow.'.format(self.task_id),
+                            category=DeprecationWarning)
+                        task_copy.post_execute(context=context)
+                    else:
+                        raise
+
                 Stats.incr('operator_successes_{}'.format(
                     self.task.__class__.__name__), 1, 1)
             self.state = State.SUCCESS
@@ -2154,8 +2169,7 @@ class BaseOperator(object):
 
     def pre_execute(self, context):
         """
-        This is triggered right before self.execute, it's mostly a hook
-        for people deriving operators.
+        This hook is triggered right before self.execute() is called.
         """
         pass
 
@@ -2168,10 +2182,11 @@ class BaseOperator(object):
         """
         raise NotImplementedError()
 
-    def post_execute(self, context):
+    def post_execute(self, context, result=None):
         """
-        This is triggered right after self.execute, it's mostly a hook
-        for people deriving operators.
+        This hook is triggered right after self.execute() is called.
+        It is passed the execution context and any results returned by the
+        operator.
         """
         pass
 

http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/4da3611c/tests/models.py
----------------------------------------------------------------------
diff --git a/tests/models.py b/tests/models.py
index 97d703e..6f4337f 100644
--- a/tests/models.py
+++ b/tests/models.py
@@ -649,3 +649,29 @@ class TaskInstanceTest(unittest.TestCase):
                                       key=key,
                                       include_prior_dates=True),
                          value)
+
+    def test_post_execute_hook(self):
+        """
+        Test that post_execute hook is called with the Operator's result.
+        The result ('error') will cause an error to be raised and trapped.
+        """
+
+        class TestError(Exception):
+            pass
+
+        class TestOperator(PythonOperator):
+            def post_execute(self, context, result):
+                if result == 'error':
+                    raise TestError('expected error.')
+
+        dag = models.DAG(dag_id='test_post_execute_dag')
+        task = TestOperator(
+            task_id='test_operator',
+            dag=dag,
+            python_callable=lambda: 'error',
+            owner='airflow',
+            start_date=datetime.datetime(2017, 2, 1))
+        ti = TI(task=task, execution_date=datetime.datetime.now())
+
+        with self.assertRaises(TestError):
+            ti.run()

Reply via email to