This is an automated email from the ASF dual-hosted git repository.
ferruzzi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 637b741f1a6 Pass Dag runs directly to Deadline references (#70714)
637b741f1a6 is described below
commit 637b741f1a60f0a059dd60f2a62bd47c0e38cd54
Author: Richard <[email protected]>
AuthorDate: Thu Oct 15 10:48:26 2026 -0600
Pass Dag runs directly to Deadline references (#70714)
* Fix custom Deadline references with optional keyword arguments
Preserve the Dag run evaluation context for custom references that accept
additional keyword arguments, so they continue to evaluate correctly.
Co-authored-by: D. Ferruzzi <[email protected]>
---
airflow-core/docs/howto/deadline-alerts.rst | 25 +-
airflow-core/newsfragments/70714.significant.rst | 1 +
airflow-core/src/airflow/models/deadline.py | 139 ++++-------
airflow-core/src/airflow/models/taskinstance.py | 25 +-
.../src/airflow/serialization/definitions/dag.py | 2 +-
.../airflow/serialization/definitions/deadline.py | 90 ++------
airflow-core/tests/unit/models/test_deadline.py | 257 ++++-----------------
.../tests/unit/models/test_deadline_alert.py | 84 ++++++-
task-sdk/src/airflow/sdk/definitions/deadline.py | 37 ++-
9 files changed, 224 insertions(+), 436 deletions(-)
diff --git a/airflow-core/docs/howto/deadline-alerts.rst
b/airflow-core/docs/howto/deadline-alerts.rst
index 76b18a99e45..538173c3505 100644
--- a/airflow-core/docs/howto/deadline-alerts.rst
+++ b/airflow-core/docs/howto/deadline-alerts.rst
@@ -466,9 +466,9 @@ Place the reference classes and the plugin that registers
them in your plugins f
class MyCustomDecoratedReference(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
+ def _evaluate_with(self, *, session: Session, dagrun) -> datetime:
+ my_datetime = my_business_logic(dagrun.logical_date)
+ return my_datetime
# You can specify when evaluate_with will be called by providing a
DeadlineReference.TYPES value.
@@ -476,14 +476,9 @@ Place the reference classes and the plugin that registers
them in your plugins f
class MyQueuedReference(BaseDeadlineReference):
"""A custom reference evaluated when Dag runs are queued."""
- # Ask for the Dag run context values supplied by Airflow; see notes
below.
- required_kwargs = {"dag_id", "run_id"}
-
- def _evaluate_with(self, *, session: Session, **kwargs) -> datetime:
- dag_id = kwargs["dag_id"]
- run_id = kwargs["run_id"]
- # Use dag_id and run_id in your calculation
- return your_datetime
+ def _evaluate_with(self, *, session: Session, dagrun) -> datetime:
+ my_datetime = my_business_logic(dagrun.queued_at)
+ return my_datetime
# Register the classes so the scheduler can resolve them when it
deserializes the Dag.
@@ -564,8 +559,8 @@ followed by a more urgent escalation if the Dag is still
running.
* **Plugin Registration**: Custom references must be listed in the
``deadline_references`` attribute
of an ``AirflowPlugin``, so the plugins directory is the natural home for
them.
* **API Server Restart**: Restart the Airflow API Server after adding or
modifying custom references.
-* **Required Parameters**: ``required_kwargs`` declares which Dag run context
values Airflow should
- forward to ``_evaluate_with()``. Only ``dag_id`` and ``run_id`` are
available; declaring anything
- else raises a ``ValueError`` when the deadline is evaluated. To configure a
reference itself, give
- it constructor fields or read from an Airflow Variable.
+* **Dag run context**: Add a ``dagrun`` parameter to ``_evaluate_with()`` to
use the Dag run being evaluated.
+* **Required Parameters** *(deprecated)*: ``required_kwargs`` declares which
parameters your reference
+ needs. It still works but is deprecated and will be removed in Airflow 4.0 —
declare the keyword-only
+ parameters your ``_evaluate_with()`` needs instead.
* **Database Access**: Use the ``session`` parameter for Airflow database
queries if needed.
diff --git a/airflow-core/newsfragments/70714.significant.rst
b/airflow-core/newsfragments/70714.significant.rst
new file mode 100644
index 00000000000..70572b6d598
--- /dev/null
+++ b/airflow-core/newsfragments/70714.significant.rst
@@ -0,0 +1 @@
+Deprecate the ``required_kwargs`` attribute for custom deadline references.
Implement ``_evaluate_with()`` with a ``dagrun`` parameter instead.
diff --git a/airflow-core/src/airflow/models/deadline.py
b/airflow-core/src/airflow/models/deadline.py
index 77e3c682e7b..b9c212dc834 100644
--- a/airflow-core/src/airflow/models/deadline.py
+++ b/airflow-core/src/airflow/models/deadline.py
@@ -17,21 +17,23 @@
from __future__ import annotations
import logging
+import warnings
from abc import ABC, abstractmethod
from collections.abc import Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any, cast
+from inspect import signature
+from typing import TYPE_CHECKING, Any, Protocol, cast
from uuid import UUID
import uuid6
-from sqlalchemy import Boolean, ForeignKey, Index, Integer, Uuid, and_, func,
inspect, select, text
-from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy import Boolean, ForeignKey, Index, Integer, Uuid, and_, func,
select, text
from sqlalchemy.orm import Mapped, mapped_column, relationship
from airflow._shared.observability.metrics import stats
from airflow._shared.timezones import timezone
from airflow.configuration import conf
+from airflow.exceptions import RemovedInAirflow4Warning
from airflow.models.base import Base
from airflow.models.callback import (
Callback,
@@ -57,6 +59,38 @@ logger = logging.getLogger(__name__)
CALLBACK_METRICS_PREFIX = "deadline_alerts"
+class DeadlineDagRunProtocol(Protocol):
+ """The Dag run interface deadline references are evaluated against."""
+
+ dag_id: str
+ logical_date: datetime | None
+ queued_at: datetime | None
+
+
+def _get_evaluation_kwargs(reference: Any, evaluator: Any, kwargs: dict[str,
Any]) -> dict[str, Any]:
+ """Return the evaluation arguments accepted by a deadline reference."""
+ required_kwargs: set[str] | None = getattr(reference, "required_kwargs",
None)
+ if required_kwargs is not None:
+ warnings.warn(
+ "required_kwargs is deprecated. Declare the keyword-only
parameters your "
+ "_evaluate_with() implementation needs instead.",
+ RemovedInAirflow4Warning,
+ stacklevel=3,
+ )
+ kwargs = {key: value for key, value in kwargs.items() if key in
required_kwargs}
+ if missing_kwargs := required_kwargs - kwargs.keys():
+ raise ValueError(
+ f"{reference.__class__.__name__} is missing required
parameters: {', '.join(missing_kwargs)}"
+ )
+ return kwargs
+
+ parameters = signature(evaluator).parameters
+ if any(param.kind is param.VAR_KEYWORD for param in parameters.values()):
+ return kwargs
+
+ return {key: value for key, value in kwargs.items() if key in parameters}
+
+
class classproperty:
"""
Decorator that converts a method with a single cls argument into a
property.
@@ -323,34 +357,18 @@ class ReferenceModels:
class BaseDeadlineReference(LoggingMixin, ABC):
"""Base class for all Deadline implementations."""
- # Set of required kwargs - subclasses should override this.
- required_kwargs: set[str] = set()
-
@classproperty
def reference_name(cls: Any) -> str:
return cls.__name__
def evaluate_with(self, *, session: Session, interval: timedelta,
**kwargs: Any) -> datetime | None:
- """Validate the provided kwargs and evaluate this deadline with
the given conditions."""
- filtered_kwargs = {k: v for k, v in kwargs.items() if k in
self.required_kwargs}
-
- if missing_kwargs := self.required_kwargs - filtered_kwargs.keys():
- raise ValueError(
- f"{self.__class__.__name__} is missing required
parameters: {', '.join(missing_kwargs)}"
- )
-
- if extra_kwargs := kwargs.keys() - filtered_kwargs.keys():
- self.log.debug(
- "%s ignoring unexpected parameters: %s",
- self.reference_name,
- ", ".join(extra_kwargs),
- )
-
- base_time = self._evaluate_with(session=session, **filtered_kwargs)
+ """Evaluate this deadline with the supplied context."""
+ evaluation_kwargs = _get_evaluation_kwargs(self,
self._evaluate_with, kwargs)
+ base_time = self._evaluate_with(session=session,
**evaluation_kwargs)
return base_time + interval if base_time is not None else None
@abstractmethod
- def _evaluate_with(self, *, session: Session, **kwargs: Any) ->
datetime | None:
+ def _evaluate_with(self, *, session: Session, dagrun:
DeadlineDagRunProtocol) -> datetime | None:
"""Must be implemented by subclasses to perform the actual
evaluation."""
raise NotImplementedError
@@ -384,7 +402,7 @@ class ReferenceModels:
_datetime: datetime
- def _evaluate_with(self, *, session: Session, **kwargs: Any) ->
datetime | None:
+ def _evaluate_with(self, *, session: Session, **kwargs) -> datetime |
None:
return self._datetime
def serialize_reference(self) -> dict:
@@ -400,23 +418,14 @@ class ReferenceModels:
class DagRunLogicalDateDeadline(BaseDeadlineReference):
"""A deadline that returns a DagRun's logical date."""
- required_kwargs = {"dag_id", "run_id"}
-
- def _evaluate_with(self, *, session: Session, **kwargs: Any) ->
datetime | None:
- from airflow.models import DagRun
-
- return _fetch_from_db(DagRun.logical_date, session=session,
**kwargs)
+ def _evaluate_with(self, *, session: Session, dagrun:
DeadlineDagRunProtocol) -> datetime | None:
+ return dagrun.logical_date
class DagRunQueuedAtDeadline(BaseDeadlineReference):
"""A deadline that returns when a DagRun was queued."""
- required_kwargs = {"dag_id", "run_id"}
-
- @provide_session
- def _evaluate_with(self, *, session: Session, **kwargs: Any) ->
datetime | None:
- from airflow.models import DagRun
-
- return _fetch_from_db(DagRun.queued_at, session=session, **kwargs)
+ def _evaluate_with(self, *, session: Session, dagrun:
DeadlineDagRunProtocol) -> datetime | None:
+ return dagrun.queued_at
@dataclass
class AverageRuntimeDeadline(BaseDeadlineReference):
@@ -425,7 +434,6 @@ class ReferenceModels:
DEFAULT_LIMIT = 10
max_runs: int
min_runs: int | None = None
- required_kwargs = {"dag_id"}
def __post_init__(self):
if self.min_runs is None:
@@ -434,10 +442,10 @@ class ReferenceModels:
raise ValueError("min_runs must be at least 1")
@provide_session
- def _evaluate_with(self, *, session: Session, **kwargs: Any) ->
datetime | None:
+ def _evaluate_with(self, *, session: Session, dagrun:
DeadlineDagRunProtocol) -> datetime | None:
from airflow.models import DagRun
- dag_id = kwargs["dag_id"]
+ dag_id = dagrun.dag_id
# Get database dialect to use appropriate time difference
calculation
dialect = get_dialect_name(session)
@@ -510,54 +518,3 @@ class ReferenceModels:
DeadlineReferenceType = ReferenceModels.BaseDeadlineReference
-
-
-@provide_session
-def _fetch_from_db(model_reference: Mapped, *, session=None, **conditions) ->
datetime | None:
- """
- Fetch a datetime value from the database using the provided model
reference and filtering conditions.
-
- For example, to fetch a TaskInstance's start_date:
- _fetch_from_db(
- TaskInstance.start_date, dag_id='example_dag',
task_id='example_task', run_id='example_run'
- )
-
- This generates SQL equivalent to:
- SELECT start_date
- FROM task_instance
- WHERE dag_id = 'example_dag'
- AND task_id = 'example_task'
- AND run_id = 'example_run'
-
- :param model_reference: SQLAlchemy Column to select (e.g.,
DagRun.logical_date, TaskInstance.start_date)
- :param conditions: Filtering conditions applied as equality comparisons in
the WHERE clause.
- Multiple conditions are combined with AND.
- :param session: SQLAlchemy session (auto-provided by decorator)
- """
- query = select(model_reference)
-
- for key, value in conditions.items():
- inspected = inspect(model_reference)
- if inspected is not None:
- query = query.where(getattr(inspected.class_, key) == value)
-
- compiled_query = query.compile(compile_kwargs={"literal_binds": True})
- pretty_query = "\n ".join(str(compiled_query).splitlines())
- logger.debug(
- "Executing query:\n %r\nAs SQL:\n %s",
- query,
- pretty_query,
- )
-
- try:
- result = session.scalar(query)
- except SQLAlchemyError:
- logger.exception("Database query failed.")
- raise
-
- if result is None:
- message = f"No matching record found in the database for query:\n
{pretty_query}"
- logger.error(message)
- raise ValueError(message)
-
- return result
diff --git a/airflow-core/src/airflow/models/taskinstance.py
b/airflow-core/src/airflow/models/taskinstance.py
index 8c3a811a17b..c90deef1b60 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -255,14 +255,11 @@ def _add_and_prime_mapped_ti(
set_committed_value(ti, "dag_run", dag_run)
-def _recalculate_dagrun_queued_at_deadlines(
- dagrun: DagRun, new_queued_at: datetime, session: Session
-) -> None:
+def _recalculate_dagrun_queued_at_deadlines(dagrun: DagRun, *, session:
Session) -> None:
"""
Recalculate deadline times for deadlines that reference dagrun.queued_at.
:param dagrun: The DagRun whose deadlines should be recalculated
- :param new_queued_at: The new queued_at timestamp to use for calculation
:param session: Database session
:meta private:
@@ -285,12 +282,9 @@ def _recalculate_dagrun_queued_at_deadlines(
from airflow.serialization.decoders import decode_deadline_alert_model,
resolve_deadline_alert_interval
for deadline, deadline_alert in results:
- # We can't use evaluate_with() since the new queued_at is not written
to the DB yet, and
- # interval is stored as JSON, so it has to be decoded rather than
passed to timedelta().
try:
- interval = resolve_deadline_alert_interval(
- decode_deadline_alert_model(deadline_alert), session=session
- )
+ decoded_alert = decode_deadline_alert_model(deadline_alert)
+ interval = resolve_deadline_alert_interval(decoded_alert,
session=session)
except (ValueError, TypeError):
# Either step can fail: the alert may hold a payload the decoder
refuses, or a
# variable-backed interval may point at a Variable that is missing
or non-numeric.
@@ -304,7 +298,16 @@ def _recalculate_dagrun_queued_at_deadlines(
)
continue
- new_deadline_time = new_queued_at + interval
+ new_deadline_time = decoded_alert.reference.evaluate_with(
+ session=session,
+ interval=interval,
+ dagrun=dagrun,
+ dag_id=dagrun.dag_id,
+ run_id=dagrun.run_id,
+ )
+
+ if new_deadline_time is None:
+ continue
log.debug(
"Recalculating deadline %s for DagRun %s.%s: old=%s, new=%s",
@@ -603,7 +606,7 @@ def clear_task_instances(
parent_context=parent_trace_context(dr.conf),
)
- _recalculate_dagrun_queued_at_deadlines(dr, dr.queued_at,
session)
+ _recalculate_dagrun_queued_at_deadlines(dr, session=session)
if was_finished:
dr.state = dag_run_state
diff --git a/airflow-core/src/airflow/serialization/definitions/dag.py
b/airflow-core/src/airflow/serialization/definitions/dag.py
index c529bf95ade..64ceeff7ccc 100644
--- a/airflow-core/src/airflow/serialization/definitions/dag.py
+++ b/airflow-core/src/airflow/serialization/definitions/dag.py
@@ -795,7 +795,7 @@ class SerializedDAG:
deadline_time =
deserialized_deadline_alert.reference.evaluate_with(
session=session,
interval=interval,
- # TODO : Pretty sure we can drop these last two;
verify after testing is complete
+ dagrun=orm_dagrun,
dag_id=self.dag_id,
run_id=run_id,
)
diff --git a/airflow-core/src/airflow/serialization/definitions/deadline.py
b/airflow-core/src/airflow/serialization/definitions/deadline.py
index 25399aabc62..c7f1f1ee111 100644
--- a/airflow-core/src/airflow/serialization/definitions/deadline.py
+++ b/airflow-core/src/airflow/serialization/definitions/deadline.py
@@ -27,7 +27,7 @@ import attrs
from sqlalchemy import select
from airflow._shared.timezones import timezone
-from airflow.models.deadline import classproperty
+from airflow.models.deadline import _get_evaluation_kwargs, classproperty
from airflow.models.variable import Variable
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.session import provide_session
@@ -39,6 +39,8 @@ if TYPE_CHECKING:
from sqlalchemy import ColumnElement
from sqlalchemy.orm import Session
+ from airflow.models.deadline import DeadlineDagRunProtocol
+
logger = logging.getLogger(__name__)
@@ -95,33 +97,18 @@ class SerializedReferenceModels:
class SerializedBaseDeadlineReference(LoggingMixin, ABC):
"""Base class for all serialized Deadline implementations."""
- required_kwargs: set[str] = set()
-
@classproperty
def reference_name(cls: Any) -> str:
return cls.__name__
def evaluate_with(self, *, session: Session, interval: timedelta,
**kwargs: Any) -> datetime | None:
- """Validate the provided kwargs and evaluate this deadline with
the given conditions."""
- filtered_kwargs = {k: v for k, v in kwargs.items() if k in
self.required_kwargs}
-
- if missing_kwargs := self.required_kwargs - filtered_kwargs.keys():
- raise ValueError(
- f"{self.__class__.__name__} is missing required
parameters: {', '.join(missing_kwargs)}"
- )
-
- if extra_kwargs := kwargs.keys() - filtered_kwargs.keys():
- self.log.debug(
- "%s ignoring unexpected parameters: %s",
- self.reference_name,
- ", ".join(extra_kwargs),
- )
-
- base_time = self._evaluate_with(session=session, **filtered_kwargs)
+ """Evaluate this deadline with the supplied context."""
+ evaluation_kwargs = _get_evaluation_kwargs(self,
self._evaluate_with, kwargs)
+ base_time = self._evaluate_with(session=session,
**evaluation_kwargs)
return base_time + interval if base_time is not None else None
@abstractmethod
- def _evaluate_with(self, *, session: Session, **kwargs: Any) ->
datetime | None:
+ def _evaluate_with(self, *, session: Session, dagrun:
DeadlineDagRunProtocol) -> datetime | None:
"""Must be implemented by subclasses to perform the actual
evaluation."""
raise NotImplementedError
@@ -168,23 +155,14 @@ class SerializedReferenceModels:
class DagRunLogicalDateDeadline(SerializedBaseDeadlineReference):
"""A deadline that returns a DagRun's logical date."""
- required_kwargs = {"dag_id", "run_id"}
-
- def _evaluate_with(self, *, session: Session, **kwargs: Any) ->
datetime | None:
- from airflow.models import DagRun
-
- return _fetch_from_db(DagRun.logical_date, session=session,
**kwargs)
+ def _evaluate_with(self, *, session: Session, dagrun:
DeadlineDagRunProtocol) -> datetime | None:
+ return dagrun.logical_date
class DagRunQueuedAtDeadline(SerializedBaseDeadlineReference):
"""A deadline that returns when a DagRun was queued."""
- required_kwargs = {"dag_id", "run_id"}
-
- @provide_session
- def _evaluate_with(self, *, session: Session, **kwargs: Any) ->
datetime | None:
- from airflow.models import DagRun
-
- return _fetch_from_db(DagRun.queued_at, session=session, **kwargs)
+ def _evaluate_with(self, *, session: Session, dagrun:
DeadlineDagRunProtocol) -> datetime | None:
+ return dagrun.queued_at
@dataclass
class AverageRuntimeDeadline(SerializedBaseDeadlineReference):
@@ -193,7 +171,6 @@ class SerializedReferenceModels:
DEFAULT_LIMIT = 10
max_runs: int
min_runs: int | None = None
- required_kwargs = {"dag_id"}
def __post_init__(self):
if self.min_runs is None:
@@ -202,13 +179,13 @@ class SerializedReferenceModels:
raise ValueError("min_runs must be at least 1")
@provide_session
- def _evaluate_with(self, *, session: Session, **kwargs: Any) ->
datetime | None:
- from sqlalchemy import func, select, text
+ def _evaluate_with(self, *, session: Session, dagrun:
DeadlineDagRunProtocol) -> datetime | None:
+ from sqlalchemy import func, text
from airflow.models import DagRun
from airflow.utils.state import DagRunState
- dag_id = kwargs["dag_id"]
+ dag_id = dagrun.dag_id
dialect = get_dialect_name(session)
@@ -279,7 +256,7 @@ class SerializedReferenceModels:
"""
Wrapper for custom deadline references.
- This class dynamically delegates to the wrapped reference for
required_kwargs and evaluation logic.
+ This class dynamically delegates to the wrapped reference's evaluation
logic.
"""
def __init__(self, inner_ref):
@@ -290,23 +267,9 @@ class SerializedReferenceModels:
return self.inner_ref.reference_name
def evaluate_with(self, *, session: Session, interval: timedelta,
**kwargs: Any) -> datetime | None:
- """Validate the provided kwargs and evaluate this deadline with
the given conditions."""
- required_kwargs: set[str] = getattr(self.inner_ref,
"required_kwargs", set())
- filtered_kwargs = {k: v for k, v in kwargs.items() if k in
required_kwargs}
-
- if missing_kwargs := required_kwargs - filtered_kwargs.keys():
- raise ValueError(
- f"{self.inner_ref.__class__.__name__} is missing required
parameters: {', '.join(missing_kwargs)}"
- )
-
- if extra_kwargs := kwargs.keys() - filtered_kwargs.keys():
- self.log.debug(
- "%s ignoring unexpected parameters: %s",
- self.reference_name,
- ", ".join(extra_kwargs),
- )
-
- deadline = self.inner_ref._evaluate_with(session=session,
**filtered_kwargs)
+ """Evaluate the wrapped custom reference with the supplied
context."""
+ evaluation_kwargs = _get_evaluation_kwargs(self.inner_ref,
self.inner_ref._evaluate_with, kwargs)
+ deadline = self.inner_ref._evaluate_with(session=session,
**evaluation_kwargs)
return deadline + interval if deadline is not None else None
def _evaluate_with(self, *, session: Session, **kwargs: Any) ->
datetime | None:
@@ -365,23 +328,6 @@ SerializedReferenceModels.TYPES.DAGRUN = (
)
-def _fetch_from_db(column, *, session: Session, dag_id: str, run_id: str) ->
datetime | None:
- """
- Fetch a datetime column from the DagRun table.
-
- :meta private:
- """
- from airflow.models import DagRun
-
- row = session.execute(
- select(column).where(DagRun.dag_id == dag_id, DagRun.run_id == run_id)
- ).one_or_none()
- if row is None:
- logger.warning("Could not find DagRun for dag_id=%s, run_id=%s",
dag_id, run_id)
- return None
- return row[0]
-
-
@attrs.define(frozen=True)
class SerializedVariableInterval:
"""Core-side serialized representation of a variable-backed deadline
interval."""
diff --git a/airflow-core/tests/unit/models/test_deadline.py
b/airflow-core/tests/unit/models/test_deadline.py
index 87c383d3e9c..6c09ff069b1 100644
--- a/airflow-core/tests/unit/models/test_deadline.py
+++ b/airflow-core/tests/unit/models/test_deadline.py
@@ -19,17 +19,17 @@ from __future__ import annotations
import re
from dataclasses import dataclass
from datetime import datetime, timedelta
+from types import SimpleNamespace
from typing import TYPE_CHECKING
from unittest import mock
import pytest
import time_machine
from sqlalchemy import select
-from sqlalchemy.exc import SQLAlchemyError
from airflow.api_fastapi.core_api.datamodels.dag_run import DAGRunResponse
from airflow.models import DagRun
-from airflow.models.deadline import Deadline, _fetch_from_db
+from airflow.models.deadline import Deadline, ReferenceModels
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.sdk import timezone
from airflow.sdk.definitions.callback import AsyncCallback, SyncCallback
@@ -42,10 +42,7 @@ from airflow.sdk.definitions.deadline import (
FixedDatetimeDeadline,
deadline_reference,
)
-from airflow.serialization.definitions.deadline import (
- SerializedReferenceModels,
- _fetch_from_db as _fetch_serialized_deadline_from_db,
-)
+from airflow.serialization.definitions.deadline import
SerializedReferenceModels
from airflow.utils.state import DagRunState
from tests_common.test_utils import db
@@ -58,15 +55,6 @@ DAG_ID = "dag_id_1"
INVALID_DAG_ID = "invalid_dag_id"
INVALID_RUN_ID = -1
-REFERENCE_TYPES = [
- pytest.param(SerializedReferenceModels.DagRunLogicalDateDeadline(),
id="logical_date"),
- pytest.param(SerializedReferenceModels.DagRunQueuedAtDeadline(),
id="queued_at"),
-
pytest.param(SerializedReferenceModels.FixedDatetimeDeadline(DEFAULT_DATE),
id="fixed_deadline"),
- pytest.param(
- SerializedReferenceModels.AverageRuntimeDeadline(max_runs=10,
min_runs=10), id="average_runtime"
- ),
-]
-
async def callback_for_deadline():
"""Used in a number of tests to confirm that Deadlines and DeadlineAlerts
function correctly."""
@@ -338,168 +326,41 @@ class TestDeadline:
@pytest.mark.db_test
-class TestCalculatedDeadlineDatabaseCalls:
+class TestCalculatedDeadlineReferences:
@staticmethod
def teardown_method():
_clean_db()
@pytest.mark.parametrize(
- ("column", "conditions", "expected_query"),
- [
- pytest.param(
- DagRun.logical_date,
- {"dag_id": DAG_ID},
- "SELECT dag_run.logical_date \nFROM dag_run \nWHERE
dag_run.dag_id = :dag_id_1",
- id="single_condition_logical_date",
- ),
- pytest.param(
- DagRun.queued_at,
- {"dag_id": DAG_ID},
- "SELECT dag_run.queued_at \nFROM dag_run \nWHERE
dag_run.dag_id = :dag_id_1",
- id="single_condition_queued_at",
- ),
- pytest.param(
- DagRun.logical_date,
- {"dag_id": DAG_ID, "state": "running"},
- "SELECT dag_run.logical_date \nFROM dag_run \nWHERE
dag_run.dag_id = :dag_id_1 AND dag_run.state = :state_1",
- id="multiple_conditions",
- ),
- ],
- )
- @mock.patch("sqlalchemy.orm.Session")
- def test_fetch_from_db_success(self, mock_session, column, conditions,
expected_query):
- """Test successful database queries."""
- mock_session.scalar.return_value = DEFAULT_DATE
-
- result = _fetch_from_db(column, session=mock_session, **conditions)
-
- assert isinstance(result, datetime)
- mock_session.scalar.assert_called_once()
-
- # Check that the correct query was constructed
- call_args = mock_session.scalar.call_args[0][0]
- assert str(call_args) == expected_query
-
- # Verify the actual parameter values
- compiled = call_args.compile()
- for key, value in conditions.items():
- # Note that SQLAlchemy appends the _1 to ensure unique template
field names
- assert compiled.params[f"{key}_1"] == value
-
- @pytest.mark.parametrize(
- ("use_valid_conditions", "scalar_side_effect", "expected_error",
"expected_message"),
- [
- pytest.param(
- False,
- mock.DEFAULT, # This will allow the call to pass through
- AttributeError,
- None,
- id="invalid_attribute",
- ),
- pytest.param(
- True,
- SQLAlchemyError("Database connection failed"),
- SQLAlchemyError,
- "Database connection failed",
- id="database_error",
- ),
- pytest.param(
- True, lambda x: None, ValueError, "No matching record found in
the database", id="no_results"
- ),
- ],
- )
- @mock.patch("sqlalchemy.orm.Session")
- def test_fetch_from_db_error_cases(
- self, mock_session, use_valid_conditions, scalar_side_effect,
expected_error, expected_message
- ):
- """Test database access error handling."""
- model_reference = DagRun.logical_date
- conditions = {"dag_id": "test_dag"} if use_valid_conditions else
{"non_existent_column": "some_value"}
-
- # Configure mock session
- mock_session.scalar.side_effect = scalar_side_effect
-
- with pytest.raises(expected_error, match=expected_message):
- _fetch_from_db(model_reference, session=mock_session, **conditions)
-
- @pytest.mark.parametrize(
- ("reference", "expected_column"),
+ ("reference", "attribute"),
[
pytest.param(
- SerializedReferenceModels.DagRunLogicalDateDeadline(),
DagRun.logical_date, id="logical_date"
- ),
- pytest.param(
- SerializedReferenceModels.DagRunQueuedAtDeadline(),
DagRun.queued_at, id="queued_at"
+ SerializedReferenceModels.DagRunLogicalDateDeadline(),
"logical_date", id="logical_date"
),
+ pytest.param(SerializedReferenceModels.DagRunQueuedAtDeadline(),
"queued_at", id="queued_at"),
pytest.param(
- SerializedReferenceModels.FixedDatetimeDeadline(DEFAULT_DATE),
None, id="fixed_deadline"
- ),
- pytest.param(
- SerializedReferenceModels.AverageRuntimeDeadline(max_runs=10,
min_runs=10),
- None,
- id="average_runtime",
+ ReferenceModels.DagRunLogicalDateDeadline(), "logical_date",
id="legacy_logical_date"
),
+ pytest.param(ReferenceModels.DagRunQueuedAtDeadline(),
"queued_at", id="legacy_queued_at"),
],
)
- def test_deadline_database_integration(self, reference, expected_column,
session):
- """
- Test database integration for all deadline types.
-
- Verifies:
- 1. Calculated deadlines call _fetch_from_db with correct column.
- 2. Fixed deadlines do not interact with database.
- 3. Intervals are added to reference times.
- """
- conditions = {"dag_id": DAG_ID, "run_id": "dagrun_1"}
+ def test_dagrun_references_use_supplied_dagrun(self, reference, attribute,
session):
+ """DagRun references use the in-memory DagRun instead of querying it
again."""
+ dagrun = SimpleNamespace(dag_id=DAG_ID, logical_date=DEFAULT_DATE,
queued_at=DEFAULT_DATE)
interval = timedelta(hours=1)
- with
mock.patch("airflow.serialization.definitions.deadline._fetch_from_db") as
mock_fetch:
- mock_fetch.return_value = DEFAULT_DATE
-
- if expected_column is not None:
- result = reference.evaluate_with(session=session,
interval=interval, **conditions)
- mock_fetch.assert_called_once_with(expected_column,
session=session, **conditions)
- elif isinstance(reference,
SerializedReferenceModels.AverageRuntimeDeadline):
- with mock.patch("airflow._shared.timezones.timezone.utcnow")
as mock_utcnow:
- mock_utcnow.return_value = DEFAULT_DATE
- # No DAG runs exist, so it should use 24-hour default
- result = reference.evaluate_with(session=session,
interval=interval, dag_id=DAG_ID)
- mock_fetch.assert_not_called()
- # Should return None when no DAG runs exist
- assert result is None
- else:
- result = reference.evaluate_with(session=session,
interval=interval)
- mock_fetch.assert_not_called()
- assert result == DEFAULT_DATE + interval
-
- def
test_serialized_fetch_from_db_returns_null_column_without_missing_dagrun_warning(
- self, session, dag_maker, caplog
- ):
- with dag_maker(DAG_ID):
- EmptyOperator(task_id="test_task")
-
- dag_maker.create_dagrun(
- logical_date=None,
- run_id="manual__null_logical_date",
- state=DagRunState.QUEUED,
- )
- session.commit()
- with caplog.at_level("WARNING",
logger="airflow.serialization.definitions.deadline"):
- result = _fetch_serialized_deadline_from_db(
- DagRun.logical_date, session=session, dag_id=DAG_ID,
run_id="manual__null_logical_date"
+ assert getattr(dagrun, attribute) == DEFAULT_DATE
+ assert (
+ reference.evaluate_with(
+ session=session,
+ interval=interval,
+ dagrun=dagrun,
+ dag_id=DAG_ID,
+ run_id="dagrun_1",
+ unexpected="ignored",
)
-
- assert result is None
- assert not any("Could not find DagRun" in record.message for record in
caplog.records)
-
- def test_serialized_fetch_from_db_logs_missing_dagrun(self, session,
caplog):
- with caplog.at_level("WARNING",
logger="airflow.serialization.definitions.deadline"):
- result = _fetch_serialized_deadline_from_db(
- DagRun.logical_date, session=session, dag_id=DAG_ID,
run_id="missing_run"
- )
-
- assert result is None
- assert any("Could not find DagRun" in record.message for record in
caplog.records)
+ == DEFAULT_DATE + interval
+ )
def test_average_runtime_with_sufficient_history(self, session, dag_maker):
"""Test AverageRuntimeDeadline when enough historical data exists."""
@@ -530,7 +391,9 @@ class TestCalculatedDeadlineDatabaseCalls:
with mock.patch("airflow._shared.timezones.timezone.utcnow") as
mock_utcnow:
mock_utcnow.return_value = DEFAULT_DATE
- result = reference.evaluate_with(session=session,
interval=interval, dag_id=DAG_ID)
+ result = reference.evaluate_with(
+ session=session, interval=interval,
dagrun=SimpleNamespace(dag_id=DAG_ID)
+ )
# Calculate expected average: sum(durations) / len(durations)
expected_avg_seconds = sum(durations) / len(durations)
@@ -567,7 +430,9 @@ class TestCalculatedDeadlineDatabaseCalls:
with mock.patch("airflow._shared.timezones.timezone.utcnow") as
mock_utcnow:
mock_utcnow.return_value = DEFAULT_DATE
- result = reference.evaluate_with(session=session,
interval=interval, dag_id=DAG_ID)
+ result = reference.evaluate_with(
+ session=session, interval=interval,
dagrun=SimpleNamespace(dag_id=DAG_ID)
+ )
# Should return None since insufficient runs
assert result is None
@@ -601,7 +466,9 @@ class TestCalculatedDeadlineDatabaseCalls:
with mock.patch("airflow._shared.timezones.timezone.utcnow") as
mock_utcnow:
mock_utcnow.return_value = DEFAULT_DATE
- result = reference.evaluate_with(session=session,
interval=interval, dag_id=DAG_ID)
+ result = reference.evaluate_with(
+ session=session, interval=interval,
dagrun=SimpleNamespace(dag_id=DAG_ID)
+ )
# Should calculate average from 3 runs
expected_avg_seconds = sum(durations) / len(durations) # 4200
seconds
@@ -614,7 +481,9 @@ class TestCalculatedDeadlineDatabaseCalls:
with mock.patch("airflow._shared.timezones.timezone.utcnow") as
mock_utcnow:
mock_utcnow.return_value = DEFAULT_DATE
- result = reference.evaluate_with(session=session,
interval=interval, dag_id=DAG_ID)
+ result = reference.evaluate_with(
+ session=session, interval=interval,
dagrun=SimpleNamespace(dag_id=DAG_ID)
+ )
assert result is None
def test_average_runtime_min_runs_validation(self):
@@ -663,7 +532,9 @@ class TestCalculatedDeadlineDatabaseCalls:
with mock.patch("airflow._shared.timezones.timezone.utcnow") as
mock_utcnow:
mock_utcnow.return_value = DEFAULT_DATE
- result = reference.evaluate_with(session=session,
interval=interval, dag_id=DAG_ID)
+ result = reference.evaluate_with(
+ session=session, interval=interval,
dagrun=SimpleNamespace(dag_id=DAG_ID)
+ )
# Average must be over the 3 successful 60s runs only (not the 36000s
failures).
expected = DEFAULT_DATE + timedelta(seconds=success_duration) +
interval
@@ -687,61 +558,15 @@ class TestCalculatedDeadlineDatabaseCalls:
session.commit()
reference =
SerializedReferenceModels.AverageRuntimeDeadline(max_runs=10, min_runs=3)
- result = reference.evaluate_with(session=session,
interval=timedelta(hours=1), dag_id=DAG_ID)
+ result = reference.evaluate_with(
+ session=session, interval=timedelta(hours=1),
dagrun=SimpleNamespace(dag_id=DAG_ID)
+ )
assert result is None
class TestDeadlineReference:
"""DeadlineReference lives in definitions/deadlines.py but properly
testing them requires DB access."""
- DEFAULT_INTERVAL = timedelta(hours=1)
- DEFAULT_ARGS = {"interval": DEFAULT_INTERVAL}
-
- @pytest.mark.parametrize("reference", REFERENCE_TYPES)
- @pytest.mark.db_test
- def test_deadline_evaluate_with(self, reference, session):
- """Test that all deadline types evaluate correctly with their required
conditions."""
- conditions = {
- "dag_id": DAG_ID,
- "run_id": "dagrun_1",
- "unexpected": "param", # Add an unexpected parameter.
- "extra": "kwarg", # Add another unexpected parameter.
- }
-
- with mock.patch.object(reference, "_evaluate_with") as mock_evaluate:
- mock_evaluate.return_value = DEFAULT_DATE
-
- if reference.required_kwargs:
- result = reference.evaluate_with(**self.DEFAULT_ARGS,
session=session, **conditions)
- else:
- result = reference.evaluate_with(**self.DEFAULT_ARGS,
session=session)
-
- # Verify only expected kwargs are passed through.
- expected_kwargs = {k: conditions[k] for k in
reference.required_kwargs if k in conditions}
- expected_kwargs["session"] = session
-
- mock_evaluate.assert_called_once_with(**expected_kwargs)
-
- assert result == DEFAULT_DATE + self.DEFAULT_INTERVAL
-
- @pytest.mark.parametrize("reference", REFERENCE_TYPES)
- @pytest.mark.db_test
- def test_deadline_missing_required_kwargs(self, reference, session):
- """Test that deadlines raise appropriate errors for missing required
parameters."""
- if reference.required_kwargs:
- with pytest.raises(
- ValueError, match=re.escape(f"{reference.__class__.__name__}
is missing required parameters:")
- ) as raised_exception:
- reference.evaluate_with(session=session, **self.DEFAULT_ARGS)
-
- assert all(substring in str(raised_exception.value) for substring
in reference.required_kwargs)
- else:
- # Let the lack of an exception here effectively assert that no
exception is raised.
- reference.evaluate_with(session=session, **self.DEFAULT_ARGS)
-
- for required_param in reference.required_kwargs:
- assert required_param in str(raised_exception.value)
-
def test_deadline_reference_creation(self):
"""Test that DeadlineReference provides consistent interface and
types."""
fixed_reference = DeadlineReference.FIXED_DATETIME(DEFAULT_DATE)
diff --git a/airflow-core/tests/unit/models/test_deadline_alert.py
b/airflow-core/tests/unit/models/test_deadline_alert.py
index 56ed2f8cc86..54a138b64a0 100644
--- a/airflow-core/tests/unit/models/test_deadline_alert.py
+++ b/airflow-core/tests/unit/models/test_deadline_alert.py
@@ -17,6 +17,7 @@
from __future__ import annotations
from datetime import timedelta
+from types import SimpleNamespace
from unittest.mock import Mock
import pytest
@@ -24,6 +25,7 @@ import time_machine
from sqlalchemy import select
from airflow._shared.timezones import timezone
+from airflow.exceptions import RemovedInAirflow4Warning
from airflow.models.deadline_alert import DeadlineAlert
from airflow.models.serialized_dag import SerializedDagModel
from airflow.sdk.definitions.deadline import BaseDeadlineReference,
DeadlineReference
@@ -218,23 +220,83 @@ class TestDeadlineAlert:
wrapper =
SerializedReferenceModels.SerializedCustomReference(inner_ref)
- wrapper.evaluate_with(
+ with pytest.warns(RemovedInAirflow4Warning, match="required_kwargs is
deprecated"):
+ wrapper.evaluate_with(
+ session=None,
+ interval=timedelta(hours=1),
+ dag_id="test_dag",
+ run_id="test_run",
+ extra_param="should_be_filtered",
+ )
+
+ inner_ref._evaluate_with.assert_called_once_with(session=None,
dag_id="test_dag", run_id="test_run")
+
+ # try calling with missing required parameters
+ with pytest.warns(RemovedInAirflow4Warning, match="required_kwargs is
deprecated"):
+ with pytest.raises(ValueError, match="missing required parameters:
run_id"):
+ wrapper.evaluate_with(
+ session=None,
+ interval=timedelta(hours=1),
+ dag_id="test_dag",
+ )
+
+ def test_serialized_custom_reference_receives_dagrun(self):
+ class DagRunReference(BaseDeadlineReference):
+ def _evaluate_with(self, *, session, dagrun):
+ return dagrun.queued_at
+
+ wrapper =
SerializedReferenceModels.SerializedCustomReference(DagRunReference())
+ dagrun = SimpleNamespace(queued_at=DEFAULT_DATE)
+
+ assert wrapper.evaluate_with(
session=None,
interval=timedelta(hours=1),
+ dagrun=dagrun,
dag_id="test_dag",
run_id="test_run",
- extra_param="should_be_filtered",
- )
+ unexpected="ignored",
+ ) == DEFAULT_DATE + timedelta(hours=1)
- inner_ref._evaluate_with.assert_called_once_with(session=None,
dag_id="test_dag", run_id="test_run")
+ def
test_serialized_custom_reference_receives_dagrun_with_var_keyword_arguments(self):
+ class DagRunReference(BaseDeadlineReference):
+ def _evaluate_with(self, *, session, dagrun, **kwargs):
+ return dagrun.queued_at
- # try calling with missing required parameters
- with pytest.raises(ValueError, match="missing required parameters:
run_id"):
- wrapper.evaluate_with(
- session=None,
- interval=timedelta(hours=1),
- dag_id="test_dag",
- )
+ wrapper =
SerializedReferenceModels.SerializedCustomReference(DagRunReference())
+ dagrun = SimpleNamespace(queued_at=DEFAULT_DATE)
+
+ assert wrapper.evaluate_with(
+ session=None,
+ interval=timedelta(hours=1),
+ dagrun=dagrun,
+ dag_id="test_dag",
+ run_id="test_run",
+ ) == DEFAULT_DATE + timedelta(hours=1)
+
+ def
test_serialized_custom_reference_var_keyword_receives_full_context(self):
+ class VarKeywordReference(BaseDeadlineReference):
+ received_kwargs: dict[str, object]
+
+ def _evaluate_with(self, *, session, **kwargs):
+ self.received_kwargs = kwargs
+ return DEFAULT_DATE
+
+ reference = VarKeywordReference()
+ wrapper =
SerializedReferenceModels.SerializedCustomReference(reference)
+ dagrun = SimpleNamespace()
+
+ assert wrapper.evaluate_with(
+ session=None,
+ interval=timedelta(hours=1),
+ dagrun=dagrun,
+ dag_id="test_dag",
+ run_id="test_run",
+ ) == DEFAULT_DATE + timedelta(hours=1)
+ assert reference.received_kwargs == {
+ "dagrun": dagrun,
+ "dag_id": "test_dag",
+ "run_id": "test_run",
+ }
def test_core_deadline_reference_treated_as_builtins(self):
"""Test that refs from airflow.models.deadline are still treated as
builtins."""
diff --git a/task-sdk/src/airflow/sdk/definitions/deadline.py
b/task-sdk/src/airflow/sdk/definitions/deadline.py
index 3ef17245803..61dbe09b57b 100644
--- a/task-sdk/src/airflow/sdk/definitions/deadline.py
+++ b/task-sdk/src/airflow/sdk/definitions/deadline.py
@@ -47,7 +47,8 @@ class BaseDeadlineReference(ABC):
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). A custom
+ ``_evaluate_with()`` with deferred Core imports (imports inside the method
body). Its
+ keyword-only ``dagrun`` parameter receives the DagRun being evaluated. A
custom
reference must be decorated with ``@deadline_reference`` and listed in the
``deadline_references``
attribute of an ``AirflowPlugin``; see
:external:doc:`howto/deadline-alerts`.
"""
@@ -212,8 +213,8 @@ class DeadlineReference:
The public interface class for all DeadlineReference options.
This class provides a unified interface for working with Deadlines,
supporting both
- calculated deadlines (which fetch values from the database) and fixed
deadlines
- (which return a predefined datetime).
+ calculated deadlines, including references to the DagRun being evaluated
and
+ historical runtime data, and fixed deadlines (which return a predefined
datetime).
------
Usage:
@@ -240,17 +241,15 @@ class DeadlineReference:
),
)
- 3. Evaluating deadlines will ignore unexpected parameters:
+ 3. Custom references receive the DagRun being evaluated:
.. code-block:: python
- # For deadlines requiring parameters:
- deadline = DeadlineReference.DAGRUN_LOGICAL_DATE
- deadline.evaluate_with(dag_id=dag.dag_id)
-
- # For deadlines with no required parameters:
- deadline = DeadlineReference.FIXED_DATETIME(datetime(2025, 5, 4))
- deadline.evaluate_with()
+ class MyDeadlineReference(BaseDeadlineReference):
+ def _evaluate_with(self, *, session, dagrun):
+ # Add your business logic here; dagrun.logical_date is
available to use.
+ my_datetime = my_business_logic(dagrun.logical_date)
+ return my_datetime
"""
class TYPES:
@@ -395,11 +394,10 @@ def deadline_reference(deadline_reference_type=None):
@deadline_reference()
class MyCustomReference(BaseDeadlineReference):
# By default, evaluate_with will be called when a new dagrun is
created.
- def _evaluate_with(self, *, session: Session, **kwargs) ->
datetime:
- # Put your business logic here (use deferred imports for Core
types)
- from airflow.models import DagRun
-
- return some_datetime
+ def _evaluate_with(self, *, session: Session, dagrun) -> datetime:
+ # Add your business logic here; dagrun.logical_date is
available to use.
+ my_datetime = my_business_logic(dagrun.logical_date)
+ return my_datetime
def serialize_reference(self) -> dict:
return {"reference_type": self.reference_name}
@@ -408,9 +406,10 @@ def deadline_reference(deadline_reference_type=None):
# Optionally, specify when it is calculated by providing a
DeadlineReference.TYPES value.
@deadline_reference(DeadlineReference.TYPES.DAGRUN_QUEUED)
class MyQueuedRef(BaseDeadlineReference):
- def _evaluate_with(self, *, session: Session, **kwargs) ->
datetime:
- # Put your business logic here
- return some_datetime
+ def _evaluate_with(self, *, session: Session, dagrun) -> datetime:
+ # Add your business logic here; dagrun.queued_at is available
to use.
+ my_datetime = my_business_logic(dagrun.queued_at)
+ return my_datetime
def serialize_reference(self) -> dict:
return {"reference_type": self.reference_name}
