[v3-4-test] Add asset ID filtering to database cleanup (#74498)

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__<table>__<timestamp>``. 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.

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,

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,
     )
 

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

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,