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

sqlcg/parsers/ansi_parser.py

@@ -0,0 +1,215 @@
+"""ANSI SQL parser implementation."""
+
+from pathlib import Path
+from typing import Any
+
+import sqlglot
+import sqlglot.expressions as exp
+
+from sqlcg.lineage.schema_resolver import SchemaResolver
+from sqlcg.parsers.base import ParsedFile, QueryNode, SqlParser, TableRef
+from sqlcg.parsers.registry import register
+from sqlcg.utils.logging import getLogger
+
+logger = getLogger(__name__)
+
+
+@register(None)  # None = ANSI/default dialect
+class AnsiParser(SqlParser):
+    """ANSI SQL parser for standard SQL dialects.
+
+    Attributes:
+        DIALECT: Set to None for ANSI/default dialect
+    """
+
+    DIALECT: str | None = None
+
+    def __init__(self, schema_resolver: SchemaResolver):
+        """Initialize ANSI parser.
+
+        Args:
+            schema_resolver: SchemaResolver instance for table/column lookups
+        """
+        super().__init__(schema_resolver)
+
+    def parse_file(self, path: Path, sql: str) -> ParsedFile:
+        """Parse SQL file and extract table/column lineage.
+
+        Args:
+            path: Path to the source file
+            sql: SQL text to parse
+
+        Returns:
+            ParsedFile with parsed statements and metadata
+        """
+        out = ParsedFile(path=path, dialect=self.DIALECT)
+
+        # Parse all statements in the file
+        try:
+            statements = sqlglot.parse(sql, dialect=self.DIALECT)
+        except Exception as exc:
+            logger.warning("Failed to parse file %s: %s", path, exc)
+            out.errors.append(f"parse_error:{exc}")
+            return out
+
+        # Process each statement
+        for stmt_index, stmt in enumerate(statements):
+            if stmt is None:
+                continue
+
+            try:
+                query_node = self._parse_statement(stmt, path, stmt_index)
+                out.statements.append(query_node)
+
+                # Track defined and referenced tables
+                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 statement %d in %s: %s", stmt_index, path, exc)
+                out.errors.append(f"statement_error:{stmt_index}:{exc}")
+
+        return out
+
+    def _parse_statement(self, stmt: Any, path: Path, stmt_index: int) -> QueryNode:
+        """Parse a single SQL statement into a QueryNode.
+
+        Args:
+            stmt: sqlglot AST node
+            path: Path to the source file
+            stmt_index: Statement index in the file
+
+        Returns:
+            QueryNode with extracted metadata
+        """
+        kind = self._classify(stmt)
+        sql = stmt.sql(dialect=self.DIALECT)
+
+        # Extract target table for CREATE/INSERT statements
+        target = None
+        if isinstance(stmt, exp.Create):
+            target = self._extract_target_table(stmt)
+        elif isinstance(stmt, exp.Insert):
+            target = self._extract_insert_target(stmt)
+
+        # Extract source tables
+        sources = []
+        ctes = {}
+        confidence = 1.0
+        parse_failed = False
+
+        # Try to extract table references using scope analysis
+        try:
+            from sqlglot.optimizer.scope import build_scope
+
+            root_scope = build_scope(stmt)
+            if root_scope:
+                sources = self._real_tables(root_scope)
+            else:
+                # Fallback to basic table extraction
+                sources = self._fallback_table_scan(stmt)
+                parse_failed = True
+        except Exception as exc:
+            logger.warning(
+                "Failed to build scope for statement %d in %s: %s", stmt_index, path, exc
+            )
+            sources = self._fallback_table_scan(stmt)
+            parse_failed = True
+
+        # Extract column lineage (currently minimal implementation)
+        column_lineage = []
+
+        # Remove duplicates while preserving order
+        sources = self._deduplicate_table_refs(sources)
+
+        return QueryNode(
+            file=path,
+            statement_index=stmt_index,
+            sql=sql,
+            kind=kind,
+            target=target,
+            sources=sources,
+            ctes=ctes,
+            column_lineage=column_lineage,
+            parse_failed=parse_failed,
+            confidence=confidence,
+            parsing_mode="sqlglot",
+        )
+
+    @staticmethod
+    def _extract_target_table(create_stmt: exp.Create) -> TableRef | None:
+        """Extract the target table from a CREATE statement.
+
+        Args:
+            create_stmt: CREATE expression
+
+        Returns:
+            TableRef or None if not a table create
+        """
+        if create_stmt.kind not in ("TABLE", "VIEW"):
+            return None
+
+        if not create_stmt.this:
+            return None
+
+        return AnsiParser._convert_table_expr_to_ref(create_stmt.this)
+
+    @staticmethod
+    def _extract_insert_target(insert_stmt: exp.Insert) -> TableRef | None:
+        """Extract the target table from an INSERT statement.
+
+        Args:
+            insert_stmt: INSERT expression
+
+        Returns:
+            TableRef or None
+        """
+        if not insert_stmt.this:
+            return None
+
+        return AnsiParser._convert_table_expr_to_ref(insert_stmt.this)
+
+    @staticmethod
+    def _fallback_table_scan(stmt: Any) -> list[TableRef]:
+        """Fallback table extraction when scope analysis fails.
+
+        Walks the AST tree to find all Table nodes.
+
+        Args:
+            stmt: SQL expression
+
+        Returns:
+            List of TableRef objects found
+        """
+        tables = []
+        seen = set()
+
+        # Use find_all to iterate through all nodes of type Table
+        for table_node in stmt.find_all(exp.Table):
+            ref = AnsiParser._convert_table_expr_to_ref(table_node)
+            if ref and ref.full_id not in seen:
+                tables.append(ref)
+                seen.add(ref.full_id)
+
+        return tables
+
+    @staticmethod
+    def _deduplicate_table_refs(tables: list[TableRef]) -> list[TableRef]:
+        """Remove duplicate table references while preserving order.
+
+        Args:
+            tables: List of TableRef objects
+
+        Returns:
+            Deduplicated list
+        """
+        seen = set()
+        result = []
+        for table in tables:
+            if table.full_id not in seen:
+                result.append(table)
+                seen.add(table.full_id)
+        return result

sqlcg/parsers/base.py

@@ -0,0 +1,414 @@
+"""Base data models and abstract parser for SQL parsing and lineage extraction."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import StrEnum
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from sqlcg.lineage.schema_resolver import SchemaResolver
+
+
+class QueryKind(StrEnum):
+    """SQL query classification types."""
+
+    SELECT = "SELECT"
+    INSERT = "INSERT"
+    UPDATE = "UPDATE"
+    DELETE = "DELETE"
+    CREATE_TABLE = "CREATE_TABLE"
+    CREATE_VIEW = "CREATE_VIEW"
+    CREATE_PROC = "CREATE_PROC"
+    MERGE = "MERGE"
+    OTHER = "OTHER"
+
+
+@dataclass(frozen=True)
+class TableRef:
+    """A reference to a table (immutable).
+
+    Attributes:
+        catalog: Optional catalog name (e.g., "my_warehouse")
+        db: Optional database/schema name (e.g., "public")
+        name: Table name (always required)
+        alias: Optional alias used in the query context
+    """
+
+    catalog: str | None = None
+    db: str | None = None
+    name: str = ""
+    alias: str | None = None
+
+    @property
+    def full_id(self) -> str:
+        """Return the fully qualified table identifier.
+
+        Includes all non-None components: catalog.db.name or db.name or name.
+        Used as the Table.qualified key in the graph.
+        """
+        parts = [p for p in (self.catalog, self.db, self.name) if p]
+        return ".".join(parts)
+
+    @property
+    def qualified(self) -> str:
+        """Alias for full_id for compatibility."""
+        return self.full_id
+
+
+@dataclass(frozen=True)
+class ColumnRef:
+    """A reference to a column (immutable).
+
+    Attributes:
+        table: The TableRef this column belongs to
+        name: Column name
+    """
+
+    table: TableRef
+    name: str = ""
+
+    @property
+    def full_id(self) -> str:
+        """Return the fully qualified column identifier.
+
+        Format: table.full_id.column_name
+        Used as the Column.id key in the graph (finding 5.4).
+        """
+        return f"{self.table.full_id}.{self.name}"
+
+
+@dataclass(frozen=True)
+class LineageEdge:
+    """A column-level lineage edge (immutable by design for hashing).
+
+    Attributes:
+        src: Source ColumnRef
+        dst: Destination ColumnRef
+        transform: Description of the transformation applied (e.g., "PASS_THROUGH", "AGGREGATION")
+        confidence: Confidence score (0.0 to 1.0)
+        query_id: Optional identifier of the query that created this edge
+    """
+
+    src: ColumnRef
+    dst: ColumnRef
+    transform: str = "UNKNOWN"
+    confidence: float = 1.0
+    query_id: str | None = None
+
+    def __hash__(self) -> int:
+        """Support hashing for use in sets/dicts."""
+        return hash((self.src, self.dst, self.transform, self.confidence, self.query_id))
+
+    def __eq__(self, other: object) -> bool:
+        """Support equality comparison."""
+        if not isinstance(other, LineageEdge):
+            return NotImplemented
+        return (
+            self.src == other.src
+            and self.dst == other.dst
+            and self.transform == other.transform
+            and self.confidence == other.confidence
+            and self.query_id == other.query_id
+        )
+
+
+@dataclass
+class QueryNode:
+    """A parsed SQL query node (mutable by design for pass-2 patching).
+
+    This class is intentionally mutable to support the two-pass resolution strategy.
+    Pass 1 extracts initial lineage; pass 2 patches QueryNode fields directly.
+    Do not freeze this class.
+
+    Note: Do not rely on QueryNode identity after pass 2; use dataclasses.replace()
+    to create a new instance if mutation semantics are a concern.
+
+    Attributes:
+        file: Path to the source file
+        statement_index: 0-based index of the statement in the file
+        sql: The original SQL text
+        kind: Query type (SELECT, INSERT, UPDATE, DELETE, CREATE_TABLE, CREATE_VIEW, etc.)
+        target: Optional TableRef for INSERT/CREATE targets
+        sources: List of source TableRefs
+        ctes: Dict of CTE names to their TableRef definitions
+        column_lineage: List of LineageEdge objects
+        parse_failed: Whether parsing failed (fallback to table-only lineage)
+        confidence: Overall confidence score for this query's lineage
+        parsing_mode: How the query was parsed (e.g., "sqlglot", "fallback", "scripting")
+    """
+
+    file: Path
+    statement_index: int
+    sql: str
+    kind: str = ""
+    target: TableRef | None = None
+    sources: list[TableRef] = field(default_factory=list)
+    ctes: dict[str, TableRef] = field(default_factory=dict)
+    column_lineage: list[LineageEdge] = field(default_factory=list)
+    parse_failed: bool = False
+    confidence: float = 1.0
+    parsing_mode: str = "sqlglot"
+
+
+@dataclass
+class ParsedFile:
+    """Result of parsing a single SQL file (mutable for aggregation).
+
+    Attributes:
+        path: Path to the source file
+        dialect: SQL dialect used for parsing
+        statements: List of QueryNode objects parsed from the file
+        defined_tables: List of TableRef for tables defined in this file
+        referenced_tables: List of TableRef for tables referenced in this file
+        errors: List of error messages encountered during parsing
+    """
+
+    path: Path
+    dialect: str | None = None
+    statements: list[QueryNode] = field(default_factory=list)
+    defined_tables: list[TableRef] = field(default_factory=list)
+    referenced_tables: list[TableRef] = field(default_factory=list)
+    errors: list[str] = field(default_factory=list)
+
+    @property
+    def path_str(self) -> str:
+        """Return the path as a string."""
+        return str(self.path)
+
+
+class SqlParser(ABC):
+    """Abstract base class for SQL parsers.
+
+    Attributes:
+        DIALECT: SQL dialect identifier (None for ANSI, "snowflake", "bigquery", etc.)
+        _schema: SchemaResolver instance for table/column lookups
+        _log: Logger instance for this parser
+    """
+
+    DIALECT: str | None = None
+
+    def __init__(self, schema_resolver: "SchemaResolver"):
+        """Initialize parser with schema resolver.
+
+        Args:
+            schema_resolver: SchemaResolver instance for table/column lookups
+        """
+        from sqlcg.utils.logging import getLogger
+
+        self._schema = schema_resolver
+        self._log = getLogger(f"{__name__}.{self.__class__.__name__}")
+
+    @abstractmethod
+    def parse_file(self, path: Path, sql: str) -> ParsedFile:
+        """Parse SQL text and return a ParsedFile with all statements.
+
+        Args:
+            path: Path to the source file
+            sql: SQL text to parse
+
+        Returns:
+            ParsedFile containing parsed statements and metadata
+        """
+        ...
+
+    def _classify(self, stmt: Any) -> str:
+        """Classify a SQL statement into a kind string.
+
+        Args:
+            stmt: sqlglot AST node
+
+        Returns:
+            One of: SELECT, INSERT, UPDATE, DELETE, CREATE_TABLE, CREATE_VIEW,
+            CREATE_PROC, MERGE, OTHER
+        """
+        import sqlglot.expressions as exp
+
+        match stmt:
+            case exp.Select():
+                return QueryKind.SELECT
+            case exp.Insert():
+                return QueryKind.INSERT
+            case exp.Update():
+                return QueryKind.UPDATE
+            case exp.Delete():
+                return QueryKind.DELETE
+            case exp.Create():
+                match stmt.kind:
+                    case "TABLE":
+                        return QueryKind.CREATE_TABLE
+                    case "VIEW":
+                        return QueryKind.CREATE_VIEW
+                    case "PROCEDURE" | "FUNCTION":
+                        return QueryKind.CREATE_PROC
+                    case _:
+                        return QueryKind.CREATE_TABLE  # Default to table
+            case exp.Merge():
+                return QueryKind.MERGE
+            case _:
+                return QueryKind.OTHER
+
+    def _real_tables(self, scope: Any) -> list[TableRef]:
+        """Return real (non-CTE) tables referenced in a scope.
+
+        Args:
+            scope: sqlglot scope object
+
+        Returns:
+            List of TableRef objects for non-CTE tables
+        """
+        if not scope:
+            return []
+
+        cte_sources = set(scope.cte_sources.keys()) if hasattr(scope, "cte_sources") else set()
+        tables = []
+
+        if hasattr(scope, "tables"):
+            # scope.tables is a list of Table expressions
+            table_list = scope.tables
+            if isinstance(table_list, list):
+                for table_expr in table_list:
+                    # Extract the table name to check if it's a CTE
+                    table_name = None
+                    if hasattr(table_expr, "name"):
+                        table_name = table_expr.name
+
+                    if table_name not in cte_sources:
+                        # Convert table expression to TableRef
+                        ref = self._convert_table_expr_to_ref(table_expr)
+                        if ref:
+                            tables.append(ref)
+
+        return tables
+
+    @staticmethod
+    def _convert_table_expr_to_ref(table_expr: Any) -> "TableRef | None":
+        """Convert a table expression to a TableRef.
+
+        Args:
+            table_expr: Table expression (e.g., Table, Identifier, Schema)
+
+        Returns:
+            TableRef or None
+        """
+        import sqlglot.expressions as exp
+
+        match table_expr:
+            case exp.Table():
+                return TableRef(
+                    catalog=table_expr.catalog,
+                    db=table_expr.db,
+                    name=table_expr.name,
+                )
+            case exp.Identifier():
+                return TableRef(name=table_expr.name)
+            case exp.Schema():
+                # Schema wraps the actual table
+                return SqlParser._convert_table_expr_to_ref(table_expr.this)
+            case _:
+                # Try to extract name from string representation
+                name = str(table_expr)
+                if name:
+                    return TableRef(name=name)
+                return None
+
+    def _extract_column_lineage(
+        self, stmt: Any, path: Path, out: ParsedFile, schema: dict
+    ) -> list[LineageEdge]:
+        """Extract column-level lineage with structured error recording.
+
+        On sqlglot.lineage failure: log WARNING, append to ParsedFile.errors,
+        emit LineageEdge with confidence=0.0, continue (do NOT raise or skip silently).
+
+        For columns not found in the schema, emits edges with reduced confidence (0.5)
+        rather than skipping silently (only applies when schema is non-empty).
+
+        Args:
+            stmt: sqlglot AST node (Select/Insert/Create)
+            path: Path to the source file
+            out: ParsedFile object to append errors to
+            schema: Schema dict from _schema.as_dict()
+
+        Returns:
+            List of LineageEdge objects
+        """
+        import sqlglot.expressions as exp
+        from sqlglot.lineage import lineage as sg_lineage
+
+        edges: list[LineageEdge] = []
+
+        # Only extract column lineage for certain statement types
+        if not isinstance(stmt, (exp.Select, exp.Insert, exp.Create)):
+            return edges
+
+        # Extract column references from SELECT list or target
+        try:
+            # Get the body of the query for lineage extraction
+            if isinstance(stmt, exp.Select):
+                body = stmt
+                # Extract output columns
+                for col_expr in stmt.expressions:
+                    col_name = col_expr.alias if col_expr.alias else str(col_expr)
+                    # Schema validation: if schema is loaded and column isn't in it,
+                    # emit a reduced-confidence edge rather than a full-confidence one.
+                    if schema:
+                        table_cols: list[str] | None = None
+                        for _scope_name, cols in schema.items():
+                            if isinstance(cols, list):
+                                table_cols = cols
+                                break
+                            elif isinstance(cols, dict):
+                                for _db, tables in cols.items():
+                                    if isinstance(tables, dict):
+                                        for _tbl, tcols in tables.items():
+                                            if col_name in tcols:
+                                                table_cols = tcols
+                                                break
+                        if table_cols is not None and col_name not in table_cols:
+                            self._log.warning(
+                                "column %s not found in schema, emitting low-confidence edge",
+                                col_name,
+                            )
+                            edges.append(
+                                LineageEdge(
+                                    src=ColumnRef(TableRef(None, None, "<unknown>"), col_name),
+                                    dst=ColumnRef(TableRef(None, None, "<unknown>"), col_name),
+                                    transform="UNKNOWN",
+                                    confidence=0.5,
+                                )
+                            )
+                            continue
+
+                    try:
+                        root = sg_lineage(col_name, body, schema=schema, dialect=self.DIALECT)
+                        if root:
+                            # Successfully extracted lineage
+                            # TODO: convert root to LineageEdge(s)
+                            pass
+                    except Exception as exc:
+                        self._log.warning(
+                            "column lineage extraction failed: file=%s col=%s error=%s",
+                            path,
+                            col_name,
+                            exc,
+                        )
+                        out.errors.append(f"col_lineage:{col_name}:{exc}")
+                        # Emit a zero-confidence placeholder edge
+                        edges.append(
+                            LineageEdge(
+                                src=ColumnRef(TableRef(None, None, "<unknown>"), col_name),
+                                dst=ColumnRef(TableRef(None, None, "<unknown>"), col_name),
+                                transform="UNKNOWN",
+                                confidence=0.0,
+                            )
+                        )
+
+        except Exception as exc:
+            self._log.warning(
+                "column lineage extraction failed for entire statement: file=%s error=%s",
+                path,
+                exc,
+            )
+            out.errors.append(f"col_lineage:statement:{exc}")
+
+        return edges

sqlcg/parsers/bigquery_parser.py

@@ -0,0 +1,77 @@
+"""BigQuery SQL parser with scripting block detection."""
+
+from pathlib import Path
+
+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__)
+
+
+@register("bigquery")
+class BigQueryParser(AnsiParser):
+    """BigQuery SQL parser.
+
+    Handles BigQuery-specific features:
+    - Scripting block detection (DECLARE, BEGIN, IF)
+    - Standard table/column lineage extraction
+    """
+
+    DIALECT: str | None = "bigquery"
+
+    def __init__(self, schema_resolver: SchemaResolver):
+        """Initialize BigQuery parser.
+
+        Args:
+            schema_resolver: SchemaResolver instance for table/column lookups
+        """
+        super().__init__(schema_resolver)
+
+    def parse_file(self, path: Path, sql: str) -> ParsedFile:
+        """Parse BigQuery 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("BigQuery scripting block detected in %s, marking as parse_failed", path)
+            # Scripting blocks are not fully parseable; mark as parse_failed
+            out = ParsedFile(path=path, dialect=self.DIALECT)
+            out.errors.append("parse_mode:scripting_block")
+            return out
+
+        # Otherwise use standard ANSI parsing with BigQuery dialect
+        return AnsiParser.parse_file(self, path, sql)  # type: ignore
+
+    def _has_scripting_block(self, sql: str) -> bool:
+        """Token-aware scripting block detection for BigQuery.
+
+        Detects DECLARE, BEGIN, IF, and other scripting keywords.
+
+        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("bigquery").tokenize(sql)  # type: ignore
+            scripting_tokens = {
+                TokenType.DECLARE,  # type: ignore
+                TokenType.BEGIN,  # type: ignore
+            }
+            return any(t.token_type in scripting_tokens for t in toks)
+        except Exception:
+            # Fallback: check for keywords in uppercase
+            sql_upper = sql.upper()
+            return any(keyword in sql_upper for keyword in ("DECLARE ", "BEGIN "))