magnus919_agent-skills/backend-engineering/scripts/n1-query-spotter.py

#!/usr/bin/env python3
"""N+1 query spotter for backend-engineering.

Scans Python source files for the classic N+1 query pattern: a database query
or ORM fetch invoked inside a loop body. When a loop runs N iterations and
each iteration issues its own query, the code makes N+1 round trips instead
of one batched query — the fix is a WHERE IN batch, a join, or eager loading.

Detection is static and heuristic: any query-like call that appears inside a
for or while loop is flagged. A call that also references the loop variable
in one of its arguments is flagged with higher confidence, because the query
is almost certainly varying per iteration.

Input: one or more file paths. With no paths, Python source is read from
stdin. Output is one line per finding, or a JSON report with --json.

Exit codes:
  0  no potential N+1 patterns found
  1  one or more potential N+1 patterns found
  2  usage or I/O error (missing file, unparseable source, bad flags)
"""

import argparse
import ast
import json
import sys
from pathlib import Path

# Attribute method names (obj.NAME(...)) treated as database query / ORM fetch calls.
QUERY_METHODS = frozenset(
    {
        "query",
        "execute",
        "fetchall",
        "fetchone",
        "fetchmany",
        "fetch",
        "first",
        "one",
        "all",
        "get",
        "filter",
        "select",
        "find",
        "save",
        "create",
        "update",
        "delete",
        "insert",
        "commit",
        "persist",
        "bulk_create",
        "bulk_update",
    }
)

# Bare function names (NAME(...)) treated as query entry points.
QUERY_NAMES = frozenset({"query", "execute", "run", "fetch", "find", "select"})

_STDIN_LABEL = "<stdin>"


class N1Finding:
    """One suspected N+1 pattern: a query-like call inside a loop."""

    def __init__(self, line, column, call_text, loop_line, loop_targets, high_confidence):
        self.line = line
        self.column = column
        self.call_text = call_text
        self.loop_line = loop_line
        self.loop_targets = sorted(loop_targets)
        self.high_confidence = high_confidence

    def to_dict(self):
        return {
            "line": self.line,
            "column": self.column,
            "call": self.call_text,
            "loop_line": self.loop_line,
            "loop_targets": self.loop_targets,
            "confidence": "high" if self.high_confidence else "possible",
        }

    def render(self, source_path):
        confidence = "high confidence" if self.high_confidence else "possible"
        target_note = f", loop target {', '.join(self.loop_targets)}" if self.loop_targets else ""
        return (
            f"{source_path}:{self.line}:{self.column}: potential N+1: "
            f"{self.call_text!r} inside loop at line {self.loop_line}{target_note} ({confidence})"
        )


class N1Scanner(ast.NodeVisitor):
    """Walks one module, flagging query-like calls that sit inside a loop."""

    def __init__(self, extra_methods=(), extra_names=()):
        self.methods = QUERY_METHODS | set(extra_methods)
        self.names = QUERY_NAMES | set(extra_names)
        self.loop_stack = []  # (loop_node, loop_target_names)
        self.findings = []

    # -- loop tracking -----------------------------------------------------

    def _target_names(self, target):
        if isinstance(target, ast.Name):
            return {target.id}
        if isinstance(target, ast.Tuple):
            return {elt.id for elt in target.elts if isinstance(elt, ast.Name)}
        return set()

    def visit_For(self, node):
        self.loop_stack.append((node, self._target_names(node.target)))
        self.generic_visit(node)
        self.loop_stack.pop()

    def visit_While(self, node):
        self.loop_stack.append((node, set()))
        self.generic_visit(node)
        self.loop_stack.pop()

    # -- call inspection ---------------------------------------------------

    def _call_name(self, node):
        """Return (name, kind) for a call's callee, or None if not matched."""
        func = node.func
        if isinstance(func, ast.Attribute):
            return func.attr, "attribute"
        if isinstance(func, ast.Name):
            return func.id, "name"
        return None

    @staticmethod
    def _names_in_node(node):
        """All identifier names reachable inside an AST node."""
        return {name.id for name in ast.walk(node) if isinstance(name, ast.Name)}

    @classmethod
    def _references_loop_var(cls, node, names):
        """True if any argument of the call references a loop variable.

        Attribute access (user.id) and nested expressions count as references,
        since a per-iteration query keyed off the loop item is the N+1 signature.
        """
        for arg in node.args:
            if cls._names_in_node(arg) & names:
                return True
        for kw in node.keywords:
            if kw.arg in names:
                return True
            if kw.value is not None and cls._names_in_node(kw.value) & names:
                return True
        return False

    def visit_Call(self, node):
        if self.loop_stack:
            matched = self._call_name(node)
            if matched:
                call_name, kind = matched
                if (kind == "attribute" and call_name in self.methods) or (
                    kind == "name" and call_name in self.names
                ):
                    loop_node, targets = self.loop_stack[-1]
                    high = bool(targets) and self._references_loop_var(node, targets)
                    self.findings.append(
                        N1Finding(
                            line=node.lineno,
                            column=getattr(node, "col_offset", 0),
                            call_text=ast.unparse(node),
                            loop_line=loop_node.lineno,
                            loop_targets=targets,
                            high_confidence=high,
                        )
                    )
        self.generic_visit(node)


def scan_source(source, source_path, extra_methods=(), extra_names=()):
    """Parse source and return the list of N1Finding objects."""
    tree = ast.parse(source, filename=source_path)
    scanner = N1Scanner(extra_methods=extra_methods, extra_names=extra_names)
    scanner.visit(tree)
    return scanner.findings


def _split_csv(value):
    return [item.strip() for item in value.split(",") if item.strip()]


def build_parser():
    parser = argparse.ArgumentParser(
        prog="n1-query-spotter.py",
        description=(
            "Spot potential N+1 query patterns (query-like calls inside loops) in "
            "Python source files. With no FILE arguments, reads source from stdin."
        ),
        epilog="Exit codes: 0 no findings, 1 findings, 2 usage or I/O error.",
    )
    parser.add_argument("files", nargs="*", metavar="FILE", help="Python source files to scan")
    parser.add_argument(
        "--json",
        action="store_true",
        help="emit a machine-readable JSON report instead of human lines",
    )
    parser.add_argument(
        "--extra-methods",
        default="",
        metavar="NAME[,NAME...]",
        help="additional query-like attribute method names to detect, e.g. 'run_query,raw'",
    )
    parser.add_argument(
        "--extra-names",
        default="",
        metavar="NAME[,NAME...]",
        help="additional query-like bare function names to detect",
    )
    return parser


def _read_input(paths, errors):
    """Yield (label, source) pairs for every input; collect IO errors."""
    if not paths:
        yield _STDIN_LABEL, sys.stdin.read()
        return
    for raw in paths:
        path = Path(raw)
        try:
            yield str(path), path.read_text(encoding="utf-8")
        except OSError as exc:
            errors.append(f"{path}: cannot read: {exc.strerror}")


def main(argv=None):
    parser = build_parser()
    args = parser.parse_args(argv)
    errors = []
    all_findings = []
    for label, source in _read_input(args.files, errors):
        try:
            findings = scan_source(
                source,
                label,
                extra_methods=_split_csv(args.extra_methods),
                extra_names=_split_csv(args.extra_names),
            )
        except SyntaxError as exc:
            errors.append(f"{label}:{exc.lineno}: cannot parse source: {exc.msg}")
            continue
        for finding in findings:
            all_findings.append((label, finding))

    if errors:
        for message in errors:
            print(f"ERROR: {message}", file=sys.stderr)
        return 2

    if args.json:
        report = {
            "findings": [{"file": label, **finding.to_dict()} for label, finding in all_findings],
            "count": len(all_findings),
        }
        print(json.dumps(report, indent=2))
    else:
        for label, finding in all_findings:
            print(finding.render(label))
        if all_findings:
            print(
                f"{len(all_findings)} potential N+1 pattern(s) found; "
                "consider batch queries, joins, or eager loading.",
                file=sys.stderr,
            )
    return 1 if all_findings else 0


if __name__ == "__main__":
    sys.exit(main())