diff --git a/airflow-core/newsfragments/71339.bugfix.rst b/airflow-core/newsfragments/71339.bugfix.rst new file mode 100644 index 0000000000000..bc8b9cfb0542a --- /dev/null +++ b/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. diff --git a/airflow-core/src/airflow/utils/db_cleanup.py b/airflow-core/src/airflow/utils/db_cleanup.py index 1a79dbc9e30ce..6e28a85a8aefd 100644 --- a/airflow-core/src/airflow/utils/db_cleanup.py +++ b/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) diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py b/airflow-core/tests/unit/utils/test_db_cleanup.py index abe11aa7c589d..ff9cf6c60efd2 100644 --- a/airflow-core/tests/unit/utils/test_db_cleanup.py +++ b/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"):