sowai-fiscal-svc/tests/shared/test_for_update_populate_existing.py

"""Ported verbatim (mechanism unchanged, only the scanned directory) from
`auto/backend/tests/shared/test_for_update_populate_existing.py` -- the
static AST guard for the exact bug `documents.service.allocate_fiscal_number`
exists to avoid: `SELECT ... FOR UPDATE` locks the row in Postgres, but if
the calling session already has that row's object in its identity map (an
earlier unlocked fetch of the same object, same session), SQLAlchemy's
default behavior hands back the CACHED object instead of repopulating it
from the freshly-(re)locked row. The Postgres lock is real; the in-memory
object the code decides against is a stale snapshot.
`.execution_options(populate_existing=True)` is the fix -- see
`fiscal_svc.documents.service.allocate_fiscal_number`'s docstring for the
full write-up (ported from the auto's own `tenants.service.
allocate_fiscal_number`).

Scans every `.py` file under `src/fiscal_svc/` (never `tests/`) with the
`ast` module -- deliberately not regex over the source text, because a
fluent call chain like
`select(...).where(...).with_for_update().execution_options(...)` routinely
breaks across several lines/parens, and a text-based check would either
miss those or need to reimplement a parser badly. Requires that every
`.with_for_update()` call has a `.execution_options(populate_existing=
True)` call somewhere in the SAME chain of method calls."""
from __future__ import annotations

import ast
from pathlib import Path

SRC_DIR = Path(__file__).resolve().parents[2] / "src" / "fiscal_svc"
REPO_DIR = SRC_DIR.parents[1]

_FIX_EXPLANATION = """
WHY: `.with_for_update()` alone locks the row in Postgres, but if this
session already has an object for that row in its identity map (e.g. an
earlier unlocked lookup of the same row, same session, done for a 404/
access check before calling into the locking code), SQLAlchemy's default
identity-map behavior returns that CACHED object instead of repopulating it
from the row `FOR UPDATE` just (re)read. The lock becomes decorative: the
code holds a real Postgres lock on a row it never actually re-reads, and
makes its decision against a stale in-memory snapshot instead.

HOW TO FIX: chain `.execution_options(populate_existing=True)` onto the
SAME query as the `.with_for_update()`, e.g.:

    query = query.with_for_update().execution_options(populate_existing=True)

or, split across lines/parens, as long as it's the same chain:

    result = await session.execute(
        select(Model)
        .where(Model.id == model_id)
        .with_for_update()
        .execution_options(populate_existing=True)
    )

If you have found a `FOR UPDATE` that must NOT populate_existing (none exist
today -- this would be unusual), do not just delete this check or special-
case your file/line here. Talk to the team about adding an explicit,
commented opt-out marker to this test first, so the next reader still gets
an explanation instead of a silently-shrinking guard.
""".strip()


def _is_populate_existing_true_kwarg(call: ast.Call) -> bool:
    """True if `call` is `.execution_options(...)` carrying a literal
    `populate_existing=True` keyword argument."""
    if not (isinstance(call.func, ast.Attribute) and call.func.attr == "execution_options"):
        return False
    return any(
        kw.arg == "populate_existing" and isinstance(kw.value, ast.Constant) and kw.value.value is True
        for kw in call.keywords
    )


def _build_parent_map(tree: ast.AST) -> dict[int, ast.AST]:
    parents: dict[int, ast.AST] = {}
    for parent in ast.walk(tree):
        for child in ast.iter_child_nodes(parent):
            parents[id(child)] = parent
    return parents


def _chain_root(node: ast.AST, parents: dict[int, ast.AST]) -> ast.AST:
    """Walk up from `node` while still inside the same fluent method-call
    chain -- i.e. while each successive parent is itself a `Call` or
    `Attribute` node, the two node shapes chaining (`a.b().c().d()`) is
    built from in the AST. This is what lets a chain split across many
    lines/parens (still a single expression to the parser) be treated as
    one unit, and stops at the first non-chain boundary (assignment,
    statement, argument to an unrelated call, ...).
    """
    current = node
    while True:
        parent = parents.get(id(current))
        if parent is None or not isinstance(parent, (ast.Call, ast.Attribute)):
            return current
        current = parent


def _find_violations(src_dir: Path) -> list[str]:
    """Returns one `path:line` string per `.with_for_update()` call site
    under `src_dir` that does not have a `.execution_options(
    populate_existing=True)` call in the same method-call chain."""
    violations: list[str] = []
    for path in sorted(src_dir.rglob("*.py")):
        source = path.read_text()
        try:
            tree = ast.parse(source, filename=str(path))
        except SyntaxError:
            continue
        parents = _build_parent_map(tree)
        for node in ast.walk(tree):
            if not (
                isinstance(node, ast.Call)
                and isinstance(node.func, ast.Attribute)
                and node.func.attr == "with_for_update"
            ):
                continue
            root = _chain_root(node, parents)
            has_populate_existing = any(
                isinstance(candidate, ast.Call) and _is_populate_existing_true_kwarg(candidate)
                for candidate in ast.walk(root)
            )
            if not has_populate_existing:
                # `path.relative_to(REPO_DIR)` when `src_dir` is actually
                # inside this repo (the real guard below); falls back to the
                # absolute path when it isn't (the synthetic-`tmp_path`
                # self-tests further down, which scan a throwaway directory
                # outside the repo entirely).
                try:
                    rel = path.relative_to(REPO_DIR)
                except ValueError:
                    rel = path
                # `end_lineno`, not `lineno`: for a multi-line chain, a
                # Call/Attribute node's `lineno` is inherited from where the
                # WHOLE expression starts (e.g. the `select(...)` at the top
                # of the chain), while `end_lineno` lands on the line the
                # `.with_for_update()` text itself is on -- the line the
                # next developer actually needs to look at.
                violations.append(f"{rel}:{node.end_lineno}")
    return violations


def test_every_with_for_update_call_site_chains_populate_existing():
    violations = _find_violations(SRC_DIR)
    assert not violations, (
        "Found `.with_for_update()` call site(s) missing "
        "`.execution_options(populate_existing=True)` in the same call "
        "chain:\n  " + "\n  ".join(violations) + "\n\n" + _FIX_EXPLANATION
    )


# --- Detection proven in BOTH directions (Global Constraints) --------------
#
# The test above only proves the guard is currently GREEN against this
# repo's real code -- on its own that's equally consistent with "the guard
# actually detects the bug" and "the guard is a no-op that always passes".
# These two exercise `_find_violations` directly against synthetic files in
# a throwaway directory, so the guard's OWN detection logic is proven both
# ways: it flags the missing fix, and it does not false-positive on the
# fixed shape (including a chain split across lines, the exact shape
# `allocate_fiscal_number`'s real call site uses).


def test_find_violations_flags_with_for_update_missing_populate_existing(tmp_path):
    (tmp_path / "offender.py").write_text(
        "async def f(session):\n"
        "    return await session.execute(\n"
        "        select(Model).where(Model.id == x).with_for_update()\n"
        "    )\n"
    )

    violations = _find_violations(tmp_path)

    assert len(violations) == 1
    assert "offender.py" in violations[0]


def test_find_violations_does_not_flag_a_correctly_fixed_multiline_chain(tmp_path):
    (tmp_path / "fixed.py").write_text(
        "async def f(session):\n"
        "    return await session.execute(\n"
        "        select(Model)\n"
        "        .where(Model.id == x)\n"
        "        .with_for_update()\n"
        "        .execution_options(populate_existing=True)\n"
        "    )\n"
    )

    assert _find_violations(tmp_path) == []