diff --git a/codegraph/analysis/lint.py b/codegraph/analysis/lint.py index 5bd9453..8eccc2e 100644 --- a/codegraph/analysis/lint.py +++ b/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 diff --git a/codegraph/parsers/python.py b/codegraph/parsers/python.py index 5e8c1d9..4117637 100644 --- a/codegraph/parsers/python.py +++ b/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, diff --git a/codegraph/parsers/typescript.py b/codegraph/parsers/typescript.py index 592e114..7182bed 100644 --- a/codegraph/parsers/typescript.py +++ b/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) diff --git a/tests/fixtures/df0_typescript/loop_calls.ts b/tests/fixtures/df0_typescript/loop_calls.ts new file mode 100644 index 0000000..8549b6d --- /dev/null +++ b/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)); +} diff --git a/tests/fixtures/lint_sample/src/batch.py b/tests/fixtures/lint_sample/src/batch.py new file mode 100644 index 0000000..3a90f32 --- /dev/null +++ b/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)") diff --git a/tests/fixtures/lint_sample/src/loops.ts b/tests/fixtures/lint_sample/src/loops.ts new file mode 100644 index 0000000..6a90692 --- /dev/null +++ b/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 { + 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 { + 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 { + await db.insert(jobs).values(item); +} diff --git a/tests/test_df0_python.py b/tests/test_df0_python.py index 75ab69c..071f693 100644 --- a/tests/test_df0_python.py +++ b/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 diff --git a/tests/test_df0_typescript.py b/tests/test_df0_typescript.py index 71b4c32..3476b91 100644 --- a/tests/test_df0_typescript.py +++ b/tests/test_df0_typescript.py @@ -196,3 +196,53 @@ def test_jsx_function_captures_params() -> None: {"name": "name", "type": "string", "default": None}, ] assert f.metadata["returns"] == "JSX.Element" + + +# ---------- loop_depth / in_loop metadata ---------- + + +@pytest.fixture(scope="module") +def loop_parsed() -> tuple[list[Node], list[Edge]]: + extractor = TypeScriptExtractor() + return extractor.parse_file(FIXTURE_DIR / "loop_calls.ts", FIXTURE_DIR) + + +def test_call_inside_for_of_has_loop_depth( + loop_parsed: tuple[list[Node], list[Edge]], +) -> None: + _, edges = loop_parsed + # doWork called inside for...of — loop_depth must be >= 1 + calls = _calls_in(edges, "doWork") + in_loop_calls = [e for e in calls if e.metadata.get("loop_depth", 0) >= 1] + assert in_loop_calls, "expected at least one doWork call with loop_depth >= 1" + + +def test_call_inside_foreach_callback_has_loop_depth( + loop_parsed: tuple[list[Node], list[Edge]], +) -> None: + _, edges = loop_parsed + # doWork called inside .forEach callback — in_loop must be True + calls = _calls_in(edges, "doWork") + forEach_in_loop = [e for e in calls if e.metadata.get("in_loop") is True] + assert forEach_in_loop, "expected at least one doWork call with in_loop=True" + + +def test_call_outside_loop_has_zero_loop_depth( + loop_parsed: tuple[list[Node], list[Edge]], +) -> None: + _, edges = loop_parsed + # doWork called in withoutLoop — loop_depth must be 0 + calls = _calls_in(edges, "doWork") + outside_calls = [e for e in calls if e.metadata.get("loop_depth", -1) == 0] + assert outside_calls, "expected at least one doWork call with loop_depth == 0" + + +def test_callee_chained_before_map_still_collected( + loop_parsed: tuple[list[Node], list[Edge]], +) -> None: + _, edges = loop_parsed + # getUsers() is the receiver of .map — its CALLS edge must not be + # dropped by the iteration-method depth handling. + calls = _calls_in(edges, "getUsers") + assert calls, "expected getUsers() call edge to survive .map handling" + assert calls[0].metadata.get("loop_depth") == 0 diff --git a/tests/test_lint.py b/tests/test_lint.py index 7e57aeb..af3f8fd 100644 --- a/tests/test_lint.py +++ b/tests/test_lint.py @@ -183,3 +183,78 @@ def test_sensitive_literal_redacts_value(tmp_path: Path) -> None: assert "hunter2" not in f.message assert "sk-test" not in f.snippet assert "sk-test" not in f.message + + +# --- db-call-in-loop ------------------------------------------------------- + + +def test_db_call_in_loop_typescript_for_of(tmp_path: Path) -> None: + repo = _sample_repo(tmp_path) + findings = [ + f for f in run_lint(repo) + if f.rule_id == "db-call-in-loop" and f.file == "src/loops.ts" + ] + # for...of loop (line 6) should fire + loop_lines = {f.line for f in findings} + assert 6 in loop_lines + + +def test_db_call_in_loop_typescript_foreach_callback(tmp_path: Path) -> None: + repo = _sample_repo(tmp_path) + findings = [ + f for f in run_lint(repo) + if f.rule_id == "db-call-in-loop" and f.file == "src/loops.ts" + ] + # forEach callback (line 13) should fire + loop_lines = {f.line for f in findings} + assert 13 in loop_lines + + +def test_db_call_in_loop_typescript_non_loop_not_flagged(tmp_path: Path) -> None: + repo = _sample_repo(tmp_path) + findings = [ + f for f in run_lint(repo) + if f.rule_id == "db-call-in-loop" and f.file == "src/loops.ts" + ] + # non-loop db.insert (line 19) must NOT fire + loop_lines = {f.line for f in findings} + assert 19 not in loop_lines + + +def test_db_call_in_loop_python(tmp_path: Path) -> None: + repo = _sample_repo(tmp_path) + findings = [ + f for f in run_lint(repo) + if f.rule_id == "db-call-in-loop" and f.file == "src/batch.py" + ] + # for loop session.execute (line 7) should fire + assert any(f.line == 7 for f in findings) + + +def test_db_call_in_loop_python_non_loop_not_flagged(tmp_path: Path) -> None: + repo = _sample_repo(tmp_path) + findings = [ + f for f in run_lint(repo) + if f.rule_id == "db-call-in-loop" and f.file == "src/batch.py" + ] + # non-loop session.execute (line 11) must NOT fire + assert not any(f.line == 11 for f in findings) + + +def test_db_call_in_loop_severity_high(tmp_path: Path) -> None: + repo = _sample_repo(tmp_path) + findings = [ + f for f in run_lint(repo) if f.rule_id == "db-call-in-loop" + ] + assert findings + assert all(f.severity == "high" for f in findings) + + +def test_db_call_in_loop_message_contains_chain(tmp_path: Path) -> None: + repo = _sample_repo(tmp_path) + findings = [ + f for f in run_lint(repo) if f.rule_id == "db-call-in-loop" + ] + assert findings + for f in findings: + assert "db." in f.message or "session." in f.message diff --git a/tests/test_lint_cli.py b/tests/test_lint_cli.py index dacf8f1..51245a6 100644 --- a/tests/test_lint_cli.py +++ b/tests/test_lint_cli.py @@ -38,7 +38,12 @@ def test_lint_markdown_output(lint_repo: Path) -> None: def test_lint_fail_on_gates_exit_code(lint_repo: Path) -> None: result = runner.invoke(app, ["lint", "--fail-on", "low"]) assert result.exit_code == 1 + # db-call-in-loop (high) fires on the sample fixture, so --fail-on high + # also exits non-zero. result = runner.invoke(app, ["lint", "--fail-on", "high"]) + assert result.exit_code == 1 + # No critical findings in the fixture → exit 0. + result = runner.invoke(app, ["lint", "--fail-on", "critical"]) assert result.exit_code == 0