This is an automated email from the ASF dual-hosted git repository.
eladkal 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 1ab249327b5 Speed up trigger cleanup callback lookups (#75228)
1ab249327b5 is described below
commit 1ab249327b50eb0bddfad3cb01976b47dfe85833
Author: Shivam Rastogi <[email protected]>
AuthorDate: Mon Nov 2 03:38:46 2026 -0800
Speed up trigger cleanup callback lookups (#75228)
---
airflow-core/docs/migrations-ref.rst | 4 +-
.../0153_3_5_0_add_index_on_callback_trigger_id.py | 71 ++++++
airflow-core/src/airflow/models/callback.py | 11 +-
airflow-core/src/airflow/utils/db.py | 2 +-
.../test_0153_add_index_on_callback_trigger_id.py | 249 +++++++++++++++++++++
airflow-core/tests/unit/models/test_trigger.py | 18 +-
6 files changed, 350 insertions(+), 5 deletions(-)
diff --git a/airflow-core/docs/migrations-ref.rst
b/airflow-core/docs/migrations-ref.rst
index 413d4b3bd64..324dda5a79f 100644
--- a/airflow-core/docs/migrations-ref.rst
+++ b/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
|
+=========================+==================+===================+==============================================================+
-| ``5fbfa42a2d01`` (head) | ``38c8df3fcb9c`` | ``3.5.0`` | Add
recorded state and owner to trigger. |
+| ``16d7301f124b`` (head) | ``5fbfa42a2d01`` | ``3.5.0`` | Add index
on callback trigger ID. |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``5fbfa42a2d01`` | ``38c8df3fcb9c`` | ``3.5.0`` | Add
recorded state and owner to trigger. |
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
| ``38c8df3fcb9c`` | ``c71065a1b14f`` | ``3.5.0`` | Add index
on dag_run (dag_id, run_after). |
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
diff --git
a/airflow-core/src/airflow/migrations/versions/0153_3_5_0_add_index_on_callback_trigger_id.py
b/airflow-core/src/airflow/migrations/versions/0153_3_5_0_add_index_on_callback_trigger_id.py
new file mode 100644
index 00000000000..f88b451c7e5
--- /dev/null
+++
b/airflow-core/src/airflow/migrations/versions/0153_3_5_0_add_index_on_callback_trigger_id.py
@@ -0,0 +1,71 @@
+#
+# 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 index on callback trigger ID.
+
+Revision ID: 16d7301f124b
+Revises: 5fbfa42a2d01
+Create Date: 2026-10-19 15:11:12.337117
+
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+from airflow.migrations.utils import disable_sqlite_fkeys
+
+revision = "16d7301f124b"
+down_revision = "5fbfa42a2d01"
+branch_labels = None
+depends_on = None
+airflow_version = "3.5.0"
+
+
+def upgrade():
+ """Add an index on callback.trigger_id."""
+ trigger_id_is_not_null = sa.column("trigger_id").is_not(None)
+ with disable_sqlite_fkeys(op):
+ with op.batch_alter_table("callback", schema=None) as batch_op:
+ batch_op.create_index(
+ "idx_callback_trigger_id",
+ ["trigger_id"],
+ unique=False,
+ postgresql_where=trigger_id_is_not_null,
+ sqlite_where=trigger_id_is_not_null,
+ )
+
+
+def downgrade():
+ """Remove the index on callback.trigger_id."""
+ with disable_sqlite_fkeys(op):
+ conn = op.get_bind()
+ with op.batch_alter_table("callback", schema=None) as batch_op:
+ if conn.dialect.name == "mysql":
+
batch_op.drop_constraint(batch_op.f("callback_trigger_id_fkey"),
type_="foreignkey")
+ batch_op.drop_index("idx_callback_trigger_id")
+ batch_op.create_foreign_key(
+ batch_op.f("callback_trigger_id_fkey"),
+ "trigger",
+ ["trigger_id"],
+ ["id"],
+ )
+ else:
+ batch_op.drop_index("idx_callback_trigger_id")
diff --git a/airflow-core/src/airflow/models/callback.py
b/airflow-core/src/airflow/models/callback.py
index 490a167d7cc..d1fa640bab6 100644
--- a/airflow-core/src/airflow/models/callback.py
+++ b/airflow-core/src/airflow/models/callback.py
@@ -26,7 +26,7 @@ from uuid import UUID
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
@@ -160,6 +160,15 @@ class Callback(Base, BaseWorkload):
trigger_id: Mapped[int | None] = mapped_column(Integer,
ForeignKey("trigger.id"), nullable=True)
trigger = relationship("Trigger", back_populates="callback", uselist=False)
+ __table_args__ = (
+ Index(
+ "idx_callback_trigger_id",
+ trigger_id,
+ postgresql_where=trigger_id.is_not(None),
+ sqlite_where=trigger_id.is_not(None),
+ ),
+ )
+
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.
diff --git a/airflow-core/src/airflow/utils/db.py
b/airflow-core/src/airflow/utils/db.py
index a0956e49b1d..5aa8350c822 100644
--- a/airflow-core/src/airflow/utils/db.py
+++ b/airflow-core/src/airflow/utils/db.py
@@ -112,7 +112,7 @@ _REVISION_HEADS_MAP: dict[str, str] = {
"3.2.0": "1d6611b6ab7c",
"3.3.0": "d2f4e1b3c5a7",
"3.4.0": "b4f1a7d2c903",
- "3.5.0": "5fbfa42a2d01",
+ "3.5.0": "16d7301f124b",
}
# Prefix used to identify tables holding data moved during migration.
diff --git
a/airflow-core/tests/unit/migrations/test_0153_add_index_on_callback_trigger_id.py
b/airflow-core/tests/unit/migrations/test_0153_add_index_on_callback_trigger_id.py
new file mode 100644
index 00000000000..3ad2269a6c7
--- /dev/null
+++
b/airflow-core/tests/unit/migrations/test_0153_add_index_on_callback_trigger_id.py
@@ -0,0 +1,249 @@
+# 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.
+"""Tests for the callback trigger index migration."""
+
+from __future__ import annotations
+
+import re
+from datetime import UTC, datetime
+from importlib import import_module
+from uuid import UUID
+
+import pytest
+import sqlalchemy as sa
+from sqlalchemy import inspect
+from sqlalchemy.exc import IntegrityError
+
+from airflow import settings
+from airflow.models.callback import Callback
+from airflow.utils.db import downgrade, upgradedb
+
+pytestmark = pytest.mark.db_test
+
+_migration =
import_module("airflow.migrations.versions.0153_3_5_0_add_index_on_callback_trigger_id")
+REVISION = _migration.revision
+PREVIOUS_REVISION = _migration.down_revision
+INDEX_NAME = "idx_callback_trigger_id"
+CALLBACK_IDS = (UUID(int=9_014_901), UUID(int=9_014_902), UUID(int=9_014_903),
UUID(int=9_014_904))
+INVALID_CALLBACK_ID = UUID(int=9_014_999)
+TEST_TIMESTAMP = datetime(2026, 10, 19, tzinfo=UTC)
+TRIGGER_CLASSPATHS = (
+ "airflow.triggers.testing.CallbackIndexTrigger1",
+ "airflow.triggers.testing.CallbackIndexTrigger2",
+)
+
+
+def _reflect_tables(engine):
+ metadata = sa.MetaData()
+ return (
+ sa.Table("trigger", metadata, autoload_with=engine),
+ sa.Table("callback", metadata, sa.Column("id", sa.Uuid(),
primary_key=True), autoload_with=engine),
+ )
+
+
+def _delete_test_rows(engine):
+ trigger, callback = _reflect_tables(engine)
+ with engine.begin() as connection:
+
connection.execute(callback.delete().where(callback.c.id.in_((*CALLBACK_IDS,
INVALID_CALLBACK_ID))))
+
connection.execute(trigger.delete().where(trigger.c.classpath.in_(TRIGGER_CLASSPATHS)))
+
+
[email protected]
+def database_at_previous_revision():
+ """Move the database to the pre-index revision, then restore the current
head."""
+ original_async_conn = settings.SQL_ALCHEMY_CONN_ASYNC
+ settings.SQL_ALCHEMY_CONN_ASYNC = ""
+ settings.dispose_orm(do_log=False)
+ settings.configure_orm(disable_connection_pool=True)
+ try:
+ downgrade(to_revision=PREVIOUS_REVISION)
+ _delete_test_rows(settings.engine)
+ yield settings.engine
+ finally:
+ try:
+ upgradedb()
+ _delete_test_rows(settings.engine)
+ finally:
+ settings.SQL_ALCHEMY_CONN_ASYNC = original_async_conn
+ settings.reconfigure_orm()
+
+
[email protected]
+def stored_rows(database_at_previous_revision):
+ trigger, callback = _reflect_tables(database_at_previous_revision)
+ with database_at_previous_revision.begin() as connection:
+ trigger_ids = tuple(
+ connection.execute(
+ trigger.insert().values(
+ classpath=classpath,
+ kwargs="{}",
+ created_date=TEST_TIMESTAMP,
+ )
+ ).inserted_primary_key[0]
+ for classpath in TRIGGER_CLASSPATHS
+ )
+ connection.execute(
+ callback.insert(),
+ [
+ {
+ "id": CALLBACK_IDS[0],
+ "type": "triggerer",
+ "fetch_method": "import_path",
+ "data": {},
+ "state": "scheduled",
+ "priority_weight": 1,
+ "created_at": TEST_TIMESTAMP,
+ "trigger_id": None,
+ },
+ {
+ "id": CALLBACK_IDS[1],
+ "type": "triggerer",
+ "fetch_method": "import_path",
+ "data": {},
+ "state": "scheduled",
+ "priority_weight": 1,
+ "created_at": TEST_TIMESTAMP,
+ "trigger_id": trigger_ids[0],
+ },
+ {
+ "id": CALLBACK_IDS[2],
+ "type": "triggerer",
+ "fetch_method": "import_path",
+ "data": {},
+ "state": "scheduled",
+ "priority_weight": 1,
+ "created_at": TEST_TIMESTAMP,
+ "trigger_id": trigger_ids[0],
+ },
+ {
+ "id": CALLBACK_IDS[3],
+ "type": "triggerer",
+ "fetch_method": "import_path",
+ "data": {},
+ "state": "scheduled",
+ "priority_weight": 1,
+ "created_at": TEST_TIMESTAMP,
+ "trigger_id": trigger_ids[1],
+ },
+ ],
+ )
+ return database_at_previous_revision, trigger_ids
+
+
+def _get_callback_rows(engine):
+ _, callback = _reflect_tables(engine)
+ with engine.connect() as connection:
+ rows = connection.execute(
+ sa.select(callback.c.id, callback.c.trigger_id)
+ .where(callback.c.id.in_(CALLBACK_IDS))
+ .order_by(callback.c.id)
+ ).all()
+ return [(row.id, row.trigger_id) for row in rows]
+
+
+def _get_callback_trigger_indexes(engine):
+ return [
+ index
+ for index in inspect(engine).get_indexes("callback")
+ if index["column_names"] and index["column_names"][0] == "trigger_id"
+ ]
+
+
+def _assert_explicit_index(engine):
+ indexes = _get_callback_trigger_indexes(engine)
+ assert [index["name"] for index in indexes] == [INDEX_NAME]
+ index = indexes[0]
+ assert index["column_names"] == ["trigger_id"]
+ assert not index["unique"]
+
+ dialect = engine.dialect.name
+ if dialect in {"postgresql", "sqlite"}:
+ predicate = index["dialect_options"][f"{dialect}_where"]
+ normalized = re.sub(r'[\s()"`]+', "", str(predicate).lower())
+ assert normalized == "trigger_idisnotnull"
+
+
+def _assert_callback_foreign_key(engine):
+ foreign_key = next(
+ constraint
+ for constraint in inspect(engine).get_foreign_keys("callback")
+ if constraint["name"] == "callback_trigger_id_fkey"
+ )
+ assert foreign_key["constrained_columns"] == ["trigger_id"]
+ assert foreign_key["referred_table"] == "trigger"
+ assert foreign_key["referred_columns"] == ["id"]
+
+
+def _assert_mysql_foreign_key_is_enforced(engine):
+ if engine.dialect.name != "mysql":
+ return
+ _, callback = _reflect_tables(engine)
+ with pytest.raises(IntegrityError), engine.begin() as connection:
+ connection.execute(
+ callback.insert().values(
+ id=INVALID_CALLBACK_ID,
+ type="triggerer",
+ fetch_method="import_path",
+ data={},
+ state="scheduled",
+ priority_weight=1,
+ created_at=TEST_TIMESTAMP,
+ trigger_id=99_999_999,
+ )
+ )
+
+
+def test_callback_model_declares_partial_trigger_id_index():
+ indexes = [index for index in Callback.__table__.indexes if index.name ==
INDEX_NAME]
+ assert len(indexes) == 1
+ index = indexes[0]
+ assert [column.name for column in index.columns] == ["trigger_id"]
+ assert not index.unique
+ assert str(index.dialect_options["postgresql"]["where"]) ==
"callback.trigger_id IS NOT NULL"
+ assert str(index.dialect_options["sqlite"]["where"]) ==
"callback.trigger_id IS NOT NULL"
+
+
[email protected]("sqlite", "postgres", "mysql")
+def
test_upgrade_downgrade_reupgrade_preserves_rows_and_foreign_key(stored_rows):
+ engine, trigger_ids = stored_rows
+ expected_rows = [
+ (CALLBACK_IDS[0], None),
+ (CALLBACK_IDS[1], trigger_ids[0]),
+ (CALLBACK_IDS[2], trigger_ids[0]),
+ (CALLBACK_IDS[3], trigger_ids[1]),
+ ]
+
+ upgradedb(to_revision=REVISION)
+ _assert_explicit_index(engine)
+ _assert_callback_foreign_key(engine)
+ assert _get_callback_rows(engine) == expected_rows
+
+ downgrade(to_revision=PREVIOUS_REVISION)
+ trigger_indexes = _get_callback_trigger_indexes(engine)
+ if engine.dialect.name == "mysql":
+ assert len(trigger_indexes) == 1
+ assert trigger_indexes[0]["name"] != INDEX_NAME
+ else:
+ assert trigger_indexes == []
+ _assert_callback_foreign_key(engine)
+ _assert_mysql_foreign_key_is_enforced(engine)
+ assert _get_callback_rows(engine) == expected_rows
+
+ upgradedb(to_revision=REVISION)
+ _assert_explicit_index(engine)
+ _assert_callback_foreign_key(engine)
+ assert _get_callback_rows(engine) == expected_rows
diff --git a/airflow-core/tests/unit/models/test_trigger.py
b/airflow-core/tests/unit/models/test_trigger.py
index 66a9fd6565f..f696e61eb0d 100644
--- a/airflow-core/tests/unit/models/test_trigger.py
+++ b/airflow-core/tests/unit/models/test_trigger.py
@@ -28,7 +28,7 @@ import pytz
import sqlalchemy as sa
import time_machine
from cryptography.fernet import Fernet
-from sqlalchemy import delete, event, func, insert, select, update
+from sqlalchemy import delete, event, func, insert, select, text, update
from sqlalchemy.exc import OperationalError
from airflow._shared.timezones import timezone
@@ -55,7 +55,7 @@ from airflow.triggers.base import (
from airflow.utils.session import create_session
from airflow.utils.state import State, TaskInstanceState
-from tests_common.test_utils.asserts import assert_queries_count
+from tests_common.test_utils.asserts import assert_queries_count,
capture_orm_selects
from tests_common.test_utils.config import conf_vars
if TYPE_CHECKING:
@@ -125,6 +125,20 @@ def
test_fetch_trigger_ids_with_non_task_associations(session):
assert results == {asset_trigger.id, callback_trigger.id}
[email protected]("sqlite")
+def test_clean_unused_callback_lookup_uses_callback_trigger_index(session):
+ with capture_orm_selects("trigger") as statements:
+ Trigger._delete_unused_batch(500, session=session)
+ plan = [
+ row[-1]
+ for statement in statements
+ for row in session.execute(text(f"EXPLAIN QUERY PLAN
{statement.removesuffix(' FOR UPDATE')}")).all()
+ ]
+
+ assert statements
+ assert any("callback" in row and "idx_callback_trigger_id" in row for row
in plan), plan
+
+
def test_fetch_assignments_maps_surviving_rows_to_their_triggerer(session):
owned = Trigger(classpath="airflow.triggers.testing.SuccessTrigger1",
kwargs={})
owned.triggerer_id = 42
