diff --git a/airflow-core/docs/howto/usage-cli.rst b/airflow-core/docs/howto/usage-cli.rst index a45b30b2ad0e1..0a8ecf576a544 100644 --- a/airflow-core/docs/howto/usage-cli.rst +++ b/airflow-core/docs/howto/usage-cli.rst @@ -259,6 +259,15 @@ You can filter cleanup to specific DAGs using ``--dag-ids`` (comma-separated lis The ``--dry-run`` output shows a ``dag_id_column`` for each table it would clean. Where that column is ``None``, the DAG filters do not apply to that table. +To clean events for specific assets, pass ``--asset-ids`` with a comma-separated list of positive integer +asset IDs. This requires explicitly selecting only ``--tables asset_event``; omitted or additional tables +are rejected. The timestamp cutoff still applies, and any Dag filters further restrict the events by +``source_dag_id``. Existing cascading deletes, including queued events, still apply. + +.. code-block:: bash + + airflow db clean --tables asset_event --asset-ids 1,2 --clean-before-timestamp '2024-01-01T00:00:00+00:00' --dry-run + By default, ``db clean`` will archive purged rows in tables of the form ``_airflow_deleted____``. Cascading attempt data may use a unique suffix instead of a timestamp. If you don't want the data preserved in this way, you may supply argument ``--skip-archive``. When you encounter an error without using ``--skip-archive``, an archive table may remain in the DB. You can use ``db drop-archived`` to remove it. ``db export-archived`` and ``db drop-archived`` also accept the pre-upgrade names ``xcom`` and ``task_instance_history`` to find older archives. diff --git a/airflow-core/src/airflow/cli/cli_config.py b/airflow-core/src/airflow/cli/cli_config.py index 579c92efdda0d..c3f9d14ce8226 100644 --- a/airflow-core/src/airflow/cli/cli_config.py +++ b/airflow-core/src/airflow/cli/cli_config.py @@ -596,6 +596,12 @@ 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_ASSET_IDS = Arg( + ("--asset-ids",), + default=None, + help="Only clean events for these comma-separated positive integer asset IDs; requires --tables asset_event", + type=lambda value: [positive_int(allow_zero=False)(item) for item in string_list_type(value)], +) ARG_DAG_IDS = Arg( ("--dag-ids",), default=None, @@ -1809,6 +1815,7 @@ class GroupCommand(NamedTuple): ARG_YES, ARG_DB_SKIP_ARCHIVE, ARG_DB_BATCH_SIZE, + ARG_ASSET_IDS, ARG_DAG_IDS, ARG_EXCLUDE_DAG_IDS, ARG_DB_ERROR_ON_CLEANUP_FAILURE, diff --git a/airflow-core/src/airflow/cli/commands/db_command.py b/airflow-core/src/airflow/cli/commands/db_command.py index f6c4403f751ef..ee027222df0ef 100644 --- a/airflow-core/src/airflow/cli/commands/db_command.py +++ b/airflow-core/src/airflow/cli/commands/db_command.py @@ -379,6 +379,7 @@ def cleanup_tables(args): batch_size=args.batch_size, dag_ids=args.dag_ids, exclude_dag_ids=args.exclude_dag_ids, + asset_ids=args.asset_ids, error_on_cleanup_failure=args.error_on_cleanup_failure, ) diff --git a/airflow-core/src/airflow/utils/db_cleanup.py b/airflow-core/src/airflow/utils/db_cleanup.py index 23ea7e41bc1a3..c466626b07f71 100644 --- a/airflow-core/src/airflow/utils/db_cleanup.py +++ b/airflow-core/src/airflow/utils/db_cleanup.py @@ -876,12 +876,15 @@ def _confirm_delete( tables: list[str], dag_ids: list[str] | None = None, exclude_dag_ids: list[str] | None = None, + asset_ids: list[int] | None = None, ) -> None: for_tables = f" for tables {tables!r}" if tables else "" for_dags = f" for the following dags: {dag_ids!r}" if dag_ids else "" excluding_dags = f" excluding the following dags: {exclude_dag_ids!r}" if exclude_dag_ids else "" + for_assets = f" for asset IDs {asset_ids!r}" if asset_ids is not None else "" question = ( - f"You have requested that we purge all data prior to {date}{for_tables}{for_dags}{excluding_dags}.\n" + f"You have requested that we purge all data prior to " + f"{date}{for_tables}{for_dags}{excluding_dags}{for_assets}.\n" f"This is irreversible. Consider backing up the tables first and / or doing a dry run " f"with option --dry-run.\n" f"Enter 'delete rows' (without quotes) to proceed." @@ -1027,6 +1030,7 @@ def run_cleanup( table_names: list[str] | None = None, dag_ids: list[str] | None = None, exclude_dag_ids: list[str] | None = None, + asset_ids: list[int] | None = None, dry_run: bool = False, verbose: bool = False, confirm: bool = True, @@ -1051,6 +1055,8 @@ def run_cleanup( :param dag_ids: Optional. List of dag ids to perform maintenance on. If list not provided, will perform maintenance on all dags. :param exclude_dag_ids: Optional. List of dag ids to exclude from maintenance. + :param asset_ids: Optional. Only clean events for these asset IDs. Requires table_names=["asset_event"]. + An empty list selects no events. :param dry_run: If true, print rows meeting deletion criteria :param verbose: If true, may provide more detailed output. :param confirm: Require user input to confirm before processing deletions. @@ -1061,12 +1067,27 @@ def run_cleanup( 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. """ + if asset_ids is not None and set(table_names or []) != {"asset_event"}: + raise SystemExit("--asset-ids requires explicitly selecting only --tables asset_event.") + clean_before_timestamp = timezone.coerce_datetime(clean_before_timestamp) # Get all tables to clean (root + dependents) effective_table_names, effective_config_dict = _effective_table_names(table_names=table_names) + if asset_ids is not None: + asset_config = effective_config_dict["asset_event"] + effective_config_dict["asset_event"] = dataclasses.replace( + asset_config, + extra_columns=[*(asset_config.extra_columns or []), "asset_id"], + extra_filters=[ + *(asset_config.extra_filters or []), + literal_column(f"{_BASE_TABLE_ALIAS}.asset_id").in_(asset_ids), + ], + ) if dry_run: print("Performing dry run for db cleanup.") + if asset_ids is not None: + print(f"Restricting cleanup to asset IDs {asset_ids!r}.") print( f"Data prior to {clean_before_timestamp} would be purged " f"from tables {effective_table_names} with the following config:\n" @@ -1078,6 +1099,7 @@ def run_cleanup( tables=sorted(effective_table_names), dag_ids=dag_ids, exclude_dag_ids=exclude_dag_ids, + asset_ids=asset_ids, ) existing_tables: set[str] = { table 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 df09941b73b5c..19fcf338c6a1f 100644 --- a/airflow-core/tests/unit/cli/commands/test_db_command.py +++ b/airflow-core/tests/unit/cli/commands/test_db_command.py @@ -865,6 +865,7 @@ def test_date_timezone_omitted(self, run_cleanup_mock, timezone): table_names=None, dag_ids=None, exclude_dag_ids=None, + asset_ids=None, dry_run=False, clean_before_timestamp=pendulum.parse(timestamp, tz=timezone), verbose=False, @@ -889,6 +890,7 @@ def test_date_timezone_supplied(self, run_cleanup_mock, timezone): table_names=None, dag_ids=None, exclude_dag_ids=None, + asset_ids=None, dry_run=False, clean_before_timestamp=pendulum.parse(timestamp), verbose=False, @@ -919,6 +921,7 @@ def test_confirm(self, run_cleanup_mock, confirm_arg, expected): table_names=None, dag_ids=None, exclude_dag_ids=None, + asset_ids=None, dry_run=False, clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"), verbose=False, @@ -949,6 +952,7 @@ def test_skip_archive(self, run_cleanup_mock, extra_arg, expected): table_names=None, dag_ids=None, exclude_dag_ids=None, + asset_ids=None, dry_run=False, clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"), verbose=False, @@ -979,6 +983,7 @@ def test_dry_run(self, run_cleanup_mock, dry_run_arg, expected): table_names=None, dag_ids=None, exclude_dag_ids=None, + asset_ids=None, dry_run=expected, clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"), verbose=False, @@ -1011,6 +1016,7 @@ def test_tables(self, run_cleanup_mock, extra_args, expected): table_names=expected, dag_ids=None, exclude_dag_ids=None, + asset_ids=None, dry_run=False, clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"), verbose=False, @@ -1041,6 +1047,7 @@ def test_verbose(self, run_cleanup_mock, extra_args, expected): table_names=None, dag_ids=None, exclude_dag_ids=None, + asset_ids=None, dry_run=False, clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"), verbose=expected, @@ -1071,6 +1078,7 @@ def test_batch_size(self, run_cleanup_mock, extra_args, expected): table_names=None, dag_ids=None, exclude_dag_ids=None, + asset_ids=None, dry_run=False, clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"), verbose=False, @@ -1104,6 +1112,7 @@ def test_dag_ids(self, run_cleanup_mock, extra_args, expected): dry_run=False, dag_ids=expected, exclude_dag_ids=None, + asset_ids=None, clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"), verbose=False, confirm=True, @@ -1112,6 +1121,79 @@ def test_dag_ids(self, run_cleanup_mock, extra_args, expected): error_on_cleanup_failure=False, ) + @patch("airflow.cli.commands.db_command.run_cleanup", autospec=True) + def test_asset_ids(self, run_cleanup_mock): + args = self.parser.parse_args( + [ + "db", + "clean", + "--clean-before-timestamp", + "2021-01-01", + "--tables", + "asset_event", + "--asset-ids", + " 1, 2 ", + "--dag-ids", + "producer", + "--exclude-dag-ids", + "excluded", + ] + ) + db_command.cleanup_tables(args) + run_cleanup_mock.assert_called_once_with( + table_names=["asset_event"], + asset_ids=[1, 2], + dag_ids=["producer"], + exclude_dag_ids=["excluded"], + clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"), + dry_run=False, + verbose=False, + confirm=True, + skip_archive=False, + batch_size=None, + error_on_cleanup_failure=False, + ) + + @pytest.mark.parametrize("asset_ids", ["", " ", ",", "1,", ",1", "1,,2", "0", "-1", "1.5", "abc", "1, 0"]) + @patch("airflow.cli.commands.db_command.run_cleanup", autospec=True) + def test_invalid_asset_ids(self, run_cleanup_mock, asset_ids, capsys): + with pytest.raises(SystemExit) as exc_info: + db_command.cleanup_tables( + self.parser.parse_args( + [ + "db", + "clean", + "--clean-before-timestamp", + "2021-01-01", + "--tables", + "asset_event", + f"--asset-ids={asset_ids}", + ] + ) + ) + assert exc_info.value.code == 2 + assert "argument --asset-ids: invalid positive int value:" in capsys.readouterr().err + run_cleanup_mock.assert_not_called() + + @pytest.mark.parametrize( + "tables", + [[], ["--tables", "log"], ["--tables", "asset_event,log"], ["--tables", "asset_event,invalid"]], + ) + @patch("airflow.utils.db_cleanup.reflect_tables", autospec=True) + @patch("airflow.utils.db_cleanup._cleanup_table", autospec=True) + @patch("airflow.utils.db_cleanup._confirm_delete", autospec=True) + def test_asset_ids_require_only_asset_event(self, confirm_mock, cleanup_mock, reflect_mock, tables): + args = self.parser.parse_args( + ["db", "clean", "--clean-before-timestamp", "2021-01-01", "--asset-ids", "1,2", *tables] + ) + with pytest.raises( + SystemExit, match="--asset-ids requires explicitly selecting only --tables asset_event" + ): + db_command.cleanup_tables(args) + confirm_mock.assert_not_called() + cleanup_mock.assert_not_called() + reflect_mock.assert_not_called() + @pytest.mark.parametrize( ("extra_args", "expected"), [(["--exclude-dag-ids", "dag1, dag2"], ["dag1", "dag2"]), ([], None)] ) @@ -1136,6 +1218,7 @@ def test_exclude_dag_ids(self, run_cleanup_mock, extra_args, expected): dry_run=False, dag_ids=None, exclude_dag_ids=expected, + asset_ids=None, clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"), verbose=False, confirm=True, @@ -1166,6 +1249,7 @@ def test_error_on_cleanup_failure(self, run_cleanup_mock, extra_args, expected): dry_run=False, dag_ids=None, exclude_dag_ids=None, + asset_ids=None, clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"), verbose=False, confirm=True, diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py b/airflow-core/tests/unit/utils/test_db_cleanup.py index 045a158797738..b7126f1acb998 100644 --- a/airflow-core/tests/unit/utils/test_db_cleanup.py +++ b/airflow-core/tests/unit/utils/test_db_cleanup.py @@ -50,6 +50,15 @@ from airflow._shared.timezones import timezone from airflow.exceptions import AirflowException from airflow.models import DagModel, DagRun, TaskInstance +from airflow.models.asset import ( + AssetAliasModel, + AssetDagRunQueue, + AssetEvent, + AssetModel, + DagScheduleAssetReference, + asset_alias_asset_event_association_table, + association_table, +) from airflow.models.dag_version import DagVersion from airflow.models.dagbundle import DagBundleModel from airflow.models.deadline import Deadline @@ -1812,6 +1821,175 @@ def test_cleanup_scopes_attempt_cascade_to_selected_dag(self, ownership_session, ) +class TestAssetEventCleanup: + CUTOFF = pendulum.datetime(2024, 1, 2, tz="UTC") + + @pytest.fixture + def asset_events(self, session): + drop_tables_with_prefix(ARCHIVE_TABLE_PREFIX) + assets = [AssetModel(name=f"asset-{i}", uri=f"test://asset-{i}") for i in range(3)] + session.add_all(assets) + session.flush() + asset_ids = [asset.id for asset in assets] + events = [ + AssetEvent(asset_id=asset_id, source_dag_id=source, timestamp=self.CUTOFF.add(days=offset)) + for asset_id in asset_ids + for source in ("producer", "other", None) + for offset in (-1, 0, 1) + ] + session.add_all(events) + session.flush() + rows = {event.id: (event.asset_id, event.source_dag_id, event.timestamp) for event in events} + session.commit() + yield asset_ids, rows + drop_tables_with_prefix(ARCHIVE_TABLE_PREFIX) + + @pytest.mark.parametrize( + ("dry_run", "skip_archive", "batch_size", "dag_ids", "exclude_dag_ids", "sources"), + [ + pytest.param(False, False, None, None, None, {"producer", "other", None}, id="archive"), + pytest.param(False, False, 2, None, None, {"producer", "other", None}, id="batched-archive"), + pytest.param(False, True, 2, None, None, {"producer", "other", None}, id="batched-skip-archive"), + pytest.param(True, False, 2, None, None, {"producer", "other", None}, id="dry-run"), + pytest.param(False, True, 2, ["producer"], None, {"producer"}, id="include-dag"), + pytest.param(False, True, 2, None, ["other"], {"producer", None}, id="exclude-dag"), + pytest.param( + False, True, 2, ["producer", "other"], ["other"], {"producer"}, id="intersect-dag-filters" + ), + ], + ) + def test_cleanup_selected_assets( + self, + asset_events, + session, + dag_maker, + capsys, + dry_run, + skip_archive, + batch_size, + dag_ids, + exclude_dag_ids, + sources, + ): + asset_ids, rows = asset_events + selected_assets = asset_ids[:2] + with dag_maker("asset_cleanup_consumer", session=session): + pass + dag_run = dag_maker.create_dagrun() + dag_run_id, dag_id = dag_run.id, dag_run.dag_id + alias = AssetAliasModel(name="cleanup-alias") + session.add(alias) + session.flush() + alias_id = alias.id + session.add_all(DagScheduleAssetReference(asset_id=asset_id, dag_id=dag_id) for asset_id in asset_ids) + session.add_all( + AssetDagRunQueue(asset_event_id=event_id, asset_id=row[0], target_dag_id=dag_id) + for event_id, row in rows.items() + ) + session.execute( + association_table.insert(), + [{"dag_run_id": dag_run_id, "event_id": event_id} for event_id in rows], + ) + session.execute( + asset_alias_asset_event_association_table.insert(), + [{"alias_id": alias_id, "event_id": event_id} for event_id in rows], + ) + session.commit() + matching = { + event_id + for event_id, (asset_id, source, timestamp) in rows.items() + if asset_id in selected_assets and source in sources and timestamp < self.CUTOFF + } + + run_cleanup( + clean_before_timestamp=self.CUTOFF, + table_names=["asset_event"], + asset_ids=selected_assets, + dag_ids=dag_ids, + exclude_dag_ids=exclude_dag_ids, + dry_run=dry_run, + skip_archive=skip_archive, + batch_size=batch_size, + confirm=False, + error_on_cleanup_failure=True, + session=session, + ) + + remaining = set(rows) if dry_run else set(rows) - matching + assert set(session.scalars(select(AssetEvent.id))) == remaining + assert set(session.scalars(select(AssetDagRunQueue.asset_event_id))) == remaining + assert set(session.scalars(select(association_table.c.event_id))) == remaining + assert set(session.scalars(select(asset_alias_asset_event_association_table.c.event_id))) == remaining + assert set(session.scalars(select(AssetModel.id))) == set(asset_ids) + assert session.scalar(select(DagModel.dag_id).where(DagModel.dag_id == dag_id)) == dag_id + assert session.scalar(select(DagRun.id).where(DagRun.id == dag_run_id)) == dag_run_id + assert session.scalar(select(AssetAliasModel.id)) == alias_id + assert set(session.scalars(select(DagScheduleAssetReference.asset_id))) == set(asset_ids) + + archives = [ + name + for name in _get_archived_table_names(["asset_event"], session) + if name.startswith(f"{ARCHIVE_TABLE_PREFIX}asset_event__") + ] + if dry_run or skip_archive: + assert archives == [] + else: + archived_ids = { + event_id + for name in archives + for event_id in session.scalars(select(sa.column("id")).select_from(sa.table(name))) + } + assert archived_ids == matching + assert len(archives) == (len(matching) // batch_size if batch_size else 1) + if dry_run: + output = capsys.readouterr().out + assert f"Restricting cleanup to asset IDs {selected_assets!r}." in output + assert f"Found {len(matching)} rows meeting deletion criteria." in output + + @pytest.mark.parametrize("empty_selection", [False, True]) + def test_asset_filter_is_request_local(self, asset_events, session, empty_selection): + asset_ids, rows = asset_events + selected = [] if empty_selection else asset_ids[:1] + run_cleanup( + clean_before_timestamp=self.CUTOFF, + table_names=["asset_event"], + asset_ids=selected, + skip_archive=True, + confirm=False, + error_on_cleanup_failure=True, + session=session, + ) + deleted = { + event_id + for event_id, (asset_id, _, timestamp) in rows.items() + if asset_id in selected and timestamp < self.CUTOFF + } + assert set(session.scalars(select(AssetEvent.id))) == set(rows) - deleted + + run_cleanup( + clean_before_timestamp=self.CUTOFF, + table_names=["asset_event"], + asset_ids=None, + skip_archive=True, + confirm=False, + error_on_cleanup_failure=True, + session=session, + ) + assert set(session.scalars(select(AssetEvent.id))) == { + event_id for event_id, (_, _, timestamp) in rows.items() if timestamp >= self.CUTOFF + } + + @patch("builtins.input", autospec=True, return_value="cancel") + def test_confirmation_includes_asset_ids(self, mock_input, capsys): + with pytest.raises(SystemExit, match="User did not confirm"): + run_cleanup( + clean_before_timestamp=self.CUTOFF, + table_names=["asset_event"], + asset_ids=[1, 2], + ) + assert "for asset IDs [1, 2]" in capsys.readouterr().out + + def create_tis(base_date, num_tis, run_type=DagRunType.SCHEDULED): from tests_common.test_utils.taskinstance import create_task_instance