AIP-86 - Add Deadline References (#50677)

ferruzzi · web-flow · commit a3295a88eb5d · 2025-06-02T11:17:44.000-07:00
* AIP-86 - Add Deadline References
diff --git a/airflow-core/src/airflow/models/deadline.py b/airflow-core/src/airflow/models/deadline.py
@@ -17,25 +17,30 @@
 from __future__ import annotations
 
 import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from datetime import datetime, timedelta
-from enum import Enum
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Any, Callable
 
 import sqlalchemy_jsonfield
 import uuid6
-from sqlalchemy import Column, ForeignKey, Index, Integer, String
+from sqlalchemy import Column, ForeignKey, Index, Integer, String, select
+from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.orm import relationship
 from sqlalchemy_utils import UUIDType
 
 from airflow.models.base import Base, StringID
 from airflow.settings import json
+from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.module_loading import import_string, is_valid_dotpath
 from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.utils.sqlalchemy import UtcDateTime
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
 
+    from airflow.sdk.definitions.deadline import DeadlineReference
+
 logger = logging.getLogger(__name__)
 
 
@@ -100,33 +105,67 @@ def add_deadline(cls, deadline: Deadline, session: Session = NEW_SESSION):
         session.add(deadline)
 
 
-class DeadlineReference(Enum):
+class ReferenceModels:
     """
-    Store the calculation methods for the various Deadline Alert triggers.
+    Store the implementations for the different Deadline References.
 
-    TODO:  PLEASE NOTE This class is a placeholder and will be expanded in the next PR.
+    After adding the implementations here, all DeadlineReferences should be added
+    to the user interface in airflow.sdk.definitions.deadline.DeadlineReference
+    """
 
-    ------
-    Usage:
-    ------
+    class BaseDeadlineReference(LoggingMixin, ABC):
+        """Base class for all Deadline implementations."""
 
-    Example use when defining a deadline in a DAG:
+        # Set of required kwargs - subclasses should override this.
+        required_kwargs: set[str] = set()
 
-    DAG(
-        dag_id='dag_with_deadline',
-        deadline=DeadlineAlert(
-            reference=DeadlineReference.DAGRUN_LOGICAL_DATE,
-            interval=timedelta(hours=1),
-            callback=hello_callback,
-        )
-    )
+        def evaluate_with(self, **kwargs: Any) -> datetime:
+            """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}
 
-    To parse the deadline reference later we will use something like:
+            if missing_kwargs := self.required_kwargs - filtered_kwargs.keys():
+                raise ValueError(
+                    f"{self.__class__.__name__} is missing required parameters: {', '.join(missing_kwargs)}"
+                )
 
-    dag.deadline.reference.evaluate_with(dag_id=dag.dag_id)
-    """
+            if extra_kwargs := kwargs.keys() - filtered_kwargs.keys():
+                self.log.debug("Ignoring unexpected parameters: %s", ", ".join(extra_kwargs))
+
+            return self._evaluate_with(**filtered_kwargs)
+
+        @abstractmethod
+        def _evaluate_with(self, **kwargs: Any) -> datetime:
+            """Must be implemented by subclasses to perform the actual evaluation."""
+            raise NotImplementedError
+
+    @dataclass
+    class FixedDatetimeDeadline(BaseDeadlineReference):
+        """A deadline that always returns a fixed datetime."""
+
+        _datetime: datetime
+
+        def _evaluate_with(self, **kwargs: Any) -> datetime:
+            return self._datetime
+
+    class DagRunLogicalDateDeadline(BaseDeadlineReference):
+        """A deadline that returns a DagRun's logical date."""
+
+        required_kwargs = {"dag_id"}
+
+        def _evaluate_with(self, **kwargs: Any) -> datetime:
+            from airflow.models import DagRun
+
+            return _fetch_from_db(DagRun.logical_date, **kwargs)
+
+    class DagRunQueuedAtDeadline(BaseDeadlineReference):
+        """A deadline that returns when a DagRun was queued."""
 
-    DAGRUN_LOGICAL_DATE = "dagrun_logical_date"
+        required_kwargs = {"dag_id"}
+
+        def _evaluate_with(self, **kwargs: Any) -> datetime:
+            from airflow.models import DagRun
+
+            return _fetch_from_db(DagRun.queued_at, **kwargs)
 
 
 class DeadlineAlert:
@@ -186,3 +225,52 @@ def serialize_deadline_alert(self):
                 "callback_kwargs": self.callback_kwargs,
             }
         )
+
+
+@provide_session
+def _fetch_from_db(model_reference: Column, session=None, **conditions) -> datetime:
+    """
+    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():
+        query = query.where(getattr(model_reference.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/tests/unit/models/test_deadline.py b/airflow-core/tests/unit/models/test_deadline.py
@@ -18,14 +18,18 @@
 
 import json
 import logging
+from datetime import datetime
+from unittest import mock
 
 import pytest
 import time_machine
 from sqlalchemy import select
+from sqlalchemy.exc import SQLAlchemyError
 
 from airflow.models import DagRun
-from airflow.models.deadline import Deadline, DeadlineAlert
+from airflow.models.deadline import Deadline, DeadlineAlert, _fetch_from_db
 from airflow.providers.standard.operators.empty import EmptyOperator
+from airflow.sdk.definitions.deadline import DeadlineReference
 from airflow.utils.state import DagRunState
 
 from tests_common.test_utils import db
@@ -182,3 +186,126 @@ def test_log_unimportable_but_properly_formatted_callback(self, caplog):
 
             assert "could not be imported" in caplog.text
             assert path == UNIMPORTABLE_DOT_PATH
+
+
+@pytest.mark.db_test
+class TestCalculatedDeadlineDatabaseCalls:
+    @staticmethod
+    def setup_method():
+        _clean_db()
+
+    @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, caplog
+    ):
+        """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 caplog.at_level(logging.ERROR):
+            with pytest.raises(expected_error, match=expected_message):
+                _fetch_from_db(model_reference, session=mock_session, **conditions)
+            if expected_message:
+                assert expected_message in caplog.text
+
+    @pytest.mark.parametrize(
+        "reference, expected_column",
+        [
+            pytest.param(DeadlineReference.DAGRUN_LOGICAL_DATE, DagRun.logical_date, id="logical_date"),
+            pytest.param(DeadlineReference.DAGRUN_QUEUED_AT, DagRun.queued_at, id="queued_at"),
+            pytest.param(DeadlineReference.FIXED_DATETIME(DEFAULT_DATE), None, id="fixed_deadline"),
+        ],
+    )
+    def test_deadline_database_integration(self, reference, expected_column):
+        """
+        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.
+        """
+        conditions = {"dag_id": DAG_ID}
+
+        with mock.patch("airflow.models.deadline._fetch_from_db") as mock_fetch:
+            mock_fetch.return_value = DEFAULT_DATE
+
+            if expected_column is not None:
+                result = reference.evaluate_with(**conditions)
+                mock_fetch.assert_called_once_with(expected_column, **conditions)
+            else:
+                result = reference.evaluate_with(**conditions)
+                mock_fetch.assert_not_called()
+
+            assert result == DEFAULT_DATE
diff --git a/task-sdk/src/airflow/sdk/definitions/deadline.py b/task-sdk/src/airflow/sdk/definitions/deadline.py
@@ -0,0 +1,66 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+
+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).
+
+    ------
+    Usage:
+    ------
+
+    1. Example deadline references:
+       fixed = DeadlineReference.FIXED_DATETIME(datetime(2025, 5, 4))
+       logical = DeadlineReference.DAGRUN_LOGICAL_DATE
+       queued = DeadlineReference.DAGRUN_QUEUED_AT
+
+    2. Using in a DAG:
+       DAG(
+           dag_id='dag_with_deadline',
+           deadline=DeadlineAlert(
+               reference=DeadlineReference.DAGRUN_LOGICAL_DATE,
+               interval=timedelta(hours=1),
+               callback=hello_callback,
+           )
+       )
+
+    3. Evaluating deadlines will ignore unexpected parameters:
+       # 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()
+    """
+
+    from airflow.models.deadline import ReferenceModels
+
+    DAGRUN_LOGICAL_DATE: ReferenceModels.BaseDeadlineReference = ReferenceModels.DagRunLogicalDateDeadline()
+    DAGRUN_QUEUED_AT: ReferenceModels.BaseDeadlineReference = ReferenceModels.DagRunQueuedAtDeadline()
+
+    @classmethod
+    def FIXED_DATETIME(cls, datetime: datetime) -> ReferenceModels.BaseDeadlineReference:
+        return cls.ReferenceModels.FixedDatetimeDeadline(datetime)
diff --git a/task-sdk/tests/task_sdk/definitions/test_deadline.py b/task-sdk/tests/task_sdk/definitions/test_deadline.py
