Dark Mode

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

AIP-86 - Implement Deadline Callback definition #53951

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking "Sign up for GitHub", you agree to our terms of service and privacy statement. We'll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ferruzzi merged 5 commits into apache:main from aws-mwaa:ramitkataria/deadlines/callback-def
Aug 6, 2025
Merged

AIP-86 - Implement Deadline Callback definition #53951

Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .pre-commit-config.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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$|
Expand Down
2 changes: 1 addition & 1 deletion airflow-core/docs/img/airflow_erd.sha256
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1 +1 @@
efbae2f1de68413e5a6f671a306e748581fe454b81e25eeb2927567f11ebd59c
625f362919679fe85b2bfb3f9e053261248ac6ec7a974eee51012e55a8105b94
264 changes: 130 additions & 134 deletions airflow-core/docs/img/airflow_erd.svg
View file
Open in desktop
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
4 changes: 3 additions & 1 deletion airflow-core/docs/migrations-ref.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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. |
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
Expand Down
57 changes: 57 additions & 0 deletions airflow-core/src/airflow/migrations/versions/0081_3_1_0_modi fy_deadline_callback_schema.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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))
1 change: 0 additions & 1 deletion airflow-core/src/airflow/models/dag.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down
58 changes: 36 additions & 22 deletions airflow-core/src/airflow/models/deadline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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


Expand Down Expand Up @@ -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"
Expand All @@ -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))

Expand All @@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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)

Expand Down
6 changes: 3 additions & 3 deletions airflow-core/src/airflow/triggers/deadline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand All @@ -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]:
Expand Down
2 changes: 1 addition & 1 deletion airflow-core/src/airflow/utils/db.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
}


Expand Down
1 change: 1 addition & 0 deletions airflow-core/src/airflow/utils/module_loading.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
9 changes: 7 additions & 2 deletions airflow-core/tests/unit/models/test_dag.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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:
...
Expand Down
6 changes: 3 additions & 3 deletions airflow-core/tests/unit/models/test_dagrun.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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:
...
Expand Down
Loading
Loading