Merge pull request #63 from smochan/feat/lint-db-call-in-loop

feat(lint): db-call-in-loop rule + loop_depth on call records

Author: smochan

codegraph/analysis/lint.py

@@ -318,10 +318,84 @@ def _check_sensitive_literal(
     )
 
 
+_TS_DB_PATTERN = r"\b(db|database)\.(select|insert|update|delete|execute|query)\b"
+_PY_DB_PATTERN = r"\b(session|cursor|db)\.(query|execute|add|commit)\b"
+
+# Higher-order iteration methods treated as implicit loop bodies.
+_TS_ITER_METHODS = frozenset({"map", "forEach", "filter", "reduce", "flatMap"})
+# Function-typed AST node types that can appear as callbacks.
+_FUNC_NODE_TYPES = frozenset({"arrow_function", "function_expression", "function"})
+
+
+def _py_chain_text(node: tree_sitter.Node, src: bytes) -> str | None:
+    """Flatten a Python attribute-access chain to dotted text.
+
+    For a ``call`` node whose ``function`` child is an ``attribute``,
+    walks the attribute chain to produce e.g. ``session.execute``.
+    Returns ``None`` when the chain doesn't bottom out at an identifier.
+    """
+    parts: list[str] = []
+    cur: tree_sitter.Node | None = node.child_by_field_name("function")
+    while cur is not None:
+        if cur.type == "attribute":
+            attr = cur.child_by_field_name("attribute")
+            if attr is not None:
+                parts.append(node_text(attr, src))
+            cur = cur.child_by_field_name("object")
+        elif cur.type in ("identifier",):
+            parts.append(node_text(cur, src))
+            return ".".join(reversed(parts))
+        else:
+            return None
+    return None
+
+
+def _check_db_call_in_loop(
+    node: tree_sitter.Node, ctx: _LintContext
+) -> LintFinding | None:
+    """DB call inside a loop body — potential N+1 query."""
+    if ctx.is_test_file or ctx.loop_depth < 1:
+        return None
+
+    if ctx.language in _TS_LANGS:
+        if node.type != "call_expression" or not _is_chain_root(node):
+            return None
+        chain = _member_chain(node, ctx.src)
+        if chain is None:
+            return None
+        base, methods = chain
+        chain_text = ".".join([base, *methods])
+        pattern = str(ctx.rule.options.get("pattern") or _TS_DB_PATTERN)
+    elif ctx.language == "python":
+        if node.type != "call":
+            return None
+        chain_text = _py_chain_text(node, ctx.src) or ""
+        if not chain_text:
+            return None
+        pattern = str(ctx.rule.options.get("pattern") or _PY_DB_PATTERN)
+    else:
+        return None
+
+    if not re.search(pattern, chain_text):
+        return None
+
+    return LintFinding(
+        rule_id=ctx.rule.id,
+        severity=ctx.rule.severity,
+        message=_format_message(
+            ctx.rule, chain=chain_text, depth=ctx.loop_depth
+        ),
+        file=ctx.rel_path,
+        line=node.start_point[0] + 1,
+        snippet=_snippet(node, ctx.src),
+    )
+
+
 _CHECK_REGISTRY: dict[str, _Checker] = {
     "console-in-prod": _check_console_in_prod,
     "unfiltered-query": _check_unfiltered_query,
     "sensitive-literal": _check_sensitive_literal,
+    "db-call-in-loop": _check_db_call_in_loop,
 }
 
 
@@ -344,6 +418,12 @@ def _check_sensitive_literal(
         severity="med",
         message="Literal assigned to `{name}` looks sensitive ({reason})",
     ),
+    LintRule(
+        id="db-call-in-loop",
+        check="db-call-in-loop",
+        severity="high",
+        message="DB call `{chain}` inside a loop (N+1 risk)",
+    ),
 ]
 
 
@@ -412,6 +492,45 @@ def _is_test_path(rel_path: str, language: str) -> bool:
     return _is_ts_test_file(rel_path)
 
 
+def _iter_children_with_depth(
+    node: tree_sitter.Node, base_depth: int, language: str, src: bytes
+) -> list[tuple[tree_sitter.Node, int]]:
+    """Return ``(child, depth)`` pairs for a node's children.
+
+    For TS/JS ``call_expression`` nodes whose member property is one of the
+    higher-order iteration methods (.map, .forEach, .filter, .reduce,
+    .flatMap), function-typed arguments receive ``base_depth + 1`` so the
+    DFS treats their bodies as implicit loop bodies.  All other children
+    keep ``base_depth``.
+    """
+    if language not in _TS_LANGS or node.type != "call_expression":
+        return [(child, base_depth) for child in reversed(node.children)]
+
+    fn = node.child_by_field_name("function")
+    if fn is None or fn.type != "member_expression":
+        return [(child, base_depth) for child in reversed(node.children)]
+
+    prop = fn.child_by_field_name("property")
+    if prop is None or node_text(prop, src) not in _TS_ITER_METHODS:
+        return [(child, base_depth) for child in reversed(node.children)]
+
+    args_node = node.child_by_field_name("arguments")
+    if args_node is None:
+        return [(child, base_depth) for child in reversed(node.children)]
+
+    result: list[tuple[tree_sitter.Node, int]] = []
+    for child in reversed(node.children):
+        if child == args_node:
+            for arg in reversed(child.children):
+                if arg.type in _FUNC_NODE_TYPES:
+                    result.append((arg, base_depth + 1))
+                else:
+                    result.append((arg, base_depth))
+        else:
+            result.append((child, base_depth))
+    return result
+
+
 def _lint_file(
     path: Path, rel_path: str, language: str, rules: list[LintRule]
 ) -> list[LintFinding]:
@@ -450,8 +569,13 @@ def _lint_file(
         child_depth = (
             loop_depth + 1 if node.type in _LOOP_NODE_TYPES else loop_depth
         )
-        for child in reversed(node.children):
-            stack.append((child, child_depth))
+        # For .map/.forEach/.filter/.reduce/.flatMap calls (TS only), treat
+        # function-typed arguments as implicit loop bodies.
+        children_with_depth = _iter_children_with_depth(
+            node, child_depth, language, src
+        )
+        for child, cd in children_with_depth:
+            stack.append((child, cd))
     return findings
 
 

codegraph/parsers/python.py

@@ -1319,6 +1319,11 @@ def _visit_nested_defs(
                     continue
             stack.extend(node.children)
 
+    # Python loop node types that increase loop depth for their children.
+    _PY_LOOP_TYPES: frozenset[str] = frozenset(
+        {"for_statement", "while_statement"}
+    )
+
     def _collect_calls(
         self,
         node: tree_sitter.Node,
@@ -1328,9 +1333,11 @@ def _collect_calls(
         edges: list[Edge],
     ) -> None:
         """Walk subtree collecting call expressions, stopping at nested defs."""
-        stack: list[tree_sitter.Node] = list(node.children)
+        stack: list[tuple[tree_sitter.Node, int]] = [
+            (child, 0) for child in node.children
+        ]
         while stack:
-            child = stack.pop()
+            child, depth = stack.pop()
             if child.type == "call":
                 func_child = child.child_by_field_name("function")
                 if func_child is None and child.children:
@@ -1352,6 +1359,8 @@ def _collect_calls(
                             "target_name": name,
                             "args": args,
                             "kwargs": kwargs,
+                            "loop_depth": depth,
+                            "in_loop": depth > 0,
                         },
                     ))
             # ``decorator`` subtrees are handled by ``_emit_decorator_calls``
@@ -1361,7 +1370,10 @@ def _collect_calls(
             if child.type not in (
                 "class_definition", "function_definition", "decorator",
             ):
-                stack.extend(child.children)
+                child_depth = (
+                    depth + 1 if child.type in self._PY_LOOP_TYPES else depth
+                )
+                stack.extend((gc, child_depth) for gc in child.children)
 
     def _emit_decorator_calls(
         self,

codegraph/parsers/typescript.py

@@ -1017,6 +1017,20 @@ def _handle_lexical_decl(
                     self._collect_calls(body, rel, func_id, src, edges)
                     self._collect_fetches(body, rel, func_id, src, nodes, edges)
 
+    # Higher-order iteration methods whose function-typed arguments are treated
+    # as implicit loop bodies for loop-depth tracking purposes.
+    _ITER_METHODS: frozenset[str] = frozenset(
+        {"map", "forEach", "filter", "reduce", "flatMap"}
+    )
+    # TS/JS loop node types that increase loop depth for their children.
+    _TS_LOOP_TYPES: frozenset[str] = frozenset(
+        {"for_statement", "for_in_statement", "while_statement", "do_statement"}
+    )
+    # Function-typed AST node types that can appear as callbacks.
+    _FUNC_ARG_TYPES: frozenset[str] = frozenset(
+        {"arrow_function", "function_expression", "function"}
+    )
+
     def _collect_calls(
         self,
         node: tree_sitter.Node,
@@ -1025,9 +1039,11 @@ def _collect_calls(
         src: bytes,
         edges: list[Edge],
     ) -> None:
-        stack: list[tree_sitter.Node] = list(node.children)
+        stack: list[tuple[tree_sitter.Node, int]] = [
+            (child, 0) for child in node.children
+        ]
         while stack:
-            child = stack.pop()
+            child, depth = stack.pop()
             if child.type == "call_expression":
                 func_child = child.child_by_field_name("function")
                 if func_child is None and child.children:
@@ -1056,9 +1072,37 @@ def _collect_calls(
                             "target_name": name,
                             "args": call_args,
                             "kwargs": call_kwargs,
+                            "loop_depth": depth,
+                            "in_loop": depth > 0,
                         },
                     ))
-            stack.extend(child.children)
+                    # Determine whether this call is an iteration method
+                    # (.map/.forEach/.filter/.reduce/.flatMap) and bump depth
+                    # for any function-typed arguments.
+                    iter_method = ""
+                    if func_child is not None and func_child.type == "member_expression":
+                        prop_node = func_child.child_by_field_name("property")
+                        if prop_node is not None:
+                            iter_method = node_text(prop_node, src)
+                    if iter_method in self._ITER_METHODS and args_node is not None:
+                        # The callee subtree may itself contain calls
+                        # (e.g. ``getUsers().map(cb)``) — keep visiting it
+                        # at the current depth.
+                        stack.extend(
+                            (gc, depth) for gc in func_child.children
+                        )
+                        for arg in args_node.children:
+                            if arg.type in self._FUNC_ARG_TYPES:
+                                stack.extend(
+                                    (gc, depth + 1) for gc in arg.children
+                                )
+                            elif arg.type not in (",", "(", ")"):
+                                stack.extend((gc, depth) for gc in arg.children)
+                        continue
+            child_depth = (
+                depth + 1 if child.type in self._TS_LOOP_TYPES else depth
+            )
+            stack.extend((gc, child_depth) for gc in child.children)
 
     # ------------------------------------------------------------------
     # DF2: HTTP call-site detection (fetch / axios / SWR / api-clients)

tests/fixtures/df0_typescript/loop_calls.ts

@@ -0,0 +1,19 @@
+// Fixture for loop_depth / in_loop metadata tests.
+
+export function withLoops(items: string[]): void {
+  for (const item of items) {
+    doWork(item);
+  }
+
+  items.forEach((x) => {
+    doWork(x);
+  });
+}
+
+export function withoutLoop(): void {
+  doWork("direct");
+}
+
+export function chainedBeforeMap(): void {
+  getUsers().map((u) => doWork(u));
+}

tests/fixtures/lint_sample/src/batch.py

@@ -0,0 +1,12 @@
+from db import session
+
+
+# for loop with session.execute — should fire db-call-in-loop
+def process_items(items):
+    for item in items:
+        session.execute("INSERT INTO log VALUES (?)", [item])
+
+
+# non-loop call — must NOT fire
+def setup():
+    session.execute("CREATE TABLE IF NOT EXISTS log (val TEXT)")

tests/fixtures/lint_sample/src/loops.ts

@@ -0,0 +1,20 @@
+import { db, jobs } from "./schema";
+
+// for...of loop — should fire db-call-in-loop
+export async function insertAll(items: string[]): Promise<void> {
+  for (const item of items) {
+    await db.insert(jobs).values(item);
+  }
+}
+
+// .forEach callback — should fire db-call-in-loop
+export async function updateAll(ids: number[]): Promise<void> {
+  ids.forEach(async (id) => {
+    await db.update(jobs).set({ processed: true }).where(id);
+  });
+}
+
+// non-loop call — must NOT fire
+export async function insertOne(item: string): Promise<void> {
+  await db.insert(jobs).values(item);
+}

tests/test_df0_python.py

@@ -185,3 +185,42 @@ def test_call_literal_kinds_captured(
     _, edges = _run_extractor_on_source(tmp_path, src)
     call = next(e for e in _calls_edges(edges) if e.metadata.get("target_name") == "f")
     assert call.metadata["args"] == [expected]
+
+
+# --- loop_depth / in_loop metadata ----------------------------------------
+
+
+def test_call_inside_for_loop_has_loop_depth(tmp_path: Path) -> None:
+    src = "def caller(items):\n    for item in items:\n        f(item)\n"
+    _, edges = _run_extractor_on_source(tmp_path, src)
+    call = next(e for e in _calls_edges(edges) if e.metadata.get("target_name") == "f")
+    assert call.metadata["loop_depth"] == 1
+    assert call.metadata["in_loop"] is True
+
+
+def test_call_inside_while_loop_has_loop_depth(tmp_path: Path) -> None:
+    src = "def caller(cond):\n    while cond:\n        f()\n"
+    _, edges = _run_extractor_on_source(tmp_path, src)
+    call = next(e for e in _calls_edges(edges) if e.metadata.get("target_name") == "f")
+    assert call.metadata["loop_depth"] == 1
+    assert call.metadata["in_loop"] is True
+
+
+def test_call_outside_loop_has_zero_loop_depth(tmp_path: Path) -> None:
+    src = "def caller():\n    f()\n"
+    _, edges = _run_extractor_on_source(tmp_path, src)
+    call = next(e for e in _calls_edges(edges) if e.metadata.get("target_name") == "f")
+    assert call.metadata["loop_depth"] == 0
+    assert call.metadata["in_loop"] is False
+
+
+def test_nested_loop_increments_depth(tmp_path: Path) -> None:
+    src = (
+        "def caller(matrix):\n"
+        "    for row in matrix:\n"
+        "        for cell in row:\n"
+        "            f(cell)\n"
+    )
+    _, edges = _run_extractor_on_source(tmp_path, src)
+    call = next(e for e in _calls_edges(edges) if e.metadata.get("target_name") == "f")
+    assert call.metadata["loop_depth"] == 2