This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new 88801da4987 [v3-4-test] Decode deadline callback kwargs where the
callback runs (#74830) (#75114)
88801da4987 is described below
commit 88801da49878a6f8d93893a0d0c9763d2bb8eabd
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Mon Oct 19 18:46:16 2026 +0200
[v3-4-test] Decode deadline callback kwargs where the callback runs
(#74830) (#75114)
* Decode deadline callback kwargs where the callback runs
Deadline Alert callback kwargs were decoded in the scheduler and the API
server whenever deadline alerts were loaded, although neither component
calls the callback. Keep the kwargs serde-encoded on the server and
decode them where they are used: the triggerer decodes AsyncCallback
kwargs with its existing pipeline, and the callback subprocess decodes
SyncCallback kwargs right before calling the user callable.
Executor callback data now carries a "kwargs_encoding" marker so the
worker only decodes encoded kwargs, callbacks created before upgrade run
unchanged, and an older Task SDK fails such callbacks instead of passing
encoded arguments.
The runtime-provided callback context is serialized strictly, as before;
only the stored callback kwargs keep their encoded form.
Generated-by: Claude Opus 5 following the guidelines at
https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions
* Render deadline callback kwargs readably in metrics and repr
Deadline callback kwargs are now kept in their encoded form, so metric tags
and the callback repr showed the raw encoding instead of the values. Show them
in a readable form again, without importing or building anything they name.
Generated-by: Claude Opus 5 following the guidelines at
https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions
* Keep deadline callback kwargs display total for non-finite intervals
A timedelta payload holding NaN made the display of callback kwargs raise,
which would escape from metric emission and the callback repr. Show such
values, and any other interval that cannot be built, in the generic form
instead.
Generated-by: Claude Opus 5 following the guidelines at
https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions
* Assert that the deadline context never reaches the triggerer, not how it
is refused
The triggerer-path deadline test pinned the AttributeError raised by
strict serialization. Assert the outcome instead (no trigger is created
and nothing in the context is constructed), so the test holds whether
the context is refused by raising or by failing the callback.
Generated-by: Claude Opus 5 following the guidelines at
https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions
* Fail only the affected trigger when its kwargs cannot be decoded
(cherry picked from commit 233f9c3fc78a3cc91dc157f5ace5e1354dce68a7)
Generated-by: Claude Opus 5
Co-authored-by: Jarek Potiuk <[email protected]>
---
airflow-core/newsfragments/74830.significant.rst | 10 ++
.../src/airflow/executors/base_executor.py | 8 +
.../src/airflow/executors/workloads/callback.py | 6 +
.../src/airflow/jobs/triggerer_job_runner.py | 6 +-
airflow-core/src/airflow/models/callback.py | 52 ++++--
airflow-core/src/airflow/models/trigger.py | 22 ++-
airflow-core/src/airflow/serialization/decoders.py | 89 +++++++--
airflow-core/src/airflow/serialization/encoders.py | 68 ++++++-
.../tests/unit/executors/test_local_executor.py | 23 +++
airflow-core/tests/unit/jobs/test_triggerer_job.py | 41 +++++
airflow-core/tests/unit/models/test_callback.py | 200 ++++++++++++++++++++-
airflow-core/tests/unit/models/test_deadline.py | 110 ++++++++++++
.../unit/serialization/test_serialized_objects.py | 86 +++++++--
generated/known_sdk_imports_in_core.txt | 2 +-
.../sdk/execution_time/callback_supervisor.py | 29 +++
.../execution_time/test_callback_supervisor.py | 107 ++++++++++-
16 files changed, 808 insertions(+), 51 deletions(-)
diff --git a/airflow-core/newsfragments/74830.significant.rst
b/airflow-core/newsfragments/74830.significant.rst
new file mode 100644
index 00000000000..d36b7312b3d
--- /dev/null
+++ b/airflow-core/newsfragments/74830.significant.rst
@@ -0,0 +1,10 @@
+Deadline Alert callback kwargs are decoded where the callback runs
+
+The keyword arguments of Deadline Alert callbacks are no longer decoded by the
scheduler. They
+are kept in their serialized form and decoded where the callback runs: by the
triggerer for an
+``AsyncCallback``, and by the callback runner on the worker for a
``SyncCallback``.
+
+Workers that run ``SyncCallback`` Deadline Alerts need a Task SDK that
includes this change, so
+upgrade ``apache-airflow`` and ``apache-airflow-task-sdk`` on workers
together. A worker with
+an upgraded ``apache-airflow`` and an older Task SDK fails such callbacks with
an error instead
+of running them. Callbacks created before the upgrade run as before.
diff --git a/airflow-core/src/airflow/executors/base_executor.py
b/airflow-core/src/airflow/executors/base_executor.py
index 0e86013e0a7..8d7d3b9b76b 100644
--- a/airflow-core/src/airflow/executors/base_executor.py
+++ b/airflow-core/src/airflow/executors/base_executor.py
@@ -912,8 +912,15 @@ class BaseExecutor(LoggingMixin):
sentry_integration=getattr(workload, "sentry_integration", ""),
)
if isinstance(workload, ExecuteCallback):
+ from airflow.executors.workloads.callback import
CALLBACK_KWARGS_ENCODING_KEY
from airflow.sdk.execution_time.callback_supervisor import
supervise_callback
+ # Only passed when set, so callbacks stored without it run on any
Task SDK, while
+ # encoded kwargs fail on a Task SDK that cannot decode them
instead of reaching
+ # the callback encoded.
+ encoding_kwargs = {}
+ if kwargs_encoding :=
workload.callback.data.get(CALLBACK_KWARGS_ENCODING_KEY):
+ encoding_kwargs["callback_kwargs_encoding"] = kwargs_encoding
return supervise_callback(
id=workload.callback.id,
callback_path=workload.callback.data.get("path", ""),
@@ -923,6 +930,7 @@ class BaseExecutor(LoggingMixin):
bundle_info=workload.bundle_info,
token=workload.token,
server=server,
+ **encoding_kwargs,
)
if isinstance(workload, TestConnection):
from airflow.sdk.execution_time.connection_test_supervisor import
supervise_connection_test
diff --git a/airflow-core/src/airflow/executors/workloads/callback.py
b/airflow-core/src/airflow/executors/workloads/callback.py
index 2f4ef1f09a8..80b0fad2cf5 100644
--- a/airflow-core/src/airflow/executors/workloads/callback.py
+++ b/airflow-core/src/airflow/executors/workloads/callback.py
@@ -47,6 +47,12 @@ class CallbackFetchMethod(str, Enum):
IMPORT_PATH = "import_path"
+# Key in ``CallbackDTO.data`` naming how ``data["kwargs"]`` is encoded.
Without it the kwargs are
+# passed to the callback as they are, as for callbacks stored before the key
was introduced.
+CALLBACK_KWARGS_ENCODING_KEY = "kwargs_encoding"
+SERDE_KWARGS_ENCODING = "serde"
+
+
class CallbackDTO(BaseModel):
"""Schema for Callback with minimal required fields needed for Executors
and Task SDK."""
diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
index 326eae953e4..0bcc790224b 100644
--- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
+++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
@@ -1515,8 +1515,10 @@ class TriggerRunner:
else:
trigger_name = f"ID {trigger_id}"
trigger_instance = trigger_class(**deserialised_kwargs)
- except TypeError as err:
- self.log.error("Trigger failed to inflate", error=err)
+ except Exception as err:
+ # Kwargs are decoded here rather than when queued, so one
trigger with undecodable
+ # kwargs (e.g. a class not in allowed_deserialization_classes)
must not stop the runner.
+ self.log.error("Trigger failed to inflate", error=err,
trigger_id=trigger_id)
self.failed_triggers.append((trigger_id, err))
continue
trigger_instance.trigger_id = trigger_id
diff --git a/airflow-core/src/airflow/models/callback.py
b/airflow-core/src/airflow/models/callback.py
index 6295b7dfa7b..360f7b8b7af 100644
--- a/airflow-core/src/airflow/models/callback.py
+++ b/airflow-core/src/airflow/models/callback.py
@@ -30,12 +30,19 @@ from sqlalchemy import ForeignKey, Integer, String, Text,
Uuid
from sqlalchemy.orm import Mapped, mapped_column, relationship
# Re-exporting as _accepts_context for backward compatibility
-from airflow._shared.module_loading import accepts_context as _accepts_context
# noqa: F401
+from airflow._shared.module_loading import (
+ accepts_context as _accepts_context, # noqa: F401
+ qualname,
+)
from airflow._shared.observability.metrics import stats
from airflow._shared.timezones import timezone
from airflow.configuration import conf
from airflow.executors.workloads import BaseWorkload
-from airflow.executors.workloads.callback import CallbackFetchMethod
+from airflow.executors.workloads.callback import (
+ CALLBACK_KWARGS_ENCODING_KEY,
+ SERDE_KWARGS_ENCODING,
+ CallbackFetchMethod,
+)
from airflow.models import Base
from airflow.models.base import StringID
from airflow.models.dagbundle import DagBundleModel
@@ -67,6 +74,12 @@ ACTIVE_STATES = frozenset((CallbackState.PENDING,
CallbackState.QUEUED, Callback
TERMINAL_STATES = frozenset((CallbackState.SUCCESS, CallbackState.FAILED))
+def _format_kwargs_for_display(kwargs: dict | None) -> Any:
+ from airflow.serialization.decoders import format_encoded_value_for_display
+
+ return format_encoded_value_for_display(kwargs) if kwargs else ""
+
+
class CallbackType(str, Enum):
"""
Types of Callbacks.
@@ -167,10 +180,15 @@ class Callback(Base, BaseWorkload):
def get_metric_info(self, status: CallbackState, result: Any, team_name:
str | None = None) -> dict:
tags = {"result": result, **self.data}
tags.pop("prefix", None)
+ tags.pop(CALLBACK_KWARGS_ENCODING_KEY, None)
if "kwargs" in tags:
+ from airflow.serialization.decoders import
format_encoded_value_for_display
+
# Remove the context (if exists) to keep the tags simple
- tags["kwargs"] = {k: v for k, v in tags["kwargs"].items() if k !=
"context"}
+ tags["kwargs"] = format_encoded_value_for_display(
+ {k: v for k, v in tags["kwargs"].items() if k != "context"}
+ )
# Metric backends (statsd, OpenTelemetry) require tag values to be
primitives.
# OTel's aggregation key is built via
``frozenset(attributes.items())``, which
@@ -237,22 +255,25 @@ class TriggererCallback(Callback):
self.data |= callback_def.serialize()
def __repr__(self):
- return f"{self.data['path']}({self.data['kwargs'] or ''}) on a
triggerer"
+ return
f"{self.data['path']}({_format_kwargs_for_display(self.data['kwargs'])}) on a
triggerer"
def queue(self, *, session: Session) -> None:
from airflow.models.trigger import Trigger
+ from airflow.serialization.encoders import
encode_queued_callback_kwargs
from airflow.triggers.callback import CallbackTrigger
team_name: str | None = None
if self.bundle_name and conf.getboolean("core", "multi_team"):
team_name = DagBundleModel.get_team_name(self.bundle_name,
session=session)
- self.trigger = Trigger.from_object(
- CallbackTrigger(
- callback_path=self.data["path"],
- callback_kwargs=self.data["kwargs"],
- queue=self.data.get("queue"),
- )
+ # The kwargs stay encoded: the triggerer decodes them when it creates
the CallbackTrigger.
+ self.trigger = Trigger.from_serialized_kwargs(
+ classpath=qualname(CallbackTrigger),
+ serialized_kwargs={
+ "callback_path": self.data["path"],
+ "callback_kwargs":
encode_queued_callback_kwargs(self.data["kwargs"] or {}),
+ },
+ queue=self.data.get("queue"),
)
self.trigger.team_name = team_name
super().queue(session=session)
@@ -290,12 +311,21 @@ class ExecutorCallback(Callback):
:param fetch_method: Method to fetch the callback at runtime
:param kwargs: Passed to parent Callback.__init__ (see base class for
details)
"""
+ from airflow.serialization.encoders import
encode_stored_callback_kwargs
+
super().__init__(**kwargs)
self.fetch_method = fetch_method
self.data |= callback_def.serialize()
+ # The kwargs are decoded by the callback runner on the worker, which
needs the marker
+ # to tell them apart from the already-decoded kwargs of callbacks
stored before it.
+ self.data["kwargs"] =
encode_stored_callback_kwargs(self.data.get("kwargs") or {})
+ self.data[CALLBACK_KWARGS_ENCODING_KEY] = SERDE_KWARGS_ENCODING
def __repr__(self):
- return f"{self.data['path']}({self.data['kwargs'] or ''}) on
{self.data.get('executor', 'default')} executor"
+ return (
+
f"{self.data['path']}({_format_kwargs_for_display(self.data['kwargs'])}) "
+ f"on {self.data.get('executor', 'default')} executor"
+ )
class DagProcessorCallback(Callback):
diff --git a/airflow-core/src/airflow/models/trigger.py
b/airflow-core/src/airflow/models/trigger.py
index 6a67e261bc3..1713f9aebb7 100644
--- a/airflow-core/src/airflow/models/trigger.py
+++ b/airflow-core/src/airflow/models/trigger.py
@@ -154,12 +154,16 @@ class Trigger(Base):
@staticmethod
def encrypt_kwargs(kwargs: dict[str, Any]) -> str:
"""Encrypt the kwargs of the trigger."""
+ from airflow.sdk.serde import serialize
+
+ return
Trigger._encrypt_serialized_kwargs(serialize(stringify_encoding_keys(kwargs)))
+
+ @staticmethod
+ def _encrypt_serialized_kwargs(serialized_kwargs: dict[str, Any]) -> str:
import json
from airflow.models.crypto import get_fernet
- from airflow.sdk.serde import serialize
- serialized_kwargs = serialize(stringify_encoding_keys(kwargs))
return
get_fernet().encrypt(json.dumps(serialized_kwargs).encode("utf-8")).decode("utf-8")
@staticmethod
@@ -203,6 +207,20 @@ class Trigger(Base):
classpath, kwargs = trigger.serialize()
return cls(classpath=classpath, kwargs=kwargs, queue=trigger.queue)
+ @classmethod
+ def from_serialized_kwargs(
+ cls, classpath: str, serialized_kwargs: dict[str, Any], queue: str |
None = None
+ ) -> Trigger:
+ """
+ Create a trigger row from kwargs that are already in their
serde-encoded form.
+
+ The kwargs are stored exactly as given and decoded by the triggerer
when it creates
+ the trigger, the same as kwargs stored by :meth:`from_object`.
+ """
+ trigger = cls(classpath=classpath, kwargs={}, queue=queue)
+ trigger.encrypted_kwargs =
cls._encrypt_serialized_kwargs(serialized_kwargs)
+ return trigger
+
@classmethod
@provide_session
def bulk_fetch(cls, ids: Iterable[int], *, session: Session = NEW_SESSION)
-> dict[int, Trigger]:
diff --git a/airflow-core/src/airflow/serialization/decoders.py
b/airflow-core/src/airflow/serialization/decoders.py
index 8c42ec0c987..d3f9986762d 100644
--- a/airflow-core/src/airflow/serialization/decoders.py
+++ b/airflow-core/src/airflow/serialization/decoders.py
@@ -19,11 +19,20 @@
from __future__ import annotations
import datetime
+import math
from typing import TYPE_CHECKING, Any, TypeVar
import dateutil.relativedelta
from airflow._shared.module_loading import import_string
+from airflow._shared.serialization import (
+ CLASSNAME,
+ DATA,
+ OLD_DATA,
+ OLD_DICT,
+ OLD_TYPE,
+ OLD_TYPE_TO_FULL_QUALNAME,
+)
from airflow.serialization.definitions.assets import (
SerializedAsset,
SerializedAssetAlias,
@@ -250,20 +259,75 @@ def _decode_deadline_interval(raw_interval: Any) ->
datetime.timedelta | Seriali
)
+def _get_encoded_callback_kwargs(raw_kwargs: Any) -> dict[str, Any]:
+ """
+ Return the encoded callback ``kwargs`` as a mapping, leaving every value
encoded.
+
+ Only the container is normalised: the legacy ``{"__type": "dict", "__var":
...}`` spelling
+ is unwrapped so the result is always a plain mapping of argument name to
encoded value.
+ """
+ from airflow.sdk.serde import _convert
+
+ if not raw_kwargs:
+ return {}
+ kwargs = _convert(raw_kwargs) if isinstance(raw_kwargs, dict) else
raw_kwargs
+ if not isinstance(kwargs, dict) or CLASSNAME in kwargs:
+ raise ValueError("Deadline callback kwargs is not a mapping.")
+ return kwargs
+
+
+_DISPLAY_SEQUENCE_CLASSNAMES = frozenset({"builtins.tuple", "builtins.set",
"builtins.frozenset"})
+
+
+def format_encoded_value_for_display(value: Any) -> Any:
+ """
+ Return a readable form of a serde-encoded value, for logs, metric tags and
reprs.
+
+ Nothing the value names is imported or instantiated: ``timedelta`` and the
built-in
+ collections are shown by value, any other class as ``classname(data)``.
Values that are
+ not encoded are returned unchanged.
+
+ :meta private:
+ """
+ if isinstance(value, list):
+ return [format_encoded_value_for_display(v) for v in value]
+ if not isinstance(value, dict):
+ return value
+ if OLD_TYPE in value and OLD_DATA in value:
+ if value[OLD_TYPE] == OLD_DICT:
+ return format_encoded_value_for_display(value[OLD_DATA])
+ classname = OLD_TYPE_TO_FULL_QUALNAME.get(value[OLD_TYPE],
value[OLD_TYPE])
+ data = value[OLD_DATA]
+ elif CLASSNAME in value:
+ classname, data = value[CLASSNAME], value.get(DATA)
+ else:
+ return {str(k): format_encoded_value_for_display(v) for k, v in
value.items()}
+
+ if (
+ classname == "datetime.timedelta"
+ and isinstance(data, (int, float))
+ and not isinstance(data, bool)
+ and math.isfinite(data)
+ ):
+ try:
+ return str(datetime.timedelta(seconds=data))
+ except (OverflowError, ValueError, TypeError):
+ pass
+ if classname in _DISPLAY_SEQUENCE_CLASSNAMES and isinstance(data, list):
+ return [format_encoded_value_for_display(v) for v in data]
+ return f"{classname}({format_encoded_value_for_display(data)})"
+
+
def _decode_deadline_callback(raw_callback: Any) -> AsyncCallback |
SyncCallback:
"""
Build the callback from its encoded form without importing what the
payload names.
- ``kwargs`` is still passed through generic deserialization, and that is a
deliberate,
- documented limit rather than an oversight. Leaving it encoded would close
a real
- residual -- a legitimate callback can carry an arbitrary allow-listed
class under its
- kwargs, which serde constructs while the outer payload looks entirely
valid -- but the
- kwargs are consumed through two different paths that use two different
encodings
- (``BaseSerialization`` in the triggerer, serde here), and deferring the
decode without
- getting both exactly right silently hands user code an encoded dict in
place of its
- argument. That residual is not specific to deadlines: it is the general
property of
- deserializing Dag-author data, shared with every other serde call site.
Closing it
- belongs with that broader work, not smuggled in here.
+ ``kwargs`` stays in its serde-encoded form. The scheduler and the API
server only route the
+ callback and never call it, so they have no use for the decoded values.
The kwargs are
+ decoded where the callback runs: by the triggerer when it loads the
``CallbackTrigger``
+ kwargs for an ``AsyncCallback``, and by the callback runner on the worker
for a
+ ``SyncCallback``.
:func:`airflow.serialization.encoders.encode_deadline_alert` relies on
+ this to re-encode a decoded alert without serializing the kwargs a second
time.
"""
from airflow.sdk.definitions.callback import (
DEADLINE_CALLBACK_TYPES,
@@ -297,10 +361,7 @@ def _decode_deadline_callback(raw_callback: Any) ->
AsyncCallback | SyncCallback
if not isinstance(path, str):
raise ValueError("Deadline callback payload has no string path.")
- from airflow.sdk.serde import deserialize
-
- raw_kwargs = data.get("kwargs") or {}
- fields: dict[str, Any] = {"kwargs": deserialize(raw_kwargs) if raw_kwargs
else {}}
+ fields: dict[str, Any] = {"kwargs":
_get_encoded_callback_kwargs(data.get("kwargs"))}
for optional in sorted(permitted_fields - {"path", "kwargs"}):
if optional in data:
value = data[optional]
diff --git a/airflow-core/src/airflow/serialization/encoders.py
b/airflow-core/src/airflow/serialization/encoders.py
index d1f801d0b53..2100cedb3b8 100644
--- a/airflow-core/src/airflow/serialization/encoders.py
+++ b/airflow-core/src/airflow/serialization/encoders.py
@@ -26,6 +26,7 @@ import attrs
import pendulum
from airflow._shared.module_loading import qualname
+from airflow._shared.serialization import CLASSNAME, DATA, VERSION
from airflow.partition_mappers.base import PartitionMapper as
CorePartitionMapper
from airflow.partition_mappers.wait_policy import WaitPolicy as CoreWaitPolicy
from airflow.partition_mappers.window import Window as CoreWindow
@@ -256,10 +257,75 @@ def encode_deadline_alert(d: DeadlineAlert |
SerializedDeadlineAlert) -> dict[st
"name": d.name,
"reference": encode_deadline_reference(d.reference),
"interval": serialize(d.interval),
- "callback": serialize(d.callback),
+ "callback": (
+ _encode_decoded_deadline_callback(d.callback)
+ if isinstance(d, SerializedDeadlineAlert)
+ else serialize(d.callback)
+ ),
}
+def _encode_decoded_deadline_callback(callback: Any) -> dict[str, Any]:
+ """
+ Re-encode the callback of a decoded deadline alert.
+
+ A decoded alert keeps the callback ``kwargs`` encoded (see
+ ``airflow.serialization.decoders._decode_deadline_callback``), so they are
written back
+ as they are rather than serialized a second time.
+ """
+ return {
+ CLASSNAME: qualname(callback),
+ VERSION: getattr(type(callback), "__version__", 0),
+ DATA: {
+ field: encode_stored_callback_kwargs(value) if field == "kwargs"
else value
+ for field, value in callback.serialize().items()
+ },
+ }
+
+
+def encode_stored_callback_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
+ """
+ Serialize the stored kwargs of a callback, keeping values that are already
serde-encoded.
+
+ Deadline callback kwargs travel in their encoded form from the serialized
Dag to the
+ component that runs the callback, and only that component decodes them.
Kwargs stored by
+ earlier versions in decoded form are serialized here, so the result is
fully encoded either
+ way. This is only for the kwargs of the callback definition; values
provided at runtime,
+ such as the callback ``context``, are serialized by
:func:`encode_queued_callback_kwargs`.
+
+ :meta private:
+ """
+ from airflow.sdk.serde import serialize
+
+ def _encode(value: Any) -> Any:
+ if isinstance(value, dict):
+ if CLASSNAME in value:
+ return value
+ return {str(k): _encode(v) for k, v in value.items()}
+ if isinstance(value, list):
+ return [_encode(v) for v in value]
+ return serialize(value)
+
+ return {str(k): _encode(v) for k, v in kwargs.items()}
+
+
+def encode_queued_callback_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
+ """
+ Serialize the kwargs of a queued callback: its stored kwargs plus the
runtime ``context``.
+
+ Only the stored kwargs keep their encoded form. The ``context`` is
provided at runtime and
+ is serialized strictly with serde, as any other value.
+
+ :meta private:
+ """
+ from airflow.sdk.serde import serialize
+
+ encoded = encode_stored_callback_kwargs({k: v for k, v in kwargs.items()
if k != "context"})
+ if "context" in kwargs:
+ encoded["context"] = serialize(kwargs["context"])
+ return encoded
+
+
_BUILTIN_DEADLINE_MODULES = (
"airflow.sdk.definitions.deadline",
"airflow.serialization.definitions.deadline",
diff --git a/airflow-core/tests/unit/executors/test_local_executor.py
b/airflow-core/tests/unit/executors/test_local_executor.py
index 0b4916c1390..08d95cf370c 100644
--- a/airflow-core/tests/unit/executors/test_local_executor.py
+++ b/airflow-core/tests/unit/executors/test_local_executor.py
@@ -1138,6 +1138,29 @@ class TestLocalExecutorCallbackSupport:
server=TestLocalExecutorCallbackSupport.TEST_SERVER,
)
+
@mock.patch("airflow.sdk.execution_time.callback_supervisor.supervise_callback",
return_value=0)
+ def
test_execute_workload_passes_kwargs_encoding_to_supervise_callback(self,
mock_supervise_callback):
+ encoded_kwargs = {
+ "interval": {"__classname__": "datetime.timedelta", "__version__":
2, "__data__": 1.0}
+ }
+ callback_data = CallbackDTO(
+ id=self.CALLBACK_UUID,
+ fetch_method=CallbackFetchMethod.IMPORT_PATH,
+ data={"path": "test.module.my_callback", "kwargs": encoded_kwargs,
"kwargs_encoding": "serde"},
+ )
+ callback_workload = workloads.ExecuteCallback(
+ callback=callback_data,
+ dag_rel_path=Path("test.py"),
+ bundle_info=BundleInfo(name="test_bundle", version="1.0"),
+ token="test_token",
+ log_path="test.log",
+ )
+
+ BaseExecutor.run_workload(callback_workload)
+
+ assert mock_supervise_callback.call_args.kwargs["callback_kwargs"] ==
encoded_kwargs
+ assert
mock_supervise_callback.call_args.kwargs["callback_kwargs_encoding"] == "serde"
+
@mock.patch(
"airflow.sdk.execution_time.callback_supervisor.supervise_callback",
side_effect=RuntimeError("Callback subprocess exited with code 1"),
diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py
b/airflow-core/tests/unit/jobs/test_triggerer_job.py
index 0573e0d3d6d..345a5859743 100644
--- a/airflow-core/tests/unit/jobs/test_triggerer_job.py
+++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py
@@ -1876,6 +1876,47 @@ class TestTriggerRunner:
assert isinstance(err, TypeError)
assert "got an unexpected keyword argument 'not_exists_arg'" in
str(err)
+ @pytest.mark.parametrize(
+ "decode_error",
+ [
+ pytest.param(KeyError("__var"), id="key_error"),
+ pytest.param(ImportError("not in
allowed_deserialization_classes"), id="import_error"),
+ ],
+ )
+ @patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs",
autospec=True)
+ @patch(
+
"airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath",
+ return_value=DateTimeTrigger,
+ )
+ @pytest.mark.asyncio
+ async def test_create_triggers_fails_only_trigger_with_undecodable_kwargs(
+ self, mock_get_trigger_by_classpath, mock_decrypt_kwargs,
decode_error, cap_structlog
+ ) -> None:
+ moment = timezone.utcnow() + datetime.timedelta(hours=1)
+
+ def decrypt(encrypted_kwargs):
+ if encrypted_kwargs == "bad":
+ raise decode_error
+ return {"moment": moment}
+
+ mock_decrypt_kwargs.side_effect = decrypt
+
+ trigger_runner = TriggerRunner()
+ trigger_runner.to_create.extend(
+ [
+ workloads.RunTrigger.model_construct(id=1, classpath="abc",
encrypted_kwargs="bad"),
+ workloads.RunTrigger.model_construct(id=2, classpath="abc",
encrypted_kwargs="good"),
+ ]
+ )
+ await trigger_runner.create_triggers()
+
+ assert list(trigger_runner.failed_triggers) == [(1, decode_error)]
+ assert list(trigger_runner.triggers) == [2]
+ assert {"event": "Trigger failed to inflate", "trigger_id": 1,
"error": decode_error} in cap_structlog
+
+ trigger_runner.triggers[2]["task"].cancel()
+ await trigger_runner.cleanup_finished_triggers()
+
@pytest.mark.asyncio
@patch("airflow.sdk.execution_time.task_runner.SUPERVISOR_COMMS",
create=True)
async def test_invalid_trigger(self, supervisor_builder):
diff --git a/airflow-core/tests/unit/models/test_callback.py
b/airflow-core/tests/unit/models/test_callback.py
index 27639f8f88f..4622616d04b 100644
--- a/airflow-core/tests/unit/models/test_callback.py
+++ b/airflow-core/tests/unit/models/test_callback.py
@@ -16,13 +16,19 @@
# under the License.
from __future__ import annotations
+import json
+from datetime import timedelta
+from pathlib import Path
from unittest.mock import patch
import pytest
+import structlog
from sqlalchemy import select
from airflow._shared.module_loading import accepts_context
+from airflow._shared.timezones import timezone
from airflow.callbacks.callback_requests import DagCallbackRequest
+from airflow.executors.workloads.callback import CALLBACK_KWARGS_ENCODING_KEY,
CallbackDTO
from airflow.models import Trigger
from airflow.models.callback import (
Callback,
@@ -32,9 +38,14 @@ from airflow.models.callback import (
ExecutorCallback,
TriggererCallback,
)
+from airflow.sdk import serde
from airflow.sdk.definitions.callback import AsyncCallback, SyncCallback
+from airflow.sdk.definitions.deadline import DeadlineAlert, DeadlineReference
+from airflow.sdk.execution_time.callback_supervisor import execute_callback
+from airflow.serialization.decoders import decode_deadline_alert,
smart_decode_trigger_kwargs
+from airflow.serialization.encoders import encode_deadline_alert
from airflow.triggers.base import TriggerEvent
-from airflow.triggers.callback import PAYLOAD_BODY_KEY, PAYLOAD_STATUS_KEY
+from airflow.triggers.callback import PAYLOAD_BODY_KEY, PAYLOAD_STATUS_KEY,
CallbackTrigger
from airflow.utils.session import create_session
from tests_common.test_utils.config import conf_vars
@@ -59,6 +70,71 @@ TEST_SYNC_CALLBACK = SyncCallback(sync_callback,
kwargs=TEST_CALLBACK_KWARGS)
TEST_DAG_ID = "test_dag"
+class _TrackedKwarg:
+ """A class named in callback kwargs, to tell where the kwargs are
decoded."""
+
+ instantiated = False
+
+ def __init__(self):
+ type(self).instantiated = True
+
+ def serialize(self):
+ return {}
+
+ @staticmethod
+ def deserialize(data, version):
+ return _TrackedKwarg()
+
+
+_TRACKED_KWARG_NAME =
f"{_TrackedKwarg.__module__}.{_TrackedKwarg.__qualname__}"
+
+ENCODED_CALLBACK_KWARGS = {
+ "current": {
+ "interval": {serde.CLASSNAME: "datetime.timedelta", serde.VERSION: 2,
serde.DATA: 3600.0},
+ "pair": {serde.CLASSNAME: "builtins.tuple", serde.VERSION: 1,
serde.DATA: [1, 2]},
+ "plain": {"nested": [1, "two"]},
+ "tracked": {serde.CLASSNAME: _TRACKED_KWARG_NAME, serde.VERSION: 0,
serde.DATA: {}},
+ },
+ "legacy": {
+ "interval": {"__type": "timedelta", "__var": 3600.0},
+ "pair": {"__type": "tuple", "__var": [1, 2]},
+ "plain": {"nested": [1, "two"]},
+ "tracked": {"__type": _TRACKED_KWARG_NAME, "__var": {}},
+ },
+}
+
+received_sync_kwargs: dict = {}
+
+
+def recording_sync_callback(**kwargs):
+ """Sync callable which records the kwargs it was called with."""
+ received_sync_kwargs.clear()
+ received_sync_kwargs.update(kwargs)
+
+
+def _decode_callback_def(callback_def, encoded_kwargs):
+ """Return ``callback_def`` as the scheduler sees it after decoding the
stored deadline alert."""
+ encoded = encode_deadline_alert(
+ DeadlineAlert(
+ reference=DeadlineReference.DAGRUN_QUEUED_AT,
interval=timedelta(hours=1), callback=callback_def
+ )
+ )
+ encoded["callback"][serde.DATA]["kwargs"] = encoded_kwargs
+ return decode_deadline_alert(encoded).callback
+
+
+def _assert_decoded(kwargs):
+ assert isinstance(kwargs.pop("tracked"), _TrackedKwarg)
+ assert kwargs == {"interval": timedelta(hours=1), "pair": (1, 2), "plain":
{"nested": [1, "two"]}}
+
+
[email protected]
+def allow_tracked_kwarg():
+ _TrackedKwarg.instantiated = False
+ with patch.object(serde, "_extra_allowed", serde._extra_allowed |
{_TRACKED_KWARG_NAME}):
+ yield
+
+
@pytest.fixture
def session():
"""Fixture that provides a SQLAlchemy session"""
@@ -117,6 +193,54 @@ class TestCallback:
"dag_id": TEST_DAG_ID,
}
+ @pytest.mark.parametrize(
+ "encoded_interval",
+ [
+ pytest.param(
+ {serde.CLASSNAME: "datetime.timedelta", serde.VERSION: 2,
serde.DATA: 3600.0}, id="current"
+ ),
+ pytest.param({"__type": "timedelta", "__var": 3600.0},
id="legacy"),
+ ],
+ )
+ @pytest.mark.usefixtures("allow_tracked_kwarg")
+ def test_metric_tags_and_repr_show_kwargs_readably(self, encoded_interval):
+ kwargs = {
+ "interval": encoded_interval,
+ "pair": {serde.CLASSNAME: "builtins.tuple", serde.VERSION: 1,
serde.DATA: [1, 2]},
+ "tracked": {serde.CLASSNAME: _TRACKED_KWARG_NAME, serde.VERSION:
0, serde.DATA: {}},
+ }
+ expected = {"interval": "1:00:00", "pair": [1, 2], "tracked":
f"{_TRACKED_KWARG_NAME}({{}})"}
+ executor_callback = ExecutorCallback(
+ SyncCallback(sync_callback, kwargs=kwargs),
fetch_method=CallbackFetchMethod.IMPORT_PATH
+ )
+ triggerer_callback = TriggererCallback(AsyncCallback(async_callback,
kwargs=kwargs))
+
+ tags = executor_callback.get_metric_info(CallbackState.SUCCESS,
"0")["tags"]
+
+ assert tags["kwargs"] == json.dumps(expected, sort_keys=True)
+ assert repr(executor_callback) ==
f"{executor_callback.data['path']}({expected}) on None executor"
+ assert repr(triggerer_callback) ==
f"{triggerer_callback.data['path']}({expected}) on a triggerer"
+ assert _TrackedKwarg.instantiated is False
+
+ @pytest.mark.parametrize(
+ ("data", "expected"),
+ [
+ pytest.param(float("nan"), "datetime.timedelta(nan)", id="nan"),
+ pytest.param(float("inf"), "datetime.timedelta(inf)", id="inf"),
+ pytest.param(float("-inf"), "datetime.timedelta(-inf)", id="-inf"),
+ pytest.param(1e300, "datetime.timedelta(1e+300)", id="too-large"),
+ pytest.param("an hour", "datetime.timedelta(an hour)",
id="not-a-number"),
+ ],
+ )
+ def test_metric_tags_and_repr_fall_back_for_unusable_intervals(self, data,
expected):
+ kwargs = {"interval": {serde.CLASSNAME: "datetime.timedelta",
serde.VERSION: 2, serde.DATA: data}}
+ callback = TriggererCallback(AsyncCallback(async_callback,
kwargs=kwargs))
+
+ tags = callback.get_metric_info(CallbackState.SUCCESS, "0")["tags"]
+
+ assert tags["kwargs"] == json.dumps({"interval": expected})
+ assert repr(callback) == f"{callback.data['path']}({ {'interval':
expected} }) on a triggerer"
+
def test_get_metric_info_includes_queue(self):
"""``queue`` is stored in ``self.data`` for a queued AsyncCallback, so
it flows into metric tags."""
queued_callback = AsyncCallback(async_callback,
kwargs=TEST_CALLBACK_KWARGS, queue="custom-queue")
@@ -197,6 +321,33 @@ class TestTriggererCallback:
assert callback.trigger.queue is None
assert callback.state == CallbackState.QUEUED
+ @pytest.mark.usefixtures("allow_tracked_kwarg")
+ @pytest.mark.parametrize("encoding", ["current", "legacy"])
+ def test_queue_leaves_kwargs_for_the_triggerer_to_decode(self, session,
encoding):
+ callback_def = _decode_callback_def(
+ AsyncCallback(async_callback, kwargs=TEST_CALLBACK_KWARGS),
ENCODED_CALLBACK_KWARGS[encoding]
+ )
+ callback = TriggererCallback(callback_def)
+ session.add(callback)
+ session.commit()
+ session.expunge_all()
+ callback = session.get(Callback, callback.id)
+ deadline_time = timezone.datetime(2025, 1, 1)
+ context = {"deadline": {"id": "1", "deadline_time": deadline_time}}
+ callback.data = {**callback.data, "kwargs":
{**callback.data["kwargs"], "context": context}}
+
+ callback.queue(session=session)
+
+ assert _TrackedKwarg.instantiated is False
+
+ # Decode the way the triggerer does before it creates the trigger.
+ trigger_kwargs =
Trigger._decrypt_kwargs(callback.trigger.encrypted_kwargs)
+ trigger = CallbackTrigger(**{k: smart_decode_trigger_kwargs(v) for k,
v in trigger_kwargs.items()})
+
+ assert trigger.callback_path == callback_def.path
+ assert trigger.callback_kwargs.pop("context") == context
+ _assert_decoded(trigger.callback_kwargs)
+
def test_queue_populates_trigger_queue(self, session):
queued_callback = AsyncCallback(async_callback,
kwargs=TEST_CALLBACK_KWARGS, queue="custom-queue")
callback = TriggererCallback(queued_callback)
@@ -323,7 +474,7 @@ class TestExecutorCallback:
retrieved = session.scalar(select(Callback).where(Callback.id ==
callback.id))
assert isinstance(retrieved, ExecutorCallback)
assert retrieved.fetch_method == CallbackFetchMethod.IMPORT_PATH
- assert retrieved.data == TEST_SYNC_CALLBACK.serialize()
+ assert retrieved.data == {**TEST_SYNC_CALLBACK.serialize(),
CALLBACK_KWARGS_ENCODING_KEY: "serde"}
assert retrieved.state == CallbackState.SCHEDULED.value
assert retrieved.output is None
assert retrieved.priority_weight == 1
@@ -337,6 +488,51 @@ class TestExecutorCallback:
callback.queue(session=session)
assert callback.state == CallbackState.QUEUED
+ @pytest.mark.usefixtures("allow_tracked_kwarg")
+ @pytest.mark.parametrize("encoding", ["current", "legacy"])
+ def test_workload_leaves_kwargs_for_the_callback_runner_to_decode(self,
session, encoding):
+ callback_def = _decode_callback_def(
+ SyncCallback(recording_sync_callback, kwargs=TEST_CALLBACK_KWARGS),
+ ENCODED_CALLBACK_KWARGS[encoding],
+ )
+ callback = ExecutorCallback(callback_def,
fetch_method=CallbackFetchMethod.IMPORT_PATH)
+ session.add(callback)
+ session.commit()
+ session.expunge_all()
+ callback = session.get(Callback, callback.id)
+
+ workload_callback = CallbackDTO.model_validate_json(
+ CallbackDTO.model_validate(callback,
from_attributes=True).model_dump_json()
+ )
+
+ assert _TrackedKwarg.instantiated is False
+ assert workload_callback.data["kwargs"] ==
ENCODED_CALLBACK_KWARGS[encoding]
+ assert workload_callback.data[CALLBACK_KWARGS_ENCODING_KEY] == "serde"
+
+ success, error = execute_callback(
+ callback_path=workload_callback.data["path"],
+ callback_kwargs=workload_callback.data["kwargs"],
+ dag_rel_path=Path("test.py"),
+ bundle_path=None,
+ log=structlog.get_logger(),
+
callback_kwargs_encoding=workload_callback.data[CALLBACK_KWARGS_ENCODING_KEY],
+ )
+
+ assert (success, error) == (True, None)
+ _assert_decoded(received_sync_kwargs)
+
+ def test_kwargs_are_stored_encoded_with_their_encoding(self):
+ callback = ExecutorCallback(
+ SyncCallback(sync_callback, kwargs={"interval":
timedelta(hours=1), "arg1": "value1"}),
+ fetch_method=CallbackFetchMethod.IMPORT_PATH,
+ )
+
+ assert callback.data["kwargs"] == serde.serialize({"interval":
timedelta(hours=1), "arg1": "value1"})
+ assert callback.data[CALLBACK_KWARGS_ENCODING_KEY] == "serde"
+ assert (
+ CALLBACK_KWARGS_ENCODING_KEY not in
callback.get_metric_info(CallbackState.SUCCESS, "ok")["tags"]
+ )
+
def test_session_get_requires_uuid_not_str(self, session):
"""Filtering the UUID id column with a plain str breaks on SQLite, so
callers must wrap with ``UUID(...)`` before querying."""
diff --git a/airflow-core/tests/unit/models/test_deadline.py
b/airflow-core/tests/unit/models/test_deadline.py
index fa009620f0b..3848a89c7c9 100644
--- a/airflow-core/tests/unit/models/test_deadline.py
+++ b/airflow-core/tests/unit/models/test_deadline.py
@@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations
+import contextlib
import re
from dataclasses import dataclass
from datetime import datetime, timedelta
@@ -63,6 +64,47 @@ async def callback_for_deadline():
TEST_CALLBACK_PATH = f"{__name__}.{callback_for_deadline.__name__}"
+
+
+class _ContextValue:
+ """A class named by a value in the Dag run conf, to tell whether that
value is rebuilt."""
+
+ instantiated = False
+
+ def __init__(self):
+ type(self).instantiated = True
+
+ def serialize(self):
+ return {}
+
+ @staticmethod
+ def deserialize(data, version):
+ return _ContextValue()
+
+
+_CONTEXT_VALUE_NAME =
f"{_ContextValue.__module__}.{_ContextValue.__qualname__}"
+_CLASSNAME_SHAPED_CONF = {"value": {"__classname__": _CONTEXT_VALUE_NAME,
"__version__": 0, "__data__": {}}}
+
+received_sync_kwargs: dict = {}
+
+
+def recording_sync_callback(**kwargs):
+ """Sync callable which records the kwargs it was called with."""
+ received_sync_kwargs.clear()
+ received_sync_kwargs.update(kwargs)
+
+
[email protected]
+def classname_shaped_conf(dagrun, session):
+ from airflow.sdk import serde
+
+ _ContextValue.instantiated = False
+ dagrun.conf = _CLASSNAME_SHAPED_CONF
+ session.flush()
+ with mock.patch.object(serde, "_extra_allowed", serde._extra_allowed |
{_CONTEXT_VALUE_NAME}):
+ yield
+
+
TEST_CALLBACK_KWARGS = {"arg1": "value1"}
TEST_ASYNC_CALLBACK = AsyncCallback(TEST_CALLBACK_PATH,
kwargs=TEST_CALLBACK_KWARGS)
TEST_SYNC_CALLBACK = SyncCallback(TEST_CALLBACK_PATH,
kwargs=TEST_CALLBACK_KWARGS)
@@ -377,6 +419,74 @@ class TestDeadline:
assert callback.data["dag_id"] == dag_id
assert callback.data["deadline_id"] == str(deadline_id)
+ @pytest.mark.db_test
+ @pytest.mark.usefixtures("classname_shaped_conf")
+ def test_handle_miss_never_rebuilds_runtime_context_for_triggerer(self,
dagrun, session):
+ """A conf value shaped like an encoded object never reaches the
triggerer to be rebuilt."""
+ from airflow.models.trigger import Trigger
+
+ deadline_orm = Deadline(
+ deadline_time=DEFAULT_DATE,
+ callback=AsyncCallback(TEST_CALLBACK_PATH, TEST_CALLBACK_KWARGS),
+ dagrun_id=dagrun.id,
+ dag_id=dagrun.dag_id,
+ deadline_alert_id=None,
+ )
+ session.add(deadline_orm)
+ session.flush()
+
+ # Strict serialization refuses the context, either by raising here or
by failing the
+ # callback without queueing it; both keep the value away from the
triggerer.
+ with contextlib.suppress(AttributeError):
+ deadline_orm.handle_miss(session)
+
+ assert deadline_orm.callback.trigger is None
+ assert session.scalars(select(Trigger)).all() == []
+ assert _ContextValue.instantiated is False
+
+ @pytest.mark.db_test
+ @pytest.mark.usefixtures("classname_shaped_conf")
+ def
test_handle_miss_passes_runtime_context_as_is_to_executor_callback(self,
dagrun, session):
+ from pathlib import Path
+
+ import structlog
+
+ from airflow.executors.workloads.callback import
CALLBACK_KWARGS_ENCODING_KEY, CallbackDTO
+ from airflow.sdk.execution_time.callback_supervisor import
execute_callback
+
+ deadline_orm = Deadline(
+ deadline_time=DEFAULT_DATE,
+ callback=SyncCallback(f"{__name__}.recording_sync_callback",
TEST_CALLBACK_KWARGS),
+ dagrun_id=dagrun.id,
+ dag_id=dagrun.dag_id,
+ deadline_alert_id=None,
+ )
+ session.add(deadline_orm)
+ session.flush()
+ deadline_id = deadline_orm.id
+
+ deadline_orm.handle_miss(session)
+ session.commit()
+ session.expunge_all()
+ callback = session.scalar(select(Deadline).where(Deadline.id ==
deadline_id)).callback
+ workload_callback = CallbackDTO.model_validate_json(
+ CallbackDTO.model_validate(callback,
from_attributes=True).model_dump_json()
+ )
+
+ success, error = execute_callback(
+ callback_path=workload_callback.data["path"],
+ callback_kwargs=workload_callback.data["kwargs"],
+ dag_rel_path=Path("test.py"),
+ bundle_path=None,
+ log=structlog.get_logger(),
+
callback_kwargs_encoding=workload_callback.data[CALLBACK_KWARGS_ENCODING_KEY],
+ )
+
+ assert (success, error) == (True, None)
+ assert _ContextValue.instantiated is False
+ assert received_sync_kwargs["context"]["dag_run"]["conf"] ==
_CLASSNAME_SHAPED_CONF
+ assert received_sync_kwargs["arg1"] == "value1"
+
@pytest.mark.db_test
class TestCalculatedDeadlineReferences:
diff --git a/airflow-core/tests/unit/serialization/test_serialized_objects.py
b/airflow-core/tests/unit/serialization/test_serialized_objects.py
index f1d44c211f0..86dcf98d80d 100644
--- a/airflow-core/tests/unit/serialization/test_serialized_objects.py
+++ b/airflow-core/tests/unit/serialization/test_serialized_objects.py
@@ -1703,25 +1703,22 @@ def
test_deadline_fields_refuse_unexpected_classes_without_constructing_them(fie
assert _DeadlineGadget.instantiated is False, "the class was constructed
before being rejected"
-def test_deadline_callback_kwargs_still_construct_nested_classes():
- """Documents a known residual: kwargs are still generically deserialized.
-
- The outer payload is a valid AsyncCallback and the class sits under
``kwargs``,
- which serde deserializes recursively. This is NOT closed here, and the
test asserts
- the current behaviour rather than the desired one so the gap is visible
and a future
- change has to update it deliberately.
-
- It is not specific to deadlines -- it is the general property of
deserializing
- Dag-author data, shared with every other serde call site. Closing it means
deferring
- the kwargs decode to the process that runs the callback, which spans two
consumption
- paths using two different encodings and belongs with that broader work.
[email protected]("encoding", ["current", "legacy"])
+def test_deadline_callback_kwargs_are_kept_encoded(encoding):
+ """Callback kwargs are not decoded with the alert: classes named there are
not constructed.
+
+ The kwargs are decoded where the callback runs, so the decoded alert holds
them exactly as
+ they were stored, and encoding the alert again writes them back unchanged.
"""
from unittest import mock
from airflow.sdk import serde
+ from airflow.sdk._shared.serialization import OLD_DATA, OLD_TYPE
gadget_name =
f"{_DeadlineGadget.__module__}.{_DeadlineGadget.__qualname__}"
_DeadlineGadget.instantiated = False
+ nested = _encode_as(gadget_name) if encoding == "current" else {OLD_TYPE:
gadget_name, OLD_DATA: {}}
+ encoded_kwargs = {"nested": nested, "items": [nested], "plain": "value"}
valid = DeadlineAlert(
reference=DeadlineReference.DAGRUN_QUEUED_AT,
@@ -1730,15 +1727,72 @@ def
test_deadline_callback_kwargs_still_construct_nested_classes():
)
serialized = BaseSerialization.serialize(valid)
callback_payload = serialized[Encoding.VAR][DeadlineAlertFields.CALLBACK]
- callback_payload[serde.DATA]["kwargs"] = {"evil": _encode_as(gadget_name)}
+ callback_payload[serde.DATA]["kwargs"] = encoded_kwargs
with mock.patch.object(serde, "_extra_allowed", serde._extra_allowed |
{gadget_name}):
+ decoded = BaseSerialization.deserialize(serialized)
+ reencoded = BaseSerialization.serialize(decoded)
+
+ assert _DeadlineGadget.instantiated is False
+ assert isinstance(decoded.callback, AsyncCallback)
+ assert decoded.callback.kwargs == encoded_kwargs
+ assert reencoded[Encoding.VAR][DeadlineAlertFields.CALLBACK] ==
callback_payload
+
+
+def test_deadline_callback_kwargs_legacy_container_is_unwrapped():
+ from airflow.sdk import serde
+ from airflow.sdk._shared.serialization import OLD_DATA, OLD_DICT, OLD_TYPE
+
+ valid = DeadlineAlert(
+ reference=DeadlineReference.DAGRUN_QUEUED_AT,
+ interval=timedelta(hours=1),
+ callback=AsyncCallback(TEST_CALLBACK_PATH,
kwargs=TEST_CALLBACK_KWARGS),
+ )
+ serialized = BaseSerialization.serialize(valid)
+ callback_payload = serialized[Encoding.VAR][DeadlineAlertFields.CALLBACK]
+ callback_payload[serde.DATA]["kwargs"] = {OLD_TYPE: OLD_DICT, OLD_DATA:
TEST_CALLBACK_KWARGS}
+
+ decoded = BaseSerialization.deserialize(serialized)
+
+ assert decoded.callback.kwargs == TEST_CALLBACK_KWARGS
+
+
[email protected]("kwargs", [["not", "a", "mapping"],
_encode_as("builtins.dict")])
+def test_deadline_callback_kwargs_must_be_a_mapping(kwargs):
+ from airflow.sdk import serde
+
+ valid = DeadlineAlert(
+ reference=DeadlineReference.DAGRUN_QUEUED_AT,
+ interval=timedelta(hours=1),
+ callback=AsyncCallback(TEST_CALLBACK_PATH,
kwargs=TEST_CALLBACK_KWARGS),
+ )
+ serialized = BaseSerialization.serialize(valid)
+
serialized[Encoding.VAR][DeadlineAlertFields.CALLBACK][serde.DATA]["kwargs"] =
kwargs
+
+ with pytest.raises(ValueError, match="kwargs is not a mapping"):
BaseSerialization.deserialize(serialized)
- assert _DeadlineGadget.instantiated is True, (
- "kwargs deserialization behaviour changed -- if this is now closed,
update the "
- "docstring on _decode_deadline_callback and this test together"
+
+def test_deadline_alert_with_live_callback_kwargs_round_trips():
+ """An alert written by the Dag processor keeps its kwargs encoded through
decode and re-encode."""
+ from airflow.sdk import serde
+
+ original = DeadlineAlert(
+ reference=DeadlineReference.DAGRUN_QUEUED_AT,
+ interval=timedelta(hours=1),
+ callback=AsyncCallback(
+ TEST_CALLBACK_PATH, kwargs={"interval": timedelta(minutes=5),
"pair": (1, 2), "plain": {"a": 1}}
+ ),
+ )
+ serialized = BaseSerialization.serialize(original)
+
+ decoded = BaseSerialization.deserialize(serialized)
+
+ assert (
+ decoded.callback.kwargs
+ ==
serialized[Encoding.VAR][DeadlineAlertFields.CALLBACK][serde.DATA]["kwargs"]
)
+ assert BaseSerialization.serialize(decoded) == serialized
@pytest.mark.parametrize("field", [DeadlineAlertFields.CALLBACK,
DeadlineAlertFields.INTERVAL])
diff --git a/generated/known_sdk_imports_in_core.txt
b/generated/known_sdk_imports_in_core.txt
index b9a217df8b6..8d103e6886c 100644
--- a/generated/known_sdk_imports_in_core.txt
+++ b/generated/known_sdk_imports_in_core.txt
@@ -27,7 +27,7 @@ airflow-core/src/airflow/serialization/decoders.py::3
airflow-core/src/airflow/serialization/definitions/baseoperator.py::1
airflow-core/src/airflow/serialization/definitions/dag.py::1
airflow-core/src/airflow/serialization/definitions/mappedoperator.py::5
-airflow-core/src/airflow/serialization/encoders.py::11
+airflow-core/src/airflow/serialization/encoders.py::13
airflow-core/src/airflow/serialization/serialized_objects.py::17
airflow-core/src/airflow/serialization/stub_arg_bindings.py::5
airflow-core/src/airflow/settings.py::1
diff --git a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py
index 01d1ecf7f38..bc38e2abb20 100644
--- a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py
@@ -86,12 +86,34 @@ CallbackToSupervisor = Annotated[
]
+SERDE_KWARGS_ENCODING = "serde"
+
+
+def decode_callback_kwargs(callback_kwargs: dict[str, Any], encoding: str |
None) -> dict[str, Any]:
+ """
+ Decode the kwargs of a callback, right before the callback is called.
+
+ ``encoding`` is ``None`` for kwargs that are passed as they are (callbacks
stored by
+ earlier versions), or ``"serde"`` for serde-encoded kwargs. ``context`` is
added by the
+ scheduler and is never encoded.
+ """
+ if encoding is None:
+ return callback_kwargs
+ if encoding != SERDE_KWARGS_ENCODING:
+ raise ValueError(f"Unsupported callback kwargs encoding: {encoding!r}")
+
+ from airflow.sdk.serde import deserialize
+
+ return {key: value if key == "context" else deserialize(value) for key,
value in callback_kwargs.items()}
+
+
def execute_callback(
callback_path: str,
callback_kwargs: dict,
dag_rel_path: os.PathLike[str],
bundle_path: os.PathLike[str] | None,
log,
+ callback_kwargs_encoding: str | None = None,
) -> tuple[bool, str | None]:
"""
Execute a callback function by importing and calling it, returning the
success state.
@@ -112,6 +134,7 @@ def execute_callback(
:param dag_rel_path: Relative path to the DAG file.
:param bundle_path: Path to the bundle file.
:param log: Logger instance for recording execution.
+ :param callback_kwargs_encoding: How ``callback_kwargs`` is encoded, see
:func:`decode_callback_kwargs`.
:return: Tuple of (success: bool, error_message: str | None)
"""
if not callback_path:
@@ -139,6 +162,7 @@ def execute_callback(
else:
module = import_module(module_path)
callback_callable = getattr(module, function_name)
+ callback_kwargs = decode_callback_kwargs(callback_kwargs,
callback_kwargs_encoding)
log.debug("Executing callback", callback_path=callback_path,
callback_kwargs=callback_kwargs)
@@ -201,6 +225,7 @@ class CallbackSubprocess(WatchedSubprocess):
id: str,
callback_path: str,
callback_kwargs: dict,
+ callback_kwargs_encoding: str | None = None,
dag_rel_path: os.PathLike[str],
bundle_info: _BundleInfoLike | None = None,
client: Client,
@@ -251,6 +276,7 @@ class CallbackSubprocess(WatchedSubprocess):
dag_rel_path=dag_rel_path,
bundle_path=bundle_path,
log=_log,
+ callback_kwargs_encoding=callback_kwargs_encoding,
)
if not success:
_log.error("Callback failed", error=error_msg)
@@ -403,6 +429,7 @@ def supervise_callback(
callback_path: str,
callback_kwargs: dict,
dag_rel_path: os.PathLike[str],
+ callback_kwargs_encoding: str | None = None,
log_path: str | None = None,
bundle_info: _BundleInfoLike | None = None,
token: str = "",
@@ -416,6 +443,7 @@ def supervise_callback(
:param callback_path: Dot-separated import path to the callback function
or class.
:param callback_kwargs: Keyword arguments to pass to the callback.
:param dag_rel_path: Relative path to the DAG file.
+ :param callback_kwargs_encoding: How ``callback_kwargs`` is encoded, see
:func:`decode_callback_kwargs`.
:param log_path: Path to write logs, if required.
:param bundle_info: When provided, the bundle's path is added to sys.path
so callbacks in Dag Bundles are importable.
:param token: Authentication token for the API client.
@@ -444,6 +472,7 @@ def supervise_callback(
id=id,
callback_path=callback_path,
callback_kwargs=callback_kwargs,
+ callback_kwargs_encoding=callback_kwargs_encoding,
dag_rel_path=dag_rel_path,
bundle_info=bundle_info,
client=client,
diff --git a/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py
b/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py
index fae31290be5..7eb5bc9c949 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py
@@ -23,6 +23,7 @@ import signal
import socket
import uuid
from dataclasses import dataclass
+from datetime import timedelta
from operator import attrgetter
from typing import Any
from unittest.mock import ANY, Mock, patch
@@ -66,6 +67,15 @@ def callback_that_raises():
raise ValueError("something went wrong")
+received_kwargs: dict[str, Any] = {}
+
+
+def callback_recording_kwargs(**kwargs):
+ """A callback that records the kwargs it was called with."""
+ received_kwargs.clear()
+ received_kwargs.update(kwargs)
+
+
class CallableClass:
"""A class that returns a callable instance (like BaseNotifier)."""
@@ -161,6 +171,77 @@ class TestExecuteCallback:
else:
assert error is None
+ @pytest.mark.parametrize(
+ "encoded_kwargs",
+ [
+ pytest.param(
+ {
+ "interval": {"__classname__": "datetime.timedelta",
"__version__": 2, "__data__": 3600.0},
+ "pair": {"__classname__": "builtins.tuple", "__version__":
1, "__data__": [1, 2]},
+ "plain": {"nested": [1, "two"]},
+ },
+ id="current",
+ ),
+ pytest.param(
+ {
+ "interval": {"__type": "timedelta", "__var": 3600.0},
+ "pair": {"__type": "tuple", "__var": [1, 2]},
+ "plain": {"nested": [1, "two"]},
+ },
+ id="legacy",
+ ),
+ ],
+ )
+ def test_execute_callback_decodes_serde_encoded_kwargs(self,
encoded_kwargs):
+ context = {"dag_run": {"dag_id": "test_dag"}, "deadline": {"id": "1"}}
+
+ success, error = execute_callback(
+ callback_path=f"{__name__}.callback_recording_kwargs",
+ callback_kwargs={**encoded_kwargs, "context": context},
+ dag_rel_path=Path("test.py"),
+ bundle_path=Path("bundle/path"),
+ log=structlog.get_logger(),
+ callback_kwargs_encoding="serde",
+ )
+
+ assert (success, error) == (True, None)
+ assert received_kwargs == {
+ "interval": timedelta(hours=1),
+ "pair": (1, 2),
+ "plain": {"nested": [1, "two"]},
+ "context": context,
+ }
+
+ def test_execute_callback_passes_kwargs_as_is_without_encoding(self):
+ kwargs = {
+ "interval": {"__classname__": "datetime.timedelta", "__version__":
2, "__data__": 3600.0},
+ "plain": {"nested": [1, "two"]},
+ }
+
+ success, error = execute_callback(
+ callback_path=f"{__name__}.callback_recording_kwargs",
+ callback_kwargs=kwargs,
+ dag_rel_path=Path("test.py"),
+ bundle_path=Path("bundle/path"),
+ log=structlog.get_logger(),
+ )
+
+ assert (success, error) == (True, None)
+ assert received_kwargs == kwargs
+
+ def test_execute_callback_fails_on_unsupported_kwargs_encoding(self):
+ success, error = execute_callback(
+ callback_path=f"{__name__}.callback_recording_kwargs",
+ callback_kwargs={"arg1": "value1"},
+ dag_rel_path=Path("test.py"),
+ bundle_path=Path("bundle/path"),
+ log=structlog.get_logger(),
+ callback_kwargs_encoding="unknown",
+ )
+
+ assert success is False
+ assert "Unsupported callback kwargs encoding" in error
+
class TestCallbackHandleRequest:
"""Verify that CallbackSubprocess._handle_request dispatches each message
type to the correct handler."""
@@ -485,9 +566,10 @@ class TestCallbackSubprocessStart:
"bundle_path": bundle_path,
}
- def test_execute_callback_receives_correct_parameters(self,
base_start_kwargs):
+ @pytest.mark.parametrize("callback_kwargs_encoding", [None, "serde"])
+ def test_execute_callback_receives_correct_parameters(self,
base_start_kwargs, callback_kwargs_encoding):
"""Test that execute_callback receives the correct parameters."""
- CallbackSubprocess.start(**base_start_kwargs)
+ CallbackSubprocess.start(**base_start_kwargs,
callback_kwargs_encoding=callback_kwargs_encoding)
self.mock_super_start.call_args.kwargs["target"]()
self.mock_super_start.assert_called_with(
@@ -500,6 +582,7 @@ class TestCallbackSubprocessStart:
callback_path=base_start_kwargs["callback_path"],
dag_rel_path=base_start_kwargs["dag_rel_path"],
log=ANY,
+ callback_kwargs_encoding=callback_kwargs_encoding,
)
def test_execute_callback_with_bundle_info_should_pass_correct_parameters(
@@ -529,6 +612,7 @@ class TestCallbackSubprocessStart:
callback_path=adjusted_kwargs["callback_path"],
dag_rel_path=adjusted_kwargs["dag_rel_path"],
log=ANY,
+ callback_kwargs_encoding=None,
)
def test_callback_supervisor_with_bundle_info_should_adjust_sys_path(
@@ -622,3 +706,22 @@ class TestSuperviseCallbackExchangesTokenFirst:
self._supervise(client)
mock_start.assert_not_called()
+
+ @pytest.mark.parametrize("callback_kwargs_encoding", [None, "serde"])
+
@patch("airflow.sdk.execution_time.callback_supervisor._make_process_nondumpable")
+ @patch.object(CallbackSubprocess, "start")
+ def test_kwargs_encoding_is_passed_to_the_subprocess(
+ self, mock_start, _nondumpable, callback_kwargs_encoding
+ ):
+ mock_start.return_value = Mock(wait=Mock(return_value=0))
+
+ supervise_callback(
+ id=self.CALLBACK_ID,
+ callback_path="does.not.matter",
+ callback_kwargs={},
+ dag_rel_path=Path("dag.py"),
+ callback_kwargs_encoding=callback_kwargs_encoding,
+ client=Mock(),
+ )
+
+ assert mock_start.call_args.kwargs["callback_kwargs_encoding"] ==
callback_kwargs_encoding
