diff --git a/airflow-core/docs/howto/deadline-alerts.rst b/airflow-core/docs/howto/deadline-alerts.rst index 76b18a99e4508..a0cc7721f9932 100644 --- a/airflow-core/docs/howto/deadline-alerts.rst +++ b/airflow-core/docs/howto/deadline-alerts.rst @@ -513,6 +513,36 @@ other reference: # Your tasks here ... +The ``DeadlineReference.`` shorthand hands you an instance constructed with no arguments, +so a parameterized reference gets one carrying its defaults. You may optionally import the class and +instantiate it yourself, for example if you want to pass values: + +.. code-block:: python + + from datetime import timedelta + from airflow.sdk import AsyncCallback, DAG, DeadlineAlert + + from my_deadline_references import MyQueuedReference + + with DAG( + dag_id="custom_reference_imported_directly", + deadline=DeadlineAlert( + reference=MyQueuedReference(), + interval=timedelta(hours=2), + callback=AsyncCallback(my_callback), + ), + ): + # Your tasks here + ... + +.. note:: + + The shorthand is keyed on the class name alone, so two custom references with the same class name + in different modules cannot both have it. When that happens, neither one gets it and a warning + naming both classes is logged; import the class directly as above instead. Both references + remain fully usable, and the scheduler still resolves either of them: only the shorthand is + withheld. + Multiple Deadline Alerts ^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/airflow-core/tests/unit/models/test_deadline.py b/airflow-core/tests/unit/models/test_deadline.py index 87c383d3e9cb8..a8a0a8198081f 100644 --- a/airflow-core/tests/unit/models/test_deadline.py +++ b/airflow-core/tests/unit/models/test_deadline.py @@ -42,6 +42,7 @@ FixedDatetimeDeadline, deadline_reference, ) +from airflow.sdk.exceptions import RemovedInAirflow4Warning from airflow.serialization.definitions.deadline import ( SerializedReferenceModels, _fetch_from_db as _fetch_serialized_deadline_from_db, @@ -766,6 +767,20 @@ def test_deadline_reference_creation(self): class TestCustomDeadlineReference: + @staticmethod + def _reference_named(name, module): + """Build a valid reference class with a chosen ``__name__`` and module path.""" + reference = type( + name, + (BaseDeadlineReference,), + { + "_evaluate_with": lambda self, **kwargs: None, + "serialize_reference": lambda self: {}, + }, + ) + reference.__module__ = module + return reference + class MyCustomRef(BaseDeadlineReference): def _evaluate_with(self, *, session: Session, **kwargs) -> datetime: return timezone.datetime(DEFAULT_DATE) @@ -790,6 +805,9 @@ def teardown_method(self): DeadlineReference.TYPES.DAGRUN_QUEUED = self.original_dagrun_queued DeadlineReference.TYPES.DAGRUN = self.original_dagrun + DeadlineReference._custom_reference_name_map.clear() + DeadlineReference._contested_reference_names.clear() + for attr in set(dir(DeadlineReference)): if attr not in self.original_deadline_attrs: delattr(DeadlineReference, attr) @@ -809,12 +827,12 @@ def teardown_method(self): pytest.param(DeadlineReference.TYPES.DAGRUN_QUEUED, id="dagrun_queued"), ], ) - def test_register_custom_reference(self, timing, reference): + def test_add_custom_reference_to_namespace(self, timing, reference): if timing is None: - result = DeadlineReference.register_custom_reference(reference) + result = DeadlineReference._add_custom_reference_to_namespace(reference) expected_timing = DeadlineReference.TYPES.DAGRUN_CREATED else: - result = DeadlineReference.register_custom_reference(reference, timing) + result = DeadlineReference._add_custom_reference_to_namespace(reference, timing) expected_timing = timing assert result is reference @@ -826,11 +844,11 @@ def test_register_custom_reference(self, timing, reference): DeadlineReference.TYPES.DAGRUN_QUEUED, DeadlineReference.TYPES.DAGRUN_CREATED ) - def test_register_custom_reference_invalid_inheritance(self): + def test_add_custom_reference_invalid_inheritance(self): with pytest.raises(ValueError, match="must inherit from BaseDeadlineReference"): - DeadlineReference.register_custom_reference(self.MyInvalidCustomRef) + DeadlineReference._add_custom_reference_to_namespace(self.MyInvalidCustomRef) - def test_register_custom_reference_invalid_timing(self): + def test_add_custom_reference_invalid_timing(self): invalid_timing = ("not", "a", "valid", "timing") with pytest.raises( @@ -840,18 +858,116 @@ def test_register_custom_reference_invalid_timing(self): f"must be a valid DeadlineReference.TYPES option." ), ): - DeadlineReference.register_custom_reference(self.MyCustomRef, invalid_timing) + DeadlineReference._add_custom_reference_to_namespace(self.MyCustomRef, invalid_timing) def test_custom_reference_discoverable_on_deadline_reference(self): # Custom references are only registered on DeadlineReference, not on ReferenceModels. # During deserialization, custom refs are discovered via __class_path in the # serialized data (using import_string), not through ReferenceModels lookup. - DeadlineReference.register_custom_reference(self.MyCustomRef) + DeadlineReference._add_custom_reference_to_namespace(self.MyCustomRef) assert hasattr(DeadlineReference, self.MyCustomRef.__name__) found_instance = getattr(DeadlineReference, self.MyCustomRef.__name__) assert isinstance(found_instance, self.MyCustomRef) + def test_reparsing_the_same_reference_rebinds_the_shorthand(self): + """ + A Dag re-parse builds a new class object for the same reference. + + The shorthand must point at the new class's instance, not keep the one belonging to the class + object the re-parse replaced. + """ + first = self._reference_named("Reparsed", "team_a.refs") + DeadlineReference._add_custom_reference_to_namespace(first) + assert DeadlineReference.Reparsed.__class__ is first + + second = self._reference_named("Reparsed", "team_a.refs") + assert second is not first + DeadlineReference._add_custom_reference_to_namespace(second) + + assert DeadlineReference.Reparsed.__class__ is second + + def test_reparsing_the_same_reference_replaces_its_types_entry(self): + """A re-parse must replace the reference's classification entry rather than append beside it.""" + first = self._reference_named("Reparsed", "team_a.refs") + DeadlineReference._add_custom_reference_to_namespace(first) + after_first = DeadlineReference.TYPES.DAGRUN_CREATED + + second = self._reference_named("Reparsed", "team_a.refs") + DeadlineReference._add_custom_reference_to_namespace(second) + + assert len(DeadlineReference.TYPES.DAGRUN_CREATED) == len(after_first) + assert second in DeadlineReference.TYPES.DAGRUN_CREATED + assert first not in DeadlineReference.TYPES.DAGRUN_CREATED + + def test_reclassifying_a_reference_moves_it_between_types(self): + """Re-registering under a different classification must not leave it in the one it left.""" + first = self._reference_named("Reclassified", "team_a.refs") + DeadlineReference._add_custom_reference_to_namespace(first) + assert first in DeadlineReference.TYPES.DAGRUN_CREATED + + second = self._reference_named("Reclassified", "team_a.refs") + DeadlineReference._add_custom_reference_to_namespace(second, DeadlineReference.TYPES.DAGRUN_QUEUED) + + assert second in DeadlineReference.TYPES.DAGRUN_QUEUED + assert second not in DeadlineReference.TYPES.DAGRUN_CREATED + assert first not in DeadlineReference.TYPES.DAGRUN_CREATED + + def test_two_modules_claiming_one_name_leaves_it_bound_to_neither(self, cap_structlog): + """ + The shorthand is withheld from both rather than awarded by import order. + + A Dag file is written against the short name, so letting the last import win means the same + Dag silently evaluates a different class after a restart. + """ + first = self._reference_named("Contested", "team_a.refs") + second = self._reference_named("Contested", "team_b.reports") + + DeadlineReference._add_custom_reference_to_namespace(first) + assert DeadlineReference.Contested.__class__ is first + + DeadlineReference._add_custom_reference_to_namespace(second) + + assert not hasattr(DeadlineReference, "Contested") + assert { + "event": re.compile(r".*team_a\.refs\.Contested.*team_b\.reports\.Contested"), + "level": "warning", + } in cap_structlog + + # Both remain valid references; only the shorthand is gone. + assert first in DeadlineReference.TYPES.DAGRUN_CREATED + assert second in DeadlineReference.TYPES.DAGRUN_CREATED + + def test_a_contested_name_is_not_rebound_by_a_later_claimant(self): + """A third claim is recorded alongside the others, and still does not get the shorthand.""" + for module in ("team_a.refs", "team_b.reports", "team_c.other"): + DeadlineReference._add_custom_reference_to_namespace(self._reference_named("Contested", module)) + + assert not hasattr(DeadlineReference, "Contested") + assert DeadlineReference._contested_reference_names["Contested"] == [ + "team_a.refs.Contested", + "team_b.reports.Contested", + "team_c.other.Contested", + ] + + def test_shadowing_a_builtin_reference_raises(self): + """A builtin must never be unbound, so this is an error rather than a contested name.""" + with pytest.raises(ValueError, match="DeadlineReference.DAGRUN_QUEUED_AT already exists"): + DeadlineReference._add_custom_reference_to_namespace( + self._reference_named("DAGRUN_QUEUED_AT", "team_a.refs") + ) + + assert DeadlineReference.DAGRUN_QUEUED_AT is not None + + def test_register_custom_reference_is_a_deprecated_shim(self): + reference = self._reference_named("ViaShim", "team_a.refs") + + with pytest.warns(RemovedInAirflow4Warning, match="register_custom_reference is deprecated"): + result = DeadlineReference.register_custom_reference(reference) + + assert result is reference + assert DeadlineReference.ViaShim.__class__ is reference + class TestDeadlineReferenceDecorator: def setup_method(self): @@ -865,6 +981,9 @@ def teardown_method(self): DeadlineReference.TYPES.DAGRUN_QUEUED = self.original_dagrun_queued DeadlineReference.TYPES.DAGRUN = self.original_dagrun + DeadlineReference._custom_reference_name_map.clear() + DeadlineReference._contested_reference_names.clear() + for attr in set(dir(DeadlineReference)): if attr not in self.original_deadline_attrs: delattr(DeadlineReference, attr) @@ -953,8 +1072,8 @@ class DecoratedCustomRef(BaseDeadlineReference): def _evaluate_with(self, *, session: Session, **kwargs) -> datetime: return timezone.datetime(DEFAULT_DATE) - @mock.patch.object(DeadlineReference, "register_custom_reference") - def test_deadline_reference_decorator_calls_register_method(self, mock_register): + @mock.patch.object(DeadlineReference, "_add_custom_reference_to_namespace") + def test_deadline_reference_decorator_calls_the_namespace_method(self, mock_add): timing = DeadlineReference.TYPES.DAGRUN_QUEUED @deadline_reference(timing) @@ -962,7 +1081,7 @@ class DecoratedCustomRef(BaseDeadlineReference): def _evaluate_with(self, *, session: Session, **kwargs) -> datetime: return timezone.datetime(DEFAULT_DATE) - mock_register.assert_called_once_with(DecoratedCustomRef, timing) + mock_add.assert_called_once_with(DecoratedCustomRef, timing) def test_deadline_reference_decorator_without_parentheses(self): @deadline_reference diff --git a/task-sdk/src/airflow/sdk/definitions/deadline.py b/task-sdk/src/airflow/sdk/definitions/deadline.py index ba923868b4495..93b8ca0f9813f 100644 --- a/task-sdk/src/airflow/sdk/definitions/deadline.py +++ b/task-sdk/src/airflow/sdk/definitions/deadline.py @@ -25,6 +25,7 @@ import attrs +from airflow.sdk._shared.module_loading import qualname from airflow.sdk.definitions.callback import AsyncCallback, Callback, SyncCallback from airflow.sdk.definitions.variable import Variable from airflow.sdk.exceptions import AirflowRuntimeError, RemovedInAirflow4Warning @@ -243,9 +244,17 @@ class TYPES: # All DagRun-related deadline types. DAGRUN: DeadlineReferenceTypes = DAGRUN_CREATED + DAGRUN_QUEUED + # The attributes above that a custom reference can be classified under. DAGRUN is their + # union rather than a classification of its own, so it is deliberately absent. + _CLASSIFICATIONS = ("DAGRUN_CREATED", "DAGRUN_QUEUED") + DAGRUN_LOGICAL_DATE: DeadlineReferenceType = DagRunLogicalDateDeadline() DAGRUN_QUEUED_AT: DeadlineReferenceType = DagRunQueuedAtDeadline() + # These map short class name to qualname(s); used to handle custom reference name collisions. + _custom_reference_name_map: dict[str, str] = {} + _contested_reference_names: dict[str, list[str]] = {} + @classmethod def AVERAGE_RUNTIME(cls, max_runs: int = 0, min_runs: int | None = None) -> DeadlineReferenceType: if max_runs == 0: @@ -269,29 +278,33 @@ def FIXED_DATETIME(cls, dt: datetime) -> DeadlineReferenceType: )() @classmethod - def register_custom_reference( + def _add_custom_reference_to_namespace( cls, reference_class: type[BaseDeadlineReference], deadline_reference_type: DeadlineReferenceTypes | None = None, ) -> type[BaseDeadlineReference]: """ - Register a custom deadline reference class for use in Dag files. + Add a custom deadline reference class to the DeadlineReference namespace for use in Dag files. This makes the reference available to Dag authors as ``DeadlineReference.`` and records when it should be evaluated. - .. warning:: + .. note:: - Registering the reference is **not** the same as registering the plugin, despite the - name of this method. This only affects the process that runs the Dag file; it does not - make the class resolvable when the scheduler deserializes the Dag. The class must - *also* be listed in the ``deadline_references`` attribute of an ``AirflowPlugin``, or - deserialization raises ``DeadlineReferenceNotRegistered``. See - :external:doc:`howto/deadline-alerts`. + Adding the reference to this namespace is **not** the same as registering the plugin. + This only affects the process that runs the Dag file; it does not make the class + resolvable when the scheduler deserializes the Dag. The class must *also* be listed in + the ``deadline_references`` attribute of an ``AirflowPlugin``, or deserialization raises + ``DeadlineReferenceNotRegistered``. See :external:doc:`howto/deadline-alerts`. + + If two classes share the same name but have different module paths, the shorthand is + withheld from both and a warning names the competing classes. In that case, users can + still import the class directly; the classes themselves remain valid references either way. :param reference_class: The custom reference class inheriting from BaseDeadlineReference - :param deadline_reference_type: A DeadlineReference.TYPES for when the deadline should be evaluated ("DAGRUN_CREATED", - "DAGRUN_QUEUED", etc.); defaults to DeadlineReference.TYPES.DAGRUN_CREATED + :param deadline_reference_type: A DeadlineReference.TYPES for when the deadline should be + evaluated ("DAGRUN_CREATED", "DAGRUN_QUEUED", etc.); + defaults to DeadlineReference.TYPES.DAGRUN_CREATED """ # Default to DAGRUN_CREATED if no deadline_reference_type specified if deadline_reference_type is None: @@ -313,25 +326,99 @@ def register_custom_reference( f"registered as a deadline reference. If it takes parameters, decorate it with " f"@dataclass and give every field a default value. Original error: {e}" ) from e - setattr(cls, reference_class.__name__, reference_instance) - logger.info("Registered DeadlineReference %s", reference_class.__name__) - - # Add to appropriate deadline_reference_type classification - if deadline_reference_type is cls.TYPES.DAGRUN_CREATED: - cls.TYPES.DAGRUN_CREATED = cls.TYPES.DAGRUN_CREATED + (reference_class,) - elif deadline_reference_type is cls.TYPES.DAGRUN_QUEUED: - cls.TYPES.DAGRUN_QUEUED = cls.TYPES.DAGRUN_QUEUED + (reference_class,) + name = reference_class.__name__ + key = qualname(reference_class) + + if name in cls._contested_reference_names: + # There is a conflict; do not register the short name. + claimed_names = cls._contested_reference_names[name] + if key not in claimed_names: + claimed_names.append(key) + logger.warning( + "DeadlineReference.%s is claimed by more than one class (%s), so it is not available " + "as a shorthand. Import the class you want directly instead.", + name, + ", ".join(sorted(claimed_names)), + ) + elif (previous := cls._custom_reference_name_map.get(name)) is not None: + if previous == key: + # Class was registered in a previous parse, refresh the value in case user code changed. + setattr(cls, name, reference_instance) + else: + # Two different classes want the same shorthand, reject BOTH. + # Without this, the import order arbitrarily decides which gets the + # name and potentially breaks a Dag file which uses the short name. + delattr(cls, name) + cls._contested_reference_names[name] = [cls._custom_reference_name_map.pop(name), key] + logger.warning( + "DeadlineReference.%s is claimed by more than one class (%s), so it is no longer " + "available as a shorthand. Import the class you want directly instead.", + name, + ", ".join(sorted((previous, key))), + ) + elif hasattr(cls, name): + # Custom reference name can not shadow built-in references. + raise ValueError( + f"{key} cannot be added to the DeadlineReference namespace as {name!r}, because " + f"DeadlineReference.{name} already exists. Rename the class." + ) + else: + # New reference class name, no collisions. + setattr(cls, name, reference_instance) + cls._custom_reference_name_map[name] = key + + if cls._custom_reference_name_map.get(name) == key: + logger.info("Added DeadlineReference.%s (%s)", name, key) + + # Resolve which classification this reference belongs in. + for candidate in cls.TYPES._CLASSIFICATIONS: + if deadline_reference_type is getattr(cls.TYPES, candidate): + classification = candidate + break else: raise ValueError( f"Invalid deadline reference type {deadline_reference_type}; " "must be a valid DeadlineReference.TYPES option." ) + # A contested class is still a valid reference: it serializes under its own qualname and + # resolves through the plugin registry, so only the shorthand is withheld. Drop any earlier + # registration of the same qualname rather than appending beside it, because a re-parse + # produces a new class object for the same reference. Both tuples are filtered so that a + # reference whose classification changes does not linger in the one it left. + for attr in cls.TYPES._CLASSIFICATIONS: + kept = tuple(existing for existing in getattr(cls.TYPES, attr) if qualname(existing) != key) + setattr(cls.TYPES, attr, kept + (reference_class,) if attr == classification else kept) + # Refresh the combined DAGRUN tuple cls.TYPES.DAGRUN = cls.TYPES.DAGRUN_CREATED + cls.TYPES.DAGRUN_QUEUED return reference_class + @classmethod + def register_custom_reference( + cls, + reference_class: type[BaseDeadlineReference], + deadline_reference_type: DeadlineReferenceTypes | None = None, + ) -> type[BaseDeadlineReference]: + """ + Add a custom deadline reference to the DeadlineReference namespace. + + .. deprecated:: 3.4.0 + Use the ``@deadline_reference`` decorator instead. Despite its name this method never + registered anything with the scheduler; see + :meth:`_add_custom_reference_to_namespace`. + """ + warnings.warn( + "DeadlineReference.register_custom_reference is deprecated and will be removed in " + "Airflow 4. It does not register anything with the scheduler; use the " + "@deadline_reference decorator, and list the class in an AirflowPlugin's " + "deadline_references to make it resolvable.", + RemovedInAirflow4Warning, + stacklevel=2, + ) + return cls._add_custom_reference_to_namespace(reference_class, deadline_reference_type) + @overload def deadline_reference( @@ -391,12 +478,12 @@ def serialize_reference(self) -> dict: """ # Used bare, without parentheses: the decorated class is passed in directly. if isinstance(deadline_reference_type, type): - return DeadlineReference.register_custom_reference(deadline_reference_type) + return DeadlineReference._add_custom_reference_to_namespace(deadline_reference_type) def decorator( reference_class: type[BaseDeadlineReference], ) -> type[BaseDeadlineReference]: - DeadlineReference.register_custom_reference(reference_class, deadline_reference_type) + DeadlineReference._add_custom_reference_to_namespace(reference_class, deadline_reference_type) return reference_class return decorator