From 9ad7c673603b0b1be7367ed3f19fc20ed3356d92 Mon Sep 17 00:00:00 2001 From: Biley Roy Date: Tue, 14 Jul 2026 21:40:50 +0530 Subject: [PATCH 1/3] fix(db_cleanup): handle NULL start_date in dag_run cleanup When start_date is NULL (e.g. runs that failed or were killed before execution began), the WHERE condition start_date < cutoff evaluates to NULL and silently skips these rows forever. Add --fallback-cleanup-on-null CLI flag that uses COALESCE(start_date, created_at) in the WHERE condition, keep_last subquery, and JOIN so that NULL start_date rows are cleaned based on their created_at timestamp while the most recent run per dag is still protected by keep_last. closes: #69872 --- airflow-core/src/airflow/cli/cli_config.py | 10 ++ .../src/airflow/cli/commands/db_command.py | 1 + airflow-core/src/airflow/utils/db_cleanup.py | 45 ++++- .../tests/unit/utils/test_db_cleanup.py | 167 ++++++++++++++++++ 4 files changed, 219 insertions(+), 4 deletions(-) diff --git a/airflow-core/src/airflow/cli/cli_config.py b/airflow-core/src/airflow/cli/cli_config.py index 46a0ebfa391f0..ad8da77f74f13 100644 --- a/airflow-core/src/airflow/cli/cli_config.py +++ b/airflow-core/src/airflow/cli/cli_config.py @@ -586,6 +586,15 @@ def string_lower_type(val): help="Command will exit with a non-zero exit code if any table cleanup failed. By default errors are suppressed and the command exits 0.", action="store_true", ) +ARG_DB_FALLBACK_CLEANUP_ON_NULL = Arg( + ("--fallback-cleanup-on-null",), + help=( + "When set, records with a NULL recency column (e.g. dag_run.start_date for runs that never " + "started) will be included in cleanup using a fallback column (e.g. created_at) to determine " + "their age. Without this flag, such records are silently skipped." + ), + action="store_true", +) ARG_DAG_IDS = Arg( ("--dag-ids",), default=None, @@ -1782,6 +1791,7 @@ class GroupCommand(NamedTuple): ARG_DAG_IDS, ARG_EXCLUDE_DAG_IDS, ARG_DB_ERROR_ON_CLEANUP_FAILURE, + ARG_DB_FALLBACK_CLEANUP_ON_NULL, ), ), ActionCommand( diff --git a/airflow-core/src/airflow/cli/commands/db_command.py b/airflow-core/src/airflow/cli/commands/db_command.py index 4c15ccc2d488d..c91f373dd770d 100644 --- a/airflow-core/src/airflow/cli/commands/db_command.py +++ b/airflow-core/src/airflow/cli/commands/db_command.py @@ -366,6 +366,7 @@ def cleanup_tables(args): dag_ids=args.dag_ids, exclude_dag_ids=args.exclude_dag_ids, error_on_cleanup_failure=args.error_on_cleanup_failure, + fallback_cleanup_on_null=args.fallback_cleanup_on_null, ) diff --git a/airflow-core/src/airflow/utils/db_cleanup.py b/airflow-core/src/airflow/utils/db_cleanup.py index 6deb4e78f6856..6187ceb5f3955 100644 --- a/airflow-core/src/airflow/utils/db_cleanup.py +++ b/airflow-core/src/airflow/utils/db_cleanup.py @@ -78,6 +78,11 @@ class _TableConfig: :param keep_last_group_by: if keeping the last record, can keep the last record for each group :param dependent_tables: list of tables which have FK relationship with this table :param extra_filters: SQLAlchemy expressions ANDed with the recency filter; referenced columns must be in ``extra_columns``. + :param fallback_recency_column_name: optional fallback date column used when ``recency_column_name`` + is NULL for a row. When the ``--fallback-cleanup-on-null`` CLI flag is enabled, the cleanup query + uses ``COALESCE(recency_column, fallback_recency_column)`` so that rows with a NULL recency value + (e.g. ``dag_run.start_date`` for runs that never started) can still be cleaned based on the + fallback column (e.g. ``created_at``). The fallback column must be listed in ``extra_columns``. :param skip_if_referenced: list of ``(referencing_table, fk_column)`` pairs whose FK points at this table's ``referenced_pk_column``. A row that is still referenced by any of these is excluded from deletion. This avoids issuing deletes that would violate an ``ON DELETE RESTRICT`` foreign key @@ -98,11 +103,15 @@ class _TableConfig: # Relying on automation here would increase complexity and reduce maintainability. dependent_tables: list[str] | None = None extra_filters: list[Any] | None = None + fallback_recency_column_name: str | None = None skip_if_referenced: list[tuple[str, str]] | None = None referenced_pk_column: str = "id" def __post_init__(self): self.recency_column = column(self.recency_column_name) + self.fallback_recency_column = ( + column(self.fallback_recency_column_name) if self.fallback_recency_column_name else None + ) if self.dag_id_column_name is None: self.dag_id_column = None self.orm_model: Base = table( @@ -154,11 +163,12 @@ def readable_config(self): table_name="dag_run", recency_column_name="start_date", dag_id_column_name="dag_id", - extra_columns=["dag_id", "run_type"], + extra_columns=["dag_id", "run_type", "created_at"], keep_last=True, keep_last_filters=[column("run_type") != DagRunType.MANUAL], keep_last_group_by=["dag_id"], dependent_tables=["task_instance", "task_state_store", "deadline"], + fallback_recency_column_name="created_at", ), _TableConfig(table_name="asset_event", recency_column_name="timestamp", dag_id_column_name="dag_id"), _TableConfig(table_name="import_error", recency_column_name="timestamp"), @@ -372,8 +382,14 @@ def _subquery_keep_last( keep_last_filters, group_by_columns, max_date_colname, + fallback_recency_column=None, + fallback_cleanup_on_null: bool = False, ): - subquery = select(*group_by_columns, func.max(recency_column).label(max_date_colname)) + if fallback_cleanup_on_null and fallback_recency_column is not None: + effective_col = func.coalesce(recency_column, fallback_recency_column) + else: + effective_col = recency_column + subquery = select(*group_by_columns, func.max(effective_col).label(max_date_colname)) if keep_last_filters is not None: for entry in keep_last_filters: @@ -415,13 +431,22 @@ def _build_query( extra_filters: list[Any] | None = None, skip_if_referenced: list[tuple[str, str]] | None = None, referenced_pk_column: str = "id", + fallback_recency_column=None, + fallback_cleanup_on_null: bool = False, **kwargs, ) -> Select: base_table_alias = "base" base_table = aliased(orm_model, name=base_table_alias) query = select(text(f"{base_table_alias}.*")).select_from(base_table) base_table_recency_col = base_table.c[recency_column.name] - conditions = [base_table_recency_col < clean_before_timestamp] + + if fallback_cleanup_on_null and fallback_recency_column is not None: + base_table_fallback_col = base_table.c[fallback_recency_column.name] + effective_recency = func.coalesce(base_table_recency_col, base_table_fallback_col) + else: + effective_recency = base_table_recency_col + + conditions = [effective_recency < clean_before_timestamp] if extra_filters: conditions.extend(extra_filters) @@ -458,12 +483,14 @@ def _build_query( keep_last_filters=keep_last_filters, group_by_columns=group_by_columns, max_date_colname=max_date_col_name, + fallback_recency_column=fallback_recency_column, + fallback_cleanup_on_null=fallback_cleanup_on_null, ) query = query.outerjoin( subquery, and_( *[base_table.c[x] == subquery.c[x] for x in keep_last_group_by], # type: ignore[attr-defined] - base_table_recency_col == column(max_date_col_name), + effective_recency == column(max_date_col_name), ), ) conditions.append(column(max_date_col_name).is_(None)) @@ -490,6 +517,8 @@ def _cleanup_table( extra_filters: list[Any] | None = None, skip_if_referenced: list[tuple[str, str]] | None = None, referenced_pk_column: str = "id", + fallback_recency_column=None, + fallback_cleanup_on_null: bool = False, **kwargs, ) -> None: print() @@ -508,6 +537,8 @@ def _cleanup_table( extra_filters=extra_filters, skip_if_referenced=skip_if_referenced, referenced_pk_column=referenced_pk_column, + fallback_recency_column=fallback_recency_column, + fallback_cleanup_on_null=fallback_cleanup_on_null, session=session, ) logger.debug("old rows query:\n%s", query.selectable.compile()) @@ -667,6 +698,7 @@ def run_cleanup( session: Session = NEW_SESSION, batch_size: int | None = None, error_on_cleanup_failure: bool = False, + fallback_cleanup_on_null: bool = False, ) -> None: """ Purges old records in airflow metadata database. @@ -693,6 +725,10 @@ def run_cleanup( :param error_on_cleanup_failure: If True, raise a RuntimeError after processing all tables if any per-table cleanup encountered an error. By default errors are suppressed, a warning summary is logged, and the command exits 0 even if some tables were not cleaned. + :param fallback_cleanup_on_null: If True, use a fallback column (e.g. ``created_at``) for + tables whose recency column (e.g. ``dag_run.start_date``) can be NULL. Without this flag, + rows with a NULL recency column are silently skipped. When enabled, the query uses + ``COALESCE(recency_column, fallback_column)`` to determine the row's age. """ clean_before_timestamp = timezone.coerce_datetime(clean_before_timestamp) @@ -728,6 +764,7 @@ def run_cleanup( skip_archive=skip_archive, session=session, batch_size=batch_size, + fallback_cleanup_on_null=fallback_cleanup_on_null, ) if ctx.failed: failed_tables.append(table_name) diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py b/airflow-core/tests/unit/utils/test_db_cleanup.py index 8271d38968b98..78aaffa8c0d17 100644 --- a/airflow-core/tests/unit/utils/test_db_cleanup.py +++ b/airflow-core/tests/unit/utils/test_db_cleanup.py @@ -1427,3 +1427,170 @@ def _dag_version_config_without_row_exclusion(): if key in config: config[key] = None return config + + +@pytest.mark.db_test +class TestDagRunNullStartDateCleanup: + """Tests for cleaning dag_run records where start_date is NULL.""" + + @pytest.fixture(autouse=True) + def clear_airflow_tables(self): + drop_tables_with_prefix("_airflow_") + + def _create_dag_runs_with_null_start_date(self, base_date): + """Create dag_runs: some with start_date set, one with start_date=None.""" + with create_session() as session: + bundle_name = "testing" + session.add(DagBundleModel(name=bundle_name)) + session.flush() + + dag_id = f"test-dag-null-sd_{uuid4()}" + dag = DAG(dag_id=dag_id) + dm = DagModel(dag_id=dag_id, bundle_name=bundle_name) + session.add(dm) + SerializedDagModel.write_dag(LazyDeserializedDAG.from_dag(dag), bundle_name=bundle_name) + + # Run with start_date set (normal run) + dag_run_normal = DagRun( + dag.dag_id, + run_id="normal_run", + run_type=DagRunType.MANUAL, + start_date=base_date, + ) + session.add(dag_run_normal) + + # Run with start_date=None (failed before starting) + dag_run_null = DagRun( + dag.dag_id, + run_id="null_start_date_run", + run_type=DagRunType.MANUAL, + start_date=None, + ) + dag_run_null.created_at = base_date + session.add(dag_run_null) + session.commit() + return dag_id + + def test_build_query_without_flag_skips_null_start_date(self): + """Without --fallback-cleanup-on-null, NULL start_date rows are invisible.""" + base_date = pendulum.DateTime(2022, 1, 1, tzinfo=pendulum.timezone("UTC")) + self._create_dag_runs_with_null_start_date(base_date) + + with create_session() as session: + # Use a future cutoff so all records are eligible by age + clean_before = base_date.add(days=365) + query = _build_query( + **config_dict["dag_run"].__dict__, + clean_before_timestamp=clean_before, + session=session, + ) + num_rows = session.scalars(select(func.count()).select_from(query.subquery())).one() + # Only the normal run (with start_date) should be found; NULL row is skipped + assert num_rows == 1 + + def test_build_query_with_flag_finds_null_start_date(self): + """With --fallback-cleanup-on-null, NULL start_date rows are found via created_at.""" + base_date = pendulum.DateTime(2022, 1, 1, tzinfo=pendulum.timezone("UTC")) + self._create_dag_runs_with_null_start_date(base_date) + + with create_session() as session: + clean_before = base_date.add(days=365) + query = _build_query( + **config_dict["dag_run"].__dict__, + clean_before_timestamp=clean_before, + fallback_cleanup_on_null=True, + session=session, + ) + num_rows = session.scalars(select(func.count()).select_from(query.subquery())).one() + # Both runs should be found: normal (via start_date) and null (via created_at fallback) + assert num_rows == 2 + + def test_cleanup_table_with_flag_deletes_null_start_date(self): + """Verify actual deletion of NULL start_date records when flag is enabled.""" + base_date = pendulum.DateTime(2022, 1, 1, tzinfo=pendulum.timezone("UTC")) + self._create_dag_runs_with_null_start_date(base_date) + + with create_session() as session: + # Verify 2 dag_runs exist before cleanup + model = config_dict["dag_run"].orm_model + assert session.scalar(select(func.count()).select_from(model)) == 2 + + clean_before = base_date.add(days=365) + _cleanup_table( + **config_dict["dag_run"].__dict__, + clean_before_timestamp=clean_before, + dry_run=False, + session=session, + fallback_cleanup_on_null=True, + ) + + # Both should be deleted (all are manual, keep_last only protects non-manual) + assert session.scalar(select(func.count()).select_from(model)) == 0 + + def _create_scheduled_runs_with_null_start_date(self, base_date): + """Create 2 scheduled dag_runs with NULL start_date and different created_at.""" + with create_session() as session: + bundle_name = "testing" + session.add(DagBundleModel(name=bundle_name)) + session.flush() + + dag_id = f"test-dag-sched-null_{uuid4()}" + dag = DAG(dag_id=dag_id) + dm = DagModel(dag_id=dag_id, bundle_name=bundle_name) + session.add(dm) + SerializedDagModel.write_dag(LazyDeserializedDAG.from_dag(dag), bundle_name=bundle_name) + + # Older scheduled run (NULL start_date) + dr_old = DagRun( + dag.dag_id, + run_id="scheduled_old", + run_type=DagRunType.SCHEDULED, + start_date=None, + ) + dr_old.created_at = base_date + session.add(dr_old) + + # Newer scheduled run (NULL start_date) — should be protected by keep_last + dr_new = DagRun( + dag.dag_id, + run_id="scheduled_new", + run_type=DagRunType.SCHEDULED, + start_date=None, + ) + dr_new.created_at = base_date.add(months=1) + session.add(dr_new) + session.commit() + return dag_id + + def test_keep_last_protects_latest_null_start_date_run(self): + """With flag, keep_last protects the newest scheduled run even when start_date is NULL.""" + base_date = pendulum.DateTime(2022, 1, 1, tzinfo=pendulum.timezone("UTC")) + self._create_scheduled_runs_with_null_start_date(base_date) + + with create_session() as session: + clean_before = base_date.add(years=1) + query = _build_query( + **config_dict["dag_run"].__dict__, + clean_before_timestamp=clean_before, + fallback_cleanup_on_null=True, + session=session, + ) + num_rows = session.scalars(select(func.count()).select_from(query.subquery())).one() + # Only the older run should be eligible; the newer one is protected by keep_last + assert num_rows == 1 + + def test_keep_last_all_null_without_flag_skips_all(self): + """Without flag, scheduled runs with NULL start_date are invisible to cleanup.""" + base_date = pendulum.DateTime(2022, 1, 1, tzinfo=pendulum.timezone("UTC")) + self._create_scheduled_runs_with_null_start_date(base_date) + + with create_session() as session: + clean_before = base_date.add(years=1) + query = _build_query( + **config_dict["dag_run"].__dict__, + clean_before_timestamp=clean_before, + session=session, + ) + num_rows = session.scalars(select(func.count()).select_from(query.subquery())).one() + # Without flag, both are invisible (NULL < cutoff → NULL → skipped) + assert num_rows == 0 From b34a9f0f28ced19fa2bc016c17b7d632a7d5597c Mon Sep 17 00:00:00 2001 From: Biley Roy Date: Thu, 16 Jul 2026 20:50:38 +0530 Subject: [PATCH 2/3] add fallback_cleanup_on_null to existing CLI test assertions --- .../tests/unit/cli/commands/test_db_command.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/airflow-core/tests/unit/cli/commands/test_db_command.py b/airflow-core/tests/unit/cli/commands/test_db_command.py index 76e391e1144ed..84f50b48fe8ef 100644 --- a/airflow-core/tests/unit/cli/commands/test_db_command.py +++ b/airflow-core/tests/unit/cli/commands/test_db_command.py @@ -815,6 +815,7 @@ def test_date_timezone_omitted(self, run_cleanup_mock, timezone): skip_archive=False, batch_size=None, error_on_cleanup_failure=False, + fallback_cleanup_on_null=False, ) @pytest.mark.parametrize("timezone", ["UTC", "Europe/Berlin", "America/Los_Angeles"]) @@ -839,6 +840,7 @@ def test_date_timezone_supplied(self, run_cleanup_mock, timezone): skip_archive=False, batch_size=None, error_on_cleanup_failure=False, + fallback_cleanup_on_null=False, ) @pytest.mark.parametrize(("confirm_arg", "expected"), [(["-y"], False), ([], True)]) @@ -869,6 +871,7 @@ def test_confirm(self, run_cleanup_mock, confirm_arg, expected): skip_archive=False, batch_size=None, error_on_cleanup_failure=False, + fallback_cleanup_on_null=False, ) @pytest.mark.parametrize(("extra_arg", "expected"), [(["--skip-archive"], True), ([], False)]) @@ -899,6 +902,7 @@ def test_skip_archive(self, run_cleanup_mock, extra_arg, expected): skip_archive=expected, batch_size=None, error_on_cleanup_failure=False, + fallback_cleanup_on_null=False, ) @pytest.mark.parametrize(("dry_run_arg", "expected"), [(["--dry-run"], True), ([], False)]) @@ -929,6 +933,7 @@ def test_dry_run(self, run_cleanup_mock, dry_run_arg, expected): skip_archive=False, batch_size=None, error_on_cleanup_failure=False, + fallback_cleanup_on_null=False, ) @pytest.mark.parametrize( @@ -961,6 +966,7 @@ def test_tables(self, run_cleanup_mock, extra_args, expected): skip_archive=False, batch_size=None, error_on_cleanup_failure=False, + fallback_cleanup_on_null=False, ) @pytest.mark.parametrize(("extra_args", "expected"), [(["--verbose"], True), ([], False)]) @@ -991,6 +997,7 @@ def test_verbose(self, run_cleanup_mock, extra_args, expected): skip_archive=False, batch_size=None, error_on_cleanup_failure=False, + fallback_cleanup_on_null=False, ) @pytest.mark.parametrize(("extra_args", "expected"), [(["--batch-size", "1234"], 1234), ([], None)]) @@ -1021,6 +1028,7 @@ def test_batch_size(self, run_cleanup_mock, extra_args, expected): skip_archive=False, batch_size=expected, error_on_cleanup_failure=False, + fallback_cleanup_on_null=False, ) @pytest.mark.parametrize( @@ -1053,6 +1061,7 @@ def test_dag_ids(self, run_cleanup_mock, extra_args, expected): skip_archive=False, batch_size=None, error_on_cleanup_failure=False, + fallback_cleanup_on_null=False, ) @pytest.mark.parametrize( @@ -1085,6 +1094,7 @@ def test_exclude_dag_ids(self, run_cleanup_mock, extra_args, expected): skip_archive=False, batch_size=None, error_on_cleanup_failure=False, + fallback_cleanup_on_null=False, ) @pytest.mark.parametrize( @@ -1115,6 +1125,7 @@ def test_error_on_cleanup_failure(self, run_cleanup_mock, extra_args, expected): skip_archive=False, batch_size=None, error_on_cleanup_failure=expected, + fallback_cleanup_on_null=False, ) @patch("airflow.cli.commands.db_command.export_archived_records") From bd28693aaa2255e05a1b0a284e54359a5317b42e Mon Sep 17 00:00:00 2001 From: Biley Roy Date: Thu, 16 Jul 2026 20:55:11 +0530 Subject: [PATCH 3/3] add CLI test for --fallback-cleanup-on-null flag --- .../unit/cli/commands/test_db_command.py | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/airflow-core/tests/unit/cli/commands/test_db_command.py b/airflow-core/tests/unit/cli/commands/test_db_command.py index 84f50b48fe8ef..67ce58834dc99 100644 --- a/airflow-core/tests/unit/cli/commands/test_db_command.py +++ b/airflow-core/tests/unit/cli/commands/test_db_command.py @@ -1128,6 +1128,37 @@ def test_error_on_cleanup_failure(self, run_cleanup_mock, extra_args, expected): fallback_cleanup_on_null=False, ) + @pytest.mark.parametrize( + ("extra_args", "expected"), [(["--fallback-cleanup-on-null"], True), ([], False)] + ) + @patch("airflow.cli.commands.db_command.run_cleanup") + def test_fallback_cleanup_on_null(self, run_cleanup_mock, extra_args, expected): + """When --fallback-cleanup-on-null is passed, fallback_cleanup_on_null should be True.""" + args = self.parser.parse_args( + [ + "db", + "clean", + "--clean-before-timestamp", + "2021-01-01", + *extra_args, + ] + ) + db_command.cleanup_tables(args) + + run_cleanup_mock.assert_called_once_with( + table_names=None, + dry_run=False, + dag_ids=None, + exclude_dag_ids=None, + clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"), + verbose=False, + confirm=True, + skip_archive=False, + batch_size=None, + error_on_cleanup_failure=False, + fallback_cleanup_on_null=expected, + ) + @patch("airflow.cli.commands.db_command.export_archived_records") @patch("airflow.cli.commands.db_command.os.path.isdir", return_value=True) def test_export_archived_records(self, os_mock, export_archived_mock):