Add new arguments to db_clean to explicitly include or exclude DAGs (#56663)

* Specify DAG IDs to include or exclude from db_clean

* Fix _confirm_delete

* Add documentation

* Fix tables without a dag_id_column, and keep last for dag_version

* Formatting, add more tests

* Fix deprecated sqlalchemy usage

Author: mattusifer

airflow-core/docs/howto/usage-cli.rst

@@ -213,6 +213,8 @@ The ``db clean`` command works by deleting from each table the records older tha
 
 You can optionally provide a list of tables to perform deletes on. If no list of tables is supplied, all tables will be included.
 
+You can filter cleanup to specific DAGs using ``--dag-ids`` (comma-separated list), or exclude specific DAGs using ``--exclude-dag-ids`` (comma-separated list). These options allow you to target or avoid cleanup for particular DAGs.
+
 You can use the ``--dry-run`` option to print the row counts in the primary tables to be cleaned.
 
 By default, ``db clean`` will archive purged rows in tables of the form ``_airflow_deleted__<table>__<timestamp>``.  If you don't want the data preserved in this way, you may supply argument ``--skip-archive``.

airflow-core/src/airflow/cli/cli_config.py

@@ -520,6 +520,18 @@ def string_lower_type(val):
         "Lower values reduce long-running locks but increase the number of batches."
     ),
 )
+ARG_DAG_IDS = Arg(
+    ("--dag-ids",),
+    default=None,
+    help="Only cleanup data related to the given dag_id",
+    type=string_list_type,
+)
+ARG_EXCLUDE_DAG_IDS = Arg(
+    ("--exclude-dag-ids",),
+    default=None,
+    help="Avoid cleaning up data related to the given dag_ids",
+    type=string_list_type,
+)
 
 # pool
 ARG_POOL_NAME = Arg(("pool",), metavar="NAME", help="Pool name")
@@ -1527,6 +1539,8 @@ class GroupCommand(NamedTuple):
             ARG_YES,
             ARG_DB_SKIP_ARCHIVE,
             ARG_DB_BATCH_SIZE,
+            ARG_DAG_IDS,
+            ARG_EXCLUDE_DAG_IDS,
         ),
     ),
     ActionCommand(

airflow-core/src/airflow/cli/commands/db_command.py

@@ -301,6 +301,8 @@ def cleanup_tables(args):
         confirm=not args.yes,
         skip_archive=args.skip_archive,
         batch_size=args.batch_size,
+        dag_ids=args.dag_ids,
+        exclude_dag_ids=args.exclude_dag_ids,
     )
 
 

airflow-core/src/airflow/utils/db_cleanup.py

@@ -79,6 +79,7 @@ class _TableConfig:
     table_name: str
     recency_column_name: str
     extra_columns: list[str] | None = None
+    dag_id_column_name: str | None = None
     keep_last: bool = False
     keep_last_filters: Any | None = None
     keep_last_group_by: Any | None = None
@@ -89,9 +90,19 @@ class _TableConfig:
 
     def __post_init__(self):
         self.recency_column = column(self.recency_column_name)
-        self.orm_model: Base = table(
-            self.table_name, *[column(x) for x in self.extra_columns or []], self.recency_column
-        )
+        if self.dag_id_column_name is None:
+            self.dag_id_column = None
+            self.orm_model: Base = table(
+                self.table_name, *[column(x) for x in self.extra_columns or []], self.recency_column
+            )
+        else:
+            self.dag_id_column = column(self.dag_id_column_name)
+            self.orm_model: Base = table(
+                self.table_name,
+                *[column(x) for x in self.extra_columns or []],
+                self.dag_id_column,
+                self.recency_column,
+            )
 
     def __lt__(self, other):
         return self.table_name < other.table_name
@@ -101,41 +112,47 @@ def readable_config(self):
         return {
             "table": self.orm_model.name,
             "recency_column": str(self.recency_column),
+            "dag_id_column": str(self.dag_id_column),
             "keep_last": self.keep_last,
             "keep_last_filters": [str(x) for x in self.keep_last_filters] if self.keep_last_filters else None,
             "keep_last_group_by": str(self.keep_last_group_by),
         }
 
 
 config_list: list[_TableConfig] = [
-    _TableConfig(table_name="job", recency_column_name="latest_heartbeat"),
+    _TableConfig(table_name="job", recency_column_name="latest_heartbeat", dag_id_column_name="dag_id"),
     _TableConfig(
         table_name="dag",
         recency_column_name="last_parsed_time",
         dependent_tables=["dag_version", "deadline"],
+        dag_id_column_name="dag_id",
     ),
     _TableConfig(
         table_name="dag_run",
         recency_column_name="start_date",
+        dag_id_column_name="dag_id",
         extra_columns=["dag_id", "run_type"],
         keep_last=True,
         keep_last_filters=[column("run_type") != DagRunType.MANUAL],
         keep_last_group_by=["dag_id"],
         dependent_tables=["task_instance", "deadline"],
     ),
-    _TableConfig(table_name="asset_event", recency_column_name="timestamp"),
+    _TableConfig(table_name="asset_event", recency_column_name="timestamp", dag_id_column_name="dag_id"),
     _TableConfig(table_name="import_error", recency_column_name="timestamp"),
-    _TableConfig(table_name="log", recency_column_name="dttm"),
-    _TableConfig(table_name="sla_miss", recency_column_name="timestamp"),
+    _TableConfig(table_name="log", recency_column_name="dttm", dag_id_column_name="dag_id"),
+    _TableConfig(table_name="sla_miss", recency_column_name="timestamp", dag_id_column_name="dag_id"),
     _TableConfig(
         table_name="task_instance",
         recency_column_name="start_date",
         dependent_tables=["task_instance_history", "xcom"],
+        dag_id_column_name="dag_id",
+    ),
+    _TableConfig(
+        table_name="task_instance_history", recency_column_name="start_date", dag_id_column_name="dag_id"
     ),
-    _TableConfig(table_name="task_instance_history", recency_column_name="start_date"),
-    _TableConfig(table_name="task_reschedule", recency_column_name="start_date"),
-    _TableConfig(table_name="xcom", recency_column_name="timestamp"),
-    _TableConfig(table_name="_xcom_archive", recency_column_name="timestamp"),
+    _TableConfig(table_name="task_reschedule", recency_column_name="start_date", dag_id_column_name="dag_id"),
+    _TableConfig(table_name="xcom", recency_column_name="timestamp", dag_id_column_name="dag_id"),
+    _TableConfig(table_name="_xcom_archive", recency_column_name="timestamp", dag_id_column_name="dag_id"),
     _TableConfig(table_name="callback_request", recency_column_name="created_at"),
     _TableConfig(table_name="celery_taskmeta", recency_column_name="date_done"),
     _TableConfig(table_name="celery_tasksetmeta", recency_column_name="date_done"),
@@ -148,8 +165,11 @@ def readable_config(self):
         table_name="dag_version",
         recency_column_name="created_at",
         dependent_tables=["task_instance", "dag_run"],
+        dag_id_column_name="dag_id",
+        keep_last=True,
+        keep_last_group_by=["dag_id"],
     ),
-    _TableConfig(table_name="deadline", recency_column_name="deadline_time"),
+    _TableConfig(table_name="deadline", recency_column_name="deadline_time", dag_id_column_name="dag_id"),
 ]
 
 # We need to have `fallback="database"` because this is executed at top level code and provider configuration
@@ -308,13 +328,25 @@ def _build_query(
     keep_last_group_by,
     clean_before_timestamp: DateTime,
     session: Session,
+    dag_id_column=None,
+    dag_ids: list[str] | None = None,
+    exclude_dag_ids: list[str] | None = None,
     **kwargs,
 ) -> Query:
     base_table_alias = "base"
     base_table = aliased(orm_model, name=base_table_alias)
     query = session.query(base_table).with_entities(text(f"{base_table_alias}.*"))
     base_table_recency_col = base_table.c[recency_column.name]
     conditions = [base_table_recency_col < clean_before_timestamp]
+
+    if (dag_ids or exclude_dag_ids) and dag_id_column is not None:
+        base_table_dag_id_col = base_table.c[dag_id_column.name]
+
+        if dag_ids:
+            conditions.append(base_table_dag_id_col.in_(dag_ids))
+        if exclude_dag_ids:
+            conditions.append(base_table_dag_id_col.not_in(exclude_dag_ids))
+
     if keep_last:
         max_date_col_name = "max_date_per_group"
         group_by_columns: list[Any] = [column(x) for x in keep_last_group_by]
@@ -345,6 +377,9 @@ def _cleanup_table(
     keep_last_filters,
     keep_last_group_by,
     clean_before_timestamp: DateTime,
+    dag_id_column=None,
+    dag_ids=None,
+    exclude_dag_ids=None,
     dry_run: bool = True,
     verbose: bool = False,
     skip_archive: bool = False,
@@ -358,6 +393,9 @@ def _cleanup_table(
     query = _build_query(
         orm_model=orm_model,
         recency_column=recency_column,
+        dag_id_column=dag_id_column,
+        dag_ids=dag_ids,
+        exclude_dag_ids=exclude_dag_ids,
         keep_last=keep_last,
         keep_last_filters=keep_last_filters,
         keep_last_group_by=keep_last_group_by,
@@ -380,10 +418,18 @@ def _cleanup_table(
     session.commit()
 
 
-def _confirm_delete(*, date: DateTime, tables: list[str]) -> None:
+def _confirm_delete(
+    *,
+    date: DateTime,
+    tables: list[str],
+    dag_ids: list[str] | None = None,
+    exclude_dag_ids: list[str] | 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 ""
     question = (
-        f"You have requested that we purge all data prior to {date}{for_tables}.\n"
+        f"You have requested that we purge all data prior to {date}{for_tables}{for_dags}{excluding_dags}.\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."
@@ -493,6 +539,8 @@ def run_cleanup(
     *,
     clean_before_timestamp: DateTime,
     table_names: list[str] | None = None,
+    dag_ids: list[str] | None = None,
+    exclude_dag_ids: list[str] | None = None,
     dry_run: bool = False,
     verbose: bool = False,
     confirm: bool = True,
@@ -513,6 +561,9 @@ def run_cleanup(
     :param clean_before_timestamp: The timestamp before which data should be purged
     :param table_names: Optional. List of table names to perform maintenance on.  If list not provided,
         will perform maintenance on all tables.
+    :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 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.
@@ -532,14 +583,21 @@ def run_cleanup(
         )
         _print_config(configs=effective_config_dict)
     if not dry_run and confirm:
-        _confirm_delete(date=clean_before_timestamp, tables=sorted(effective_table_names))
+        _confirm_delete(
+            date=clean_before_timestamp,
+            tables=sorted(effective_table_names),
+            dag_ids=dag_ids,
+            exclude_dag_ids=exclude_dag_ids,
+        )
     existing_tables = reflect_tables(tables=None, session=session).tables
 
     for table_name, table_config in effective_config_dict.items():
         if table_name in existing_tables:
             with _suppress_with_logging(table_name, session):
                 _cleanup_table(
                     clean_before_timestamp=clean_before_timestamp,
+                    dag_ids=dag_ids,
+                    exclude_dag_ids=exclude_dag_ids,
                     dry_run=dry_run,
                     verbose=verbose,
                     **table_config.__dict__,

airflow-core/tests/unit/cli/commands/test_db_command.py

@@ -701,6 +701,8 @@ def test_date_timezone_omitted(self, run_cleanup_mock, timezone):
             db_command.cleanup_tables(args)
         run_cleanup_mock.assert_called_once_with(
             table_names=None,
+            dag_ids=None,
+            exclude_dag_ids=None,
             dry_run=False,
             clean_before_timestamp=pendulum.parse(timestamp, tz=timezone),
             verbose=False,
@@ -722,6 +724,8 @@ def test_date_timezone_supplied(self, run_cleanup_mock, timezone):
 
         run_cleanup_mock.assert_called_once_with(
             table_names=None,
+            dag_ids=None,
+            exclude_dag_ids=None,
             dry_run=False,
             clean_before_timestamp=pendulum.parse(timestamp),
             verbose=False,
@@ -749,6 +753,8 @@ def test_confirm(self, run_cleanup_mock, confirm_arg, expected):
 
         run_cleanup_mock.assert_called_once_with(
             table_names=None,
+            dag_ids=None,
+            exclude_dag_ids=None,
             dry_run=False,
             clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"),
             verbose=False,
@@ -776,6 +782,8 @@ def test_skip_archive(self, run_cleanup_mock, extra_arg, expected):
 
         run_cleanup_mock.assert_called_once_with(
             table_names=None,
+            dag_ids=None,
+            exclude_dag_ids=None,
             dry_run=False,
             clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"),
             verbose=False,
@@ -803,6 +811,8 @@ def test_dry_run(self, run_cleanup_mock, dry_run_arg, expected):
 
         run_cleanup_mock.assert_called_once_with(
             table_names=None,
+            dag_ids=None,
+            exclude_dag_ids=None,
             dry_run=expected,
             clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"),
             verbose=False,
@@ -832,6 +842,8 @@ def test_tables(self, run_cleanup_mock, extra_args, expected):
 
         run_cleanup_mock.assert_called_once_with(
             table_names=expected,
+            dag_ids=None,
+            exclude_dag_ids=None,
             dry_run=False,
             clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"),
             verbose=False,
@@ -859,6 +871,8 @@ def test_verbose(self, run_cleanup_mock, extra_args, expected):
 
         run_cleanup_mock.assert_called_once_with(
             table_names=None,
+            dag_ids=None,
+            exclude_dag_ids=None,
             dry_run=False,
             clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"),
             verbose=expected,
@@ -886,6 +900,8 @@ def test_batch_size(self, run_cleanup_mock, extra_args, expected):
 
         run_cleanup_mock.assert_called_once_with(
             table_names=None,
+            dag_ids=None,
+            exclude_dag_ids=None,
             dry_run=False,
             clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"),
             verbose=False,
@@ -894,6 +910,68 @@ def test_batch_size(self, run_cleanup_mock, extra_args, expected):
             batch_size=expected,
         )
 
+    @pytest.mark.parametrize(
+        ("extra_args", "expected"), [(["--dag-ids", "dag1, dag2"], ["dag1", "dag2"]), ([], None)]
+    )
+    @patch("airflow.cli.commands.db_command.run_cleanup")
+    def test_dag_ids(self, run_cleanup_mock, extra_args, expected):
+        """
+        When dag_ids are included in the args then dag_ids should be passed in, or None otherwise
+        """
+        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=expected,
+            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,
+        )
+
+    @pytest.mark.parametrize(
+        ("extra_args", "expected"), [(["--exclude-dag-ids", "dag1, dag2"], ["dag1", "dag2"]), ([], None)]
+    )
+    @patch("airflow.cli.commands.db_command.run_cleanup")
+    def test_exclude_dag_ids(self, run_cleanup_mock, extra_args, expected):
+        """
+        When exclude_dag_ids are included in the args then exclude_dag_ids should be passed in, or None otherwise
+        """
+        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=expected,
+            clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"),
+            verbose=False,
+            confirm=True,
+            skip_archive=False,
+            batch_size=None,
+        )
+
     @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):