edgesql 0.1.0__py3-none-any.whl

edgesql/__init__.py

@@ -0,0 +1,251 @@
+"""EdgeSQL. Callers depend on Query, Filter, Preview, and SqlCraftError."""
+
+from __future__ import annotations
+
+from sqlglot import exp
+
+from edgesql import preview as preview_mod
+from edgesql import transforms as xforms
+from edgesql.build import col, param, select_tree
+from edgesql.dialects import get as get_dialect
+from edgesql.errors import SqlCraftError, fail
+from edgesql.inspect import build_product_ir, columns_of, ctes_of, filters_of, params_of, tables_of
+from edgesql.parse import parse_statement
+from edgesql.scope import cte_names
+
+__all__ = ["Query", "SqlCraftError", "col", "param", "select"]
+
+
+class Query:
+    """A parsed statement. Transforms return a new Query."""
+
+    def __init__(
+        self,
+        tree: exp.Expression,
+        *,
+        dialect: str | None,
+        schema: dict | None,
+        mode: str,
+        ops: list | None = None,
+        zero_hit: bool = False,
+        focus: str | None = None,
+        meta: dict | None = None,
+    ) -> None:
+        self._tree = tree
+        self._dialect = dialect
+        self._schema = None if schema is None else {key: list(value) for key, value in schema.items()}
+        self._mode = mode
+        self._ops = list(ops or [])
+        self._zero_hit = zero_hit
+        self._focus = focus
+        self._meta = dict(meta or {})
+        self.audit = None
+
+    @classmethod
+    def parse(cls, sql: str, *, dialect: str, schema: dict | None = None, mode: str = "library") -> Query:
+        tree, meta = parse_statement(sql, dialect=dialect, schema=schema, mode=mode)
+        return cls(tree, dialect=dialect, schema=schema, mode=mode, meta=meta)
+
+    def to_sql(self, dialect: str | None = None, pretty: bool = False) -> str:
+        target = self._dialect if dialect is None else dialect
+        if target not in {"postgres", "databricks"}:
+            raise fail(
+                "render",
+                "UNKNOWN_DIALECT",
+                "Name the dialect on emit.",
+                hint="Use `postgres` or `databricks`.",
+            )
+        if self._dialect is not None and target != self._dialect:
+            raise fail(
+                "render",
+                "CROSS_DIALECT",
+                f"Cannot emit {target} for a query parsed as {self._dialect}.",
+                hint="Use `migrate(target)`, then preview, then apply.",
+            )
+        module = get_dialect(target)
+        if self._dialect is None:
+            reason = module.reject(self._tree)
+            if reason:
+                raise fail(
+                    "render",
+                    "DIALECT_NODE",
+                    reason,
+                    hint="Use the dialect that accepts the clause, or rewrite it.",
+                )
+        sql = module.finish(self._tree.sql(dialect=target, pretty=pretty), self._meta)
+        if pretty:
+            return sql
+        return _stable_canonical(sql, target)
+
+    def tables(self) -> list[str]:
+        return tables_of(self._tree)
+
+    def ctes(self) -> list[str]:
+        return ctes_of(self._tree)
+
+    def columns(self):
+        return columns_of(self._tree, self._schema)
+
+    def filters(self):
+        return filters_of(self._tree, self._schema, self._dialect)
+
+    def params(self) -> list[dict]:
+        return params_of(self._tree)
+
+    def to_tree(self) -> dict:
+        document, _paths = build_product_ir(self._tree, self._dialect)
+        return document
+
+    def cte(self, name: str) -> Query:
+        if name not in cte_names(self._tree):
+            raise fail(
+                "transform",
+                "UNKNOWN_SCOPE",
+                f"CTE {name} is not in this query.",
+                hint=_cte_hint(self._tree),
+            )
+        return self._spawn(self._tree, focus=f"cte:{name}")
+
+    def rename_column(self, old: str, new: str, *, table: str, keep_output_names: bool = True) -> Query:
+        tree, hits = xforms.rename_column(
+            self._tree,
+            self._schema,
+            old,
+            new,
+            table=table,
+            keep_output_names=keep_output_names,
+            focus=self._focus,
+        )
+        return self._spawn(tree, op="rename_column", zero_hit=hits == 0)
+
+    def and_where(self, predicate, *, scope: str | None = None) -> Query:
+        chosen = scope or self._focus or "root"
+        tree = xforms.and_where(self._tree, predicate, scope=chosen, dialect=self._dialect)
+        return self._spawn(tree, op="and_where")
+
+    def where(self, predicate) -> Query:
+        return self.and_where(predicate)
+
+    def limit(self, value: int) -> Query:
+        tree = self._tree.copy()
+        limited = tree.limit(value, copy=False)
+        return self._spawn(limited if limited is not None else tree, op="limit")
+
+    def remove_filter(self, filt) -> Query:
+        self._require_focus_scope(getattr(filt, "scope", None))
+        tree = xforms.remove_filter(self._tree, filt)
+        return self._spawn(tree, op="remove_filter")
+
+    def replace_filter(self, filt, predicate) -> Query:
+        self._require_focus_scope(getattr(filt, "scope", None))
+        tree = xforms.replace_filter(self._tree, filt, predicate, self._dialect)
+        return self._spawn(tree, op="replace_filter")
+
+    def with_cte(self, name: str, query: Query) -> Query:
+        tree = xforms.with_cte(self._tree, name, query._tree, query._dialect, self._dialect)
+        spawned = self._spawn(tree, op="with_cte")
+        if spawned._dialect is None and query._dialect is not None:
+            spawned._dialect = query._dialect
+        return spawned
+
+    def rename_cte(self, old: str, new: str) -> Query:
+        tree = xforms.rename_cte(self._tree, old, new)
+        focus = self._focus
+        if focus == f"cte:{old}":
+            focus = f"cte:{new}"
+        return self._spawn(tree, op="rename_cte", focus=focus)
+
+    def drop_cte(self, name: str) -> Query:
+        tree = xforms.drop_cte(self._tree, name)
+        focus = None if self._focus == f"cte:{name}" else self._focus
+        return self._spawn(tree, op="drop_cte", focus=focus)
+
+    def inline_cte(self, name: str) -> Query:
+        raise fail(
+            "transform",
+            "NOT_IN_V1",
+            "inline_cte is not in v1.",
+            hint="The method is a fast-follow.",
+        )
+
+    def set_params(self, patches: list[dict]) -> Query:
+        tree = xforms.set_params(self._tree, self._dialect, patches)
+        return self._spawn(tree, op="set_params")
+
+    def rewrite(self, fn) -> Query:
+        raise fail(
+            "transform",
+            "NOT_IN_V1",
+            "rewrite is not in v1.",
+            hint="The method is a fast-follow.",
+        )
+
+    def migrate(self, target: str):
+        if target not in {"postgres", "databricks"}:
+            raise fail(
+                "render",
+                "UNKNOWN_DIALECT",
+                f"Dialect {target} is not known.",
+                hint="Use `postgres` or `databricks`.",
+            )
+        return preview_mod.blocked_preview(self, ["MIGRATE_STUB"], ops=["migrate"])
+
+    def preview(self, fn=None):
+        return preview_mod.build_preview(self, fn)
+
+    def _spawn(self, tree, *, op: str | None = None, zero_hit: bool = False, focus: str | None = None) -> Query:
+        if focus is None and op != "drop_cte":
+            focus = self._focus
+        return Query(
+            tree,
+            dialect=self._dialect,
+            schema=self._schema,
+            mode=self._mode,
+            ops=[*self._ops, *([op] if op else [])],
+            zero_hit=self._zero_hit or zero_hit,
+            focus=focus,
+            meta=self._meta,
+        )
+
+    def _require_focus_scope(self, scope_name: str | None) -> None:
+        if self._focus and scope_name and scope_name != self._focus:
+            raise fail(
+                "transform",
+                "UNKNOWN_SCOPE",
+                f"Scope {scope_name} is outside {self._focus}.",
+                hint=_cte_hint(self._tree),
+            )
+
+
+def select(table: str) -> Query:
+    return Query(select_tree(table), dialect=None, schema=None, mode="library")
+
+
+def _cte_hint(tree) -> str:
+    names = cte_names(tree)
+    listed = ", ".join(names) if names else "(none)"
+    return f"Known CTE names: {listed}."
+
+
+def _stable_canonical(sql: str, dialect: str) -> str:
+    """Emit the print that a second parse reproduces.
+
+    SQLGlot sometimes parks a comment on a node that the next parse drops.
+    Canonical SQL is the print that survives that second parse.
+    """
+    module = get_dialect(dialect)
+    again = _reprint(sql, dialect, module)
+    if again is None or again == sql:
+        return sql
+    third = _reprint(again, dialect, module)
+    if third == again:
+        return again
+    return sql
+
+
+def _reprint(sql: str, dialect: str, module) -> str | None:
+    try:
+        tree, meta = parse_statement(sql, dialect=dialect, schema=None, mode="library")
+    except SqlCraftError:
+        return None
+    return module.finish(tree.sql(dialect=dialect, pretty=False), meta)

edgesql/build.py

@@ -0,0 +1,91 @@
+"""Expression builders. They never paste a user string into SQL text."""
+
+from __future__ import annotations
+
+from sqlglot import exp
+
+from edgesql.errors import fail
+
+
+class BindBuilder:
+    def __init__(self, name: str) -> None:
+        self.name = name
+
+
+class ColumnBuilder:
+    def __init__(self, name: str) -> None:
+        self.name = name
+
+    def eq(self, other):
+        return self._compare(exp.EQ, other)
+
+    def ne(self, other):
+        return self._compare(exp.NEQ, other)
+
+    def gt(self, other):
+        return self._compare(exp.GT, other)
+
+    def gte(self, other):
+        return self._compare(exp.GTE, other)
+
+    def lt(self, other):
+        return self._compare(exp.LT, other)
+
+    def lte(self, other):
+        return self._compare(exp.LTE, other)
+
+    def _compare(self, kind, other):
+        return kind(this=exp.column(self.name), expression=literal(other))
+
+
+def col(name: str) -> ColumnBuilder:
+    return ColumnBuilder(name)
+
+
+def param(name: str) -> BindBuilder:
+    return BindBuilder(name)
+
+
+def literal(value):
+    if isinstance(value, exp.Expression):
+        return value
+    if isinstance(value, BindBuilder):
+        return exp.Placeholder(this=value.name)
+    if isinstance(value, ColumnBuilder):
+        return exp.column(value.name)
+    if value is None:
+        return exp.Null()
+    if isinstance(value, bool):
+        return exp.Boolean(this=value)
+    if isinstance(value, int):
+        return exp.Literal.number(value)
+    if isinstance(value, float):
+        return exp.Literal.number(value)
+    if isinstance(value, str):
+        return exp.Literal.string(value)
+    raise fail(
+        "transform",
+        "STRUCTURE_PATCH",
+        "Lane A can patch the right-hand value only.",
+        hint="Patch the right-hand value only.",
+    )
+
+
+def select_tree(table: str) -> exp.Select:
+    return exp.Select(expressions=[exp.Star()]).from_(table_expr(table))
+
+
+def table_expr(name: str) -> exp.Table:
+    parts = [exp.to_identifier(part) for part in name.split(".") if part]
+    if len(parts) == 1:
+        return exp.Table(this=parts[0])
+    if len(parts) == 2:
+        return exp.Table(this=parts[1], db=parts[0])
+    if len(parts) == 3:
+        return exp.Table(this=parts[2], db=parts[1], catalog=parts[0])
+    raise fail(
+        "parse",
+        "UNKNOWN_SOURCE",
+        f"Table name {name} has too many parts.",
+        hint="Check aliases and CTE names.",
+    )

edgesql/dialects/__init__.py

@@ -0,0 +1,14 @@
+"""Dialect registry. A new dialect adds a module and one entry here."""
+
+from __future__ import annotations
+
+from edgesql.dialects import databricks, postgres
+
+DIALECTS = {
+    "postgres": postgres,
+    "databricks": databricks,
+}
+
+
+def get(name: str):
+    return DIALECTS.get(name)

edgesql/dialects/databricks.py

@@ -0,0 +1,67 @@
+"""Databricks facts. Marker style, quotes, and nodes this dialect rejects."""
+
+from __future__ import annotations
+
+from sqlglot import exp
+
+from edgesql.dialects.maxrec import insert_max_recursion, split_max_recursion
+
+NAME = "databricks"
+QUOTES = "`"
+MARKERS = (":name", "?")
+ALLOWS_MAX_RECURSION = True
+
+_HINT = "Use the dialect that accepts the clause, or rewrite it."
+
+
+def prepare(sql: str) -> tuple[str, dict]:
+    stripped, level = split_max_recursion(sql)
+    meta = {"max_recursion": level} if level is not None else {}
+    return stripped, meta
+
+
+def finish(sql: str, meta: dict) -> str:
+    level = (meta or {}).get("max_recursion")
+    if level is None:
+        return sql
+    return insert_max_recursion(sql, int(level))
+
+
+def reject(tree: exp.Expression) -> str | None:
+    for distinct in tree.find_all(exp.Distinct):
+        if distinct.args.get("on") is not None:
+            return (
+                "DISTINCT ON is accepted on Postgres, not on Databricks. "
+                "Use the same dialect, or migrate, preview, and apply."
+            )
+    if list(tree.find_all(exp.Fetch)):
+        return "FETCH is accepted on Postgres, not on Databricks. Use the same dialect, or migrate, preview, and apply."
+    for with_ in tree.find_all(exp.With):
+        if not with_.args.get("recursive"):
+            continue
+        for cte in with_.expressions:
+            if _recursive_step_is_plain_union(cte.this):
+                return "A recursive Databricks CTE must use UNION ALL."
+    return None
+
+
+def _recursive_step_is_plain_union(body: exp.Expression) -> bool:
+    node = body.this if isinstance(body, exp.Subquery) else body
+    if not isinstance(node, exp.SetOperation):
+        return False
+    return not (isinstance(node, exp.Union) and node.args.get("distinct") is False)
+
+
+def mixed_markers(tree: exp.Expression) -> str | None:
+    named = False
+    unnamed = False
+    for node in tree.walk():
+        if not isinstance(node, exp.Placeholder):
+            continue
+        if node.args.get("jdbc") or node.args.get("this") is None:
+            unnamed = True
+        else:
+            named = True
+    if named and unnamed:
+        return "This statement mixes :name and ? markers."
+    return None

edgesql/dialects/maxrec.py

@@ -0,0 +1,211 @@
+"""Keep the Databricks MAX RECURSION LEVEL clause that SQLGlot does not model."""
+
+from __future__ import annotations
+
+
+def _starts_ident(text: str, index: int) -> bool:
+    if index <= 0:
+        return True
+    return not (text[index - 1].isalnum() or text[index - 1] == "_")
+
+
+def _keyword_at(text: str, index: int, word: str) -> bool:
+    end = index + len(word)
+    if end > len(text) or not _starts_ident(text, index):
+        return False
+    if text[index:end].lower() != word:
+        return False
+    return not (end < len(text) and (text[end].isalnum() or text[end] == "_"))
+
+
+def _skip_ignorable(text: str, index: int) -> int:
+    n = len(text)
+    while index < n:
+        if text[index].isspace():
+            index += 1
+            continue
+        if text.startswith("--", index):
+            nl = text.find("\n", index)
+            index = n if nl < 0 else nl + 1
+            continue
+        if text.startswith("/*", index):
+            index = _skip_block(text, index + 2, copy=False)[0]
+            continue
+        break
+    return index
+
+
+def _skip_block(text: str, index: int, *, copy: bool) -> tuple[int, str]:
+    """Skip a block comment. Databricks allows nested block comments."""
+    n = len(text)
+    depth = 1
+    out = ["/*"] if copy else []
+    while index < n and depth:
+        if text.startswith("/*", index):
+            depth += 1
+            if copy:
+                out.append("/*")
+            index += 2
+            continue
+        if text.startswith("*/", index):
+            depth -= 1
+            if copy:
+                out.append("*/")
+            index += 2
+            continue
+        if copy:
+            out.append(text[index])
+        index += 1
+    return index, "".join(out)
+
+
+def split_max_recursion(sql: str) -> tuple[str, int | None]:
+    """Remove one MAX RECURSION LEVEL clause. Return the SQL and the level."""
+    n = len(sql)
+    index = 0
+    out: list[str] = []
+    level: int | None = None
+    quote: str | None = None
+    while index < n:
+        if quote:
+            out.append(sql[index])
+            if sql[index] == quote:
+                if quote == "'" and index + 1 < n and sql[index + 1] == "'":
+                    out.append(sql[index + 1])
+                    index += 2
+                    continue
+                quote = None
+            index += 1
+            continue
+        if sql.startswith("--", index):
+            nl = sql.find("\n", index)
+            end = n if nl < 0 else nl
+            out.append(sql[index:end])
+            index = end
+            continue
+        if sql.startswith("/*", index):
+            index, chunk = _skip_block(sql, index + 2, copy=True)
+            out.append(chunk)
+            continue
+        if sql[index] in {"'", '"', "`"}:
+            quote = sql[index]
+            out.append(sql[index])
+            index += 1
+            continue
+        if _keyword_at(sql, index, "max"):
+            j = _skip_ignorable(sql, index + 3)
+            if _keyword_at(sql, j, "recursion"):
+                k = _skip_ignorable(sql, j + len("recursion"))
+                if _keyword_at(sql, k, "level"):
+                    m = _skip_ignorable(sql, k + len("level"))
+                    number, end = _read_int(sql, m)
+                    if number is not None:
+                        level = number
+                        index = end
+                        continue
+        out.append(sql[index])
+        index += 1
+    return "".join(out), level
+
+
+def _read_int(text: str, index: int) -> tuple[int | None, int]:
+    if index >= len(text) or not text[index].isdigit():
+        return None, index
+    end = index
+    while end < len(text) and text[end].isdigit():
+        end += 1
+    return int(text[index:end]), end
+
+
+def _match_paren(text: str, open_at: int) -> int:
+    """Return the index just after the `)` that matches `text[open_at] == '('`."""
+    n = len(text)
+    depth = 0
+    index = open_at
+    quote: str | None = None
+    while index < n:
+        if quote:
+            if text[index] == quote:
+                if quote == "'" and index + 1 < n and text[index + 1] == "'":
+                    index += 2
+                    continue
+                quote = None
+            index += 1
+            continue
+        if text.startswith("--", index):
+            nl = text.find("\n", index)
+            index = n if nl < 0 else nl + 1
+            continue
+        if text.startswith("/*", index):
+            index = _skip_block(text, index + 2, copy=False)[0]
+            continue
+        char = text[index]
+        if char in {"'", '"', "`"}:
+            quote = char
+            index += 1
+            continue
+        if char == "(":
+            depth += 1
+        elif char == ")":
+            depth -= 1
+            if depth == 0:
+                return index + 1
+        index += 1
+    return n
+
+
+def insert_max_recursion(sql: str, level: int) -> str:
+    """Put the clause after the WITH list and before the main query."""
+    index = _skip_ignorable(sql, 0)
+    if not _keyword_at(sql, index, "with"):
+        return f"MAX RECURSION LEVEL {level} {sql}"
+    index = _skip_ignorable(sql, index + len("with"))
+    if _keyword_at(sql, index, "recursive"):
+        index += len("recursive")
+    guard = 0
+    while guard < 100:
+        guard += 1
+        as_at = _find_keyword(sql, index, "as")
+        if as_at < 0:
+            break
+        after = _skip_ignorable(sql, as_at + 2)
+        if after >= len(sql) or sql[after] != "(":
+            index = as_at + 2
+            continue
+        close = _match_paren(sql, after)
+        nxt = _skip_ignorable(sql, close)
+        if nxt < len(sql) and sql[nxt] == ",":
+            index = nxt + 1
+            continue
+        clause = f" MAX RECURSION LEVEL {level}"
+        return sql[:close] + clause + sql[close:]
+    return sql
+
+
+def _find_keyword(text: str, index: int, word: str) -> int:
+    n = len(text)
+    quote: str | None = None
+    while index < n:
+        if quote:
+            if text[index] == quote:
+                if quote == "'" and index + 1 < n and text[index + 1] == "'":
+                    index += 2
+                    continue
+                quote = None
+            index += 1
+            continue
+        if text.startswith("--", index):
+            nl = text.find("\n", index)
+            index = n if nl < 0 else nl + 1
+            continue
+        if text.startswith("/*", index):
+            index = _skip_block(text, index + 2, copy=False)[0]
+            continue
+        if text[index] in {"'", '"', "`"}:
+            quote = text[index]
+            index += 1
+            continue
+        if _keyword_at(text, index, word):
+            return index
+        index += 1
+    return -1

edgesql/dialects/postgres.py

@@ -0,0 +1,53 @@
+"""Postgres facts. Marker style, quotes, and nodes this dialect rejects."""
+
+from __future__ import annotations
+
+from sqlglot import exp
+
+from edgesql.dialects.maxrec import split_max_recursion
+from edgesql.errors import fail
+
+NAME = "postgres"
+QUOTES = '"'
+MARKERS = ("$n", "%(name)s")
+ALLOWS_MAX_RECURSION = False
+
+_QUALIFY = "QUALIFY is accepted on Databricks, not on Postgres. Use the same dialect, or migrate, preview, and apply."
+_HINT = "Use the dialect that accepts the clause, or rewrite it."
+
+
+def prepare(sql: str) -> tuple[str, dict]:
+    stripped, level = split_max_recursion(sql)
+    if level is not None:
+        raise fail(
+            "parse",
+            "DIALECT_NODE",
+            "MAX RECURSION LEVEL is accepted on Databricks, not on Postgres. "
+            "Use the same dialect, or migrate, preview, and apply.",
+            hint=_HINT,
+        )
+    return stripped, {}
+
+
+def finish(sql: str, meta: dict) -> str:
+    return sql
+
+
+def reject(tree: exp.Expression) -> str | None:
+    if list(tree.find_all(exp.Qualify)):
+        return _QUALIFY
+    for join in tree.find_all(exp.Join):
+        kind = str(join.args.get("kind") or "").upper()
+        if kind == "SEMI":
+            return (
+                "SEMI is accepted on Databricks, not on Postgres. Use the same dialect, or migrate, preview, and apply."
+            )
+        if kind == "ANTI":
+            return (
+                "ANTI is accepted on Databricks, not on Postgres. Use the same dialect, or migrate, preview, and apply."
+            )
+    return None
+
+
+def mixed_markers(tree: exp.Expression) -> str | None:
+    return None