ss-tools/backend/src/services/sql_table_extractor.py

# #region SqlTableExtractorModule [C:2] [TYPE Module] [SEMANTICS sql, jinja, table, extraction, parsing]
# @BRIEF Three-phase SQL+Jinja table name extractor for virtual datasets as per FR-003.
#   Phase 1: Detect Jinja spans vs SQL spans
#   Phase 2: In Jinja spans, extract "schema.table" from string values
#   Phase 3: In SQL spans, regex pattern + sqlparse filter to reject string literal false positives
# @LAYER Service
# @RELATION DEPENDS_ON -> [EXT:Library:sqlparse]
# @INVARIANT Only exact schema.table matches (case-insensitive); unqualified references are NOT matched.
# @INVARIANT Returns a set[str] of fully-qualified table names (lowercased for case-insensitive matching).

from collections.abc import Iterable
import re

import sqlparse
from sqlparse.sql import Token, TokenList
from sqlparse.tokens import Literal

# ── Phase 1: Jinja span detection ───────────────────────────

# Minimal Jinja token detection: {% ... %}, {{ ... }}, {# ... #}
_JINJA_BLOCK_RE = re.compile(r"{%-?\s*.*?\s*-?%}", re.DOTALL)
_JINJA_EXPR_RE = re.compile(r"{{-?\s*.*?\s*-?}}", re.DOTALL)
_JINJA_COMMENT_RE = re.compile(r"{#.*?#}", re.DOTALL)

# ── Phase 3: schema.table regex (case-insensitive) ───────────
_SCHEMA_TABLE_RE = re.compile(
    r"""
    (?<!['"`\w])                     # not preceded by string delimiter or word char
    \b                                # word boundary
    (?:                               # begin: non-capturing group for full match
        (?:"([a-zA-Z_][\w]*)")?       # optional quoted schema
        (?:([a-zA-Z_][\w]*))          # schema (unquoted)
        \.
        (?:([a-zA-Z_][\w]*))          # table
    )
    \b                                # word boundary
    (?!['"`])                         # not followed by a string delimiter
    """,
    re.VERBOSE | re.IGNORECASE,
)


# #region detect_jinja_spans [C:2] [TYPE Function]
# @BRIEF Phase 1: Split raw SQL text into Jinja spans and SQL spans.
# @PRE raw_sql is a string (possibly empty).
# @POST Returns a list of (span_type: str, text: str) tuples.
#   span_type is "jinja" or "sql".
def detect_jinja_spans(raw_sql: str) -> list[tuple[str, str]]:
    """Split raw SQL+Jinja text into alternating Jinja and SQL spans.

    Phase 1 detects Jinja blocks ({%%}, {{}}, {##}) and returns the
    remaining text as SQL spans.
    """
    if not raw_sql or not raw_sql.strip():
        return [("sql", raw_sql or "")]

    spans: list[tuple[str, str]] = []
    # Collect all Jinja match positions
    jinja_matches: list[tuple[int, int]] = []
    for pattern in (_JINJA_BLOCK_RE, _JINJA_EXPR_RE, _JINJA_COMMENT_RE):
        for m in pattern.finditer(raw_sql):
            jinja_matches.append((m.start(), m.end()))

    # Merge overlapping Jinja spans
    if jinja_matches:
        jinja_matches.sort()
        merged: list[tuple[int, int]] = [jinja_matches[0]]
        for start, end in jinja_matches[1:]:
            if start <= merged[-1][1]:
                merged[-1] = (merged[-1][0], max(merged[-1][1], end))
            else:
                merged.append((start, end))

        # Build alternating sql/jinja spans
        cursor = 0
        for start, end in merged:
            if cursor < start:
                sql_chunk = raw_sql[cursor:start].strip()
                if sql_chunk:
                    spans.append(("sql", sql_chunk))
            jinja_chunk = raw_sql[start:end].strip()
            if jinja_chunk:
                spans.append(("jinja", jinja_chunk))
            cursor = end
        if cursor < len(raw_sql):
            remaining = raw_sql[cursor:].strip()
            if remaining:
                spans.append(("sql", remaining))
    else:
        spans.append(("sql", raw_sql.strip()))

    return spans
# #endregion detect_jinja_spans


# #region extract_tables_from_jinja [C:2] [TYPE Function]
# @BRIEF Phase 2: Extract "schema.table" references from Jinja string values.
#   Looks for patterns like "raw.sales" inside Jinja text.
# @PRE jinja_text is a string from a Jinja span.
# @POST Returns a set of lowercased schema.table strings.
def extract_tables_from_jinja(jinja_text: str) -> set[str]:
    """Extract ``schema.table`` references from within Jinja template blocks.

    Looks for double-quoted or single-quoted string values that match
    the ``schema.table`` pattern, e.g. ``"raw.sales"`` or ``'raw.inventory'``.
    """
    tables: set[str] = set()
    # Match string values: "schema.table" or 'schema.table'
    string_pattern = re.compile(
        r"""["']([a-zA-Z_][\w]*\.[a-zA-Z_][\w]*)["']"""
    )
    for m in string_pattern.finditer(jinja_text):
        tables.add(m.group(1).lower())
    return tables
# #endregion extract_tables_from_jinja


# #region is_string_literal [C:1] [TYPE Function]
# @BRIEF Check if a sqlparse Token is a string literal (not a schema.table identifier).
# @PRE token is a sqlparse Token.
# @POST Returns True if the token is a string/Single/Literal.String.
def is_string_literal(token: Token) -> bool:
    """Return True if token is a string literal type in sqlparse."""
    return token.ttype is Literal.String.Single or token.ttype is Literal.String
# #endregion is_string_literal


# #region extract_tables_from_sql_span [C:2] [TYPE Function]
# @BRIEF Phase 3: Extract schema.table references from SQL text using regex + sqlparse filtering.
#   Uses regex to find all schema.table candidates, then sqlparse to filter out
#   false positives inside string literals.
# @PRE sql_text is a string from a SQL span (non-Jinja).
# @POST Returns a set of lowercased schema.table strings.
def extract_tables_from_sql_span(sql_text: str) -> set[str]:
    """Extract ``schema.table`` references from plain SQL text.

    Uses regex to find all ``schema.table`` candidates, then sqlparse
    to reject matches that fall inside string literals.
    """
    tables: set[str] = set()
    raw_matches = _SCHEMA_TABLE_RE.findall(sql_text)
    if not raw_matches:
        return tables

    # Use sqlparse to identify string literal positions
    parsed = sqlparse.parse(sql_text)
    string_literal_ranges: list[tuple[int, int]] = []

    def walk_tokens(tokens: Iterable[Token], base_offset: int = 0) -> None:
        offset = base_offset
        for token in tokens:
            if isinstance(token, TokenList):
                walk_tokens(token.flatten(), offset)
            else:
                ttype = token.ttype
                val = token.value
                if is_string_literal(token):
                    string_literal_ranges.append(
                        (offset, offset + len(val))
                    )
                offset += len(val)

    for stmt in parsed:
        if stmt is None:
            continue
        walk_tokens(stmt.flatten(), base_offset=0)

    def is_in_string(pos: int) -> bool:
        for s_start, s_end in string_literal_ranges:
            if s_start <= pos <= s_end:
                return True
        return False

    # Re-scan with full match positions to filter
    for m in _SCHEMA_TABLE_RE.finditer(sql_text):
        if not is_in_string(m.start()):
            # Extract schema and table from match groups
            schema = m.group(2) or m.group(1) or ""
            table = m.group(3) or ""
            # Filter out column aliases (single-char schemas like "a.id", "o.id")
            if len(schema) < 2 and schema.isalpha() and schema.islower():
                continue
            if schema and table:
                tables.add(f"{schema.lower()}.{table.lower()}")

    return tables
# #endregion extract_tables_from_sql_span


# #region extract_tables_from_sql [C:2] [TYPE Function]
# @BRIEF Three-phase entry point: extract all schema.table references from raw SQL+Jinja text.
# @PRE raw_sql is a string (possibly empty).
# @POST Returns a set of lowercased fully-qualified table names (schema.table format).
# @RELATION CALLS -> [detect_jinja_spans]
# @RELATION CALLS -> [extract_tables_from_jinja]
# @RELATION CALLS -> [extract_tables_from_sql_span]
def extract_tables_from_sql(raw_sql: str) -> set[str]:
    """Extract all ``schema.table`` references from raw SQL + Jinja text.

    Three-phase approach per FR-003:
    1. Detect Jinja spans vs SQL spans
    2. In Jinja spans, extract ``schema.table`` from quoted string values
    3. In SQL spans, regex + sqlparse filter to reject string literal false positives

    Returns a set of lowercased ``schema.table`` strings for case-insensitive matching.

    Example:
        >>> extract_tables_from_sql("SELECT * FROM raw.sales JOIN raw.inventory ON ...")
        {'raw.sales', 'raw.inventory'}
    """
    if not raw_sql or not raw_sql.strip():
        return set()

    all_tables: set[str] = set()
    spans = detect_jinja_spans(raw_sql)

    for span_type, text in spans:
        if span_type == "jinja":
            all_tables.update(extract_tables_from_jinja(text))
        else:
            all_tables.update(extract_tables_from_sql_span(text))

    return all_tables
# #endregion extract_tables_from_sql

# #endregion SqlTableExtractorModule