Ferruzzi/deadlines/summit fixes/duplicate reference names

airflow-core/docs/howto/deadline-alerts.rst

@@ -513,6 +513,36 @@ other reference:
         # Your tasks here
         ...
 
+The ``DeadlineReference.<ClassName>`` 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
 ^^^^^^^^^^^^^^^^^^^^^^^^
 

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,16 +1072,16 @@ 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)
         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