o-nikolas commented on code in PR #73299:
URL: https://github.com/apache/airflow/pull/73299#discussion_r4557830058
##########
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)
Review Comment:
When a name gets contested, a Dag that uses the shorthand now fails to
import with a bare `AttributeError: type object 'DeadlineReference' has no
attribute 'MyRef'`, and the warning that explains why ends up in the Dag
processor log rather than in the import error the author is looking at. Would
it be worth putting something in place of the `delattr` that raises an
`AttributeError` carrying the same message? A tiny descriptor whose `__get__`
raises would do it, and `hasattr` still returns False, so your tests should
hold. Happy either way.
##########
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:
Could this note say what the Dag author actually sees? Something like "a Dag
that uses the shorthand will fail to import with an `AttributeError`" would
make it much easier to connect that import error back to this note.
--
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]
