sql-code-graph 0.2.1__py3-none-any.whl

sqlcg/indexer/indexer.py

@@ -0,0 +1,317 @@
+"""Main indexer orchestrating parsing and graph persistence."""
+
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import TimeoutError as FuturesTimeout
+from pathlib import Path
+
+from sqlcg.core.graph_db import GraphBackend
+from sqlcg.core.queries import STALE_VIEWS_QUERY
+from sqlcg.core.schema import NodeLabel, RelType
+from sqlcg.indexer.walker import walk_sql_files
+from sqlcg.lineage.aggregator import CrossFileAggregator
+from sqlcg.lineage.schema_resolver import SchemaResolver
+from sqlcg.parsers.base import ParsedFile
+from sqlcg.parsers.registry import get_parser
+from sqlcg.utils.ignore import load_ignore_spec
+from sqlcg.utils.logging import getLogger
+
+logger = getLogger(__name__)
+
+
+class Indexer:
+    """Orchestrates SQL file parsing and graph persistence."""
+
+    def index_repo(
+        self,
+        path: Path,
+        dialect: str | None,
+        db: GraphBackend,
+        dbt_manifest: Path | None = None,
+        timeout_per_file: int = 30,
+        use_git: bool = True,
+    ) -> dict:
+        """Full two-pass index. Returns summary dict.
+
+        Args:
+            path: Root directory to index
+            dialect: SQL dialect (None for ANSI)
+            db: GraphBackend instance
+            dbt_manifest: Optional path to dbt manifest.json
+            timeout_per_file: Timeout in seconds per file (0 = no timeout)
+            use_git: When True (default), use git ls-files to restrict
+                indexing to tracked files; falls back to rglob when git
+                is unavailable or the directory is not a git repository.
+
+        Returns:
+            Dict with keys: files_parsed, parse_errors, tables_found, lineage_edges_created
+        """
+        spec = load_ignore_spec(path)
+        schema_resolver = SchemaResolver(dialect=dialect)
+        parser = get_parser(dialect, schema_resolver)
+        aggregator = CrossFileAggregator()
+
+        files = list(walk_sql_files(path, spec, use_git=use_git))
+        pass1_results: list[ParsedFile] = []
+        parse_errors = 0
+
+        # Pass 1: parse all files
+        for file_path in files:
+            try:
+                sql = file_path.read_text(encoding="utf-8")
+                parsed = self._index_single_file(parser, file_path, sql, timeout_per_file)
+                aggregator.register_pass1(parsed)
+                pass1_results.append(parsed)
+                parse_errors += len(parsed.errors)
+            except KeyboardInterrupt:
+                logger.info("SIGINT received — flushing progress")
+                self._upsert_all(pass1_results, db)
+                raise
+            except Exception as exc:
+                logger.warning("Failed to parse %s: %s", file_path, exc)
+                parse_errors += 1
+
+        # Optional: load dbt manifest
+        if dbt_manifest:
+            from sqlcg.indexer.dbt_adapter import load_dbt_manifest
+
+            load_dbt_manifest(dbt_manifest, schema_resolver)
+
+        # Pass 2: resolve cross-file references
+        pass2_results: list[ParsedFile] = []
+        for parsed in pass1_results:
+            try:
+                resolved = aggregator.resolve_pass2(parser, parsed)
+                pass2_results.append(resolved)
+            except Exception as exc:
+                logger.warning("resolve_pass2 failed for %s: %s", parsed.path, exc)
+                pass2_results.append(parsed)
+
+        # Upsert all results
+        tables_found = 0
+        lineage_edges = 0
+        for parsed in pass2_results:
+            counts = self._upsert_parsed_file(parsed, db)
+            tables_found += counts["tables"]
+            lineage_edges += counts["edges"]
+
+        return {
+            "files_parsed": len(pass2_results),
+            "parse_errors": parse_errors,
+            "tables_found": tables_found,
+            "lineage_edges_created": lineage_edges,
+        }
+
+    def reindex_file(self, file_path: str, db: GraphBackend, dialect: str | None) -> None:
+        """Re-index a single file and its dependent views.
+
+        Args:
+            file_path: Path to the file to re-index
+            db: GraphBackend instance
+            dialect: SQL dialect (None for ANSI)
+        """
+        stale_views = db.run_read(STALE_VIEWS_QUERY, {"path": file_path})
+
+        with db.transaction():
+            db.delete_nodes_for_file(file_path)
+            schema_resolver = SchemaResolver(dialect=dialect)
+            parser = get_parser(dialect, schema_resolver)
+            sql = Path(file_path).read_text(encoding="utf-8")
+            parsed = parser.parse_file(Path(file_path), sql)
+            self._upsert_parsed_file(parsed, db)
+
+        for row in stale_views:
+            self._reindex_view_definition(row["view_name"], db, dialect)
+
+    def _index_single_file(self, parser, path: Path, sql: str, timeout: int) -> ParsedFile:
+        """Parse one file, with optional timeout.
+
+        Args:
+            parser: SqlParser instance
+            path: Path to the file
+            sql: SQL text
+            timeout: Timeout in seconds (0 = no timeout)
+
+        Returns:
+            ParsedFile with parse_failed flag set if timeout occurs
+        """
+        if timeout <= 0:
+            return parser.parse_file(path, sql)
+
+        with ThreadPoolExecutor(max_workers=1) as ex:
+            future = ex.submit(parser.parse_file, path, sql)
+            try:
+                return future.result(timeout=timeout)
+            except FuturesTimeout:
+                logger.warning("Timeout parsing %s (>%ds) — skipping", path, timeout)
+                out = ParsedFile(path=path, dialect=parser.DIALECT)
+                out.errors.append(f"timeout:{timeout}s")
+                return out
+
+    def _upsert_parsed_file(self, parsed: ParsedFile, db: GraphBackend) -> dict:
+        """Map ParsedFile → graph nodes/edges.
+
+        Args:
+            parsed: ParsedFile to upsert
+            db: GraphBackend instance
+
+        Returns:
+            Dict with keys: tables, edges
+        """
+        counts = {"tables": 0, "edges": 0}
+
+        # Upsert File node
+        db.upsert_node(
+            NodeLabel.FILE,
+            parsed.path_str,
+            {
+                "path": parsed.path_str,
+                "dialect": parsed.dialect or "",
+            },
+        )
+
+        # Upsert defined tables
+        for table in parsed.defined_tables:
+            db.upsert_node(
+                NodeLabel.TABLE,
+                table.full_id,
+                {
+                    "qualified": table.full_id,
+                    "name": table.name,
+                    "catalog": table.catalog or "",
+                    "db": table.db or "",
+                    "kind": "TABLE",
+                    "defined_in_file": parsed.path_str,
+                },
+            )
+            db.upsert_edge(
+                NodeLabel.TABLE,
+                table.full_id,
+                NodeLabel.FILE,
+                parsed.path_str,
+                RelType.DEFINED_IN,
+                {},
+            )
+            counts["tables"] += 1
+
+        # Upsert query nodes
+        for i, stmt in enumerate(parsed.statements):
+            query_id = f"{parsed.path_str}:{i}"
+            db.upsert_node(
+                NodeLabel.QUERY,
+                query_id,
+                {
+                    "id": query_id,
+                    "file_path": parsed.path_str,
+                    "statement_index": i,
+                    "sql": stmt.sql[:500],
+                    "kind": stmt.kind,
+                    "target_table": stmt.target.full_id if stmt.target else "",
+                    "parse_failed": stmt.parse_failed,
+                    "confidence": stmt.confidence,
+                    "parsing_mode": stmt.parsing_mode,
+                },
+            )
+            db.upsert_edge(
+                NodeLabel.QUERY,
+                query_id,
+                NodeLabel.FILE,
+                parsed.path_str,
+                RelType.QUERY_DEFINED_IN,
+                {},
+            )
+
+            # Source table edges
+            for src_table in stmt.sources:
+                db.upsert_node(
+                    NodeLabel.TABLE,
+                    src_table.full_id,
+                    {
+                        "qualified": src_table.full_id,
+                        "name": src_table.name,
+                        "catalog": src_table.catalog or "",
+                        "db": src_table.db or "",
+                        "kind": "TABLE",
+                        "defined_in_file": "",
+                    },
+                )
+                db.upsert_edge(
+                    NodeLabel.QUERY,
+                    query_id,
+                    NodeLabel.TABLE,
+                    src_table.full_id,
+                    RelType.SELECTS_FROM,
+                    {},
+                )
+
+            # Column lineage edges
+            for edge in stmt.column_lineage:
+                src_id = edge.src.full_id
+                dst_id = edge.dst.full_id
+                db.upsert_node(
+                    NodeLabel.COLUMN,
+                    src_id,
+                    {
+                        "id": src_id,
+                        "col_name": edge.src.name,
+                        "table_qualified": edge.src.table.full_id,
+                        "catalog": edge.src.table.catalog or "",
+                        "db": edge.src.table.db or "",
+                        "table_name": edge.src.table.name,
+                    },
+                )
+                db.upsert_node(
+                    NodeLabel.COLUMN,
+                    dst_id,
+                    {
+                        "id": dst_id,
+                        "col_name": edge.dst.name,
+                        "table_qualified": edge.dst.table.full_id,
+                        "catalog": edge.dst.table.catalog or "",
+                        "db": edge.dst.table.db or "",
+                        "table_name": edge.dst.table.name,
+                    },
+                )
+                db.upsert_edge(
+                    NodeLabel.COLUMN,
+                    src_id,
+                    NodeLabel.COLUMN,
+                    dst_id,
+                    RelType.COLUMN_LINEAGE,
+                    {
+                        "transform": edge.transform,
+                        "confidence": edge.confidence,
+                        "query_id": query_id,
+                    },
+                )
+                counts["edges"] += 1
+
+        return counts
+
+    def _upsert_all(self, results: list[ParsedFile], db: GraphBackend) -> None:
+        """Upsert all parsed files.
+
+        Args:
+            results: List of ParsedFile objects
+            db: GraphBackend instance
+        """
+        for parsed in results:
+            self._upsert_parsed_file(parsed, db)
+
+    def _reindex_view_definition(
+        self, view_name: str, db: GraphBackend, dialect: str | None
+    ) -> None:
+        """Re-index the file that defines a view.
+
+        Args:
+            view_name: Qualified view name
+            db: GraphBackend instance
+            dialect: SQL dialect
+        """
+        query = (
+            f"MATCH (t:{NodeLabel.TABLE} {{qualified: $name}})"
+            f"-[:{RelType.DEFINED_IN}]->(f:{NodeLabel.FILE}) "
+            "RETURN f.path AS path"
+        )
+        result = db.run_read(query, {"name": view_name})
+        for row in result:
+            self.reindex_file(row["path"], db, dialect)

sqlcg/indexer/walker.py

@@ -0,0 +1,55 @@
+"""SQL file walker with ignore pattern support."""
+
+import subprocess
+from collections.abc import Iterator
+from pathlib import Path
+
+import pathspec
+
+from sqlcg.utils.ignore import is_ignored
+
+
+def _git_sql_files(root: Path) -> list[Path] | None:
+    """Return tracked .sql files via git ls-files, or None if git unavailable."""
+    try:
+        result = subprocess.run(
+            ["git", "ls-files", "--cached"],
+            cwd=root,
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        return [root / f for f in result.stdout.splitlines() if f.endswith(".sql")]
+    except (subprocess.CalledProcessError, OSError):
+        return None
+
+
+def walk_sql_files(
+    root: Path, spec: pathspec.PathSpec, use_git: bool = True
+) -> Iterator[Path]:
+    """Walk directory tree and yield SQL files not matching ignore patterns.
+
+    When use_git=True (default) and git is available, only tracked files are
+    returned — this prevents flooding from build artefacts, node_modules, and
+    other untracked directories. Falls back to rglob when git is unavailable
+    or the directory is not a git repository.
+
+    Args:
+        root: Root directory to walk
+        spec: PathSpec object with ignore patterns
+        use_git: Use git ls-files instead of rglob (default True)
+
+    Yields:
+        Path objects for .sql files not matching ignore patterns
+    """
+    if use_git:
+        git_files = _git_sql_files(root)
+        if git_files is not None:
+            for path in git_files:
+                if path.exists() and not is_ignored(path, root, spec):
+                    yield path
+            return
+
+    for path in root.rglob("*.sql"):
+        if not is_ignored(path, root, spec):
+            yield path

sqlcg/indexer/watcher.py

@@ -0,0 +1,195 @@
+"""File system watcher for SQL file changes."""
+
+import subprocess
+import threading
+import time
+from pathlib import Path
+
+import pathspec
+from watchdog.events import FileSystemEventHandler
+
+from sqlcg.utils.ignore import is_ignored
+from sqlcg.utils.logging import getLogger
+
+logger = getLogger(__name__)
+
+
+class SqlFileEventHandler(FileSystemEventHandler):
+    """Watchdog event handler for SQL file changes."""
+
+    def __init__(self, job_manager, db, ignore_spec: pathspec.PathSpec, root: Path, indexer=None):
+        """Initialize the event handler.
+
+        Args:
+            job_manager: WatchJobManager instance
+            db: GraphBackend instance
+            ignore_spec: PathSpec with ignore patterns
+            root: Root directory being watched
+            indexer: Indexer instance (used by BranchMonitor)
+        """
+        super().__init__()
+        self._jobs = job_manager
+        self._db = db
+        self._spec = ignore_spec
+        self._root = root
+        # Create and start BranchMonitor if indexer is provided
+        self._branch_monitor: BranchMonitor | None = None
+        if indexer is not None:
+            self._branch_monitor = BranchMonitor(root, job_manager, indexer, db)
+            self._branch_monitor.start()
+
+    def _is_sql(self, path: str | bytes) -> bool:
+        """Check if a path is a non-ignored SQL file.
+
+        Args:
+            path: File path to check (str or bytes from watchdog)
+
+        Returns:
+            True if the path is a .sql file not matching ignore patterns
+        """
+        if isinstance(path, bytes):
+            path = path.decode("utf-8")
+        path_obj = Path(path)
+        return path_obj.suffix == ".sql" and not is_ignored(path_obj, self._root, self._spec)
+
+    def on_modified(self, event):
+        """Handle file modification events.
+
+        Args:
+            event: Watchdog event object
+        """
+        if not event.is_directory and self._is_sql(event.src_path):
+            self._jobs.schedule(event.src_path)
+
+    def on_created(self, event):
+        """Handle file creation events.
+
+        Args:
+            event: Watchdog event object
+        """
+        if not event.is_directory and self._is_sql(event.src_path):
+            self._jobs.schedule(event.src_path)
+
+    def on_moved(self, event):
+        """Handle file move events.
+
+        Args:
+            event: Watchdog event object
+        """
+        if not event.is_directory and self._is_sql(event.dest_path):
+            self._jobs.schedule(event.dest_path)
+
+    def on_deleted(self, event):
+        """Handle file deletion events.
+
+        Args:
+            event: Watchdog event object
+        """
+        if not event.is_directory and self._is_sql(event.src_path):
+            self._db.delete_nodes_for_file(event.src_path)
+
+
+class BranchMonitor(threading.Thread):
+    """Background thread that detects branch changes and triggers full resyncs.
+
+    Polls `git rev-parse --abbrev-ref HEAD` every 2 seconds. When the branch
+    changes, pauses the job manager, runs a full reindex, then resumes and
+    drains queued file events.
+    """
+
+    def __init__(
+        self, watched_path: Path, job_manager, indexer, db, _poll_interval: float = 2.0
+    ):
+        """Initialize the branch monitor.
+
+        Args:
+            watched_path: Path being watched (used to find git root)
+            job_manager: WatchJobManager instance
+            indexer: Indexer instance
+            db: GraphBackend instance
+            _poll_interval: Polling interval in seconds (for testing)
+        """
+        # daemon=False ensures that if index_repo() is in-flight when shutdown is requested,
+        # the process will wait (via join(timeout=5) in watch.py) up to 5 seconds before exiting.
+        # This avoids data loss from killing an in-progress resync.
+        super().__init__(daemon=False)
+        self._watched_path = watched_path
+        self._job_manager = job_manager
+        self._indexer = indexer
+        self._db = db
+        self._stop_event = threading.Event()
+        self._current_branch: str | None = None
+        self._poll_interval = _poll_interval
+
+    def run(self) -> None:
+        """Poll git branch and trigger resync on change."""
+        while not self._stop_event.is_set():
+            try:
+                branch = self._get_current_branch()
+                if branch is not None and branch != self._current_branch:
+                    logger.debug(
+                        "Branch change detected: %s -> %s", self._current_branch, branch
+                    )
+                    self._current_branch = branch
+                    self._on_branch_change()
+            except subprocess.CalledProcessError:
+                # Not a git repo or git not available
+                logger.debug("Could not get current branch (not a git repo or git unavailable)")
+                self._stop_event.set()
+                break
+            except Exception as exc:
+                logger.debug("BranchMonitor error: %s", exc)
+
+            # Sleep in small increments to allow quick shutdown
+            for _ in range(int(self._poll_interval * 10)):
+                if self._stop_event.is_set():
+                    break
+                time.sleep(0.1)
+
+    def _get_current_branch(self) -> str | None:
+        """Get the current git branch name.
+
+        Returns:
+            Branch name, or None if git command fails
+
+        Raises:
+            subprocess.CalledProcessError if git command fails
+        """
+        result = subprocess.run(
+            ["git", "rev-parse", "--abbrev-ref", "HEAD"],
+            cwd=str(self._watched_path),
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        return result.stdout.strip()
+
+    def _on_branch_change(self) -> None:
+        """Handle branch change: pause, resync, resume, drain queue."""
+        # Pause new file events
+        self._job_manager.set_paused(True)
+
+        # Cancel pending file timers
+        self._job_manager.cancel_all()
+
+        # Run full resync
+        try:
+            self._indexer.index_repo(self._watched_path, dialect=None, db=self._db)
+        except Exception as exc:
+            logger.error("Branch change resync failed: %s", exc)
+        finally:
+            # Resume and drain queued events
+            self._job_manager.set_paused(False)
+            self._job_manager.drain_queued()
+
+    def stop(self) -> None:
+        """Signal the thread to stop."""
+        self._stop_event.set()
+
+    def join(self, timeout: float | None = None) -> None:
+        """Wait for the thread to stop.
+
+        Args:
+            timeout: Maximum time to wait in seconds
+        """
+        super().join(timeout=timeout)

sqlcg/lineage/__init__.py

@@ -0,0 +1 @@
+"""Lineage resolution and analysis module."""

sqlcg/lineage/aggregator.py

@@ -0,0 +1,58 @@
+"""Cross-file lineage aggregator for two-pass resolution."""
+
+from sqlcg.parsers.base import ParsedFile
+from sqlcg.utils.logging import getLogger
+
+logger = getLogger(__name__)
+
+
+class CrossFileAggregator:
+    """Aggregates parsed files for cross-file resolution.
+
+    Pass 1: register all parsed files and build view/table source mappings.
+    Pass 2: re-parse files with cross-file schema context.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the aggregator."""
+        # Maps table.full_id -> ParsedFile that defines it
+        self.sources: dict[str, ParsedFile] = {}
+
+    def register_pass1(self, parsed: ParsedFile) -> None:
+        """Register a pass-1 result and build view/table source map.
+
+        Args:
+            parsed: ParsedFile from pass 1
+        """
+        for table in parsed.defined_tables:
+            self.sources[table.full_id] = parsed
+
+    def resolve_pass2(self, parser, parsed: ParsedFile) -> ParsedFile:
+        """Re-parse with cross-file schema context.
+
+        Args:
+            parser: SqlParser instance
+            parsed: ParsedFile from pass 1
+
+        Returns:
+            ParsedFile from pass 2 with resolved cross-file references,
+            or the pass-1 result if the file cannot be re-read.
+
+        Raises:
+            No exceptions are raised; file read errors are logged as WARNING
+            and the pass-1 result is returned unchanged.
+        """
+        # Register view sources for schema resolution
+        parser._schema.add_view_sources(self.sources)
+
+        try:
+            sql = parsed.path.read_text(encoding="utf-8")
+        except (FileNotFoundError, OSError) as exc:
+            logger.warning(
+                "resolve_pass2: cannot re-read %s (%s) — returning pass-1 result",
+                parsed.path,
+                exc,
+            )
+            return parsed
+
+        return parser.parse_file(parsed.path, sql)