This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new b9ae12ac432 Add asset ID filtering to database cleanup (#74498)
b9ae12ac432 is described below
commit b9ae12ac432a31910a7acb56b27c93ef84e7a5cd
Author: Hussein Awala <[email protected]>
AuthorDate: Sat Oct 24 21:39:30 2026 +0200
Add asset ID filtering to database cleanup (#74498)
Administrators need to retain event history independently for different
assets without cleaning every asset in the table. Restricting the new filter to
an explicit asset_event selection avoids unintentionally cleaning unrelated
data.
---
airflow-core/docs/howto/usage-cli.rst | 9 ++
airflow-core/src/airflow/cli/cli_config.py | 7 +
.../src/airflow/cli/commands/db_command.py | 1 +
airflow-core/src/airflow/utils/db_cleanup.py | 24 ++-
.../tests/unit/cli/commands/test_db_command.py | 84 ++++++++++
airflow-core/tests/unit/utils/test_db_cleanup.py | 178 +++++++++++++++++++++
6 files changed, 302 insertions(+), 1 deletion(-)
diff --git a/airflow-core/docs/howto/usage-cli.rst
b/airflow-core/docs/howto/usage-cli.rst
index a45b30b2ad0..0a8ecf576a5 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__<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.
diff --git a/airflow-core/src/airflow/cli/cli_config.py
b/airflow-core/src/airflow/cli/cli_config.py
index 961a295ddee..6f9be99f3c2 100644
--- a/airflow-core/src/airflow/cli/cli_config.py
+++ b/airflow-core/src/airflow/cli/cli_config.py
@@ -596,6 +596,12 @@ ARG_DB_ERROR_ON_CLEANUP_FAILURE = Arg(
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,
@@ -1826,6 +1832,7 @@ DB_COMMANDS = (
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 f6c4403f751..ee027222df0 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 ca81c3a503f..a3cd6e6679e 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 df09941b73b..19fcf338c6a 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 @@ class TestCLIDBClean:
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 @@ class TestCLIDBClean:
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 @@ class TestCLIDBClean:
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 @@ class TestCLIDBClean:
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 @@ class TestCLIDBClean:
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 @@ class TestCLIDBClean:
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 @@ class TestCLIDBClean:
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 @@ class TestCLIDBClean:
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 @@ class TestCLIDBClean:
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 @@ class TestCLIDBClean:
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 @@ class TestCLIDBClean:
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 @@ class TestCLIDBClean:
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 045a1587977..b7126f1acb9 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 import DAG, settings
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 @@ class TestDBCleanup:
)
+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
