AIP-86 - Implement Deadline Callback definition

.pre-commit-config.yaml

@@ -1681,6 +1681,7 @@ repos:
           ^airflow-core/src/airflow/models/baseoperator\.py$|
           ^airflow-core/src/airflow/models/connection\.py$|
           ^airflow-core/src/airflow/models/dag\.py$|
+          ^airflow-core/src/airflow/models/deadline\.py$|
           ^airflow-core/src/airflow/models/dagbag\.py$|
           ^airflow-core/src/airflow/models/dagrun\.py$|
           ^airflow-core/src/airflow/models/mappedoperator\.py$|

airflow-core/docs/img/airflow_erd.sha256

@@ -1 +1 @@
-efbae2f1de68413e5a6f671a306e748581fe454b81e25eeb2927567f11ebd59c
+625f362919679fe85b2bfb3f9e053261248ac6ec7a974eee51012e55a8105b94

airflow-core/docs/migrations-ref.rst

@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | Revision ID             | Revises ID       | Airflow Version   | Description                                                  |
 +=========================+==================+===================+==============================================================+
-| ``3bda03debd04`` (head) | ``f56f68b9e02f`` | ``3.1.0``         | Add url template and template params to DagBundleModel.      |
+| ``808787349f22`` (head) | ``3bda03debd04`` | ``3.1.0``         | Modify deadline's callback schema.                           |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``3bda03debd04``        | ``f56f68b9e02f`` | ``3.1.0``         | Add url template and template params to DagBundleModel.      |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | ``f56f68b9e02f``        | ``09fa89ba1710`` | ``3.1.0``         | Add callback_state to deadline.                              |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+

airflow-core/src/airflow/migrations/versions/0081_3_1_0_modify_deadline_callback_schema.py

@@ -0,0 +1,57 @@
+#
+# 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.
+
+"""
+Modify deadline's callback schema.
+
+Revision ID: 808787349f22
+Revises: 3bda03debd04
+Create Date: 2025-07-31 19:35:53.150465
+
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+import sqlalchemy_jsonfield
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "808787349f22"
+down_revision = "3bda03debd04"
+branch_labels = None
+depends_on = None
+airflow_version = "3.1.0"
+
+
+def upgrade():
+    """Replace deadline table's string callback and JSON callback_kwargs with JSON callback."""
+    with op.batch_alter_table("deadline", schema=None) as batch_op:
+        batch_op.drop_column("callback")
+        batch_op.drop_column("callback_kwargs")
+        batch_op.add_column(sa.Column("callback", sqlalchemy_jsonfield.jsonfield.JSONField(), nullable=False))
+
+
+def downgrade():
+    """Replace deadline table's JSON callback with string callback and JSON callback_kwargs."""
+    with op.batch_alter_table("deadline", schema=None) as batch_op:
+        batch_op.drop_column("callback")
+        batch_op.add_column(
+            sa.Column("callback_kwargs", sqlalchemy_jsonfield.jsonfield.JSONField(), nullable=True)
+        )
+        batch_op.add_column(sa.Column("callback", sa.String(length=500), nullable=False))

airflow-core/src/airflow/models/dag.py

@@ -1605,7 +1605,6 @@ def create_dagrun(
                         run_id=run_id,
                     ),
                     callback=self.deadline.callback,
-                    callback_kwargs=self.deadline.callback_kwargs or {},
                     dag_id=self.dag_id,
                     dagrun_id=orm_dagrun.id,
                 )

airflow-core/src/airflow/models/deadline.py

@@ -21,7 +21,8 @@
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 from enum import Enum
-from typing import TYPE_CHECKING, Any
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, cast
 
 import sqlalchemy_jsonfield
 import uuid6
@@ -33,6 +34,7 @@
 from airflow._shared.timezones import timezone
 from airflow.models import Trigger
 from airflow.models.base import Base, StringID
+from airflow.serialization.serde import deserialize, serialize
 from airflow.settings import json
 from airflow.triggers.deadline import PAYLOAD_STATUS_KEY, DeadlineCallbackTrigger
 from airflow.utils.log.logging_mixin import LoggingMixin
@@ -42,6 +44,7 @@
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
 
+    from airflow.sdk.definitions.deadline import Callback
     from airflow.triggers.base import TriggerEvent
 
 
@@ -76,7 +79,11 @@ def __get__(self, instance, cls=None):
 
 
 class DeadlineCallbackState(str, Enum):
-    """All possible states of deadline callbacks."""
+    """
+    All possible states of deadline callbacks once the deadline is missed.
+
+    `None` state implies that the deadline is pending (`deadline_time` hasn't passed yet).
+    """
 
     QUEUED = "queued"
     SUCCESS = "success"
@@ -96,10 +103,8 @@ class Deadline(Base):
 
     # The time after which the Deadline has passed and the callback should be triggered.
     deadline_time = Column(UtcDateTime, nullable=False)
-    # The Callback to be called when the Deadline has passed.
-    callback = Column(String(500), nullable=False)
-    # Serialized kwargs to pass to the callback.
-    callback_kwargs = Column(sqlalchemy_jsonfield.JSONField(json=json))
+    # The (serialized) callback to be called when the Deadline has passed.
+    _callback = Column("callback", sqlalchemy_jsonfield.JSONField(json=json), nullable=False)
     # The state of the deadline callback
     callback_state = Column(String(20))
 
@@ -114,15 +119,13 @@ class Deadline(Base):
     def __init__(
         self,
         deadline_time: datetime,
-        callback: str,
-        callback_kwargs: dict | None = None,
+        callback: Callback,
         dag_id: str | None = None,
         dagrun_id: int | None = None,
     ):
         super().__init__()
         self.deadline_time = deadline_time
-        self.callback = callback
-        self.callback_kwargs = callback_kwargs
+        self._callback = serialize(callback)
         self.dag_id = dag_id
         self.dagrun_id = dagrun_id
 
@@ -136,11 +139,10 @@ def _determine_resource() -> tuple[str, str]:
             return "Unknown", ""
 
         resource_type, resource_details = _determine_resource()
-        callback_kwargs = json.dumps(self.callback_kwargs) if self.callback_kwargs else ""
 
         return (
             f"[{resource_type} Deadline] {resource_details} needed by "
-            f"{self.deadline_time} or run: {self.callback}({callback_kwargs})"
+            f"{self.deadline_time} or run: {self.callback.path}({self.callback.kwargs or ''})"
         )
 
     @classmethod
@@ -193,18 +195,30 @@ def prune_deadlines(cls, *, session: Session, conditions: dict[Column, Any]) ->
 
         return deleted_count
 
+    @cached_property
+    def callback(self) -> Callback:
+        return cast("Callback", deserialize(self._callback))
+
     def handle_miss(self, session: Session):
-        """Handle a missed deadline by creating a trigger to run the callback."""
-        # TODO: check to see if the callback is meant to run in triggerer or executor. For now, the code below assumes it's for the triggerer
-        callback_trigger = DeadlineCallbackTrigger(
-            callback_path=self.callback,
-            callback_kwargs=self.callback_kwargs,
-        )
+        """Handle a missed deadline by running the callback in the appropriate host and updating the `callback_state`."""
+        from airflow.sdk.definitions.deadline import AsyncCallback, SyncCallback
+
+        if isinstance(self.callback, AsyncCallback):
+            callback_trigger = DeadlineCallbackTrigger(
+                callback_path=self.callback.path,
+                callback_kwargs=self.callback.kwargs,
+            )
+            trigger_orm = Trigger.from_object(callback_trigger)
+            session.add(trigger_orm)
+            session.flush()
+            self.trigger = trigger_orm
+
+        elif isinstance(self.callback, SyncCallback):
+            raise NotImplementedError("SyncCallback is currently not supported")
+
+        else:
+            raise TypeError("Unknown Callback type")
 
-        trigger_orm = Trigger.from_object(callback_trigger)
-        session.add(trigger_orm)
-        session.flush()
-        self.trigger_id = trigger_orm.id
         self.callback_state = DeadlineCallbackState.QUEUED
         session.add(self)
 

airflow-core/src/airflow/triggers/deadline.py

@@ -22,7 +22,7 @@
 from typing import Any
 
 from airflow.triggers.base import BaseTrigger, TriggerEvent
-from airflow.utils.module_loading import import_string
+from airflow.utils.module_loading import import_string, qualname
 
 log = logging.getLogger(__name__)
 
@@ -40,8 +40,8 @@ def __init__(self, callback_path: str, callback_kwargs: dict[str, Any] | None =
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         return (
-            f"{type(self).__module__}.{type(self).__qualname__}",
-            {"callback_path": self.callback_path, "callback_kwargs": self.callback_kwargs},
+            qualname(self),
+            {attr: getattr(self, attr) for attr in ("callback_path", "callback_kwargs")},
         )
 
     async def run(self) -> AsyncIterator[TriggerEvent]:

airflow-core/src/airflow/utils/db.py

@@ -93,7 +93,7 @@ class MappedClassProtocol(Protocol):
     "2.10.3": "5f2621c13b39",
     "3.0.0": "29ce7909c52b",
     "3.0.3": "fe199e1abd77",
-    "3.1.0": "3bda03debd04",
+    "3.1.0": "808787349f22",
 }
 
 

airflow-core/src/airflow/utils/module_loading.py

@@ -53,6 +53,7 @@ def import_string(dotted_path: str):
 
     Raise ImportError if the import failed.
     """
+    # TODO: Add support for nested classes. Currently, it only works for top-level classes.
     try:
         module_path, class_name = dotted_path.rsplit(".", 1)
     except ValueError:

airflow-core/tests/unit/models/test_dag.py

@@ -70,7 +70,7 @@
 from airflow.sdk.definitions._internal.contextmanager import TaskGroupContext
 from airflow.sdk.definitions._internal.templater import NativeEnvironment, SandboxedEnvironment
 from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetAll, AssetAny
-from airflow.sdk.definitions.deadline import DeadlineAlert, DeadlineReference
+from airflow.sdk.definitions.deadline import AsyncCallback, DeadlineAlert, DeadlineReference
 from airflow.sdk.definitions.param import Param
 from airflow.timetables.base import DagRunInfo, DataInterval, TimeRestriction, Timetable
 from airflow.timetables.simple import (
@@ -112,6 +112,11 @@
 repo_root = Path(__file__).parents[2]
 
 
+async def empty_callback_for_deadline():
+    """Used in a number of tests to confirm that Deadlines and DeadlineAlerts function correctly."""
+    pass
+
+
 @pytest.fixture
 def clear_dags():
     clear_db_dags()
@@ -2024,7 +2029,7 @@ def test_dagrun_deadline(self, reference_type, reference_column, dag_maker, sess
             deadline=DeadlineAlert(
                 reference=reference_type,
                 interval=interval,
-                callback=print,
+                callback=AsyncCallback(empty_callback_for_deadline),
             ),
         ) as dag:
             ...

airflow-core/tests/unit/models/test_dagrun.py

@@ -44,7 +44,7 @@
 from airflow.providers.standard.operators.empty import EmptyOperator
 from airflow.providers.standard.operators.python import PythonOperator, ShortCircuitOperator
 from airflow.sdk import BaseOperator, setup, task, task_group, teardown
-from airflow.sdk.definitions.deadline import DeadlineAlert, DeadlineReference
+from airflow.sdk.definitions.deadline import AsyncCallback, DeadlineAlert, DeadlineReference
 from airflow.serialization.serialized_objects import SerializedDAG
 from airflow.stats import Stats
 from airflow.triggers.base import StartTriggerArgs
@@ -69,7 +69,7 @@
 DEFAULT_DATE = pendulum.instance(_DEFAULT_DATE)
 
 
-def test_callback_for_deadline():
+async def empty_callback_for_deadline():
     """Used in a number of tests to confirm that Deadlines and DeadlineAlerts function correctly."""
     pass
 
@@ -1294,7 +1294,7 @@ def on_success_callable(context):
             deadline=DeadlineAlert(
                 reference=DeadlineReference.FIXED_DATETIME(future_date),
                 interval=datetime.timedelta(hours=1),
-                callback=test_callback_for_deadline,
+                callback=AsyncCallback(empty_callback_for_deadline),
             ),
         ) as dag:
             ...