ferruzzi commented on code in PR #57222:
URL: https://github.com/apache/airflow/pull/57222#discussion_r2466703757
##########
airflow-core/docs/howto/deadline-alerts.rst:
##########
@@ -328,24 +328,95 @@ Custom References
^^^^^^^^^^^^^^^^^
While the built-in references should cover most use cases, and more will be
released over time, you
-can create custom references by implementing a class that inherits from
DeadlineReference. This may
-be useful if you have calendar integrations or other sources that you want to
use as a reference.
+can create custom references. This may be useful if you have calendar
integrations or other sources
+that you want to use as a reference. You can create custom references by
implementing a class that
+inherits from BaseDeadlineReference, give it am _evaluate_with() method, and
register it. There are
+two ways to accomplish this. The recommended way is to use the
``@deadline_reference`` decorator
+but for more complicated implementations, the ``register_custom_reference()``
method is available.
+
+**Recommended: Using the decorator**
.. code-block:: python
- class CustomReference(DeadlineReference):
- """A deadline reference that uses a custom data source."""
+ from airflow._shared.timezones.timezone import datetime
+ from airflow.models.deadline import ReferenceModels
+ from sqlalchemy.orm import Session
+
+ from airflow.sdk.definitions.deadline import DeadlineReference,
deadline_reference
+
+
+ # By default, the evaluate_with method will be executed when the dagrun is
created.
+ @deadline_reference()
+ class MyCustomDecoratedReference(ReferenceModels.BaseDeadlineReference):
+ """A custom reference evaluated when DAG runs are created."""
+
+ def _evaluate_with(self, *, session: Session, **kwargs) -> datetime:
+ # Add your business logic here
+ return your_datetime
+
+
+ # You can specify when evaluate_with will be called by providing a
DeadlineReference.TYPES value.
+ @deadline_reference(DeadlineReference.TYPES.DAGRUN_QUEUED)
+ class MyQueuedReference(ReferenceModels.BaseDeadlineReference):
+ """A custom reference evaluated when DAG runs are queued."""
+
+ required_kwargs = {"custom_param"}
+
+ def _evaluate_with(self, *, session: Session, **kwargs) -> datetime:
+ custom_value = kwargs["custom_param"]
+ # Use custom_value in your calculation
+ return your_datetime
+
+**Alternative: Manual Registration**
+
+For advanced use cases requiring conditional or dynamic registration, you may
wish use the registration method directly.
+In this case, the plugin file will look something like this:
+
+.. code-block:: python
+
+ from sqlalchemy.orm import Session
+
+ from airflow.models.deadline import ReferenceModels
+ from airflow.sdk.definitions.deadline import DeadlineReference
- # Define any required parameters for your reference
- required_kwargs = {"custom_id"}
+ class MyManualReference(ReferenceModels.BaseDeadlineReference):
def _evaluate_with(self, *, session: Session, **kwargs) -> datetime:
- """
- Evaluate the reference time using the provided session and kwargs.
-
- The session parameter can be used for database queries, and kwargs
- will contain any required parameters defined in required_kwargs.
- """
- custom_id = kwargs["custom_id"]
- # Your custom logic here to determine the reference time
+ # Add your business logic here
return your_datetime
+
+
+ # Register with specific timing based on configuration
+ timing = (
+ DeadlineReference.TYPES.DAGRUN_QUEUED if use_queued_timing else
DeadlineReference.TYPES.DAGRUN_CREATED
+ )
+ DeadlineReference.register_custom_reference(MyManualReference, timing)
+
+**Using Custom References in DAGs**
+
+Once registered, use your custom references in DAG definitions like any other
reference:
Review Comment:
Nah. With how that is dragging on, we can't just NOT merge anything while
they argue over capitalization of a word. It can get fixed along with all the
others once we have a decision.
--
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]
