ferruzzi commented on code in PR #57222:
URL: https://github.com/apache/airflow/pull/57222#discussion_r2466825643
##########
task-sdk/src/airflow/sdk/definitions/deadline.py:
##########
@@ -321,3 +326,71 @@ def FIXED_DATETIME(cls, datetime: datetime) ->
DeadlineReferenceType:
(DeadlineReferenceType,),
{"_evaluate_with": lambda self, **kwargs: datetime.now()},
)()
+
+ @classmethod
+ def register_custom_reference(
+ cls,
+ reference_class: type[ReferenceModels.BaseDeadlineReference],
+ timing: DeadlineReferenceTuple | None = None,
+ ):
+ """
+ Register a custom deadline reference class.
+
+ :param reference_class: The custom reference class inheriting from
BaseDeadlineReference
+ :param timing: A DeadlineReference.TYPES for when the deadline should
be evaluated ("DAGRUN_CREATED",
+ "DAGRUN_QUEUED", etc.); defaults to
DeadlineReference.TYPES.DAGRUN_CREATED
+ """
+ from airflow.models.deadline import ReferenceModels
+
+ # Default to DAGRUN_CREATED if no timing specified
+ if timing is None:
+ timing = cls.TYPES.DAGRUN_CREATED
+
+ # Validate the reference class inherits from BaseDeadlineReference
+ if not issubclass(reference_class,
ReferenceModels.BaseDeadlineReference):
+ raise ValueError(f"{reference_class.__name__} must inherit from
BaseDeadlineReference")
+
+ # Register the new reference with ReferenceModels and
DeadlineReference for discoverability
+ setattr(ReferenceModels, reference_class.__name__, reference_class)
+ setattr(cls, reference_class.__name__, reference_class())
+ logger.info("Registered DeadlineReference %s",
reference_class.__name__)
Review Comment:
Yeah, I'm only seeing it on the initial plugin load (when I restart the API
Server in Breeze)
--
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]
