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

sqlcg/parsers/postgres_parser.py

@@ -0,0 +1,27 @@
+"""PostgreSQL SQL parser."""
+
+from sqlcg.lineage.schema_resolver import SchemaResolver
+from sqlcg.parsers.ansi_parser import AnsiParser
+from sqlcg.parsers.registry import register
+from sqlcg.utils.logging import getLogger
+
+logger = getLogger(__name__)
+
+
+@register("postgres")
+class PostgresParser(AnsiParser):
+    """PostgreSQL SQL parser.
+
+    Uses standard ANSI parsing with Postgres dialect for v1.
+    No special handling for scripting blocks in v1.
+    """
+
+    DIALECT: str | None = "postgres"
+
+    def __init__(self, schema_resolver: SchemaResolver):
+        """Initialize Postgres parser.
+
+        Args:
+            schema_resolver: SchemaResolver instance for table/column lookups
+        """
+        super().__init__(schema_resolver)

sqlcg/parsers/registry.py

@@ -0,0 +1,46 @@
+"""Parser registry and factory for dialect-specific SQL parsers."""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from sqlcg.lineage.schema_resolver import SchemaResolver
+    from sqlcg.parsers.base import SqlParser
+
+# Global registry of dialect -> parser class mapping
+PARSERS: dict[str | None, type["SqlParser"]] = {}
+
+
+def register(dialect: str | None):
+    """Decorator to register a parser class for a dialect.
+
+    Args:
+        dialect: SQL dialect identifier (None for ANSI, "snowflake", etc.)
+
+    Returns:
+        Decorator function
+    """
+
+    def decorator(cls: type["SqlParser"]) -> type["SqlParser"]:
+        PARSERS[dialect] = cls
+        return cls
+
+    return decorator
+
+
+def get_parser(dialect: str | None, schema_resolver: "SchemaResolver") -> "SqlParser":
+    """Get a parser instance for the given dialect.
+
+    Args:
+        dialect: SQL dialect identifier (None for ANSI, "snowflake", etc.)
+        schema_resolver: SchemaResolver instance for table/column lookups
+
+    Returns:
+        SqlParser instance for the given dialect
+
+    Raises:
+        ValueError: If no parser is registered for the dialect
+    """
+    cls = PARSERS.get(dialect) or PARSERS.get(None)
+    if cls is None:
+        raise ValueError(f"No parser registered for dialect {dialect!r}")
+    return cls(schema_resolver)

sqlcg/parsers/snowflake_parser.py

@@ -0,0 +1,148 @@
+"""Snowflake SQL parser with scripting block detection and DML extraction."""
+
+import re
+from pathlib import Path
+from typing import Any
+
+import sqlglot
+
+from sqlcg.lineage.schema_resolver import SchemaResolver
+from sqlcg.parsers.ansi_parser import AnsiParser
+from sqlcg.parsers.base import ParsedFile
+from sqlcg.parsers.registry import register
+from sqlcg.utils.logging import getLogger
+
+logger = getLogger(__name__)
+
+# Regex for detecting scripting blocks (BEGIN/IF/LOOP)
+# Used as fallback when tokenization fails
+_SCRIPTING_BLOCK = re.compile(r"\bBEGIN\b", re.IGNORECASE)
+
+# Regex for extracting DML statements from scripting blocks.
+# Does not handle ';' inside string literals — tokenizer-based extraction deferred to v2.
+_EMBEDDED_DML = re.compile(
+    r"(SELECT\s+.+?(?=;|\Z)|INSERT\s+INTO.+?(?=;|\Z)|UPDATE\s+.+?(?=;|\Z)|DELETE\s+.+?(?=;|\Z))",
+    re.DOTALL | re.IGNORECASE | re.MULTILINE,
+)
+
+
+@register("snowflake")
+class SnowflakeParser(AnsiParser):
+    """Snowflake SQL parser with scripting block handling.
+
+    Handles Snowflake-specific features:
+    - Token-aware scripting block detection (avoids false-positives)
+    - DML extraction from scripting blocks
+    - Colon-qualified identifiers (Gap 1)
+    - LATERAL FLATTEN operations (Gap 2)
+    - Dynamic identifiers (Gap 3)
+    """
+
+    DIALECT: str | None = "snowflake"
+
+    def __init__(self, schema_resolver: SchemaResolver):
+        """Initialize Snowflake parser.
+
+        Args:
+            schema_resolver: SchemaResolver instance for table/column lookups
+        """
+        super().__init__(schema_resolver)
+
+    def parse_file(self, path: Path, sql: str) -> ParsedFile:
+        """Parse Snowflake SQL file with scripting block detection.
+
+        Args:
+            path: Path to the source file
+            sql: SQL text to parse
+
+        Returns:
+            ParsedFile with parsed statements and metadata
+        """
+        # Check for scripting blocks
+        if self._has_scripting_block(sql):
+            logger.info("Snowflake scripting block detected in %s, using DML extraction", path)
+            return self._parse_scripting_file(path, sql)
+
+        # Otherwise use standard ANSI parsing with Snowflake dialect
+        return AnsiParser.parse_file(self, path, sql)  # type: ignore
+
+    def _has_scripting_block(self, sql: str) -> bool:
+        """Token-aware BEGIN detection — avoids false-positives on string literals and comments.
+
+        Args:
+            sql: SQL text to check
+
+        Returns:
+            True if a scripting block is detected
+        """
+        try:
+            from sqlglot.tokens import Tokenizer, TokenType  # type: ignore
+
+            toks = Tokenizer.from_dialect("snowflake").tokenize(sql)  # type: ignore
+            return any(t.token_type == TokenType.BEGIN for t in toks)  # type: ignore
+        except Exception:
+            # Fallback to regex if tokenization fails
+            return bool(_SCRIPTING_BLOCK.search(sql))
+
+    def _parse_scripting_file(self, path: Path, sql: str) -> ParsedFile:
+        """Parse a Snowflake file with scripting blocks using DML extraction.
+
+        Args:
+            path: Path to the source file
+            sql: SQL text to parse
+
+        Returns:
+            ParsedFile with extracted DML statements
+        """
+        out = ParsedFile(path=path, dialect=self.DIALECT)
+        out.errors.append("parse_mode:scripting_block")
+
+        # Extract DML statements using regex
+        dml_matches = _EMBEDDED_DML.finditer(sql)
+        stmt_index = 0
+
+        for match in dml_matches:
+            dml_sql = match.group(1).strip()
+            if not dml_sql:
+                continue
+
+            try:
+                # Try to parse the extracted DML
+                statements = sqlglot.parse(dml_sql, dialect=self.DIALECT)
+                for stmt in statements:
+                    if stmt is None:
+                        continue
+
+                    try:
+                        # Call parent's _parse_statement method
+                        query_node: Any = AnsiParser._parse_statement(  # type: ignore
+                            self, stmt, path, stmt_index
+                        )
+                        # Mark as parse_failed since we're in scripting mode
+                        query_node.parse_failed = True
+                        query_node.confidence = 0.3
+                        query_node.parsing_mode = "scripting"
+                        out.statements.append(query_node)
+                        stmt_index += 1
+
+                        # Track table references
+                        if query_node.kind in ("CREATE_TABLE", "CREATE_VIEW"):
+                            if query_node.target:
+                                out.defined_tables.append(query_node.target)
+                        out.referenced_tables.extend(query_node.sources)
+
+                    except Exception as exc:
+                        logger.warning(
+                            "Failed to process extracted DML statement %d in %s: %s",
+                            stmt_index,
+                            path,
+                            exc,
+                        )
+                        out.errors.append(f"statement_error:{stmt_index}:{exc}")
+                        stmt_index += 1
+
+            except Exception as exc:
+                logger.warning("Failed to parse extracted DML from %s: %s", path, exc)
+                out.errors.append(f"dml_extraction_error:{exc}")
+
+        return out

sqlcg/parsers/tsql_parser.py

@@ -0,0 +1,27 @@
+"""T-SQL (Microsoft SQL Server) parser."""
+
+from sqlcg.lineage.schema_resolver import SchemaResolver
+from sqlcg.parsers.ansi_parser import AnsiParser
+from sqlcg.parsers.registry import register
+from sqlcg.utils.logging import getLogger
+
+logger = getLogger(__name__)
+
+
+@register("tsql")
+class TsqlParser(AnsiParser):
+    """T-SQL (Microsoft SQL Server) parser.
+
+    Uses standard ANSI parsing with T-SQL dialect for v1.
+    No special handling for scripting blocks in v1.
+    """
+
+    DIALECT: str | None = "tsql"
+
+    def __init__(self, schema_resolver: SchemaResolver):
+        """Initialize T-SQL parser.
+
+        Args:
+            schema_resolver: SchemaResolver instance for table/column lookups
+        """
+        super().__init__(schema_resolver)

sqlcg/server/__init__.py

@@ -0,0 +1 @@
+"""MCP server module for sqlcg."""

sqlcg/server/exceptions.py

@@ -0,0 +1,20 @@
+"""Exceptions raised by MCP server tools."""
+
+
+class NotIndexedError(RuntimeError):
+    """Raised when graph has no indexed repos.
+
+    This error indicates that no repositories have been indexed yet.
+    Users should run `sqlcg index <path>` first to populate the graph.
+    """
+
+    pass
+
+
+class InvalidColumnRefError(ValueError):
+    """Raised for invalid column reference format.
+
+    Expected format: "table.column" or "catalog.db.table.column".
+    """
+
+    pass

sqlcg/server/models.py

@@ -0,0 +1,83 @@
+"""Pydantic models for MCP tool return types."""
+
+from pydantic import BaseModel, Field
+
+
+class LineageNode(BaseModel):
+    """Node in a lineage graph."""
+
+    name: str = Field(..., description="Name of the node (table or column)")
+    kind: str = Field(..., description="Kind of node (table, column, query, etc.)")
+    file: str | None = Field(None, description="Source file path, if applicable")
+    confidence: float | None = Field(None, description="Confidence score 0.0-1.0")
+
+
+class LineageResult(BaseModel):
+    """Result of trace_column_lineage query."""
+
+    column: str = Field(..., description="Column reference (table.column)")
+    lineage: list[LineageNode] = Field(
+        default_factory=list, description="List of nodes in the lineage"
+    )
+
+
+class TableUsage(BaseModel):
+    """Usage of a table in a query."""
+
+    query_file: str = Field(..., description="File path where query is defined")
+    sql: str | None = Field(None, description="SQL of the query")
+    kind: str | None = Field(None, description="Kind of query (SELECT, INSERT, etc.)")
+
+
+class TableUsageResult(BaseModel):
+    """Result of find_table_usages query."""
+
+    table: str = Field(..., description="Table name")
+    usages: list[TableUsage] = Field(default_factory=list, description="List of usages")
+
+
+class DependencyNode(BaseModel):
+    """Node in a dependency graph."""
+
+    name: str = Field(..., description="Name of the node")
+    kind: str = Field(..., description="Kind of node (table, column, etc.)")
+
+
+class DependencyResult(BaseModel):
+    """Result of dependency traversal queries."""
+
+    root: str = Field(..., description="Root column or table")
+    nodes: list[DependencyNode] = Field(default_factory=list, description="List of dependent nodes")
+
+
+class SqlPatternMatch(BaseModel):
+    """Match for a SQL pattern search."""
+
+    file: str = Field(..., description="File path containing the match")
+    sql: str = Field(..., description="SQL text of the match")
+    kind: str | None = Field(None, description="Kind of statement")
+
+
+class SqlPatternResult(BaseModel):
+    """Result of search_sql_pattern query."""
+
+    pattern: str = Field(..., description="Pattern searched for")
+    matches: list[SqlPatternMatch] = Field(
+        default_factory=list, description="List of matching queries"
+    )
+
+
+class DialectRepo(BaseModel):
+    """Repository with dialect information."""
+
+    path: str = Field(..., description="Repository path")
+    name: str | None = Field(None, description="Repository name")
+    dialects: list[str] = Field(default_factory=list, description="Dialects used in this repo")
+
+
+class DialectRepoResult(BaseModel):
+    """Result of list_dialects_and_repos query."""
+
+    repos: list[DialectRepo] = Field(
+        default_factory=list, description="List of indexed repositories"
+    )

sqlcg/server/server.py

@@ -0,0 +1,57 @@
+"""MCP server for SQL Code Graph.
+
+Exposes FastMCP tools for lineage queries, pattern search, and indexing.
+MCP protocol uses stdout for message transport, so this module redirects
+stdout to stderr to prevent user logs from corrupting the protocol stream.
+"""
+
+import sys
+
+from dotenv import load_dotenv
+from mcp.server import FastMCP
+
+from sqlcg.utils.logging import getLogger
+
+logger = getLogger(__name__)
+
+
+def _configure_mcp_logging() -> None:
+    """Redirect sys.stdout to sys.stderr to protect MCP protocol.
+
+    MCP uses stdout for JSON-RPC messages. Any user print() or log output
+    to stdout would corrupt the protocol. This function must be called before
+    mcp.run() and before any code that might print to stdout.
+    """
+    sys.stdout = sys.stderr
+
+
+# Protect stdout before importing FastMCP (which may emit output during import)
+_configure_mcp_logging()
+
+# Create FastMCP instance at module scope so tools.py can import and register with it
+mcp = FastMCP("SQL Code Graph")
+
+
+def main(db_path: str | None = None) -> None:
+    """Start the MCP server.
+
+    Args:
+        db_path: Path to KùzuDB database. If None, uses SQLCG_DB_PATH env var
+                or ~/.sqlcg/graph.db (via get_db_path in tools module).
+
+    Raises:
+        RuntimeError: If tools fail to initialize or FastMCP server fails.
+    """
+    load_dotenv()
+
+    # Import tools module to trigger tool registration via @mcp.tool() decorators
+    import sqlcg.server.tools
+
+    # Initialize the backend singleton used by all tools
+    sqlcg.server.tools.init_backend(db_path)
+
+    # Run the MCP server event loop, ensuring backend is closed on shutdown
+    try:
+        mcp.run()
+    finally:
+        sqlcg.server.tools.shutdown_backend()