ferruzzi commented on code in PR #73299:
URL: https://github.com/apache/airflow/pull/73299#discussion_r4568230377
##########
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:
Sounds like a good idea, I'll see what I can do.
--
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]
