Stop airflow db clean deleting triggers that are still in use

airflow-core/newsfragments/71339.bugfix.rst

@@ -0,0 +1 @@
+Stop ``airflow db clean`` deleting ``trigger`` rows that a deferred task instance, an asset watcher, or a callback still references. Such rows were removed by age alone, which cascade-deleted the referencing task instance or watcher without archiving it, or failed the foreign key and stopped the command. ``trigger`` is also a reserved word on MySQL, where the unquoted archive DDL made the table impossible to clean.

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

@@ -51,6 +51,7 @@
     from pendulum import DateTime
     from sqlalchemy import Select
     from sqlalchemy.orm import Session
+    from sqlalchemy.sql.compiler import IdentifierPreparer
 
     from airflow.models import Base
 
@@ -78,6 +79,19 @@ def _format_table_name(schema: str | None, table: str) -> str:
     return table
 
 
+def _format_quoted_table_name(preparer: IdentifierPreparer, schema: str | None, table: str) -> str:
+    """
+    Format a fully qualified table name for interpolation into raw DDL.
+
+    Table names reach raw DDL as text rather than as bound identifiers, so a name that is a
+    reserved word in the target dialect is a syntax error there. ``trigger`` is reserved on
+    MySQL, which makes the whole table uncleanable without quoting.
+    """
+    if schema:
+        return f"{preparer.quote(schema)}.{preparer.quote(table)}"
+    return preparer.quote(table)
+
+
 @dataclasses.dataclass
 class _TableConfig:
     """
@@ -211,7 +225,18 @@ def readable_config(self):
     _TableConfig(
         table_name="trigger",
         recency_column_name="created_date",
+        extra_columns=["id"],
         dependent_tables=["task_instance"],
+        # The triggerer already drops every unreferenced trigger each loop, so an old row that
+        # survives is almost always still in use. task_instance.trigger_id and
+        # asset_watcher.trigger_id are ON DELETE CASCADE, so deleting one takes a deferred task
+        # instance or an event-driven watcher with it, neither of which is archived.
+        # callback.trigger_id has no delete rule, so the same delete fails the foreign key.
+        skip_if_referenced=[
+            ("task_instance", "trigger_id"),
+            ("asset_watcher", "trigger_id"),
+            ("callback", "trigger_id"),
+        ],
     ),
     _TableConfig(
         table_name="dag_version",
@@ -308,10 +333,8 @@ def _do_delete(
         # using bulk delete
         # create a new table and copy the rows there
         timestamp_str = re.sub(r"[^\d]", "", timezone.utcnow().isoformat())[:14]
-        target_table_name = _format_table_name(
-            orm_model.schema,
-            f"{ARCHIVE_TABLE_PREFIX}{orm_model.name}__{timestamp_str}{suffix}",
-        )
+        target_bare_name = f"{ARCHIVE_TABLE_PREFIX}{orm_model.name}__{timestamp_str}{suffix}"
+        target_table_name = _format_table_name(orm_model.schema, target_bare_name)
         print(f"Moving data to table {target_table_name}")
         target_table = None
         # Lets the ``finally`` cleanup below tell the failure path (don't let a
@@ -323,7 +346,14 @@ def _do_delete(
             if dialect_name == "mysql":
                 # MySQL with replication needs this split into two queries, so just do it for all MySQL
                 # ERROR 1786 (HY000): Statement violates GTID consistency: CREATE TABLE ... SELECT.
-                session.execute(text(f"CREATE TABLE {target_table_name} LIKE {source_table_name}"))
+                preparer = bind.dialect.identifier_preparer
+                session.execute(
+                    text(
+                        f"CREATE TABLE "
+                        f"{_format_quoted_table_name(preparer, orm_model.schema, target_bare_name)} "
+                        f"LIKE {_format_quoted_table_name(preparer, orm_model.schema, orm_model.name)}"
+                    )
+                )
                 metadata = reflect_tables([target_table_name], session)
                 target_table = metadata.tables[target_table_name]
                 insert_stm = target_table.insert().from_select(target_table.c, limited_query)

airflow-core/tests/unit/utils/test_db_cleanup.py

@@ -37,10 +37,13 @@
 from airflow._shared.timezones import timezone
 from airflow.exceptions import AirflowException
 from airflow.models import DagModel, DagRun, TaskInstance
+from airflow.models.asset import AssetModel, AssetWatcherModel
+from airflow.models.callback import Callback, CallbackFetchMethod, CallbackType
 from airflow.models.dag_version import DagVersion
 from airflow.models.dagbundle import DagBundleModel
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.models.task_state_store import TaskStateStoreModel
+from airflow.models.trigger import Trigger
 from airflow.providers.standard.operators.python import PythonOperator
 from airflow.serialization.serialized_objects import LazyDeserializedDAG
 from airflow.utils.db_cleanup import (
@@ -59,13 +62,16 @@
     run_cleanup,
 )
 from airflow.utils.session import create_session
+from airflow.utils.state import TaskInstanceState
 from airflow.utils.types import DagRunType
 
 from tests_common.test_utils.db import (
     clear_db_assets,
+    clear_db_callbacks,
     clear_db_dag_bundles,
     clear_db_dags,
     clear_db_runs,
+    clear_db_triggers,
     drop_tables_with_prefix,
 )
 from tests_common.test_utils.taskinstance import create_task_instance
@@ -87,6 +93,22 @@ def clean_database():
     clear_db_dag_bundles()
 
 
+@pytest.fixture
+def clear_callbacks_and_triggers():
+    """Isolate callback and trigger rows.
+
+    ``callback.trigger_id`` has no delete rule, so a leftover callback row blocks the shared
+    teardown from clearing triggers. Callbacks must go first for the same reason.
+    """
+    clear_db_callbacks()
+    clear_db_triggers()
+
+    yield
+
+    clear_db_callbacks()
+    clear_db_triggers()
+
+
 class TestDBCleanup:
     @pytest.fixture(autouse=True)
     def clear_airflow_tables(self):
@@ -551,6 +573,118 @@ def test_dag_version_cleanup_skips_versions_pinned_by_task_instance(self):
         assert latest_id in remaining  # kept by keep_last
         assert orphan_id not in remaining  # old and unreferenced -> pruned
 
+    @pytest.mark.usefixtures("clear_callbacks_and_triggers")
+    def test_trigger_cleanup_skips_triggers_still_referenced(self):
+        """db clean must not purge trigger rows that live objects still point at.
+
+        ``task_instance.trigger_id`` and ``asset_watcher.trigger_id`` are ``ON DELETE CASCADE``,
+        so purging a referenced trigger destroys a deferred task instance or an event-driven
+        watcher without archiving it. ``callback.trigger_id`` has no delete rule, so the same
+        purge fails the foreign key and aborts the run.
+        """
+        base_date = pendulum.DateTime(2022, 1, 1, tzinfo=pendulum.timezone("UTC"))
+        dag_id = f"test_dag_{uuid4()}"
+
+        with create_session() as session:
+            triggers = [
+                Trigger(classpath=f"airflow.triggers.testing.SuccessTrigger{i}", kwargs={}) for i in range(4)
+            ]
+            for trigger in triggers:
+                trigger.created_date = base_date
+            session.add_all(triggers)
+            session.flush()
+            ti_trigger, watcher_trigger, callback_trigger, orphan_trigger = triggers
+            trigger_ids = [trigger.id for trigger in triggers]
+
+            dag_run = DagRun(dag_id, run_id="run-1", run_type=DagRunType.MANUAL, start_date=base_date)
+            deferred_ti = create_task_instance(
+                PythonOperator(task_id="deferred-task", python_callable=print),
+                run_id=dag_run.run_id,
+                dag_version_id=None,
+            )
+            deferred_ti.dag_id = dag_id
+            deferred_ti.start_date = base_date
+            deferred_ti.state = TaskInstanceState.DEFERRED
+            deferred_ti.trigger_id = ti_trigger.id
+            session.add_all([dag_run, deferred_ti])
+
+            asset = AssetModel(name=f"asset-{uuid4()}", uri=f"s3://bucket/{uuid4()}", group="asset")
+            session.add(asset)
+            session.flush()
+            session.add(AssetWatcherModel(name="watcher", asset_id=asset.id, trigger_id=watcher_trigger.id))
+            session.execute(
+                Callback.__table__.insert().values(
+                    id=uuid4(),
+                    type=CallbackType.TRIGGERER.value,
+                    fetch_method=CallbackFetchMethod.IMPORT_PATH.value,
+                    data={},
+                    priority_weight=1,
+                    created_at=base_date,
+                    trigger_id=callback_trigger.id,
+                )
+            )
+            session.commit()
+            deferred_ti_id = deferred_ti.id
+
+            _cleanup_table(
+                **config_dict["trigger"].__dict__,
+                clean_before_timestamp=base_date.add(days=10),
+                dry_run=False,
+                session=session,
+                table_names=["trigger"],
+                skip_archive=True,
+            )
+
+            remaining = set(session.scalars(select(Trigger.id).where(Trigger.id.in_(trigger_ids))).all())
+            surviving_ti = session.get(TaskInstance, deferred_ti_id)
+            watcher_count = session.scalar(
+                select(func.count())
+                .select_from(AssetWatcherModel)
+                .where(AssetWatcherModel.trigger_id == watcher_trigger.id)
+            )
+
+        assert ti_trigger.id in remaining
+        assert watcher_trigger.id in remaining
+        assert callback_trigger.id in remaining
+        assert orphan_trigger.id not in remaining  # old and unreferenced -> pruned
+        assert surviving_ti is not None, "cascade deleted a deferred task instance that was still running"
+        assert watcher_count == 1
+
+    @pytest.mark.usefixtures("clear_callbacks_and_triggers")
+    def test_trigger_cleanup_archives_orphan_trigger(self):
+        """An unreferenced trigger is archived and purged on every backend.
+
+        ``trigger`` is a reserved word on MySQL, so the archive DDL must quote it; unquoted it is
+        a syntax error that leaves the table permanently uncleanable there.
+        """
+        base_date = pendulum.DateTime(2022, 1, 1, tzinfo=pendulum.timezone("UTC"))
+
+        with create_session() as session:
+            orphan = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={})
+            orphan.created_date = base_date
+            session.add(orphan)
+            session.commit()
+            orphan_id = orphan.id
+
+            _cleanup_table(
+                **config_dict["trigger"].__dict__,
+                clean_before_timestamp=base_date.add(days=10),
+                dry_run=False,
+                session=session,
+                table_names=["trigger"],
+                skip_archive=False,
+            )
+            session.commit()
+
+            remaining = session.scalar(
+                select(func.count()).select_from(Trigger).where(Trigger.id == orphan_id)
+            )
+            archives = _get_archived_table_names(["trigger"], session)
+
+        assert remaining == 0
+        assert len(archives) == 1
+        assert archives[0].startswith(f"{ARCHIVE_TABLE_PREFIX}trigger__")
+
     def test_table_config_skip_if_referenced_requires_pk_column(self):
         """A misconfigured skip_if_referenced (pk not in columns) must fail fast at construction."""
         with pytest.raises(ValueError, match="referenced_pk_column"):