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

sqlcg/core/neo4j_backend.py

@@ -0,0 +1,195 @@
+"""Neo4j implementation of GraphBackend."""
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import Any
+
+from sqlcg.core.graph_db import GraphBackend
+from sqlcg.core.queries import (
+    DELETE_COLUMNS_FOR_FILE,
+    DELETE_FILE,
+    DELETE_QUERIES_FOR_FILE,
+    DELETE_TABLES_FOR_FILE,
+)
+from sqlcg.core.schema import NODE_COLUMN, NODE_FILE, NODE_QUERY, NODE_REPO, NODE_TABLE
+from sqlcg.utils.logging import getLogger
+
+logger = getLogger(__name__)
+
+try:
+    from neo4j import GraphDatabase as _GraphDatabase
+
+    GraphDatabase = _GraphDatabase
+    NEO4J_AVAILABLE = True
+except ImportError:
+    GraphDatabase = None  # type: ignore[assignment,misc]
+    NEO4J_AVAILABLE = False
+
+
+class Neo4jBackend(GraphBackend):
+    """Neo4j implementation of the graph database backend."""
+
+    def __init__(self, uri: str, user: str, password: str):
+        """Initialize Neo4j backend.
+
+        Args:
+            uri: Neo4j connection URI (e.g., "bolt://localhost:7687")
+            user: Neo4j username
+            password: Neo4j password
+
+        Raises:
+            ImportError: If the neo4j package is not installed
+        """
+        if not NEO4J_AVAILABLE:
+            raise ImportError(
+                "neo4j package is not installed. "
+                "Install it with: pip install 'sql-code-graph[neo4j]'"
+            )
+
+        self._driver = GraphDatabase.driver(uri, auth=(user, password))
+        self._session = self._driver.session()
+
+    def init_schema(self) -> None:
+        """Initialize the database schema if not already present.
+
+        Creates indexes and constraints for efficient querying.
+        """
+        # IF NOT EXISTS already ensures idempotency; APOC utilities add no safety benefit here.
+        indexes = [
+            f"CREATE INDEX idx_repo_path IF NOT EXISTS FOR (r:{NODE_REPO}) ON (r.path)",
+            f"CREATE INDEX idx_file_path IF NOT EXISTS FOR (f:{NODE_FILE}) ON (f.path)",
+            f"CREATE INDEX idx_table_qualified IF NOT EXISTS FOR (t:{NODE_TABLE}) ON (t.qualified)",
+            f"CREATE INDEX idx_column_id IF NOT EXISTS FOR (c:{NODE_COLUMN}) ON (c.id)",
+            f"CREATE INDEX idx_query_id IF NOT EXISTS FOR (q:{NODE_QUERY}) ON (q.id)",
+        ]
+
+        for index_query in indexes:
+            try:
+                self._session.run(index_query)
+                logger.debug(f"Created index: {index_query[:50]}...")
+            except Exception as e:
+                logger.warning(f"Index creation skipped: {e}")
+
+    def upsert_node(self, label: str, key: str, properties: dict[str, Any]) -> None:
+        """Upsert a node with the given label and properties."""
+        # Validate property keys to prevent Cypher injection
+        self._validate_props(properties)
+
+        pk_field = self._pk_field(label)
+        query = f"MERGE (n:{label} {{{pk_field}: $key}}) SET n += $props"
+        try:
+            self._session.run(query, {"key": key, "props": properties})
+        except Exception as e:
+            logger.error(f"upsert_node failed: {label} {key}: {e}")
+            raise
+
+    def upsert_edge(
+        self,
+        src_label: str,
+        src_key: str,
+        dst_label: str,
+        dst_key: str,
+        rel_type: str,
+        properties: dict[str, Any],
+    ) -> None:
+        """Upsert a relationship between two nodes."""
+        # Validate property keys to prevent Cypher injection
+        self._validate_props(properties)
+
+        src_pk = self._pk_field(src_label)
+        dst_pk = self._pk_field(dst_label)
+        query = (
+            f"MATCH (src:{src_label} {{{src_pk}: $src_key}})"
+            f" MATCH (dst:{dst_label} {{{dst_pk}: $dst_key}})"
+            f" MERGE (src)-[r:{rel_type}]->(dst)"
+            " SET r += $props"
+        )
+        try:
+            self._session.run(query, {"src_key": src_key, "dst_key": dst_key, "props": properties})
+        except Exception as e:
+            logger.error(f"upsert_edge failed: {src_label} -> {rel_type} -> {dst_label}: {e}")
+            raise
+
+    def run_read(self, query: str, params: dict[str, Any]) -> list[dict[str, Any]]:
+        """Execute a read-only query and return results."""
+        try:
+            result = self._session.run(query, params)
+            rows = [dict(record) for record in result]
+            return rows
+        except Exception as e:
+            logger.error(f"run_read failed: {e}")
+            raise
+
+    def run_write(self, query: str, params: dict[str, Any]) -> None:
+        """Execute a write query (mutation)."""
+        try:
+            self._session.run(query, params)
+        except Exception as e:
+            logger.error(f"run_write failed: {e}")
+            raise
+
+    def delete_nodes_for_file(self, file_path: str) -> None:
+        """Delete all nodes and relationships associated with a file."""
+        params = {"path": file_path}
+        try:
+            # Step A: Delete SqlColumn nodes for tables defined in this file
+            self._session.run(DELETE_COLUMNS_FOR_FILE, params)
+            # Step B: Delete SqlQuery nodes
+            self._session.run(DELETE_QUERIES_FOR_FILE, params)
+            # Step C: Delete SqlTable nodes defined in this file
+            self._session.run(DELETE_TABLES_FOR_FILE, params)
+            # Step D: Delete the File node itself
+            self._session.run(DELETE_FILE, params)
+            logger.debug(f"Deleted all nodes for {file_path}")
+        except Exception as e:
+            logger.error(f"delete_nodes_for_file failed for {file_path}: {e}")
+            raise
+
+    def get_schema_version(self) -> str | None:
+        """Get the stored schema version from the database.
+
+        Returns:
+            The schema version string, or None if not set.
+        """
+        try:
+            result = self.run_read(
+                "MATCH (v:SchemaVersion) RETURN v.version AS version LIMIT 1", {}
+            )
+            return result[0]["version"] if result else None
+        except Exception as e:
+            logger.warning(f"Failed to read schema version: {e}")
+            return None
+
+    def close(self) -> None:
+        """Close the database connection."""
+        try:
+            self._session.close()
+            self._driver.close()
+            logger.debug("Neo4jBackend connection closed")
+        except Exception as e:
+            logger.error(f"Error closing Neo4jBackend: {e}")
+            raise
+
+    @contextmanager
+    def transaction(self) -> Iterator["Neo4jBackend"]:
+        """Context manager for Neo4j transactions.
+
+        Creates a fresh session per transaction to avoid issues with shared
+        long-lived sessions that may be closed externally.
+
+        Yields:
+            self (the Neo4jBackend instance)
+
+        Raises:
+            Any exception raised in the context triggers ROLLBACK.
+        """
+        session = self._driver.session()
+        tx = session.begin_transaction()
+        try:
+            yield self
+            tx.commit()
+        except Exception:
+            tx.rollback()
+            raise
+        finally:
+            session.close()

sqlcg/core/queries.py

@@ -0,0 +1,82 @@
+"""Centralized Cypher query strings for graph operations."""
+
+from sqlcg.core.schema import NodeLabel, RelType
+
+# Scope is bounded by exact path match; APOC procedures are not required.
+# Delete Column nodes for tables defined in a file
+DELETE_COLUMNS_FOR_FILE = (
+    f"MATCH (f:{NodeLabel.FILE} {{path: $path}})"
+    f"<-[:{RelType.DEFINED_IN}]-(t:{NodeLabel.TABLE})"
+    f"-[:{RelType.HAS_COLUMN}]->(c:{NodeLabel.COLUMN})"
+    " DETACH DELETE c"
+)
+
+# Delete Query nodes and their edges
+DELETE_QUERIES_FOR_FILE = (
+    f"MATCH (f:{NodeLabel.FILE} {{path: $path}})"
+    f"<-[:{RelType.QUERY_DEFINED_IN}]-(q:{NodeLabel.QUERY})"
+    " DETACH DELETE q"
+)
+
+# Delete Table nodes defined in a file
+DELETE_TABLES_FOR_FILE = (
+    f"MATCH (f:{NodeLabel.FILE} {{path: $path}})"
+    f"<-[:{RelType.DEFINED_IN}]-(t:{NodeLabel.TABLE})"
+    " DETACH DELETE t"
+)
+
+# Delete the File node itself
+DELETE_FILE = f"MATCH (f:{NodeLabel.FILE} {{path: $path}}) DETACH DELETE f"
+
+# Find views that depend on tables defined in a file
+STALE_VIEWS_QUERY = (
+    f"MATCH (f:{NodeLabel.FILE} {{path: $path}})"
+    f"<-[:{RelType.DEFINED_IN}]-(t:{NodeLabel.TABLE})"
+    f"<-[:{RelType.SELECTS_FROM}]-(q:{NodeLabel.QUERY})"
+    f"-[:{RelType.DECLARES}]->(v:{NodeLabel.TABLE} {{kind: 'VIEW'}})"
+    " RETURN DISTINCT v.qualified AS view_name"
+)
+
+# Get all files in a repo by path prefix
+INDEX_REPO_FILES_QUERY = (
+    "MATCH (f:File) WHERE f.path STARTS WITH $repo_prefix RETURN f.path AS path"
+)
+
+# Trace upstream lineage of a column
+TRACE_COLUMN_LINEAGE_QUERY = (
+    "MATCH (dst:SqlColumn {id: $id})<-[:COLUMN_LINEAGE]-(src:SqlColumn) "
+    "RETURN src.id AS id, src.col_name AS col_name"
+)
+
+# Find table usages in queries
+FIND_TABLE_USAGES_QUERY = (
+    "MATCH (t:SqlTable {name: $name})<-[:SELECTS_FROM]-(q:SqlQuery)"
+    "-[:QUERY_DEFINED_IN]->(f:File) "
+    "RETURN f.path AS file, q.sql AS sql, q.kind AS kind"
+)
+
+# Get downstream column dependencies
+GET_DOWNSTREAM_DEPENDENCIES_QUERY = (
+    "MATCH (src:SqlColumn {id: $id})-[:COLUMN_LINEAGE]->(dst:SqlColumn) "
+    "RETURN dst.id AS id, dst.col_name AS col_name"
+)
+
+# Get upstream column dependencies
+GET_UPSTREAM_DEPENDENCIES_QUERY = (
+    "MATCH (dst:SqlColumn {id: $id})<-[:COLUMN_LINEAGE]-(src:SqlColumn) "
+    "RETURN src.id AS id, src.col_name AS col_name"
+)
+
+# Search SQL patterns in indexed queries
+SEARCH_SQL_PATTERN_QUERY = (
+    "MATCH (q:SqlQuery)-[:QUERY_DEFINED_IN]->(f:File) "
+    "WHERE contains(q.sql, $query) "
+    "RETURN f.path AS file, q.sql AS sql, q.kind AS kind "
+    "LIMIT $limit"
+)
+
+# List dialects and repos
+LIST_DIALECTS_AND_REPOS_QUERY = (
+    "MATCH (r:Repo)<-[:BELONGS_TO]-(f:File) "
+    "RETURN r.path AS path, r.name AS name, collect(DISTINCT f.dialect) AS dialects"
+)

sqlcg/core/schema.cypher

@@ -0,0 +1,104 @@
+-- Repo node: one per indexed repository
+CREATE NODE TABLE Repo (
+    path STRING PRIMARY KEY,
+    name STRING
+);
+
+-- File node: one per .sql file
+CREATE NODE TABLE File (
+    path STRING PRIMARY KEY,
+    repo_path STRING,
+    sha STRING,
+    dialect STRING
+);
+
+-- Table node: one per unique table reference
+CREATE NODE TABLE SqlTable (
+    qualified STRING PRIMARY KEY,
+    catalog STRING,
+    db STRING,
+    name STRING,
+    kind STRING,
+    defined_in_file STRING
+);
+
+-- Column node: one per unique column reference
+CREATE NODE TABLE SqlColumn (
+    id STRING PRIMARY KEY,
+    catalog STRING,
+    db STRING,
+    table_name STRING,
+    col_name STRING,
+    table_qualified STRING
+);
+
+-- Query node: one per SQL statement parsed
+CREATE NODE TABLE SqlQuery (
+    id STRING PRIMARY KEY,
+    file_path STRING,
+    statement_index INT64,
+    sql STRING,
+    kind STRING,
+    target_table STRING,
+    parse_failed BOOLEAN,
+    confidence FLOAT,
+    parsing_mode STRING
+);
+
+-- File -> Repo: file belongs to this repository
+CREATE REL TABLE BELONGS_TO (
+    FROM File TO Repo
+);
+
+-- File -> Table: table is defined in this file
+CREATE REL TABLE DEFINED_IN (
+    FROM SqlTable TO File
+);
+
+-- Query -> File: query is defined in this file
+CREATE REL TABLE QUERY_DEFINED_IN (
+    FROM SqlQuery TO File
+);
+
+-- Table -> Column: table has this column
+CREATE REL TABLE HAS_COLUMN (
+    FROM SqlTable TO SqlColumn
+);
+
+-- Query -> Table: query selects from table
+CREATE REL TABLE SELECTS_FROM (
+    FROM SqlQuery TO SqlTable
+);
+
+-- Query -> Table: query inserts into table
+CREATE REL TABLE INSERTS_INTO (
+    FROM SqlQuery TO SqlTable
+);
+
+-- Query -> Table: query deletes from table
+CREATE REL TABLE DELETES_FROM (
+    FROM SqlQuery TO SqlTable
+);
+
+-- Query -> Table: query updates table
+CREATE REL TABLE UPDATES (
+    FROM SqlQuery TO SqlTable
+);
+
+-- Column -> Column: lineage relationship
+CREATE REL TABLE COLUMN_LINEAGE (
+    FROM SqlColumn TO SqlColumn,
+    transform STRING,
+    confidence FLOAT,
+    query_id STRING
+);
+
+-- Query -> Table: query declares/creates this table
+CREATE REL TABLE DECLARES (
+    FROM SqlQuery TO SqlTable
+);
+
+-- Schema version tracking
+CREATE NODE TABLE SchemaVersion (
+    version STRING PRIMARY KEY
+);

sqlcg/core/schema.py

@@ -0,0 +1,48 @@
+"""KùzuDB schema definition for sqlcg graph."""
+
+from enum import StrEnum
+from importlib.resources import files
+
+SCHEMA_VERSION = "1"
+
+
+class NodeLabel(StrEnum):
+    REPO = "Repo"
+    FILE = "File"
+    TABLE = "SqlTable"
+    COLUMN = "SqlColumn"
+    QUERY = "SqlQuery"
+    SCHEMA_VERSION = "SchemaVersion"
+
+
+class RelType(StrEnum):
+    BELONGS_TO = "BELONGS_TO"
+    DEFINED_IN = "DEFINED_IN"
+    QUERY_DEFINED_IN = "QUERY_DEFINED_IN"
+    HAS_COLUMN = "HAS_COLUMN"
+    SELECTS_FROM = "SELECTS_FROM"
+    INSERTS_INTO = "INSERTS_INTO"
+    DELETES_FROM = "DELETES_FROM"
+    UPDATES = "UPDATES"
+    COLUMN_LINEAGE = "COLUMN_LINEAGE"
+    DECLARES = "DECLARES"
+
+
+# Backward-compatible aliases
+NODE_REPO = NodeLabel.REPO
+NODE_FILE = NodeLabel.FILE
+NODE_TABLE = NodeLabel.TABLE
+NODE_COLUMN = NodeLabel.COLUMN
+NODE_QUERY = NodeLabel.QUERY
+NODE_SCHEMA_VERSION = NodeLabel.SCHEMA_VERSION
+
+REL_DEFINED_IN = RelType.DEFINED_IN
+REL_HAS_COLUMN = RelType.HAS_COLUMN
+REL_SELECTS_FROM = RelType.SELECTS_FROM
+REL_INSERTS_INTO = RelType.INSERTS_INTO
+REL_DELETES_FROM = RelType.DELETES_FROM
+REL_UPDATES = RelType.UPDATES
+REL_COLUMN_LINEAGE = RelType.COLUMN_LINEAGE
+REL_DECLARES = RelType.DECLARES
+
+SCHEMA_DDL: str = files("sqlcg.core").joinpath("schema.cypher").read_text(encoding="utf-8")

sqlcg/indexer/__init__.py

@@ -0,0 +1 @@
+"""Indexer module for walking and parsing SQL files."""

sqlcg/indexer/dbt_adapter.py

@@ -0,0 +1,23 @@
+"""dbt manifest adapter for schema resolution."""
+
+from pathlib import Path
+
+from sqlcg.lineage.schema_resolver import SchemaResolver
+from sqlcg.utils.logging import getLogger
+
+logger = getLogger(__name__)
+
+
+def load_dbt_manifest(manifest_path: Path, schema_resolver: SchemaResolver) -> None:
+    """Load dbt manifest and register table schemas.
+
+    Errors are logged, not raised.
+
+    Args:
+        manifest_path: Path to dbt manifest.json
+        schema_resolver: SchemaResolver instance to populate
+    """
+    try:
+        schema_resolver.add_dbt_manifest(manifest_path)
+    except Exception as exc:
+        logger.warning("Failed to load dbt manifest %s: %s", manifest_path, exc)