Stop Guessing About N+1 Queries: Add a Budget Test

Before I optimize a batch operation, I want to see one small table: input size next to SQL statement count. It tells me whether the database work has a fixed shape or grows one row at a time.

That check exposed a problem in a bulk preview I worked on. A synthetic case with 500 valid rows issued 500 duplicate lookups. It was an N+1 shape inside a batch: the number of cursor.execute calls grew with the input instead of staying constant.

Timing alone would have been a weak guard for this. Test machines vary, database caches warm up, and a threshold that passes today can fail for unrelated reasons tomorrow. A statement count is less ambitious, but it captures the defect directly. The preview should use no duplicate query for an empty or entirely invalid batch and one for any non-empty valid batch.

Fixture	Valid identities	Earlier execute calls	Batched execute calls
all invalid	0	0	0
one valid row	1	1	1
synthetic batch	500	500	1

Count the narrow operation

I count at the database adapter boundary, close to execute(). The wrapper below is intentionally small. It forwards everything to the real cursor while recording calls made by the code under test:

class QueryCounter:
    def __init__(self):
        self.executions = 0


class CountingCursor:
    def __init__(self, cursor, counter):
        self._cursor = cursor
        self._counter = counter

    def execute(self, statement, params=None):
        self._counter.executions += 1
        return self._cursor.execute(statement, params)

    def __getattr__(self, name):
        return getattr(self._cursor, name)


class CountingDatabase:
    def __init__(self, connection, counter):
        self._connection = connection
        self._counter = counter

    def cursor(self):
        return CountingCursor(
            self._connection.cursor(),
            self._counter,
        )

This is generic test instrumentation, not the application's production database class. For async drivers, ORMs, or connection pools, I would put the counter at the equivalent statement-execution hook.

Scope matters. Fixture inserts, transaction setup, and cleanup are real SQL, yet they are not queries made by the preview function. I prepare the fixed database state first, then create or reset the counter immediately before calling the unit I want to measure.

stored_rows = seed_synthetic_rows(real_connection)

counter = QueryCounter()
counted_db = CountingDatabase(real_connection, counter)
preview = build_preview(counted_db, incoming_rows)

assert counter.executions == 1

If the application mixes several kinds of SQL inside build_preview, a single total can be too blunt. In that case I tag or normalize statements and count only the duplicate-lookup signature. The test should name the boundary it protects rather than silently including whatever setup happens to run nearby.

Make the budget depend on valid work

A hard-coded == 1 assertion misses one useful case. When every incoming row is invalid, the duplicate checker should return its default answers without contacting PostgreSQL. My budget is therefore based on whether the parsed batch contains at least one valid identity:

@pytest.mark.parametrize(
    "incoming, expected, expected_calls",
    [
        (all_invalid_rows(), default_outputs(), 0),
        (one_valid_row(), one_row_output(), 1),
        (synthetic_preview_case(rows=500), expected_500_outputs(), 1),
    ],
)
def test_duplicate_preview_query_budget(
    real_connection, incoming, expected, expected_calls
):
    seed_expected_matches(real_connection, incoming)

    counter = QueryCounter()
    db = CountingDatabase(real_connection, counter)
    actual = preview_duplicates(db, user_id=7, rows=incoming)

    assert actual == expected
    assert counter.executions == expected_calls

The names in this example are illustrative. The important pairing is the two assertions. The output assertion prevents a developer from satisfying the budget by skipping the lookup. The count assertion prevents a correct-looking implementation from quietly putting the query back inside the loop.

For mixed valid and invalid input, I calculate the budget from the validated identities passed to duplicate checking, not from the raw file length. That keeps parsing failures from changing what the database test is supposed to prove.

Let the old code fail first

I prefer to run the budget test against the existing implementation before changing it. In this case, the 500-row test reported 500 executions. That established three things at once: the instrumentation was attached to the intended path, the fixture actually reached the lookup, and the proposed budget would catch the current behavior.

After batching, the same fixed database state produced the same ordered duplicate answers and existing IDs with one cursor.execute call. This is evidence about calls made by the duplicate-check function, not whole-page latency.

Varying the batch size can make the diagnosis easier to read during development:

for size in (1, 10, 100, 500):
    count = count_preview_queries(make_valid_rows(size))
    print(size, count)

The old shape prints a count that tracks size. The budget test does not need to keep this exploratory loop once the cause is known. One representative non-empty case plus the all-invalid case can protect the contract with less test time.

Keep correctness beside the count

Query budgets are good at detecting multiplicative database work. They do not prove that a batched query uses the right composite identity, preserves repeated inputs, or returns results in the original order. Those need ordinary behavior assertions and, when PostgreSQL-specific features are involved, a real-database test.

I keep the counter test focused on the number of statements and compare the full returned structure in the same fixed state. Separate cases cover a stored match, a non-match, repeated input, and an invalid row that retains its output position. This is enough to make a cheap shortcut fail visibly without turning the count test into a second implementation.

There is also a concurrency difference worth noting. Under PostgreSQL READ COMMITTED, several per-row statements can observe commits between rows; one batched statement gets one statement snapshot. The fixed-state comparison does not cover concurrent writes.

The actual PostgreSQL array query and positional mapping are covered in my implementation case study. The diagnostic habit is simpler: count first, make the current loop fail a budget test, then keep that test after the repair.

I keep this query budget beside the behavior assertions, so a future per-row lookup fails with its observed execution count.