diff --git a/airflow-core/docs/howto/deadline-alerts.rst b/airflow-core/docs/howto/deadline-alerts.rst index ac07047f34587..7eb9ec76cf2d0 100644 --- a/airflow-core/docs/howto/deadline-alerts.rst +++ b/airflow-core/docs/howto/deadline-alerts.rst @@ -431,8 +431,7 @@ choose a different time. from sqlalchemy.orm import Session - from airflow.sdk import DeadlineReference - from airflow.sdk.definitions.deadline import BaseDeadlineReference, deadline_reference + from airflow.sdk import BaseDeadlineReference, DeadlineReference, deadline_reference from airflow.sdk.timezone import datetime diff --git a/task-sdk/docs/api.rst b/task-sdk/docs/api.rst index 51aa275b90bb0..2c9bbe3f4db80 100644 --- a/task-sdk/docs/api.rst +++ b/task-sdk/docs/api.rst @@ -116,6 +116,10 @@ Deadline Alerts .. autoclass:: airflow.sdk.DeadlineReference +.. autoclass:: airflow.sdk.BaseDeadlineReference + +.. autofunction:: airflow.sdk.deadline_reference + Connections & Variables ----------------------- .. autoapiclass:: airflow.sdk.Connection diff --git a/task-sdk/src/airflow/sdk/__init__.py b/task-sdk/src/airflow/sdk/__init__.py index 807019eade80e..6cff747fccc94 100644 --- a/task-sdk/src/airflow/sdk/__init__.py +++ b/task-sdk/src/airflow/sdk/__init__.py @@ -31,6 +31,7 @@ "AsyncCallback", "BaseAsyncOperator", "BaseBranchOperator", + "BaseDeadlineReference", "BaseHook", "BaseNotifier", "BaseOperator", @@ -105,6 +106,7 @@ "conf", "cross_downstream", "dag", + "deadline_reference", "get_current_context", "get_parsing_context", "literal", @@ -151,7 +153,12 @@ from airflow.sdk.definitions.connection import Connection from airflow.sdk.definitions.context import Context, get_current_context, get_parsing_context from airflow.sdk.definitions.dag import DAG, dag - from airflow.sdk.definitions.deadline import DeadlineAlert, DeadlineReference + from airflow.sdk.definitions.deadline import ( + BaseDeadlineReference, + DeadlineAlert, + DeadlineReference, + deadline_reference, + ) from airflow.sdk.definitions.decorators import result, setup, task, teardown from airflow.sdk.definitions.decorators.task_group import task_group from airflow.sdk.definitions.edges import EdgeModifier, Label @@ -234,6 +241,7 @@ "AsyncCallback": ".definitions.callback", "BaseAsyncOperator": ".bases.operator", "BaseBranchOperator": ".bases.branch", + "BaseDeadlineReference": ".definitions.deadline", "BaseHook": ".bases.hook", "BaseNotifier": ".bases.notifier", "BaseOperator": ".bases.operator", @@ -308,6 +316,7 @@ "conf": ".configuration", "cross_downstream": ".bases.operator", "dag": ".definitions.dag", + "deadline_reference": ".definitions.deadline", "NEVER_EXPIRE": ".execution_time.context", "get_current_context": ".definitions.context", "get_parsing_context": ".definitions.context", diff --git a/task-sdk/src/airflow/sdk/definitions/deadline.py b/task-sdk/src/airflow/sdk/definitions/deadline.py index ab01a9213aa5f..6b895c2164e23 100644 --- a/task-sdk/src/airflow/sdk/definitions/deadline.py +++ b/task-sdk/src/airflow/sdk/definitions/deadline.py @@ -42,11 +42,11 @@ class BaseDeadlineReference(ABC): """ Base class for all Deadline Reference implementations. - This is a lightweight SDK class for DAG authoring. It only handles serialization. - The actual evaluation logic (_evaluate_with) is in Core's SerializedReferenceModels. + This is a lightweight SDK class for Dag authoring. It only handles serialization. + The actual evaluation logic (``_evaluate_with``) is in Core's ``SerializedReferenceModels``. For custom deadline references, users should inherit from this class and implement - _evaluate_with() with deferred Core imports (imports inside the method body). + ``_evaluate_with()`` with deferred Core imports (imports inside the method body). """ @property