CLI: add --batch-size option to "airflow db clean"

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

@@ -483,6 +483,15 @@ def string_lower_type(val):
     type=positive_int(allow_zero=False),
     help="Wait time between retries in seconds",
 )
+ARG_DB_BATCH_SIZE = Arg(
+    ("--batch-size",),
+    default=None,
+    type=positive_int(allow_zero=False),
+    help=(
+        "Maximum number of rows to delete or archive in a single transaction.\n"
+        "Lower values reduce long-running locks but increase the number of batches."
+    ),
+)
 
 # pool
 ARG_POOL_NAME = Arg(("pool",), metavar="NAME", help="Pool name")
@@ -1452,6 +1461,7 @@ class GroupCommand(NamedTuple):
             ARG_VERBOSE,
             ARG_YES,
             ARG_DB_SKIP_ARCHIVE,
+            ARG_DB_BATCH_SIZE,
         ),
     ),
     ActionCommand(

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

@@ -283,6 +283,7 @@ def cleanup_tables(args):
         verbose=args.verbose,
         confirm=not args.yes,
         skip_archive=args.skip_archive,
+        batch_size=args.batch_size,
     )
 
 

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

@@ -169,61 +169,80 @@ def _dump_table_to_file(*, target_table: str, file_path: str, export_format: str
         raise AirflowException(f"Export format {export_format} is not supported.")
 
 
-def _do_delete(*, query: Query, orm_model: Base, skip_archive: bool, session: Session) -> None:
+def _do_delete(
+    *, query: Query, orm_model: Base, skip_archive: bool, session: Session, batch_size: int | None
+) -> None:
+    import itertools
     import re
 
-    print("Performing Delete...")
-    # using bulk delete
-    # create a new table and copy the rows there
-    timestamp_str = re.sub(r"[^\d]", "", timezone.utcnow().isoformat())[:14]
-    target_table_name = f"{ARCHIVE_TABLE_PREFIX}{orm_model.name}__{timestamp_str}"
-    print(f"Moving data to table {target_table_name}")
     bind = session.get_bind()
     dialect_name = bind.dialect.name
-    target_table = None
-    try:
-        if dialect_name == "mysql":
-            # MySQL with replication needs this split into two queries, so just do it for all MySQL
-            # ERROR 1786 (HY000): Statement violates GTID consistency: CREATE TABLE ... SELECT.
-            session.execute(text(f"CREATE TABLE {target_table_name} LIKE {orm_model.name}"))
-            metadata = reflect_tables([target_table_name], session)
-            target_table = metadata.tables[target_table_name]
-            insert_stm = target_table.insert().from_select(target_table.c, query)
-            logger.debug("insert statement:\n%s", insert_stm.compile())
-            session.execute(insert_stm)
+    batch_counter = itertools.count(1)
+
+    while True:
+        limited_query = query.limit(batch_size) if batch_size else query
+        if limited_query.count() == 0:  # nothing left to delete
+            break
+
+        batch_no = next(batch_counter)
+        suffix = f"__b{batch_no}" if batch_size else ""
+
+        if batch_size:
+            print(f"Performing Delete (batch {batch_no}, max {batch_size} rows)...")
         else:
-            stmt = CreateTableAs(target_table_name, query.selectable)
-            logger.debug("ctas query:\n%s", stmt.compile())
-            session.execute(stmt)
-        session.commit()
-
-        # delete the rows from the old table
-        metadata = reflect_tables([orm_model.name, target_table_name], session)
-        source_table = metadata.tables[orm_model.name]
-        target_table = metadata.tables[target_table_name]
-        logger.debug("rows moved; purging from %s", source_table.name)
-        if dialect_name == "sqlite":
-            pk_cols = source_table.primary_key.columns
-            delete = source_table.delete().where(
-                tuple_(*pk_cols).in_(
-                    select(*[target_table.c[x.name] for x in source_table.primary_key.columns])
+            print("Performing Delete...")
+
+        # using bulk delete
+        # create a new table and copy the rows there
+        timestamp_str = re.sub(r"[^\d]", "", timezone.utcnow().isoformat())[:14]
+        target_table_name = f"{ARCHIVE_TABLE_PREFIX}{orm_model.name}__{timestamp_str}{suffix}"
+        print(f"Moving data to table {target_table_name}")
+        target_table = None
+
+        try:
+            if dialect_name == "mysql":
+                # MySQL with replication needs this split into two queries, so just do it for all MySQL
+                # ERROR 1786 (HY000): Statement violates GTID consistency: CREATE TABLE ... SELECT.
+                session.execute(text(f"CREATE TABLE {target_table_name} LIKE {orm_model.name}"))
+                metadata = reflect_tables([target_table_name], session)
+                target_table = metadata.tables[target_table_name]
+                insert_stm = target_table.insert().from_select(target_table.c, limited_query)
+                logger.debug("insert statement:\n%s", insert_stm.compile())
+                session.execute(insert_stm)
+            else:
+                stmt = CreateTableAs(target_table_name, limited_query.selectable)
+                logger.debug("ctas query:\n%s", stmt.compile())
+                session.execute(stmt)
+            session.commit()
+
+            # delete the rows from the old table
+            metadata = reflect_tables([orm_model.name, target_table_name], session)
+            source_table = metadata.tables[orm_model.name]
+            target_table = metadata.tables[target_table_name]
+            logger.debug("rows moved; purging from %s", source_table.name)
+            if dialect_name == "sqlite":
+                pk_cols = source_table.primary_key.columns
+                delete = source_table.delete().where(
+                    tuple_(*pk_cols).in_(
+                        select(*[target_table.c[x.name] for x in source_table.primary_key.columns])
+                    )
                 )
-            )
-        else:
-            delete = source_table.delete().where(
-                and_(col == target_table.c[col.name] for col in source_table.primary_key.columns)
-            )
-        logger.debug("delete statement:\n%s", delete.compile())
-        session.execute(delete)
-        session.commit()
-    except BaseException as e:
-        raise e
-    finally:
-        if target_table is not None and skip_archive:
-            bind = session.get_bind()
-            target_table.drop(bind=bind)
+            else:
+                delete = source_table.delete().where(
+                    and_(col == target_table.c[col.name] for col in source_table.primary_key.columns)
+                )
+            logger.debug("delete statement:\n%s", delete.compile())
+            session.execute(delete)
             session.commit()
-            print("Finished Performing Delete")
+        except BaseException as e:
+            raise e
+        finally:
+            if target_table is not None and skip_archive:
+                bind = session.get_bind()
+                target_table.drop(bind=bind)
+                session.commit()
+
+    print("Finished Performing Delete")
 
 
 def _subquery_keep_last(
@@ -306,6 +325,7 @@ def _cleanup_table(
     verbose: bool = False,
     skip_archive: bool = False,
     session: Session,
+    batch_size: int | None = None,
     **kwargs,
 ) -> None:
     print()
@@ -325,7 +345,13 @@ def _cleanup_table(
     num_rows = _check_for_rows(query=query, print_rows=False)
 
     if num_rows and not dry_run:
-        _do_delete(query=query, orm_model=orm_model, skip_archive=skip_archive, session=session)
+        _do_delete(
+            query=query,
+            orm_model=orm_model,
+            skip_archive=skip_archive,
+            session=session,
+            batch_size=batch_size,
+        )
 
     session.commit()
 
@@ -432,6 +458,7 @@ def run_cleanup(
     confirm: bool = True,
     skip_archive: bool = False,
     session: Session = NEW_SESSION,
+    batch_size: int | None = None,
 ) -> None:
     """
     Purges old records in airflow metadata database.
@@ -474,6 +501,7 @@ def run_cleanup(
                     **table_config.__dict__,
                     skip_archive=skip_archive,
                     session=session,
+                    batch_size=batch_size,
                 )
                 session.commit()
         else:

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

@@ -376,6 +376,7 @@ def test_date_timezone_omitted(self, run_cleanup_mock, timezone):
             verbose=False,
             confirm=False,
             skip_archive=False,
+            batch_size=None,
         )
 
     @pytest.mark.parametrize("timezone", ["UTC", "Europe/Berlin", "America/Los_Angeles"])
@@ -396,6 +397,7 @@ def test_date_timezone_supplied(self, run_cleanup_mock, timezone):
             verbose=False,
             confirm=False,
             skip_archive=False,
+            batch_size=None,
         )
 
     @pytest.mark.parametrize("confirm_arg, expected", [(["-y"], False), ([], True)])
@@ -422,6 +424,7 @@ def test_confirm(self, run_cleanup_mock, confirm_arg, expected):
             verbose=False,
             confirm=expected,
             skip_archive=False,
+            batch_size=None,
         )
 
     @pytest.mark.parametrize("extra_arg, expected", [(["--skip-archive"], True), ([], False)])
@@ -448,6 +451,7 @@ def test_skip_archive(self, run_cleanup_mock, extra_arg, expected):
             verbose=False,
             confirm=True,
             skip_archive=expected,
+            batch_size=None,
         )
 
     @pytest.mark.parametrize("dry_run_arg, expected", [(["--dry-run"], True), ([], False)])
@@ -474,6 +478,7 @@ def test_dry_run(self, run_cleanup_mock, dry_run_arg, expected):
             verbose=False,
             confirm=True,
             skip_archive=False,
+            batch_size=None,
         )
 
     @pytest.mark.parametrize(
@@ -502,6 +507,7 @@ def test_tables(self, run_cleanup_mock, extra_args, expected):
             verbose=False,
             confirm=True,
             skip_archive=False,
+            batch_size=None,
         )
 
     @pytest.mark.parametrize("extra_args, expected", [(["--verbose"], True), ([], False)])
@@ -528,6 +534,34 @@ def test_verbose(self, run_cleanup_mock, extra_args, expected):
             verbose=expected,
             confirm=True,
             skip_archive=False,
+            batch_size=None,
+        )
+
+    @pytest.mark.parametrize("extra_args, expected", [(["--batch-size", "1234"], 1234), ([], None)])
+    @patch("airflow.cli.commands.db_command.run_cleanup")
+    def test_batch_size(self, run_cleanup_mock, extra_args, expected):
+        """
+        batch_size should be forwarded to run_cleanup with correct type.
+        """
+        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,
+            clean_before_timestamp=pendulum.parse("2021-01-01 00:00:00Z"),
+            verbose=False,
+            confirm=True,
+            skip_archive=False,
+            batch_size=expected,
         )
 
     @patch("airflow.cli.commands.db_command.export_archived_records")

airflow-core/tests/unit/utils/test_db_cleanup.py

@@ -123,6 +123,20 @@ def test_run_cleanup_skip_archive(self, cleanup_table_mock, kwargs, should_skip)
         )
         assert cleanup_table_mock.call_args.kwargs["skip_archive"] is should_skip
 
+    @patch("airflow.utils.db_cleanup._cleanup_table")
+    def test_run_cleanup_batch_size_propagation(self, cleanup_table_mock):
+        """Ensure batch_size is forwarded from run_cleanup to _cleanup_table."""
+        run_cleanup(
+            clean_before_timestamp=None,
+            table_names=["log"],
+            dry_run=None,
+            verbose=None,
+            confirm=False,
+            batch_size=1234,
+        )
+        cleanup_table_mock.assert_called_once()
+        assert cleanup_table_mock.call_args.kwargs["batch_size"] == 1234
+
     @pytest.mark.parametrize(
         "table_names",
         [