Store executor callback DagRun references in the database

Author: fat-catTW

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                                                  |
 +=========================+==================+===================+==============================================================+
-| ``e5a91c7f42b3`` (head) | ``ca8499dc1004`` | ``3.4.0``         | Add language column to dag_code.                             |
+| ``0f4c9a2b8e1d`` (head) | ``e5a91c7f42b3`` | ``3.4.0``         | Add dagrun_id foreign key to callback.                       |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``e5a91c7f42b3``        | ``ca8499dc1004`` | ``3.4.0``         | Add language column to dag_code.                             |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | ``ca8499dc1004``        | ``a61f0c9d2b47`` | ``3.4.0``         | Add source_reference to import_error.                        |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -1331,6 +1331,7 @@ def _enqueue_executor_callbacks(self, session: Session) -> None:
 
         pending_callbacks = session.scalars(
             select(ExecutorCallback)
+            .options(selectinload(ExecutorCallback.dag_run))
             .where(ExecutorCallback.type == CallbackType.EXECUTOR)
             .where(ExecutorCallback.state == CallbackState.PENDING)
             .order_by(ExecutorCallback.priority_weight.desc())
@@ -1350,25 +1351,19 @@ def _enqueue_executor_callbacks(self, session: Session) -> None:
                     # Can't happen since we queried ExecutorCallback, but satisfies mypy.
                     continue
 
-                # TODO: Add dagrun_id as a proper ORM foreign key on the callback table instead of storing in data dict.
-                #       This would eliminate this reconstruction step. For now, all ExecutorCallbacks
-                #       are expected to have dag_run_id set in their data dict (e.g., by Deadline.handle_miss).
-                if not isinstance(callback.data, dict) or "dag_run_id" not in callback.data:
-                    self.log.error(
-                        "ExecutorCallback %s is missing required 'dag_run_id' in data dict. "
-                        "This indicates a bug in callback creation. Skipping callback.",
-                        callback.id,
+                if callback.dagrun_id is None:
+                    self.log.warning(
+                        "Executor callback is missing dagrun_id.",
+                        callback_id=callback.id,
                     )
                     continue
 
-                dag_run_id = callback.data["dag_run_id"]
-                dag_run = session.get(DagRun, dag_run_id)
-
+                dag_run = callback.dag_run
                 if dag_run is None:
                     self.log.warning(
-                        "Could not find DagRun with id=%s for callback %s. DagRun may have been deleted.",
-                        dag_run_id,
-                        callback.id,
+                        "Could not find DagRun for executor callback. DagRun may have been deleted.",
+                        callback_id=callback.id,
+                        dagrun_id=callback.dagrun_id,
                     )
                     continue
 

airflow-core/src/airflow/migrations/versions/0141_3_4_0_add_dagrun_id_to_callback.py

@@ -0,0 +1,165 @@
+# 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.
+
+"""
+Add dagrun_id foreign key to callback.
+
+Revision ID: 0f4c9a2b8e1d
+Revises: e5a91c7f42b3
+Create Date: 2026-08-06 00:00:00.000000
+
+"""
+
+from __future__ import annotations
+
+import json
+
+import sqlalchemy as sa
+from alembic import context, op
+
+from airflow.configuration import conf
+
+# revision identifiers, used by Alembic.
+revision = "0f4c9a2b8e1d"
+down_revision = "e5a91c7f42b3"
+branch_labels = None
+depends_on = None
+airflow_version = "3.4.0"
+
+
+def _deserialize_data(data):
+    if isinstance(data, str):
+        try:
+            data = json.loads(data)
+        except json.JSONDecodeError:
+            return {}
+    return data.get("__var", data) if isinstance(data, dict) else {}
+
+
+def _extract_dagrun_id(data) -> int | None:
+    dagrun_id = _deserialize_data(data).get("dag_run_id")
+    if dagrun_id is None:
+        return None
+    try:
+        return int(dagrun_id)
+    except (TypeError, ValueError):
+        return None
+
+
+def _backfill_dagrun_id_from_deadline(conn, batch_size: int) -> None:
+    callback = sa.table(
+        "callback",
+        sa.column("id", sa.Uuid()),
+        sa.column("type", sa.String(20)),
+        sa.column("dagrun_id", sa.Integer()),
+    )
+    deadline = sa.table(
+        "deadline",
+        sa.column("callback_id", sa.Uuid()),
+        sa.column("dagrun_id", sa.Integer()),
+    )
+    dag_run = sa.table("dag_run", sa.column("id", sa.Integer()))
+
+    last_id = None
+    while True:
+        query = (
+            sa.select(callback.c.id, deadline.c.dagrun_id)
+            .select_from(
+                callback.join(deadline, callback.c.id == deadline.c.callback_id).join(
+                    dag_run, deadline.c.dagrun_id == dag_run.c.id
+                )
+            )
+            .where(callback.c.type == "executor", callback.c.dagrun_id.is_(None))
+            .order_by(callback.c.id)
+            .limit(batch_size)
+        )
+        if last_id is not None:
+            query = query.where(callback.c.id > last_id)
+
+        rows = conn.execute(query).fetchall()
+        if not rows:
+            return
+
+        for callback_id, dagrun_id in rows:
+            conn.execute(callback.update().where(callback.c.id == callback_id).values(dagrun_id=dagrun_id))
+
+        last_id = rows[-1].id
+
+
+def _backfill_dagrun_id(conn, batch_size: int) -> None:
+    callback = sa.table(
+        "callback",
+        sa.column("id", sa.Uuid()),
+        sa.column("type", sa.String(20)),
+        sa.column("data", sa.Text()),
+        sa.column("dagrun_id", sa.Integer()),
+    )
+    dag_run = sa.table("dag_run", sa.column("id", sa.Integer()))
+
+    last_id = None
+    while True:
+        query = (
+            sa.select(callback.c.id, callback.c.data)
+            .where(callback.c.type == "executor", callback.c.dagrun_id.is_(None))
+            .order_by(callback.c.id)
+            .limit(batch_size)
+        )
+        if last_id is not None:
+            query = query.where(callback.c.id > last_id)
+
+        rows = conn.execute(query).fetchall()
+        if not rows:
+            return
+
+        for callback_id, data in rows:
+            dagrun_id = _extract_dagrun_id(data)
+            if dagrun_id is None:
+                continue
+
+            dagrun_exists = conn.execute(sa.select(dag_run.c.id).where(dag_run.c.id == dagrun_id)).first()
+            if dagrun_exists:
+                conn.execute(
+                    callback.update().where(callback.c.id == callback_id).values(dagrun_id=dagrun_id)
+                )
+
+        last_id = rows[-1].id
+
+
+def upgrade():
+    """Add callback.dagrun_id and backfill it from legacy callback.data."""
+    with op.batch_alter_table("callback") as batch_op:
+        batch_op.add_column(sa.Column("dagrun_id", sa.Integer(), nullable=True))
+
+    if not context.is_offline_mode():
+        conn = op.get_bind()
+        batch_size = conf.getint("database", "migration_batch_size")
+        _backfill_dagrun_id_from_deadline(conn, batch_size)
+        _backfill_dagrun_id(conn, batch_size)
+
+    with op.batch_alter_table("callback") as batch_op:
+        batch_op.create_index("callback_dagrun_id_idx", ["dagrun_id"], unique=False)
+        batch_op.create_foreign_key(
+            batch_op.f("callback_dagrun_id_fkey"), "dag_run", ["dagrun_id"], ["id"], ondelete="CASCADE"
+        )
+
+
+def downgrade():
+    """Remove callback.dagrun_id."""
+    with op.batch_alter_table("callback") as batch_op:
+        batch_op.drop_constraint(batch_op.f("callback_dagrun_id_fkey"), type_="foreignkey")
+        batch_op.drop_index("callback_dagrun_id_idx")
+        batch_op.drop_column("dagrun_id")

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

@@ -26,7 +26,7 @@
 
 import structlog
 import uuid6
-from sqlalchemy import ForeignKey, Integer, String, Text, Uuid
+from sqlalchemy import ForeignKey, Index, Integer, String, Text, Uuid
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 # Re-exporting as _accepts_context for backward compatibility
@@ -143,10 +143,17 @@ class Callback(Base, BaseWorkload):
     # Creation time of the callback
     created_at: Mapped[datetime] = mapped_column(UtcDateTime, default=timezone.utcnow, nullable=False)
 
+    dagrun_id: Mapped[int | None] = mapped_column(
+        Integer, ForeignKey("dag_run.id", ondelete="CASCADE"), nullable=True
+    )
+    dag_run = relationship("DagRun", back_populates="callbacks")
+
     # Used for callbacks of type CallbackType.TRIGGERER
     trigger_id: Mapped[int | None] = mapped_column(Integer, ForeignKey("trigger.id"), nullable=True)
     trigger = relationship("Trigger", back_populates="callback", uselist=False)
 
+    __table_args__ = (Index("callback_dagrun_id_idx", dagrun_id, unique=False),)
+
     def __init__(self, priority_weight: int = 1, prefix: str = "", **kwargs):
         """
         Initialize a Callback. This is the base class so it shouldn't usually need to be initialized.

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

@@ -385,6 +385,13 @@ class DagRun(Base, LoggingMixin):
         cascade="all, delete, delete-orphan",
     )
 
+    callbacks = relationship(
+        "Callback",
+        back_populates="dag_run",
+        uselist=True,
+        passive_deletes=True,
+    )
+
     created_dag_version = relationship("DagVersion", uselist=False, passive_deletes=True)
     """
     The dag version that was active when the dag run was created, if available.

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

@@ -265,9 +265,9 @@ def callback_data_with_context():
         elif isinstance(self.callback, ExecutorCallback):
             data = callback_data_with_context()
             data["deadline_id"] = str(self.id)
-            data["dag_run_id"] = str(self.dagrun.id)
             data["dag_id"] = self.dagrun.dag_id
             self.callback.data = data
+            self.callback.dagrun_id = self.dagrun.id
 
             self.callback.state = CallbackState.PENDING
             session.add(self.callback)

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

@@ -117,7 +117,7 @@ class MappedClassProtocol(Protocol):
     "3.1.8": "509b94a1042d",
     "3.2.0": "1d6611b6ab7c",
     "3.3.0": "d2f4e1b3c5a7",
-    "3.4.0": "e5a91c7f42b3",
+    "3.4.0": "0f4c9a2b8e1d",
 }
 
 # Prefix used to identify tables holding data moved during migration.

airflow-core/tests/unit/jobs/test_scheduler_job.py

@@ -64,6 +64,7 @@
 from airflow.executors.executor_utils import ExecutorName
 from airflow.executors.local_executor import LocalExecutor
 from airflow.executors.workloads import WorkloadType
+from airflow.executors.workloads.callback import CallbackFetchMethod
 from airflow.jobs.job import Job, run_job
 from airflow.jobs.scheduler_job_runner import SCHEDULER_DAG_CACHE_SIZE, SchedulerJobRunner
 from airflow.listeners.listener import get_listener_manager
@@ -837,7 +838,7 @@ def create_callback_in_state(state: CallbackState):
                 deadline_alert_id=None,
             ).callback
             callback.state = state
-            callback.data["dag_run_id"] = dag_run.id
+            callback.dagrun_id = dag_run.id
             callback.data["dag_id"] = dag_run.dag_id
             return callback
 
@@ -867,6 +868,43 @@ def create_callback_in_state(state: CallbackState):
         assert session.get(ExecutorCallback, queued_callback.id).state == CallbackState.QUEUED
         assert session.get(ExecutorCallback, running_callback.id).state == CallbackState.RUNNING
 
+    @pytest.mark.parametrize(
+        ("dagrun_id", "expected_event"),
+        [
+            (None, "Executor callback is missing dagrun_id."),
+            (123, "Could not find DagRun for executor callback. DagRun may have been deleted."),
+        ],
+    )
+    def test_enqueue_executor_callbacks_logs_missing_dagrun_reference(
+        self, dagrun_id, expected_event, caplog
+    ):
+        def test_callback():
+            pass
+
+        callback = ExecutorCallback(
+            SyncCallback(test_callback),
+            fetch_method=CallbackFetchMethod.IMPORT_PATH,
+            dag_id="test_callback_missing_dagrun_reference",
+        )
+        callback.id = uuid4()
+        callback.state = CallbackState.PENDING
+        callback.dagrun_id = dagrun_id
+
+        session = MagicMock()
+        session.scalars.return_value.all.return_value = [callback]
+
+        executor = MockExecutor()
+        executor.queue_workload = MagicMock()
+        scheduler_job = Job()
+        self.job_runner = SchedulerJobRunner(job=scheduler_job, executors=[executor])
+        self.job_runner._executor_to_workloads = MagicMock(return_value={executor: [callback]})
+
+        self.job_runner._enqueue_executor_callbacks(session)
+
+        assert {"event": expected_event, "callback_id": callback.id} in caplog
+        executor.queue_workload.assert_not_called()
+        assert callback.state == CallbackState.PENDING
+
     @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest")
     @mock.patch("airflow._shared.observability.metrics.stats._get_backend")
     def test_process_executor_events_with_callback(