From c12ed91b1579dbb370132d064748f29bb5a156cb Mon Sep 17 00:00:00 2001
From: "safae.hariri" <safae.hariri@datadoghq.com>
Date: Thu, 8 Oct 2026 11:18:11 +0200
Subject: [PATCH] Reduce db clean batch probe overhead

Large cleanups repeat the batch probe many times, and counting every row up to the batch size adds unnecessary database work.
---
 airflow-core/src/airflow/utils/db_cleanup.py  |  4 +---
 .../tests/unit/utils/test_db_cleanup.py       | 20 +++++++++++++++++++
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/airflow-core/src/airflow/utils/db_cleanup.py b/airflow-core/src/airflow/utils/db_cleanup.py
index 92cebbd236e6b..59596fc73522c 100644
--- a/airflow-core/src/airflow/utils/db_cleanup.py
+++ b/airflow-core/src/airflow/utils/db_cleanup.py
@@ -490,9 +490,7 @@ def _do_delete(
 
     while True:
         limited_query = query.limit(batch_size) if batch_size else query
-        if (
-            session.scalars(select(func.count()).select_from(limited_query.subquery())).one() == 0
-        ):  # nothing left to delete
+        if not session.scalars(select(limited_query.exists())).one():  # nothing left to delete
             break
 
         batch_no = next(batch_counter)
diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py b/airflow-core/tests/unit/utils/test_db_cleanup.py
index ddf7dc5ae4bab..00752da4da276 100644
--- a/airflow-core/tests/unit/utils/test_db_cleanup.py
+++ b/airflow-core/tests/unit/utils/test_db_cleanup.py
@@ -903,6 +903,26 @@ def test_table_config_dag_id_scope_requires_fk_column(self):
                 # "dagrun_id" intentionally omitted from extra_columns
             )
 
+    def test_do_delete_checks_batch_existence_without_counting_rows(self):
+        session = MagicMock(spec=Session)
+        session.get_bind.return_value.dialect.name = "sqlite"
+        session.scalars.return_value.one.return_value = False
+        _, source_table, _, query = _build_do_delete_test_objects()
+
+        _do_delete(
+            query=query,
+            orm_model=source_table,
+            skip_archive=False,
+            session=session,
+            batch_size=100,
+        )
+
+        existence_query = session.scalars.call_args.args[0]
+        compiled_query = str(existence_query.compile()).lower()
+        assert "exists" in compiled_query
+        assert "count(" not in compiled_query
+        session.execute.assert_not_called()
+
     def test_do_delete_rolls_back_before_drop_on_failure(self):
         session = MagicMock(spec=Session)
         session.get_bind.return_value.dialect.name = "mysql"