pretensor 0.1.0__py3-none-any.whl

pretensor/enrichment/analyze/extract_python.py

@@ -0,0 +1,196 @@
+"""Stdlib ``ast`` extractor for SQL-bearing string literals in Python source."""
+
+from __future__ import annotations
+
+import ast
+import io
+import logging
+import re
+import tokenize
+from dataclasses import dataclass
+from pathlib import Path
+
+__all__ = ["SqlCandidate", "extract_sql_candidates"]
+
+logger = logging.getLogger(__name__)
+
+_SQL_IDENTIFIER_NAMES: frozenset[str] = frozenset({"query", "sql", "stmt", "statement"})
+_SQL_IDENTIFIER_SUFFIXES: tuple[str, ...] = ("_query", "_sql", "_stmt", "_statement")
+_SQL_IDENTIFIER_PREFIXES: tuple[str, ...] = ("query_", "sql_", "stmt_", "statement_")
+_SQL_METHOD_NAMES: frozenset[str] = frozenset(
+    {"execute", "executemany", "query", "raw", "exec_driver_sql"}
+)
+_NOQA_MARKER = "noqa: pretensor-analyze"
+
+
+def _is_sql_identifier(name: str) -> bool:
+    """True for variable names that conventionally hold SQL text.
+
+    Exact names (``sql``, ``query``, …) plus affixed forms (``merge_sql``,
+    ``orders_query``, ``sql_fetch_users``). The is-this-SQL classifier still
+    gates every candidate, so a ``graphql_query`` holding non-SQL text is
+    rejected downstream rather than here.
+    """
+    lowered = name.lower()
+    return (
+        lowered in _SQL_IDENTIFIER_NAMES
+        or lowered.endswith(_SQL_IDENTIFIER_SUFFIXES)
+        or lowered.startswith(_SQL_IDENTIFIER_PREFIXES)
+    )
+
+
+_SQL_FSTRING_PREFIX_RE = re.compile(
+    r"^\s*(SELECT|INSERT|UPDATE|DELETE|MERGE|WITH)\b", re.IGNORECASE
+)
+
+
+@dataclass
+class SqlCandidate:
+    """A SQL string literal candidate extracted from Python source."""
+
+    text: str
+    confidence_bucket: (
+        str  # "high" | "medium" | "low" — updated by pipeline for fstrings
+    )
+    line_start: int
+    line_end: int
+    file_path: Path
+    symbol: str
+    kind: str  # extraction context: "assignment" | "call" | "fstring"
+
+
+def extract_sql_candidates(file_path: Path) -> list[SqlCandidate]:
+    """Extract SQL-bearing string literals from a Python source file.
+
+    Uses syntactic context only — whether the string is assigned to a SQL-named
+    variable or passed as the first argument to a SQL-execution method. SQL
+    content classification is handled separately by ``classify.py``.
+    """
+    try:
+        source = file_path.read_text(encoding="utf-8", errors="replace")
+    except OSError as exc:
+        logger.debug("extract_python: cannot read %s: %s", file_path, exc)
+        return []
+
+    noqa_lines = _collect_noqa_lines(source)
+
+    try:
+        tree = ast.parse(source, filename=str(file_path))
+    except SyntaxError as exc:
+        logger.debug("extract_python: syntax error in %s: %s", file_path, exc)
+        return []
+
+    visitor = _SqlVisitor(file_path, noqa_lines)
+    visitor.visit(tree)
+    return visitor.candidates
+
+
+def _collect_noqa_lines(source: str) -> set[int]:
+    """Return the set of line numbers (1-indexed) carrying a pretensor-analyze noqa marker."""
+    noqa: set[int] = set()
+    try:
+        tokens = tokenize.generate_tokens(io.StringIO(source).readline)
+        for tok_type, tok_val, (srow, _scol), *_ in tokens:
+            if tok_type == tokenize.COMMENT and _NOQA_MARKER in tok_val:
+                noqa.add(srow)
+    except tokenize.TokenError:
+        pass
+    return noqa
+
+
+class _SqlVisitor(ast.NodeVisitor):
+    """AST visitor that collects SQL-bearing string literals by syntactic context."""
+
+    def __init__(self, file_path: Path, noqa_lines: set[int]) -> None:
+        self.file_path = file_path
+        self.noqa_lines = noqa_lines
+        self.candidates: list[SqlCandidate] = []
+
+    def _is_noqa(self, line_start: int) -> bool:
+        return line_start in self.noqa_lines or (line_start - 1) in self.noqa_lines
+
+    def _emit_constant(
+        self,
+        node: ast.Constant,
+        *,
+        symbol: str,
+        kind: str,
+        confidence_bucket: str = "high",
+    ) -> None:
+        if not isinstance(node.value, str):
+            return
+        text = node.value
+        if not text.strip():
+            return
+        line_start = node.lineno
+        line_end = getattr(node, "end_lineno", node.lineno)
+        if self._is_noqa(line_start):
+            return
+        self.candidates.append(
+            SqlCandidate(
+                text=text,
+                confidence_bucket=confidence_bucket,
+                line_start=line_start,
+                line_end=line_end,
+                file_path=self.file_path,
+                symbol=symbol,
+                kind=kind,
+            )
+        )
+
+    def _emit_fstring(self, node: ast.JoinedStr, *, symbol: str) -> None:
+        """Emit a candidate for an f-string whose first element is a SQL-keyword prefix."""
+        if not node.values:
+            return
+        first = node.values[0]
+        if not isinstance(first, ast.Constant) or not isinstance(first.value, str):
+            return
+        prefix = first.value
+        if not _SQL_FSTRING_PREFIX_RE.match(prefix):
+            return
+        line_start = node.lineno
+        line_end = getattr(node, "end_lineno", node.lineno)
+        if self._is_noqa(line_start):
+            return
+        self.candidates.append(
+            SqlCandidate(
+                text=prefix,
+                confidence_bucket="low",
+                line_start=line_start,
+                line_end=line_end,
+                file_path=self.file_path,
+                symbol=symbol,
+                kind="fstring",
+            )
+        )
+
+    def _handle_value(self, value: ast.expr, *, symbol: str, kind: str) -> None:
+        if isinstance(value, ast.Constant):
+            self._emit_constant(value, symbol=symbol, kind=kind)
+        elif isinstance(value, ast.JoinedStr):
+            self._emit_fstring(value, symbol=symbol)
+
+    def visit_Assign(self, node: ast.Assign) -> None:
+        for target in node.targets:
+            if isinstance(target, ast.Name) and _is_sql_identifier(target.id):
+                self._handle_value(node.value, symbol=target.id, kind="assignment")
+                break
+        self.generic_visit(node)
+
+    def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
+        if (
+            isinstance(node.target, ast.Name)
+            and _is_sql_identifier(node.target.id)
+            and node.value is not None
+        ):
+            self._handle_value(node.value, symbol=node.target.id, kind="assignment")
+        self.generic_visit(node)
+
+    def visit_Call(self, node: ast.Call) -> None:
+        if (
+            isinstance(node.func, ast.Attribute)
+            and node.func.attr.lower() in _SQL_METHOD_NAMES
+            and node.args
+        ):
+            self._handle_value(node.args[0], symbol=node.func.attr, kind="call")
+        self.generic_visit(node)

pretensor/enrichment/analyze/parse.py

@@ -0,0 +1,141 @@
+"""Multi-dialect sqlglot retry wrapper for SQL table reference extraction."""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import re
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Iterator
+
+from pretensor.connectors.lineage_sqlglot import dml_write_targets, table_refs_from_sql
+
+__all__ = [
+    "ParsedSql",
+    "normalize_placeholders",
+    "normalize_sql",
+    "parse_sql",
+    "quiet_sqlglot",
+]
+
+logger = logging.getLogger(__name__)
+
+_DIALECTS = ["", "postgres", "mysql", "bigquery", "snowflake", "trino", "spark"]
+
+# DBAPI-style bind placeholders that break sqlglot's parser when left in place.
+# Replaced with NULL before parsing so parameterized statements (the norm in
+# scanned application code) keep their full structure — most importantly the
+# INSERT/UPDATE node, which carries the write-target classification.
+_PLACEHOLDER_RES: tuple[re.Pattern[str], ...] = (
+    re.compile(r"%\([A-Za-z_][A-Za-z0-9_]*\)s"),  # psycopg named: %(name)s
+    re.compile(r"%s"),  # psycopg positional: %s
+    re.compile(r"(?<![:\w]):[A-Za-z_][A-Za-z0-9_]*"),  # named binds :name, not ::casts
+    re.compile(r"(?<=[\s(=,<>])\?(?=[\s),;]|$)"),  # qmark style: ?
+)
+
+
+# MySQL/SQLite REPLACE INTO is a delete-then-insert write to the target table,
+# but sqlglot cannot parse it in any dialect (it falls back to an opaque
+# Command with no table nodes). Rewriting the head to INSERT INTO preserves
+# the write-target classification; the fingerprint is computed from the
+# original text so identity is unaffected.
+_REPLACE_INTO_RE = re.compile(
+    r"^(\s*)REPLACE\s+(?:LOW_PRIORITY\s+|DELAYED\s+)?(?:INTO\s+)?",
+    re.IGNORECASE,
+)
+
+
+def normalize_placeholders(sql: str) -> str:
+    """Replace DBAPI bind placeholders (``%s``, ``%(name)s``, ``:name``, ``?``)
+    with ``NULL`` so sqlglot can parse parameterized statements. ``::`` casts and
+    string contents that don't look like bind markers are left alone."""
+    for pattern in _PLACEHOLDER_RES:
+        sql = pattern.sub("NULL", sql)
+    return sql
+
+
+def normalize_sql(sql: str) -> str:
+    """Full pre-parse normalization: placeholder substitution plus the
+    ``REPLACE INTO`` → ``INSERT INTO`` rewrite."""
+    return _REPLACE_INTO_RE.sub(r"\1INSERT INTO ", normalize_placeholders(sql))
+
+
+# Characters that never appear in a correctly parsed identifier part. A dialect
+# that "succeeds" by tokenizing another dialect's quoting into the name itself
+# (e.g. the generic dialect turning `my-project.sales.orders` into a literal
+# backtick) must not win the retry — its refs are garbage and the right
+# dialect would never get tried.
+_INVALID_IDENT_CHARS = ("`", '"', "'", ";", "\n")
+
+
+def _refs_look_valid(refs: list[tuple[str, str]]) -> bool:
+    for schema, table in refs:
+        if not table:
+            return False
+        for part in (schema, table):
+            if any(ch in part for ch in _INVALID_IDENT_CHARS):
+                return False
+        if "." in table:
+            return False
+    return True
+
+
+@contextmanager
+def quiet_sqlglot() -> Iterator[None]:
+    """Silence sqlglot's own error logging during trial parses.
+
+    The multi-dialect retry intentionally feeds sqlglot statements it may not
+    parse; without this, every failed attempt prints an ERROR line to the
+    user's console.
+    """
+    sqlglot_logger = logging.getLogger("sqlglot")
+    previous = sqlglot_logger.level
+    sqlglot_logger.setLevel(logging.CRITICAL)
+    try:
+        yield
+    finally:
+        sqlglot_logger.setLevel(previous)
+
+
+@dataclass
+class ParsedSql:
+    """Extracted table references and fingerprint from a SQL string."""
+
+    table_refs: list[tuple[str, str]] = field(default_factory=list)
+    write_targets: list[tuple[str, str]] = field(default_factory=list)
+    dialect_used: str = ""
+    fingerprint: str = ""
+
+
+def parse_sql(sql: str, *, default_schema: str = "public") -> ParsedSql:
+    """Extract table references via multi-dialect retry.
+
+    Tries each dialect in ``_DIALECTS`` until one returns non-empty,
+    plausible ``table_refs`` — refs whose names contain quoting characters or
+    dots are treated as a failed parse so the retry can reach the dialect that
+    understands the statement's quoting (BigQuery/MySQL backticks).
+    Always returns a ``ParsedSql`` (``table_refs`` may be empty if all dialects fail).
+    The ``fingerprint`` is ``sha256(stripped+lowercased sql)[:16]`` and is always
+    computed from the original text, before normalization.
+    """
+    fingerprint = hashlib.sha256(sql.strip().lower().encode()).hexdigest()[:16]
+    normalized = normalize_sql(sql)
+
+    with quiet_sqlglot():
+        for dialect in _DIALECTS:
+            refs = table_refs_from_sql(
+                normalized, dialect=dialect, default_schema=default_schema
+            )
+            if refs and _refs_look_valid(refs):
+                targets = dml_write_targets(
+                    normalized, dialect=dialect, default_schema=default_schema
+                )
+                return ParsedSql(
+                    table_refs=refs,
+                    write_targets=targets,
+                    dialect_used=dialect,
+                    fingerprint=fingerprint,
+                )
+
+    return ParsedSql(fingerprint=fingerprint)

pretensor/enrichment/analyze/pipeline.py

@@ -0,0 +1,195 @@
+"""Orchestrate the analyze enrichment pipeline (walker → extractor → writer)."""
+
+from __future__ import annotations
+
+import logging
+import time
+import uuid
+from collections.abc import Sequence
+from pathlib import Path
+
+from pretensor.core.store import KuzuStore
+from pretensor.enrichment.analyze.classify import CONFIDENCE_FLOAT, classify_sql
+from pretensor.enrichment.analyze.extract_python import extract_sql_candidates
+from pretensor.enrichment.analyze.parse import parse_sql
+from pretensor.enrichment.analyze.summary import AnalyzeSummary
+from pretensor.enrichment.analyze.walker import walk_repo
+from pretensor.enrichment.analyze.writers import (
+    ParsedCandidate,
+    write_consumers_for_run,
+)
+from pretensor.observability import log_timed_operation
+
+__all__ = ["EmptyGraphError", "run_analyze_enrichment"]
+
+logger = logging.getLogger(__name__)
+
+_FSTRING_CONFIDENCE_BUCKET = "low"
+
+
+class EmptyGraphError(ValueError):
+    """The graph has no ``SchemaTable`` rows for the requested connection."""
+
+
+def run_analyze_enrichment(
+    repo_path: Path,
+    store: KuzuStore,
+    connection_name: str,
+    *,
+    service_name: str | None = None,
+    includes: Sequence[str] = (),
+    excludes: Sequence[str] = (),
+    max_file_bytes: int = 1_000_000,
+    min_confidence: float = 0.6,
+    default_schema: str = "public",
+    scan_run_id: str | None = None,
+    dry_run: bool = False,
+) -> AnalyzeSummary:
+    """Walk ``repo_path``, extract SQL from Python files, and write graph nodes/edges.
+
+    Hard-errors when the graph has no ``SchemaTable`` rows for ``connection_name``
+    (the user must run ``pretensor index --connection <name>`` first).
+
+    Args:
+        repo_path: Root of the repository to scan.
+        store: Open Kuzu graph store.
+        connection_name: Pretensor connection name; only tables indexed under this
+            connection are considered valid CONSUMES targets.
+        service_name: Label for the scanned service (defaults to ``repo_path.name``).
+        includes: Gitignore-style glob patterns limiting which files are scanned.
+        excludes: Gitignore-style glob patterns excluding files from scanning.
+        max_file_bytes: Skip files larger than this (default 1 MB).
+        min_confidence: Drop candidates below this confidence level (default 0.6,
+            which excludes ``low``-confidence f-string prefix candidates).
+        default_schema: Schema assumed for unqualified table references in the
+            scanned SQL (default ``"public"``; pass the connection's actual
+            default schema for non-Postgres dialects).
+        scan_run_id: Idempotency key for this run; stale consumers from other run IDs
+            are swept at the end. Defaults to a random 8-hex-char string.
+
+    Returns:
+        ``AnalyzeSummary`` with counts and timing.
+    """
+    _guard_non_empty_graph(store, connection_name)
+
+    effective_service = service_name or repo_path.name
+    effective_run_id = scan_run_id or uuid.uuid4().hex[:8]
+
+    started = time.perf_counter()
+    files_scanned = 0
+    sql_candidates_found = 0
+    sql_candidates_failed = 0
+    parsed_candidates: list[ParsedCandidate] = []
+
+    with log_timed_operation(
+        logger,
+        event="analyze.walk",
+        connection_name=connection_name,
+        repo_path=str(repo_path),
+    ):
+        file_list = list(
+            walk_repo(
+                repo_path,
+                includes=includes,
+                excludes=excludes,
+                max_file_bytes=max_file_bytes,
+            )
+        )
+
+    with log_timed_operation(
+        logger,
+        event="analyze.extract",
+        connection_name=connection_name,
+        file_count=len(file_list),
+    ):
+        for file_path, language in file_list:
+            files_scanned += 1
+            if language != "python":
+                continue
+
+            raw_candidates = extract_sql_candidates(file_path)
+            for cand in raw_candidates:
+                sql_candidates_found += 1
+
+                bucket, conf = classify_sql(cand.text)
+                if not bucket:
+                    continue
+
+                if cand.kind == "fstring":
+                    bucket = _FSTRING_CONFIDENCE_BUCKET
+                    conf = CONFIDENCE_FLOAT[bucket]
+                cand.confidence_bucket = bucket
+
+                if conf < min_confidence:
+                    continue
+
+                try:
+                    parsed = parse_sql(cand.text, default_schema=default_schema)
+                except Exception as exc:
+                    logger.debug(
+                        "analyze: parse error for candidate in %s: %s", file_path, exc
+                    )
+                    sql_candidates_failed += 1
+                    continue
+
+                if not parsed.table_refs and not parsed.write_targets:
+                    continue
+
+                try:
+                    rel_path = str(file_path.relative_to(repo_path))
+                except ValueError:
+                    rel_path = file_path.name
+
+                parsed_candidates.append(
+                    ParsedCandidate(
+                        candidate=cand,
+                        parsed=parsed,
+                        rel_path=rel_path,
+                        language=language,
+                    )
+                )
+
+    with log_timed_operation(
+        logger,
+        event="analyze.write",
+        connection_name=connection_name,
+        candidate_count=len(parsed_candidates),
+    ):
+        write_summary = write_consumers_for_run(
+            store,
+            connection_name,
+            parsed_candidates,
+            effective_run_id,
+            service_name=effective_service,
+            min_confidence=min_confidence,
+            dry_run=dry_run,
+        )
+
+    duration_ms = (time.perf_counter() - started) * 1000
+
+    return AnalyzeSummary(
+        files_scanned=files_scanned,
+        sql_candidates_found=sql_candidates_found,
+        sql_candidates_parsed=len(parsed_candidates),
+        sql_candidates_failed=sql_candidates_failed,
+        consumers_written=write_summary.consumers_written,
+        edges_written=write_summary.edges_written,
+        cross_connection_dropped=write_summary.cross_connection_dropped,
+        scan_run_id=effective_run_id,
+        duration_ms=duration_ms,
+        rows=write_summary.rows,
+    )
+
+
+def _guard_non_empty_graph(store: KuzuStore, connection_name: str) -> None:
+    """Raise ``EmptyGraphError`` when no ``SchemaTable`` rows exist for the connection."""
+    rows = store.query_all_rows(
+        "MATCH (t:SchemaTable {connection_name: $cn}) RETURN count(*) LIMIT 1",
+        {"cn": connection_name},
+    )
+    count = int(rows[0][0]) if rows and rows[0][0] is not None else 0
+    if count == 0:
+        raise EmptyGraphError(
+            f"No SchemaTable nodes found for connection '{connection_name}'. "
+            f"Index the database first: pretensor index --connection {connection_name}"
+        )

pretensor/enrichment/analyze/summary.py

@@ -0,0 +1,38 @@
+"""Frozen summary dataclasses for the analyze enrichment pipeline."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+__all__ = ["AnalyzeSummary", "ConsumerRow"]
+
+
+@dataclass(frozen=True, slots=True)
+class ConsumerRow:
+    """One planned/written CONSUMES edge, for CLI rendering (no raw SQL)."""
+
+    service_name: str
+    table: str  # "schema.table" as resolved
+    op: str  # "read" | "write"
+    file_path: str  # repo-relative
+    line_start: int
+    line_end: int
+    confidence: float
+
+
+@dataclass(frozen=True, slots=True)
+class AnalyzeSummary:
+    """Aggregate counts returned by ``run_analyze_enrichment``."""
+
+    files_scanned: int
+    sql_candidates_found: int
+    sql_candidates_parsed: int
+    sql_candidates_failed: int
+    consumers_written: int
+    edges_written: int
+    cross_connection_dropped: int
+    scan_run_id: str
+    duration_ms: float
+    # Planned (dry-run) or written CONSUMES rows, for CLI rendering. Defaulted so
+    # existing keyword construction and the frozen field set stay compatible.
+    rows: tuple[ConsumerRow, ...] = field(default=())

pretensor/enrichment/analyze/walker.py

@@ -0,0 +1,98 @@
+"""Gitignore-aware repository file walker for the analyze enrichment pipeline."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator, Sequence
+from pathlib import Path
+
+import pathspec
+
+__all__ = ["walk_repo", "SKIP_DIRS"]
+
+logger = logging.getLogger(__name__)
+
+SKIP_DIRS: frozenset[str] = frozenset(
+    {"node_modules", ".venv", "dist", "build", ".tox", ".git"}
+)
+
+_LANG_MAP: dict[str, str] = {".py": "python"}
+
+
+def walk_repo(
+    repo_path: Path,
+    *,
+    includes: Sequence[str] = (),
+    excludes: Sequence[str] = (),
+    max_file_bytes: int = 1_000_000,
+) -> Iterator[tuple[Path, str]]:
+    """Yield ``(absolute_file_path, language)`` for each candidate source file.
+
+    Respects ``.gitignore`` in ``repo_path``, prunes ``SKIP_DIRS``, skips files
+    exceeding ``max_file_bytes``, and filters by ``includes``/``excludes`` glob
+    patterns when provided.
+    """
+    ignore_spec = _load_gitignore_spec(repo_path)
+    include_spec = (
+        pathspec.PathSpec.from_lines("gitignore", includes) if includes else None
+    )
+    exclude_spec = (
+        pathspec.PathSpec.from_lines("gitignore", excludes) if excludes else None
+    )
+
+    for file_path in _iter_files(repo_path):
+        rel = file_path.relative_to(repo_path)
+        rel_str = str(rel)
+
+        if ignore_spec and ignore_spec.match_file(rel_str):
+            continue
+
+        try:
+            size = file_path.stat().st_size
+        except OSError:
+            continue
+        if size > max_file_bytes:
+            logger.debug("walker: skipping large file %s (%d bytes)", rel, size)
+            continue
+
+        lang = _LANG_MAP.get(file_path.suffix)
+        if lang is None:
+            continue
+
+        if include_spec and not include_spec.match_file(rel_str):
+            continue
+        if exclude_spec and exclude_spec.match_file(rel_str):
+            continue
+
+        yield file_path, lang
+
+
+def _iter_files(directory: Path) -> Iterator[Path]:
+    """Recursively yield files, pruning SKIP_DIRS directories early."""
+    try:
+        entries = list(directory.iterdir())
+    except PermissionError:
+        return
+    for entry in entries:
+        if entry.is_symlink():
+            # Never follow symlinks: a link cycle would recurse forever, and a
+            # link out of the repo would break repo-relative path invariants.
+            continue
+        if entry.is_dir():
+            if entry.name in SKIP_DIRS:
+                continue
+            yield from _iter_files(entry)
+        elif entry.is_file():
+            yield entry
+
+
+def _load_gitignore_spec(repo_root: Path) -> pathspec.PathSpec | None:
+    """Load ``.gitignore`` from the repo root into a PathSpec (best-effort)."""
+    gitignore = repo_root / ".gitignore"
+    if not gitignore.is_file():
+        return None
+    try:
+        lines = gitignore.read_text(encoding="utf-8", errors="replace").splitlines()
+    except OSError:
+        return None
+    return pathspec.PathSpec.from_lines("gitignore", lines)