ripple-sql 0.1.0__py3-none-any.whl

ripple/engine/tsql_tvf.py

@@ -0,0 +1,653 @@
+"""T-SQL table-valued function lineage.
+
+CREATE FUNCTION ... RETURNS @out TABLE (...) mints no model from the
+generic script path: T-SQL bodies separate statements by newline, not
+";", so sqlglot reads the whole file as one Command. This module finds
+the shape in the raw text, chops the body into statements on keyword
+boundaries, and synthesizes the @out fills into one analyzable
+SELECT/UNION: positional INSERT-SELECT mapping, if-guarded fills
+unioned, local table-variable hops inlined to the first external
+boundary. A table-variable read with no visible fill (a TVP parameter,
+or a variable filled by dynamic SQL) stays cited as written with its @
+prefix; dynamic SQL stays invisible (ceds_generate, holdout round 12).
+"""
+
+from __future__ import annotations
+
+import re
+
+import sqlglot
+from sqlglot import exp
+from sqlglot.errors import ErrorLevel
+
+from ripple.engine.safe_gen import safe_sql
+
+_RETURNS_TABLE = re.compile(r"\s*RETURNS\s+(?P<var>@[\w$]+)\s+(?:as\s+)?TABLE\s*\(", re.I)
+# AS is optional in practice: ceds' fnSplit writes RETURNS @List TABLE (...) BEGIN
+_AS_BEGIN = re.compile(r"\s*(?:WITH\s+SCHEMABINDING\s+)?(?:AS\s*)?BEGIN\b", re.I)
+_WORD = re.compile(r"[@#]?[A-Za-z_][\w$]*")
+_IDENT_ATOM = re.compile(r"\[([^\]]+)\]|\"([^\"]+)\"|([A-Za-z_][\w$]*)")
+
+_CONSTRAINT_STARTERS = {
+    "unique",
+    "primary",
+    "constraint",
+    "check",
+    "foreign",
+    "index",
+    "key",
+    "clustered",
+    "nonclustered",
+}
+
+_STATEMENT_STARTERS = {
+    "declare",
+    "insert",
+    "update",
+    "delete",
+    "set",
+    "if",
+    "while",
+    "return",
+    "merge",
+    "print",
+    "break",
+    "continue",
+    "begin",
+    "end",
+    "else",
+    "go",
+    "exec",
+    "execute",
+}
+
+_SELECT_CONTINUATIONS = {"union", "all", "except", "intersect", "as", "then", "else", "(", "="}
+
+
+def _skip_atom(s: str, i: int) -> int | None:
+    """End index of the string/comment/bracket atom starting at i, else None."""
+    c = s[i]
+    if c == "'":
+        j = i + 1
+        while j < len(s):
+            if s[j] == "'":
+                if s[j + 1 : j + 2] == "'":
+                    j += 2
+                    continue
+                return j + 1
+            j += 1
+        return len(s)
+    if c == "[":
+        j = i + 1
+        while j < len(s):
+            if s[j] == "]":
+                if s[j + 1 : j + 2] == "]":
+                    j += 2
+                    continue
+                return j + 1
+            j += 1
+        return len(s)
+    if c == '"':
+        j = s.find('"', i + 1)
+        return len(s) if j == -1 else j + 1
+    if c == "-" and s[i + 1 : i + 2] == "-":
+        j = s.find("\n", i)
+        return len(s) if j == -1 else j + 1
+    if c == "/" and s[i + 1 : i + 2] == "*":
+        j = s.find("*/", i)
+        return len(s) if j == -1 else j + 2
+    return None
+
+
+def _skip_ws_comments(s: str, i: int) -> int:
+    """First index at or after i that is neither whitespace nor a comment."""
+    n = len(s)
+    while i < n:
+        c = s[i]
+        if c.isspace():
+            i += 1
+            continue
+        if (c == "-" and s[i + 1 : i + 2] == "-") or (c == "/" and s[i + 1 : i + 2] == "*"):
+            i = _skip_atom(s, i)
+            continue
+        return i
+    return i
+
+
+def _match_word(s: str, i: int, word: str) -> int | None:
+    m = _WORD.match(s, i)
+    if m is not None and m.group(0).lower() == word:
+        return m.end()
+    return None
+
+
+def _find_tvf_head(sql: str, pos: int) -> tuple[int, int, str] | None:
+    """The next CREATE [OR ALTER] FUNCTION <name> outside strings and
+    comments: (head start, index past the name, written name).
+
+    A raw regex matched commented-out function text, minted a phantom
+    model, and excised the comment; and it refused the legal comment
+    between CREATE and FUNCTION (cycle-13 review, F2/F9). This scan skips
+    string/comment/bracket atoms and allows comments between tokens."""
+    i = pos
+    n = len(sql)
+    while i < n:
+        j = _skip_atom(sql, i)
+        if j is not None:
+            i = j
+            continue
+        m = _WORD.match(sql, i)
+        if m is None:
+            i += 1
+            continue
+        head_start = m.start()
+        i = m.end()
+        if m.group(0).lower() != "create":
+            continue
+        k = _skip_ws_comments(sql, i)
+        end = _match_word(sql, k, "or")
+        if end is not None:
+            end = _match_word(sql, _skip_ws_comments(sql, end), "alter")
+            if end is None:
+                continue
+            k = _skip_ws_comments(sql, end)
+        end = _match_word(sql, k, "function")
+        if end is None:
+            continue
+        k = _skip_ws_comments(sql, end)
+        parts: list[str] = []
+        name_end = k
+        while True:
+            atom = _IDENT_ATOM.match(sql, k)
+            if atom is None:
+                parts = []
+                break
+            parts.append(next(g for g in atom.groups() if g is not None))
+            name_end = atom.end()
+            after = _skip_ws_comments(sql, name_end)
+            if after < n and sql[after] == ".":
+                k = _skip_ws_comments(sql, after + 1)
+                continue
+            break
+        if not parts:
+            continue
+        return head_start, name_end, ".".join(parts)
+    return None
+
+
+def _balanced_paren(s: str, open_idx: int) -> int | None:
+    """Index just past the ')' matching s[open_idx] == '(', else None."""
+    depth = 0
+    i = open_idx
+    while i < len(s):
+        j = _skip_atom(s, i)
+        if j is not None:
+            i = j
+            continue
+        c = s[i]
+        if c == "(":
+            depth += 1
+        elif c == ")":
+            depth -= 1
+            if depth == 0:
+                return i + 1
+        i += 1
+    return None
+
+
+def _body_end(s: str, start: int) -> tuple[int, int]:
+    """(body end, region end) for a body whose opening BEGIN is consumed.
+
+    Counts BEGIN and CASE as openers so an inner block's or CASE
+    expression's END never closes the function early. Unbalanced text
+    ends at the string's end."""
+    depth = 1
+    i = start
+    while i < len(s):
+        j = _skip_atom(s, i)
+        if j is not None:
+            i = j
+            continue
+        m = _WORD.match(s, i)
+        if m is None:
+            i += 1
+            continue
+        word = m.group(0).lower()
+        if word in ("begin", "case"):
+            depth += 1
+        elif word == "end":
+            depth -= 1
+            if depth == 0:
+                return m.start(), m.end()
+        i = m.end()
+    return len(s), len(s)
+
+
+def _column_names(cols_text: str) -> list[str]:
+    """Declared column names, in order, from a table-type parenthesis body.
+
+    Constraint entries (unique clustered (...), primary key) are skipped;
+    a constraint glued to the last column without a comma still yields
+    that column's name (the ceds spelling)."""
+    segments: list[str] = []
+    depth = 0
+    start = 0
+    i = 0
+    while i < len(cols_text):
+        j = _skip_atom(cols_text, i)
+        if j is not None:
+            i = j
+            continue
+        c = cols_text[i]
+        if c == "(":
+            depth += 1
+        elif c == ")":
+            depth -= 1
+        elif c == "," and depth == 0:
+            segments.append(cols_text[start:i])
+            start = i + 1
+        i += 1
+    segments.append(cols_text[start:])
+    names: list[str] = []
+    for seg in segments:
+        m = _IDENT_ATOM.search(seg)
+        if m is None:
+            continue
+        name = next(g for g in m.groups() if g is not None)
+        if name.lower() in _CONSTRAINT_STARTERS:
+            continue
+        names.append(name)
+    return names
+
+
+def _chop_body(body: str) -> list[str]:
+    """Split a semicolonless T-SQL body into parseable statement chunks.
+
+    A statement-starter keyword at paren depth 0 (outside any CASE
+    expression) opens a new chunk. SELECT is a starter unless it
+    continues an INSERT/WITH that has no select yet, or follows a set
+    operator. A chunk that still will not parse is skipped by the
+    caller, never guessed at."""
+    chunks: list[str] = []
+    start = 0
+    i = 0
+    n = len(body)
+    paren = 0
+    case_depth = 0
+    chunk_head: str | None = None
+    chunk_has_select = False
+    chunk_has_set = False
+    last_word = ""
+
+    def cut(at: int) -> None:
+        nonlocal start
+        if body[start:at].strip():
+            chunks.append(body[start:at])
+        start = at
+
+    while i < n:
+        j = _skip_atom(body, i)
+        if j is not None:
+            i = j
+            continue
+        c = body[i]
+        if c == "(":
+            paren += 1
+            last_word = "("
+            i += 1
+            continue
+        if c == ")":
+            paren = max(0, paren - 1)
+            last_word = ")"
+            i += 1
+            continue
+        if c == ";":
+            cut(i + 1)
+            chunk_head = None
+            chunk_has_select = False
+            chunk_has_set = False
+            last_word = ""
+            i += 1
+            continue
+        m = _WORD.match(body, i)
+        if m is None:
+            i += 1
+            continue
+        word = m.group(0).lower()
+        if paren == 0:
+            if word == "case":
+                case_depth += 1
+            elif word == "end" and case_depth > 0:
+                case_depth -= 1
+            elif case_depth == 0:
+                split = word in _STATEMENT_STARTERS
+                if word == "select":
+                    continues = (
+                        chunk_head in ("insert", "with") and not chunk_has_select
+                    ) or last_word in _SELECT_CONTINUATIONS
+                    split = not continues
+                elif word == "insert" and chunk_head == "with" and not chunk_has_select:
+                    # the INSERT consuming an open WITH chunk continues it;
+                    # cutting there discarded the CTE (cycle-13 review, F10)
+                    split = False
+                elif word == "set" and chunk_head == "update" and not chunk_has_set:
+                    # an UPDATE's own SET clause; cutting there severed the
+                    # update from its assignments (cycle-13 review, F4)
+                    split = False
+                if split:
+                    cut(i)
+                    chunk_head = word
+                    chunk_has_select = False
+                    chunk_has_set = False
+                elif chunk_head is None:
+                    chunk_head = word
+                if word == "select":
+                    chunk_has_select = True
+                elif word == "set" and chunk_head == "update":
+                    chunk_has_set = True
+        last_word = word
+        i = m.end()
+    cut(n)
+    return chunks
+
+
+def _parse_chunks(chunks: list[str]) -> list[exp.Expression]:
+    from ripple.engine.sql_script import _flatten_blocks
+
+    stmts: list[exp.Expression] = []
+    for chunk in chunks:
+        text = chunk.strip()
+        if not text:
+            continue
+        try:
+            parsed = sqlglot.parse_one(text, dialect="tsql", error_level=ErrorLevel.IGNORE)
+        except Exception:
+            continue
+        if parsed is not None:
+            stmts.append(parsed)
+    return _flatten_blocks(stmts, "tsql")
+
+
+def _tablevar_name(node) -> str | None:
+    """ "@out" for a Table (or Schema-wrapped Table) over a Parameter."""
+    if isinstance(node, exp.Schema):
+        node = node.this
+    if isinstance(node, exp.Table) and isinstance(node.this, exp.Parameter):
+        name = node.this.name
+        return f"@{name}".lower() if name else None
+    return None
+
+
+def _is_star_item(item: exp.Expression) -> bool:
+    return isinstance(item, exp.Star) or (
+        isinstance(item, exp.Column) and isinstance(item.this, exp.Star)
+    )
+
+
+class _TvfBody:
+    def __init__(self, stmts: list[exp.Expression], out_var: str, out_columns: list[str]):
+        self.declared: dict[str, list[str]] = {out_var: out_columns}
+        # fills carry their statement position: a read of a table variable
+        # inlines only fills from EARLIER statements, because a later fill
+        # cannot retroactively reach an earlier read (cycle-13 review, F5)
+        self.fills: dict[str, list[tuple[int, list[str] | None, exp.Expression]]] = {}
+        self.poisoned: dict[str, set[str]] = {}
+        self.opaque: set[str] = set()
+        for pos, stmt in enumerate(stmts):
+            if isinstance(stmt, exp.Declare):
+                for item in stmt.expressions:
+                    kind = item.args.get("kind")
+                    if not isinstance(kind, exp.Schema):
+                        continue
+                    cols = [c.name for c in kind.expressions if isinstance(c, exp.ColumnDef)]
+                    for param in item.this if isinstance(item.this, list) else [item.this]:
+                        if isinstance(param, exp.Parameter) and param.name:
+                            self.declared[f"@{param.name}".lower()] = cols
+            elif isinstance(stmt, exp.Insert):
+                var = _tablevar_name(stmt.this)
+                if var is None:
+                    continue
+                idents = None
+                if isinstance(stmt.this, exp.Schema):
+                    names = [c.name for c in stmt.this.expressions if isinstance(c, exp.Identifier)]
+                    if len(names) == len(stmt.this.expressions):
+                        idents = names
+                inner = stmt.expression
+                while isinstance(inner, (exp.Subquery, exp.Paren)):
+                    inner = inner.this
+                if isinstance(inner, (exp.Select, exp.Union)):
+                    with_clause = stmt.args.get("with_") or stmt.args.get("with")
+                    if (
+                        with_clause is not None
+                        and isinstance(inner, exp.Select)
+                        and not (inner.args.get("with_") or inner.args.get("with"))
+                    ):
+                        # the chopper keeps WITH ... INSERT together; the CTE
+                        # must ride on the fill or its reads dangle (F10)
+                        inner = inner.copy()
+                        inner.set("with_", with_clause.copy())
+                    self.fills.setdefault(var, []).append((pos, idents, inner))
+            elif isinstance(stmt, exp.Update):
+                self._add_update(pos, stmt)
+
+    def _add_update(self, pos: int, stmt: exp.Update) -> None:
+        """UPDATE @var SET col = expr is one more fill for col, its sources
+        unioning with earlier fills (cycle-13 review, F4). A SET shape that
+        cannot be synthesized degrades the touched columns to no value
+        edges instead of leaving a stale insert source as the answer."""
+        var = _tablevar_name(stmt.this)
+        if var is None:
+            return
+        items: list[exp.Alias] = []
+        cols: list[str] = []
+        unknown_touch = False
+        for eq in stmt.expressions:
+            if isinstance(eq, exp.EQ) and isinstance(eq.this, exp.Parameter):
+                continue  # a scalar-variable assignment, not a column fill
+            if (
+                isinstance(eq, exp.EQ)
+                and isinstance(eq.this, exp.Column)
+                and not eq.this.table
+                and eq.expression is not None
+            ):
+                cols.append(eq.this.name)
+                items.append(
+                    exp.Alias(this=eq.expression.copy(), alias=exp.to_identifier(eq.this.name))
+                )
+            else:
+                unknown_touch = True
+        if unknown_touch:
+            self._poison(var, None)
+            return
+        if not items:
+            return
+        select = exp.Select(expressions=items)
+        upd_from = stmt.args.get("from_") or stmt.args.get("from")
+        if upd_from is not None:
+            select.set("from_", upd_from.copy())
+        self.fills.setdefault(var, []).append((pos, cols, select))
+
+    def _poison(self, var: str, cols: list[str] | None) -> None:
+        if cols:
+            self.poisoned.setdefault(var, set()).update(c.lower() for c in cols)
+            return
+        declared = self.declared.get(var)
+        if declared:
+            self.poisoned.setdefault(var, set()).update(c.lower() for c in declared)
+        else:
+            self.opaque.add(var)
+
+    def synthesize(
+        self, var: str, seen: frozenset[str], before: int | None = None
+    ) -> exp.Expression | None:
+        if var in self.opaque:
+            return None
+        entries = [e for e in self.fills.get(var, []) if before is None or e[0] < before]
+        if not entries:
+            return None
+        branches: list[exp.Expression] = []
+        for pos, idents, select in entries:
+            node = self._map_fill(select, idents or self.declared.get(var, []))
+            if node is None:
+                continue
+            self._inline_tablevars(node, seen | {var}, pos)
+            branches.append(node)
+        if not branches:
+            return None
+        dead = self.poisoned.get(var)
+        if dead:
+            for branch in branches:
+                _null_out_columns(branch, dead)
+        combined = branches[0]
+        for branch in branches[1:]:
+            combined = exp.union(combined, branch, distinct=False, copy=False)
+        return combined
+
+    def _map_fill(self, select: exp.Expression, idents: list[str]) -> exp.Expression | None:
+        """The fill re-aliased to the insert's (or declared) column list,
+        pairwise. Refuses instead of keeping the projection as written: a
+        kept fill left source names as the function's output schema and
+        minted confident edges to nonexistent return columns (cycle-13
+        review, F3). Union branches map individually. A lone star expands
+        positionally only through a locally declared table variable whose
+        columns are known; any other star refuses."""
+        if isinstance(select, exp.Union):
+            left = self._map_fill(select.this, idents)
+            right = self._map_fill(select.expression, idents)
+            if left is None or right is None:
+                return None
+            out = select.copy()
+            out.set("this", left)
+            out.set("expression", right)
+            return out
+        if not isinstance(select, exp.Select) or not idents:
+            return None
+        work = select.copy()
+        if any(_is_star_item(s) for s in work.selects) and not self._expand_star(work):
+            return None
+        if len(work.selects) != len(idents):
+            return None
+        realiased = []
+        for item, ident in zip(work.selects, idents, strict=True):
+            base = item.this if isinstance(item, exp.Alias) else item
+            realiased.append(exp.Alias(this=base, alias=exp.to_identifier(ident)))
+        work.set("expressions", realiased)
+        return work
+
+    def _expand_star(self, select: exp.Select) -> bool:
+        if len(select.selects) != 1:
+            return False
+        from_clause = select.args.get("from_") or select.args.get("from")
+        if from_clause is None or select.args.get("joins"):
+            return False
+        tables = list(from_clause.find_all(exp.Table))
+        if len(tables) != 1:
+            return False
+        var = _tablevar_name(tables[0])
+        cols = self.declared.get(var) if var else None
+        if not cols:
+            return False
+        select.set("expressions", [exp.column(c) for c in cols])
+        return True
+
+    def _inline_tablevars(self, tree: exp.Expression, seen: frozenset[str], pos: int) -> None:
+        """Replace reads of locally filled table variables with their own
+        synthesized derivation from fills before this statement; a variable
+        with no reachable fill (a TVP parameter, dynamic SQL, a fill that
+        only happens later) stays cited as written."""
+        for table in list(tree.find_all(exp.Table)):
+            var = _tablevar_name(table)
+            if var is None or var in seen or var not in self.fills:
+                continue
+            sub = self.synthesize(var, seen, before=pos)
+            if sub is None:
+                continue
+            alias = table.alias or var.lstrip("@")
+            table.replace(
+                exp.Subquery(this=sub, alias=exp.TableAlias(this=exp.to_identifier(alias)))
+            )
+
+
+def _null_out_columns(node: exp.Expression, dead: set[str]) -> None:
+    if isinstance(node, exp.Union):
+        _null_out_columns(node.this, dead)
+        _null_out_columns(node.expression, dead)
+        return
+    if not isinstance(node, exp.Select):
+        return
+    for item in list(node.expressions):
+        if isinstance(item, exp.Alias) and item.alias and item.alias.lower() in dead:
+            item.replace(exp.Alias(this=exp.Null(), alias=exp.to_identifier(item.alias)))
+
+
+def expand_table_functions(sql: str, warnings: list[str] | None = None):
+    """(remaining sql, models): one ScriptModel per RETURNS @table function,
+    named exactly as the file writes it, its region removed from the text."""
+    from ripple.engine.sql_script import ScriptModel, _referenced_tables
+
+    models: list[ScriptModel] = []
+    regions: list[tuple[int, int]] = []
+    pos = 0
+    while True:
+        head = _find_tvf_head(sql, pos)
+        if head is None:
+            break
+        head_start, name_end, written_name = head
+        pos = name_end
+        i = _skip_ws_comments(sql, name_end)
+        if i < len(sql) and sql[i] == "(":
+            closed = _balanced_paren(sql, i)
+            if closed is None:
+                continue
+            i = closed
+        returns = _RETURNS_TABLE.match(sql, _skip_ws_comments(sql, i))
+        if returns is None:
+            continue  # a scalar function keeps the generic path
+        cols_close = _balanced_paren(sql, returns.end() - 1)
+        if cols_close is None:
+            continue
+        as_begin = _AS_BEGIN.match(sql, _skip_ws_comments(sql, cols_close))
+        if as_begin is None:
+            continue
+        body_end, region_end = _body_end(sql, as_begin.end())
+
+        out_var = returns.group("var").lower()
+        out_columns = _column_names(sql[returns.end() : cols_close - 1])
+        body = _TvfBody(
+            _parse_chunks(_chop_body(sql[as_begin.end() : body_end])), out_var, out_columns
+        )
+        synthesized = body.synthesize(out_var, frozenset())
+        derived_sql = safe_sql(synthesized, "tsql") if synthesized is not None else None
+        if derived_sql is None:
+            # return columns stay addressable even when no fill is traceable
+            if synthesized is not None and warnings is not None:
+                warnings.append(
+                    f"FUNCTION {written_name}: its synthesized body could not be "
+                    "regenerated; return columns are schema-only"
+                )
+            models.append(
+                ScriptModel(
+                    name=written_name, columns=out_columns, is_schema=True, is_function=True
+                )
+            )
+        else:
+            models.append(
+                ScriptModel(
+                    name=written_name,
+                    sql=derived_sql,
+                    parents=_referenced_tables(synthesized),
+                    columns=out_columns,
+                    is_function=True,
+                )
+            )
+        regions.append((head_start, region_end))
+        pos = region_end
+    if regions:
+        parts = []
+        prev = 0
+        for start, end in regions:
+            parts.append(sql[prev:start])
+            prev = end
+        parts.append(sql[prev:])
+        sql = "".join(parts)
+    return sql, models