This is an automated email from the ASF dual-hosted git repository.
potiuk 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 1c3252e5805 Add `callback` and `source` to context to differentiate
between Dag and task callbacks (#76110)
1c3252e5805 is described below
commit 1c3252e5805028ee94efb497c7ff53ec8dbdea5b
Author: IndiTheo <[email protected]>
AuthorDate: Mon Nov 2 01:52:21 2026 +0200
Add `callback` and `source` to context to differentiate between Dag and
task callbacks (#76110)
* Add callback_source to context to differentiate DAG vs task callbacks
Currently a callback defined at the DAG level receives a context that
is nearly identical to a task-level callback context, making it
impossible to tell which level triggered the call.
This adds a `CallbackSource` enum ("dag" / "task") and a `CallbackMeta`
named tuple, and injects `context["callback"] = CallbackMeta(source=...)`
in every code path that invokes callbacks:
- task_runner._run_task_state_change_callbacks (in-worker task path)
- processor._execute_dag_callbacks (DAG processor path)
- processor._execute_task_callbacks (DAG processor path)
- DagRun.handle_dag_callback (dag.test() path)
Both `CallbackSource` and `CallbackMeta` are exported from `airflow.sdk`.
Closes: #61119
Made-with: Cursor
* Document CallbackSource and CallbackMeta in task-sdk API reference
Adds the new public callback metadata types to api.rst so
test_docs_inventory_matches_public_api passes in CI.
* Add callback meta structs to context
Fix test to use RuntimeError instead of AirflowException.
Add docs for callback source in context.
Add noqa SDK001 to processor.py imports.
Update known_sdk_imports_in_core.txt for dagrun.py.
* Fixed dag spelling
* Fixed imports
* Converted newsfragment to significant
* Drop the newsfragment for the additive callback context key
Significant newsfragments are only for changes users must act on before
upgrading; this adds a new context key, which the release notes pick up from
the PR.
Generated-by: Claude Opus 5
---------
Co-authored-by: vasanthrpjan1-boop <[email protected]>
Co-authored-by: IndiTheo <[email protected]>
Co-authored-by: Jarek Potiuk <[email protected]>
---
.../logging-monitoring/callbacks.rst | 2 ++
.../src/airflow/dag_processing/processor.py | 6 ++++
airflow-core/src/airflow/models/dagrun.py | 2 ++
.../tests/unit/dag_processing/test_processor.py | 9 ++++++
.../check_template_context_variable_in_sync.py | 4 +--
task-sdk/docs/api.rst | 4 +++
task-sdk/src/airflow/sdk/__init__.py | 12 +++++++-
task-sdk/src/airflow/sdk/__init__.pyi | 2 ++
task-sdk/src/airflow/sdk/definitions/context.py | 26 +++++++++++++++++
.../src/airflow/sdk/execution_time/task_runner.py | 3 ++
.../task_sdk/execution_time/test_task_runner.py | 34 ++++++++++++++++++++++
11 files changed, 101 insertions(+), 3 deletions(-)
diff --git
a/airflow-core/docs/administration-and-deployment/logging-monitoring/callbacks.rst
b/airflow-core/docs/administration-and-deployment/logging-monitoring/callbacks.rst
index 9b15c8f6c09..3e6372cf455 100644
---
a/airflow-core/docs/administration-and-deployment/logging-monitoring/callbacks.rst
+++
b/airflow-core/docs/administration-and-deployment/logging-monitoring/callbacks.rst
@@ -76,6 +76,8 @@ Context Mapping
A context mapping that contains runtime information about a task instance is
passed to every callback.
Full list of variables available in ``context`` are in :doc:`docs
<../../templates-ref>` and `code
<https://github.com/apache/airflow/blob/main/task-sdk/src/airflow/sdk/definitions/context.py>`_.
+Every callback context also includes a ``callback`` key whose ``source`` field
tells you whether the callback ran at the Dag level or the task level.
+
Dag Callbacks
^^^^^^^^^^^^^
diff --git a/airflow-core/src/airflow/dag_processing/processor.py
b/airflow-core/src/airflow/dag_processing/processor.py
index 4be5623d71d..a44e43623d3 100644
--- a/airflow-core/src/airflow/dag_processing/processor.py
+++ b/airflow-core/src/airflow/dag_processing/processor.py
@@ -376,6 +376,7 @@ def _execute_callbacks(
def _execute_dag_callbacks(dagbag: DagBag, request: DagCallbackRequest, log:
FilteringBoundLogger) -> None:
from airflow.sdk.api.datamodels._generated import TIRunContext
+ from airflow.sdk.definitions.context import CallbackMeta, CallbackSource
# noqa: SDK001
dag, _ = _get_dag_with_task(dagbag, request.dag_id)
callbacks = dag.on_failure_callback if request.is_failure_callback else
dag.on_success_callback
@@ -415,6 +416,8 @@ def _execute_dag_callbacks(dagbag: DagBag, request:
DagCallbackRequest, log: Fil
context = runtime_ti.get_template_context()
context["reason"] = request.msg
+ context["callback"] = CallbackMeta(source=CallbackSource.DAG)
+
for callback in callbacks:
log.info(
"Executing on_%s dag callback",
@@ -441,6 +444,8 @@ def _execute_dag_callbacks(dagbag: DagBag, request:
DagCallbackRequest, log: Fil
def _execute_task_callbacks(dagbag: DagBag, request: TaskCallbackRequest, log:
FilteringBoundLogger) -> None:
+ from airflow.sdk.definitions.context import CallbackMeta, CallbackSource
# noqa: SDK001
+
if not request.is_failure_callback:
log.warning(
"Task callback requested but is not a failure callback",
@@ -487,6 +492,7 @@ def _execute_task_callbacks(dagbag: DagBag, request:
TaskCallbackRequest, log: F
task=task,
)
context = runtime_ti.get_template_context()
+ context["callback"] = CallbackMeta(source=CallbackSource.TASK)
def get_callback_representation(callback):
with contextlib.suppress(AttributeError):
diff --git a/airflow-core/src/airflow/models/dagrun.py
b/airflow-core/src/airflow/models/dagrun.py
index 6b8236805f1..71b25c9b4b8 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -99,6 +99,7 @@ from airflow.models.dynamic_region import (
)
from airflow.models.taskinstance import TaskInstance as TI,
_add_and_prime_mapped_ti
from airflow.models.tasklog import LogTemplate
+from airflow.sdk.definitions.context import CallbackMeta, CallbackSource #
noqa: SDK001
from airflow.serialization.definitions.deadline import
SerializedReferenceModels
from airflow.serialization.definitions.notset import NOTSET, ArgNotSet,
is_arg_set
from airflow.ti_deps.dep_context import DepContext
@@ -1823,6 +1824,7 @@ class DagRun(Base, LoggingMixin):
}
context["reason"] = reason
+ context["callback"] = CallbackMeta(source=CallbackSource.DAG)
callbacks = dag.on_success_callback if success else
dag.on_failure_callback
if not callbacks:
diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py
b/airflow-core/tests/unit/dag_processing/test_processor.py
index 154cb7640ce..c9c4ab19e08 100644
--- a/airflow-core/tests/unit/dag_processing/test_processor.py
+++ b/airflow-core/tests/unit/dag_processing/test_processor.py
@@ -1078,6 +1078,9 @@ class TestExecuteDagCallbacks:
# Check that we have template context variables from
RuntimeTaskInstance
assert "ts" in context_received
assert "params" in context_received
+ from airflow.sdk.definitions.context import CallbackSource
+
+ assert context_received["callback"].source is CallbackSource.DAG
def
test_execute_dag_callbacks_missing_last_ti_task_falls_back_to_minimal_context(self,
spy_agency):
"""A removed last_ti task must not drop the callback; it runs with the
minimal context."""
@@ -1189,6 +1192,9 @@ class TestExecuteDagCallbacks:
# Should not have template context variables
assert "ts" not in context_received
assert "params" not in context_received
+ from airflow.sdk.definitions.context import CallbackSource
+
+ assert context_received["callback"].source is CallbackSource.DAG
def test_execute_dag_callbacks_success_callback(self, spy_agency):
"""Test _execute_dag_callbacks executes success callback with
context_from_server"""
@@ -1626,6 +1632,9 @@ class TestExecuteTaskCallbacks:
assert context_received is not None
assert context_received["dag"] == dag
assert "ti" in context_received
+ from airflow.sdk.definitions.context import CallbackSource
+
+ assert context_received["callback"].source is CallbackSource.TASK
def test_execute_task_callbacks_retry_callback(self, spy_agency):
"""Test _execute_task_callbacks executes retry callbacks"""
diff --git a/scripts/ci/prek/check_template_context_variable_in_sync.py
b/scripts/ci/prek/check_template_context_variable_in_sync.py
index 88bc3cbed3a..86a5e8cccf7 100755
--- a/scripts/ci/prek/check_template_context_variable_in_sync.py
+++ b/scripts/ci/prek/check_template_context_variable_in_sync.py
@@ -163,8 +163,8 @@ def _compare_keys(retn_keys: set[str], hint_keys: set[str],
docs_keys: set[str])
retn_keys.add("test_mode")
# Only present in callbacks. Not listed in templates-ref (that doc is for
task execution).
- retn_keys.update(("exception", "reason", "try_number"))
- docs_keys.update(("exception", "reason", "try_number"))
+ retn_keys.update(("callback", "exception", "reason", "try_number"))
+ docs_keys.update(("callback", "exception", "reason", "try_number"))
# Airflow 3 added:
retn_keys.add("task_reschedule_count")
diff --git a/task-sdk/docs/api.rst b/task-sdk/docs/api.rst
index 824fde2638e..e7f854a7785 100644
--- a/task-sdk/docs/api.rst
+++ b/task-sdk/docs/api.rst
@@ -112,6 +112,10 @@ Callbacks
.. autoclass:: airflow.sdk.SyncCallback
+.. autoclass:: airflow.sdk.CallbackSource
+
+.. autoclass:: airflow.sdk.CallbackMeta
+
Deadline Alerts
---------------
.. autoclass:: airflow.sdk.DeadlineAlert
diff --git a/task-sdk/src/airflow/sdk/__init__.py
b/task-sdk/src/airflow/sdk/__init__.py
index 46418ceeaf1..1c4d7150201 100644
--- a/task-sdk/src/airflow/sdk/__init__.py
+++ b/task-sdk/src/airflow/sdk/__init__.py
@@ -41,6 +41,8 @@ __all__ = [
"BaseSensorOperator",
"BaseXCom",
"BranchMixIn",
+ "CallbackMeta",
+ "CallbackSource",
"ChainMapper",
"ChainRetryPolicy",
"Connection",
@@ -156,7 +158,13 @@ if TYPE_CHECKING:
from airflow.sdk.definitions.asset.metadata import Metadata
from airflow.sdk.definitions.callback import AsyncCallback, SyncCallback
from airflow.sdk.definitions.connection import Connection
- from airflow.sdk.definitions.context import Context, get_current_context,
get_parsing_context
+ from airflow.sdk.definitions.context import (
+ CallbackMeta as CallbackMeta,
+ CallbackSource as CallbackSource,
+ Context,
+ get_current_context,
+ get_parsing_context,
+ )
from airflow.sdk.definitions.dag import DAG, dag
from airflow.sdk.definitions.deadline import (
BaseDeadlineReference,
@@ -262,6 +270,8 @@ __lazy_imports: dict[str, tuple[str, str]] = {
"BaseSensorOperator": (".bases.sensor", "BaseSensorOperator"),
"BaseXCom": (".bases.xcom", "BaseXCom"),
"BranchMixIn": (".bases.branch", "BranchMixIn"),
+ "CallbackMeta": (".definitions.context", "CallbackMeta"),
+ "CallbackSource": (".definitions.context", "CallbackSource"),
"ChainMapper": (".definitions.partition_mappers.chain", "ChainMapper"),
"ChainRetryPolicy": (".definitions.retry_policy", "ChainRetryPolicy"),
"Connection": (".definitions.connection", "Connection"),
diff --git a/task-sdk/src/airflow/sdk/__init__.pyi
b/task-sdk/src/airflow/sdk/__init__.pyi
index 221d3eb48a6..e35b2cb05a5 100644
--- a/task-sdk/src/airflow/sdk/__init__.pyi
+++ b/task-sdk/src/airflow/sdk/__init__.pyi
@@ -56,6 +56,8 @@ from airflow.sdk.definitions.asset.metadata import Metadata
as Metadata
from airflow.sdk.definitions.callback import AsyncCallback, SyncCallback
from airflow.sdk.definitions.connection import Connection as Connection
from airflow.sdk.definitions.context import (
+ CallbackMeta as CallbackMeta,
+ CallbackSource as CallbackSource,
Context as Context,
get_current_context as get_current_context,
get_parsing_context as get_parsing_context,
diff --git a/task-sdk/src/airflow/sdk/definitions/context.py
b/task-sdk/src/airflow/sdk/definitions/context.py
index 6f85470cc68..df305bcde0e 100644
--- a/task-sdk/src/airflow/sdk/definitions/context.py
+++ b/task-sdk/src/airflow/sdk/definitions/context.py
@@ -18,10 +18,35 @@
from __future__ import annotations
import copy
+import enum
import os
from collections.abc import MutableMapping
from typing import TYPE_CHECKING, Any, NamedTuple, NotRequired, TypedDict, cast
+
+class CallbackSource(str, enum.Enum):
+ """Whether a callback was triggered at the Dag level or the task level."""
+
+ DAG = "dag"
+ TASK = "task"
+
+ def __str__(self) -> str:
+ return self.value
+
+
+class CallbackMeta(NamedTuple):
+ """
+ Metadata about the callback that is being executed.
+
+ Present in the callback context as ``context["callback"]``.
+
+ ``source`` tells whether the callback was defined on the Dag or on a task.
+ More fields may be added in the future without breaking existing callbacks.
+ """
+
+ source: CallbackSource
+
+
if TYPE_CHECKING:
import jinja2
from pendulum import DateTime
@@ -45,6 +70,7 @@ if TYPE_CHECKING:
class Context(TypedDict, total=False):
"""Jinja2 template context for task rendering."""
+ callback: NotRequired[CallbackMeta]
conn: Any
dag: DAG
dag_run: DagRunProtocol
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 3160f2d01db..caccb6d737c 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -2222,6 +2222,9 @@ def _run_task_state_change_callbacks(
context: Context,
log: Logger,
) -> None:
+ from airflow.sdk.definitions.context import CallbackMeta, CallbackSource
+
+ context["callback"] = CallbackMeta(source=CallbackSource.TASK)
callback: Callable[[Context], None]
for i, callback in enumerate(getattr(task, kind)):
try:
diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
index 387543a0bff..ea1b4afd964 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
@@ -5980,6 +5980,40 @@ class TestTaskRunnerCallsCallbacks:
)
assert callback_data["duration"] >= 0, f"duration should be >= 0, got
{callback_data['duration']}"
+ @pytest.mark.parametrize(
+ ("callback_kind", "expected_state"),
+ [
+ pytest.param("on_success_callback", TaskInstanceState.SUCCESS,
id="success"),
+ pytest.param("on_failure_callback", TaskInstanceState.FAILED,
id="failure"),
+ ],
+ )
+ def test_task_callback_context_has_callback_meta(self, create_runtime_ti,
callback_kind, expected_state):
+ """Task-level callbacks should receive callback.source=TASK in the
context."""
+ from airflow.sdk.definitions.context import CallbackSource
+
+ captured = {}
+
+ def capture_callback(context):
+ captured["callback"] = context.get("callback")
+
+ succeeds = callback_kind == "on_success_callback"
+
+ class CustomOperator(BaseOperator):
+ def execute(self, context):
+ if succeeds:
+ return "ok"
+ raise RuntimeError("boom")
+
+ task = CustomOperator(task_id="cb_task", **{callback_kind:
capture_callback})
+ runtime_ti = create_runtime_ti(dag_id="dag", task=task)
+ log = mock.MagicMock()
+ context = runtime_ti.get_template_context()
+ state, _, error = run(runtime_ti, context, log)
+ finalize(runtime_ti, state, context, log, error)
+
+ assert state == expected_state
+ assert captured["callback"].source is CallbackSource.TASK
+
def test_task_runner_both_callbacks_have_timing_info(self,
create_runtime_ti):
"""Test that both success and failure callbacks receive accurate
timing information."""
import time
