Make db clean --skip-archive delete rows directly

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

@@ -268,7 +268,7 @@ are rejected. The timestamp cutoff still applies, and any Dag filters further re
 
     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``.
+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 supply ``--skip-archive``, the command deletes rows directly from the source tables without creating archive tables.
 
 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

@@ -550,7 +550,7 @@ def string_lower_type(val):
 )
 ARG_DB_SKIP_ARCHIVE = Arg(
     ("--skip-archive",),
-    help="Don't preserve purged records in an archive table.",
+    help="Delete purged records directly without creating an archive table.",
     action="store_true",
 )
 ARG_DB_EXPORT_FORMAT = Arg(

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

@@ -491,6 +491,7 @@ def _do_delete(
     dialect_name = bind.dialect.name
     batch_counter = itertools.count(1)
     source_table_name = _format_table_name(orm_model.schema, orm_model.name)
+    direct_source_table: Table | None = None
     attempt_archives: dict[str, Table] = {}
     attempt_archive_suffix = uuid4().hex[:16]
 
@@ -507,6 +508,26 @@ def _do_delete(
         else:
             print("Performing Delete...")
 
+        if skip_archive:
+            try:
+                if direct_source_table is None:
+                    metadata = reflect_tables([source_table_name], session)
+                    direct_source_table = metadata.tables[source_table_name]
+                _delete_directly(
+                    limited_query=limited_query,
+                    source_table=direct_source_table,
+                    session=session,
+                    dialect_name=dialect_name,
+                )
+                session.commit()
+            except BaseException:
+                with suppress(Exception):
+                    session.rollback()
+                raise
+            # The direct DELETE embeds limited_query, so skip_if_referenced is re-evaluated in
+            # that statement. The short-delete warning below is only needed after an archive copy.
+            continue
+
         # using bulk delete
         # create a new table and copy the rows there
         timestamp_str = re.sub(r"[^\d]", "", timezone.utcnow().isoformat())[:14]
@@ -515,11 +536,6 @@ def _do_delete(
             f"{ARCHIVE_TABLE_PREFIX}{orm_model.name}__{timestamp_str}{suffix}",
         )
         print(f"Moving data to table {target_table_name}")
-        target_table = None
-        # Lets the ``finally`` cleanup below tell the failure path (don't let a
-        # cleanup error mask the original) from the success path (a cleanup error
-        # is a real problem and must propagate).
-        error_raised = False
 
         try:
             if dialect_name == "mysql":
@@ -605,44 +621,49 @@ def _do_delete(
                     )
 
         except BaseException:
-            error_raised = True
-            # Roll back the failed transaction so its locks are released before
-            # the archive table is dropped in the ``finally`` block below.
-            # ``rollback()`` itself can raise (e.g. the connection died); suppress
-            # it so it does not shadow the original error being re-raised.
             with suppress(Exception):
                 session.rollback()
             raise
-        finally:
-            if target_table is not None and skip_archive:
-                # Drop the archive table on the session's own connection. Binding
-                # the drop to ``session.get_bind()`` (the Engine) would check out a
-                # *second* pooled connection, and on MySQL its ``DROP TABLE`` blocks
-                # indefinitely on the metadata lock still held by this session's
-                # open transaction when the DELETE above failed -- the ``db clean``
-                # hang reported in #66177.
-                try:
-                    target_table.drop(bind=session.connection())
-                    session.commit()
-                except Exception:
-                    # If we are already unwinding from a delete failure, a cleanup
-                    # error here must not replace the original exception (Python
-                    # makes a ``finally``-raised error the top-level one). Log and
-                    # let the original delete error keep propagating. On the success
-                    # path (no delete error), a drop/commit failure is a real
-                    # problem, so re-raise it.
-                    if not error_raised:
-                        raise
-                    logger.warning(
-                        "Failed to drop archive table %s while cleaning up after a "
-                        "delete failure; propagating the original delete error instead.",
-                        target_table_name,
-                        exc_info=True,
-                    )
 
     print("Finished Performing Delete")
 
 
+def _delete_directly(
+    *,
+    limited_query: Select,
+    source_table: Table,
+    session: Session,
+    dialect_name: str,
+) -> None:
+    pk_cols = list(source_table.primary_key.columns)
+    pk_query = limited_query.with_only_columns(
+        *[literal_column(f"{_BASE_TABLE_ALIAS}.{col.name}").label(col.name) for col in pk_cols]
+    )
+    rows_to_delete = pk_query.subquery("rows_to_delete")
+
+    if dialect_name == "mysql":
+        delete = source_table.delete().where(
+            select(literal(1))
+            .select_from(rows_to_delete)
+            .where(and_(*[source_table.c[col.name] == rows_to_delete.c[col.name] for col in pk_cols]))
+            .exists()
+        )
+    elif len(pk_cols) == 1:
+        pk_col = pk_cols[0]
+        delete = source_table.delete().where(
+            source_table.c[pk_col.name].in_(select(rows_to_delete.c[pk_col.name]))
+        )
+    else:
+        delete = source_table.delete().where(
+            tuple_(*[source_table.c[col.name] for col in pk_cols]).in_(
+                select(*[rows_to_delete.c[col.name] for col in pk_cols])
+            )
+        )
+
+    logger.debug("direct delete statement:\n%s", delete.compile())
+    session.execute(delete)
+
+
 def _subquery_keep_last(
     *,
     recency_column,