You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2025/04/08 07:58:06 UTC
(airflow) branch main updated: Call class-level pre_ and post_execute (#48908)
This is an automated email from the ASF dual-hosted git repository.
kaxilnaik pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 95917f1153f Call class-level pre_ and post_execute (#48908)
95917f1153f is described below
commit 95917f1153ff2e487c5fb9777c16821e70f2dadc
Author: Tzu-ping Chung <ur...@gmail.com>
AuthorDate: Tue Apr 8 15:58:00 2025 +0800
Call class-level pre_ and post_execute (#48908)
Closes #48387
---
airflow-core/src/airflow/models/baseoperator.py | 30 ----------------------
airflow-core/src/airflow/models/taskinstance.py | 27 +++++++++++++++++--
.../tests/unit/models/test_baseoperator.py | 17 ------------
.../src/airflow/sdk/execution_time/task_runner.py | 4 +++
4 files changed, 29 insertions(+), 49 deletions(-)
diff --git a/airflow-core/src/airflow/models/baseoperator.py b/airflow-core/src/airflow/models/baseoperator.py
index c674c030aea..0d0f8b92e2d 100644
--- a/airflow-core/src/airflow/models/baseoperator.py
+++ b/airflow-core/src/airflow/models/baseoperator.py
@@ -337,19 +337,6 @@ class BaseOperator(TaskSDKBaseOperator):
extended/overridden by subclasses.
"""
- def pre_execute(self, context: Any):
- """Execute right before self.execute() is called."""
- if self._pre_execute_hook is None:
- return
- from airflow.sdk.execution_time.callback_runner import create_executable_runner
- from airflow.sdk.execution_time.context import context_get_outlet_events
-
- create_executable_runner(
- self._pre_execute_hook,
- context_get_outlet_events(context),
- logger=self.log,
- ).run(context)
-
def execute(self, context: Context) -> Any:
"""
Derive when creating an operator.
@@ -360,23 +347,6 @@ class BaseOperator(TaskSDKBaseOperator):
"""
raise NotImplementedError()
- def post_execute(self, context: Any, result: Any = None):
- """
- Execute right after self.execute() is called.
-
- It is passed the execution context and any results returned by the operator.
- """
- if self._post_execute_hook is None:
- return
- from airflow.sdk.execution_time.callback_runner import create_executable_runner
- from airflow.sdk.execution_time.context import context_get_outlet_events
-
- create_executable_runner(
- self._post_execute_hook,
- context_get_outlet_events(context),
- logger=self.log,
- ).run(context, result)
-
@provide_session
def clear(
self,
diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py
index e86294cacd5..e202991ad63 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -2779,6 +2779,9 @@ class TaskInstance(Base, LoggingMixin):
def _execute_task_with_callbacks(self, context: Context, test_mode: bool = False, *, session: Session):
"""Prepare Task for Execution."""
+ from airflow.sdk.execution_time.callback_runner import create_executable_runner
+ from airflow.sdk.execution_time.context import context_get_outlet_events
+
if TYPE_CHECKING:
assert self.task
@@ -2843,7 +2846,17 @@ class TaskInstance(Base, LoggingMixin):
)
# Run pre_execute callback
- self.task.pre_execute(context=context)
+ if self.task._pre_execute_hook:
+ create_executable_runner(
+ self.task._pre_execute_hook,
+ context_get_outlet_events(context),
+ logger=self.log,
+ ).run(context)
+ create_executable_runner(
+ self.task.pre_execute,
+ context_get_outlet_events(context),
+ logger=self.log,
+ ).run(context)
# Run on_execute callback
self._run_execute_callback(context, self.task)
@@ -2877,7 +2890,17 @@ class TaskInstance(Base, LoggingMixin):
self._rendered_map_index = _render_map_index(context, jinja_env=jinja_env)
# Run post_execute callback
- self.task.post_execute(context=context, result=result)
+ if self.task._post_execute_hook:
+ create_executable_runner(
+ self.task._post_execute_hook,
+ context_get_outlet_events(context),
+ logger=self.log,
+ ).run(context, result)
+ create_executable_runner(
+ self.task.post_execute,
+ context_get_outlet_events(context),
+ logger=self.log,
+ ).run(context, result)
Stats.incr(f"operator_successes_{self.task.task_type}", tags=self.stats_tags)
# Same metric with tagging
diff --git a/airflow-core/tests/unit/models/test_baseoperator.py b/airflow-core/tests/unit/models/test_baseoperator.py
index 6a2e5268586..b4f7197583e 100644
--- a/airflow-core/tests/unit/models/test_baseoperator.py
+++ b/airflow-core/tests/unit/models/test_baseoperator.py
@@ -20,7 +20,6 @@ from __future__ import annotations
import copy
from collections import defaultdict
from datetime import datetime
-from unittest import mock
import pytest
@@ -113,22 +112,6 @@ class TestBaseOperator:
except Exception as e:
pytest.fail(f"Exception raised: {e}")
- def test_pre_execute_hook(self):
- hook = mock.MagicMock()
-
- op = BaseOperator(task_id="test_task", pre_execute=hook)
- op_copy = op.prepare_for_execution()
- op_copy.pre_execute({})
- assert hook.called
-
- def test_post_execute_hook(self):
- hook = mock.MagicMock()
-
- op = BaseOperator(task_id="test_task", post_execute=hook)
- op_copy = op.prepare_for_execution()
- op_copy.post_execute({})
- assert hook.called
-
def test_task_naive_datetime(self):
naive_datetime = DEFAULT_DATE.replace(tzinfo=None)
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index 40861cb4d5e..227d78e53b4 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -970,6 +970,8 @@ def _execute_task(context: Context, ti: RuntimeTaskInstance, log: Logger):
if (pre_execute_hook := task._pre_execute_hook) is not None:
create_executable_runner(pre_execute_hook, outlet_events, logger=log).run(context)
+ if getattr(pre_execute_hook := task.pre_execute, "__func__", None) is not BaseOperator.pre_execute:
+ create_executable_runner(pre_execute_hook, outlet_events, logger=log).run(context)
_run_task_state_change_callbacks(task, "on_execute_callback", context, log)
@@ -993,6 +995,8 @@ def _execute_task(context: Context, ti: RuntimeTaskInstance, log: Logger):
if (post_execute_hook := task._post_execute_hook) is not None:
create_executable_runner(post_execute_hook, outlet_events, logger=log).run(context, result)
+ if getattr(post_execute_hook := task.post_execute, "__func__", None) is not BaseOperator.post_execute:
+ create_executable_runner(post_execute_hook, outlet_events, logger=log).run(context)
return result