[Tech Debt] Improve Unit Test DB Isolation

[ferruzzi]

Body

This is a direct follow-up to #68374 with all the same reasoning: We have a lot of unit tests which start by cleaning the DB in the setup phase. This is a clear indication that something is leaking artifacts into the database causing other tests to fail. All unit test fixtures which create a row should own that row and clean it up. A blanket bd cleanup in the unit test's teardown should be an exception when there is no fixture, not the norm. No unit test should require pre-cleaning in setup in order to avoid another test's garbage. As before, I'll get to this eventually, but I'm putting up an Issue in case someone wants to get to it before I do.

After @seanghaeli pointed out other spots we didn't clean up, I set Claude on the task of finding any remaining unit tests that match the pre-cleaning pattern and it found A LOT. The numbers below are unverified Claude output, but should serve as a solid launching point.

---

  Tables cleared in setup position, by frequency:
  
   91  clear_db_runs               7  clear_db_pools
   53  clear_db_dags               7  clear_db_logs
   26  clear_db_assets             7  clear_db_import_errors
   18  clear_db_connections        6  clear_db_deadline
   17  clear_db_serialized_dags    5  clear_db_backfills
   14  clear_db_jobs               4  clear_db_apdr / pakl / callbacks / triggers
    9  clear_db_variables          3  clear_db_revoked_tokens / xcom
    9  clear_db_dag_bundles        2  clear_db_deadline_alert
    8  clear_db_connection_tests   1  clear_db_dag_warnings / dag_code / task_reschedule
    8  clear_db_teams

---

Notably, dag_maker.cleanup() has the same gap for assets that it had for bundles before #69093: it deletes AssetEvent but not the AssetModel / AssetActive / AssetAliasModel rows a Dag's schedule creates. Fixing it the same way (track what was created, delete when unreferenced, FK-safe order) should account for the ~26 clear_db_assets pre-cleans on its own.

Proposal

I'm using that term very loosely, these are just Thoughts. Cleaning up the codebase is not very useful without enforcing the change, so some ideas in addition to the actual clean up:

1. A new enforcer: I don't know if this is a pre-commit or a ruff check or what this actually looks like, but we should be preventing tests from calling db_clean in the setup. If a new test needs that to pass, something is already wrong and we're just masking an existing issue.

2. A new "leak detector": This might be in the form of a fixture which snapshots the pre- and post-test table row counts and reports if there is any variation? Maybe include it in the existing "this test needs a db" fixture?

3. Several tests create Job rows by hand and nothing cleans them up; we might want a job_maker fixture akin to dag_maker which tracks those creations and cleans up on teardown? Similarly, it doesn't look like anything owns tracking and cleaning up Connections, Variables, or Pools. maybe a similar pattern can be employed there?

We can either add the enforcer ahead of time with an allowlist of existing cases which we clean as we go, or we need to clean before that lands. I kind of like the former as it doubles as a tracker for remaining work and closes the floodgate so we're not stuck making another pass later to catch the issues that were merged while we were working.

Triage Rules:

How do you tell a deletable site from a load-bearing one? I see three cases:

• Already covered by dag_maker (it deletes DagRun, TaskInstance, DagVersion, XCom, DagModel, TaskMap, plus its own bundles, and SerializedDagModel / DeadlineAlert cascade off dag_version.id and serialized_dag.id). Delete the pre-clean, unless the test bypasses dag_maker via DagBag, example Dags, or a hand-built DagModel.
• Genuine fixture gap (assets today). Fix the fixture first, then delete.
• Nothing owns the table (Job, Log, Pool, Variable, Connection, ConnectionTest, RevokedToken, Backfill). Load-bearing until an owning fixture exists; do not delete yet.

---

Once we have the framework in place to prevent regression, here is what Claude reports for the current violations that need to be fixed. Included here for a sense of scope only, I have not validated the numbers and naturally this will drift between now and then anyway:

┌───────────────────────────────────────────────────┬───────┬───────┬────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐
│ Batch                                             │ Sites │ Files │ Notes                                                                                                                  │
├───────────────────────────────────────────────────┼───────┼───────┼────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ api_fastapi/execution_api/versions/               │ 42    │ 12    │ Almost entirely clear_db_runs. Most mechanical batch; good starter.                                                    │
├───────────────────────────────────────────────────┼───────┼───────┼────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ api_fastapi/core_api/routes/ui/                   │ 39    │ 10    │ Includes the apdr / pakl partition tables.                                                                             │
├───────────────────────────────────────────────────┼───────┼───────┼────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ api_fastapi/core_api/routes/public/               │ 34    │ 14    │ Blocked on Pool / Variable / Connection / Log ownership.                                                               │
├───────────────────────────────────────────────────┼───────┼───────┼────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ api_fastapi misc (common/, auth/, conftest.py,    │ 10    │ 5     │                                                                                                                        │
│ core_api/test_app.py)                             │       │       │                                                                                                                        │
├───────────────────────────────────────────────────┼───────┼───────┼────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ models/                                           │ 46    │ 18    │ Blocked on the dag_maker asset fix. Also has one leftover bundle/team site from #68374.                                │
├───────────────────────────────────────────────────┼───────┼───────┼────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ cli/commands/                                     │ 24    │ 9     │ Heavy on connections and variables; blocked on Connection / Variable ownership.                                        │
├───────────────────────────────────────────────────┼───────┼───────┼────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ dag_processing/                                   │ 15    │ 3     │ Includes two leftover bundle/team sites from #68374.                                                                   │
├───────────────────────────────────────────────────┼───────┼───────┼────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ jobs/                                             │ 12    │ 2     │ test_scheduler_job.py + test_triggerer_job.py. Blocked on Job ownership.                                               │
├───────────────────────────────────────────────────┼───────┼───────┼────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ utils/                                            │ 10    │ 4     │                                                                                                                        │
├───────────────────────────────────────────────────┼───────┼───────┼────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ core misc: always/, api/, assets/, core/, state/, │ 19    │ 10    │ One combined batch. Includes one leftover team site in ti_deps/.                                                       │
│ ti_deps/, timetables/                             │       │       │                                                                                                                        │
├───────────────────────────────────────────────────┼───────┼───────┼────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ providers: standard 8, celery 6, snowflake 3,     │ 19    │ 8     │ Version-gated behind AIRFLOW_V_3_0_PLUS in places, so this batch carries provider compat rules the core batches do     │
│ common.io 2                                       │       │       │ not. Also has the last leftover bundle site from #68374.                                                               │
├───────────────────────────────────────────────────┼───────┼───────┼────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┤
│ Total                                             │ 270   │ 95    │                                                                                                                        │
└───────────────────────────────────────────────────┴───────┴───────┴────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘

Committer

• I acknowledge that I am a maintainer/committer of the Apache Airflow project.

[shivaam]

@ferruzzi I am going to start working on this. I think I will break it down into smaller PRs instead of one big pr so that we dont accidentally reduce coverage.

[shivaam]

@ferruzzi I am thinking what if we take a snapshot of all the db tables before we run the test, then after the test, we can count the number of rows again. Then we run a helper function that will reset the tables where the count changed. Just an idea. I will benchmark it

[ferruzzi]

@shivaam Smaller PRs is 100% the right call, I just wanted to get the big picture down "on paper" while I was thinking about it. Your idea sounds reasonable, but benchmark it for sure, I have no idea what that translates to in runtime and memory. I don't think we want to be adding minutes to the CI runtime, but if it's milliseconds, then yeah, that sounds reasonable.

It might also make sense for the enforcer and the detector to share an allowlist as well, which I didn't really consider yesterday. If we just create the detector right now, it's going to go off on all the existing issues. Flip side, if we do it first and the detector has an optional false-by-default flag which outputs its findings to the logs, then that acts as a triage for the tests to fix: "Here are all the tests right now that violate the policy" would help us split the actual issues from the tests that defensively pre-cleaned.

So I guess there are arguments for either order. The detector does sound more fun though, so i don't blame you for taking that one first. :p

[shivaam]

@ferruzzi Loved the idea of a detector only model and sharing the list with enforcer. I benchmarked a detector prototype looking at 6 tables and it added 1ms per test which is acceptable. I think a narrow detector monitoring a handful of tables look practical.

My suggested next steps:

• Keep the enforcer as a separate small PR.
• Add an optional detector that reports leaks but does not fail tests. I have the code locally for this right now.
• Run tests and find out any unit tests that are leaking using the detector as well.
• Use your analysis and the report to fix through small PRs

[ferruzzi]

Agreed on the separate PR for the enforcer for sure. Plus, if we do the detector first, we can use that to generate the starting allowlist for the enforcer, which is nice.

I think narrowing the detector may be a mistake though, these are the ones we know about, doesn't mean they are the only ones. Maybe build it to monitor everything and trim it down if timing is a concern? As you say, 1ms is nothing. devel-common/src/tests_common/test_utils/db.py has 26 clear_db_* helpers hitting almost every table in the database, maybe we can generate a list of what those helpers hit to use as the scope somehow so we don't have to manually update the table list?

Did you benchmark on six individual queries, or as a single batch? I suspect the bulk of that 1ms (even as small as that is) was taken by the data roundtrip and not the actual query, but since we only need row counts, if you batch it all into one query that might be more efficient. I'm not great with DB stuff, but something along the lines of this (untested pseudo-code) might actually just check the entire db just as fast without bothering with a list of tables if we're already planning to hit almost all of them anyway:

counts = []
for table in Base.metadata.sorted_tables:
    counts.append(select(literal(table.name), func.count()).select_from(table))
snapshot = session.execute(union_all(*counts)).all()

And one last thought I have not put any research into: are there any cases where we WANT per-test leaks? Maybe have a look and see if we ave any per-class fixtures that persist across a few tests but then do actually clean up after themselves at the end of the suite.

Another point about this whole approach is that basing it off row counts isn't foolproof. If I add one and remove another, this will clear despite my monkeying... so maybe there's a better way, or maybe row count is "good enough", or maybe we can get this going and expand it with a per-row hash later or something.... As I said, I haven't really thought this through all the way, there are possibly pitfalls still out there to consider...

[shivaam]

And one last thought I have not put any research into: are there any cases where we WANT per-test leaks? Maybe have a look and see if we ave any per-class fixtures that persist across a few tests but then do actually clean up after themselves at the end of the suite.

Agreed that detector-first could help validate the enforcer baseline, but I don’t think the order is blocking. We can update the enforcer allowlist if the detector finds cases the initial scan missed in that PR.

One correction to my earlier explanation, the detector compares primary keys, not just row counts. This catches replacing one row with another even when the total count stays the same.
I expanded it to all 57 loaded core tables and tested again. Across 500 tests, the overhead was about 15 ms per test which is still acceptable.

I also checked class, module, and session fixtures that intentionally share rows across tests. Their rows appear in before snapshot as well, so they are not reported as leaks. A test leaving data for a later test would be reported, which is desirable because tests should not depend on execution order.

We will keep the detector opt-in for now and it can help us triage existing problems and preventing leaks.

[ferruzzi]

I think this sounds like a plan. 👍

[abhijeets25012-tech]

i would like to take it

[shivaam]

i would like to take it

Hello! I have already created a cr and benchmarking orher approach for detrctor which will be out soon as well