seanghaeli commented on code in PR #73299:
URL: https://github.com/apache/airflow/pull/73299#discussion_r4567485859
##########
airflow-core/docs/howto/deadline-alerts.rst:
##########
@@ -511,6 +511,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.
Review Comment:
Related: the plugin example above is saved as `deadline_references.py`, but
its import a few lines down (line 523) uses a different module name. It should
be `from deadline_references import MyQueuedReference`.
##########
task-sdk/src/airflow/sdk/definitions/deadline.py:
##########
@@ -351,27 +365,100 @@ 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,)
- reference_class.evaluation_timing = DAGRUN_CREATED_TIMING
- elif deadline_reference_type is cls.TYPES.DAGRUN_QUEUED:
- cls.TYPES.DAGRUN_QUEUED = cls.TYPES.DAGRUN_QUEUED +
(reference_class,)
- reference_class.evaluation_timing = DAGRUN_QUEUED_TIMING
+ name = reference_class.__name__
+ key = qualname(reference_class)
+
+ 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."
)
+ if name in cls._contested_reference_names:
+ 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):
+ # Raise rather than contest the name: contesting calls delattr,
which would unbind a
+ # built-in for every Dag that uses it.
+ raise ValueError(
+ f"{key} cannot be added to the DeadlineReference namespace as
{name!r}, because "
+ f"DeadlineReference.{name} already exists. Rename the class."
+ )
+ else:
+ 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)
+
+ # 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)
Review Comment:
My agent flagged this one: this now swaps in a new tuple for both
classifications on every registration, and the type check above compares by
identity, so a held `TYPES` value goes stale. For example, a plugin file that
does `QUEUED = DeadlineReference.TYPES.DAGRUN_QUEUED` up top, defines a
default-timing reference, and then uses `@deadline_reference(QUEUED)` now fails
with `Invalid deadline reference type`. That works on main (checked in Breeze).
Only replacing the tuple that actually changed fixes it:
```suggestion
for attr in cls.TYPES._CLASSIFICATIONS:
current = getattr(cls.TYPES, attr)
kept = tuple(existing for existing in current if
qualname(existing) != key)
if attr == classification:
setattr(cls.TYPES, attr, kept + (reference_class,))
elif len(kept) != len(current):
setattr(cls.TYPES, attr, kept)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
