sqlleaf 0.0.1__py3-none-any.whl

sqlleaf/models/node/base.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import typing as t
+from dataclasses import asdict
+
+from sqlglot import exp
+
+from sqlleaf import util
+from sqlleaf.models.context import GeneratorContext, PositionContext
+from sqlleaf.models.query import Q, QueryHolder
+
+
+class NodeAttributes:
+    KIND: str = ""
+    WITH_POSITIONS: bool = False
+
+    def __init__(
+        self,
+        gen_ctx: GeneratorContext,
+        pos_ctx: PositionContext,
+        name: str | None = "",
+    ):
+        self.expr: exp.Expr = gen_ctx.expr
+        self._data_type: exp.DataType | None = gen_ctx.data_type
+        self.data_type: str = str(self._data_type) if self._data_type else ""
+
+        self.name: str = name
+        self.kind: str = self.KIND
+        self.with_positions: bool = self.WITH_POSITIONS
+        self.ctx: PositionContext = pos_ctx
+
+    # Allows the class to be used a networkx node
+    def __hash__(self) -> int:
+        return hash(self.full_name)
+
+    def get_data_type(self) -> exp.DataType:
+        assert self._data_type is not None
+        return self._data_type
+
+    def wrap(self, name: str) -> str:
+        return f"{self.kind}[{name}]"
+
+    def fields(self) -> dict[str, str]:
+        return {"type": self.data_type}
+
+    def friendly_fields(self) -> dict[str, str]:
+        return {}
+
+    def get_name(self) -> str:
+        return self.name
+
+    def _build_name(self, fields_dict: dict[str, str], with_positions: bool = False) -> str:
+        """Build a formatted name with fields."""
+        fields_dict = fields_dict.copy()
+        main_parts = [f"name={self.get_name()}"]
+
+        type_val = fields_dict.pop("type", None)
+        if type_val:
+            main_parts.append(f"type={type_val}")
+
+        props = " ".join([f"{k}={v}" for k, v in fields_dict.items() if v])
+        if props:
+            main_parts.append(f"properties=[{props}]")
+
+        content = " ".join(main_parts)
+        if with_positions:
+            content = f"{content} position=[{self.ctx.as_str()}]"
+
+        return self.wrap(content)
+
+    @property
+    def full_name(self) -> str:
+        return self._build_name(self.fields(), with_positions=self.with_positions)
+
+    @property
+    def friendly_name(self) -> str:
+        name = self.get_name()
+        fields = self.friendly_fields()
+
+        parts = []
+        if name:
+            parts.append(name)
+
+        for k, v in fields.items():
+            if v:
+                parts.append(f"{k}={v}")
+
+        content = " ".join(parts)
+        return self.wrap(content)
+
+    @property
+    def id(self) -> str:
+        name = "node:" + util.short_sha256_hash(self.full_name)
+        return name
+
+    def to_dict(self) -> dict:
+        return {
+            "id": self.id,
+            "kind": self.kind,
+            "name": self.name,
+            "data_type": self.data_type,
+            **self.fields(),
+            "position": asdict(self.ctx) if self.with_positions else {},
+        }
+
+
+class EdgeAttributes:
+    def __init__(
+        self,
+        parent: NodeAttributes,
+        child: NodeAttributes,
+        query: Q,
+        path_idx: int,
+    ):
+        self.parent = parent
+        self.child = child
+        self.query = query
+        self.path_idx = path_idx  # The sequence number if a duplicate path
+
+    @property
+    def id(self) -> str:
+        edge_id = ":".join([
+            str(s)
+            for s in [
+                self.parent.id,
+                self.child.id,
+                self.path_idx,
+            ]
+        ])
+        return "edge:" + util.short_sha256_hash(edge_id)
+
+    def to_dict(self) -> dict:
+        result = {
+            "id": self.id,
+            "source": self.parent.id,
+            "target": self.child.id,
+            "query": self.query.id,
+            "path_index": self.path_idx,
+        }
+        return result
+
+
+class GraphAttributes:
+    def __init__(self):
+        self.queries: t.List[QueryHolder] = []
+
+    def add_query_to_graph(self, holder: QueryHolder):
+        self.queries.append(holder)

sqlleaf/models/node/column.py

@@ -0,0 +1,345 @@
+from __future__ import annotations
+
+import logging
+import typing as t
+
+from sqlglot import exp
+from sqlglot.optimizer import Scope
+from sqlglot.optimizer.scope import ScopeType
+
+from sqlleaf import exception
+from sqlleaf.models.context import GeneratorContext, PositionContext
+from sqlleaf.models.node import NodeAttributes
+from sqlleaf.typing import TableOrScopeType, TableSubtype, TableType
+
+logger = logging.getLogger("sqlleaf")
+
+
+class ColumnNode(NodeAttributes):
+    KIND = "column"
+
+    def __init__(
+        self,
+        catalog: str,
+        schema: str,
+        table: str,
+        column: str,
+        gen_ctx: GeneratorContext,
+        pos_ctx: PositionContext,
+        source: TableOrScopeType | exp.Select | exp.Subquery | exp.Lateral | None = None,
+    ):
+        super().__init__(gen_ctx, pos_ctx, name=column)
+        self.catalog = catalog
+        self.schema = schema
+        self.table = table
+        self.member = ""
+
+        self.parent_kind: str = ""
+        self.parent_subkind: str = ""
+        self.source_scope: TableOrScopeType | None = None
+
+        self.set_table_properties(catalog, schema, table, gen_ctx, source=source)
+        if source is not None and not isinstance(source, (exp.Subquery, Scope)):
+            self._apply_rename(source, gen_ctx.query.dialect)
+
+        # TODO: new algorithm
+        # if table_type == "cte":
+        #     self.member = gen_ctx.node.recursive_cte_member_kind
+
+    def to_dict(self):
+        d = super().to_dict()
+        d.update({
+            "catalog": self.catalog,
+            "schema": self.schema,
+            "table": self.table,
+        })
+        return d
+
+    def _apply_rename(self, source: exp.Table | exp.Select | exp.Lateral, dialect: str) -> None:
+        """
+        Change the column's source table to be its fully qualified name, not its alias,
+        so that the ColumnNode is provided complete information.
+        """
+        column = t.cast(exp.Column, self.expr)
+        _c = column.copy()
+
+        if isinstance(source, exp.Table):
+            if source.catalog:
+                column.set("catalog", exp.to_identifier(source.catalog))
+            if source.db:
+                column.set("db", exp.to_identifier(source.db))
+            if source.name:
+                if dialect == "snowflake":
+                    if source.this.args.get("quoted", False):  # exp.Identifier
+                        column.set("table", exp.to_identifier(source.name))
+                else:
+                    column.set("table", exp.to_identifier(source.name))
+            if _c != column:
+                logger.debug(f"Renamed node {_c.sql()} to {column.sql()}")
+
+            self.expr = column
+            self.catalog = column.catalog
+            self.schema = column.db
+            self.table = column.table
+
+    def set_table_properties(
+        self,
+        catalog: str,
+        schema: str,
+        table: str,
+        gen_ctx: GeneratorContext,
+        source: TableOrScopeType | exp.Select | exp.Subquery | exp.Lateral | None = None,
+    ) -> None:
+        """
+        Figure out the table's kind (view, table, etc) and its subkind (temp, recursive, etc) by
+        inspecting the query's scope and the table mapping.
+        """
+        scope = gen_ctx.scope
+
+        # logical_source is what we use to determine CTE/UDTF etc.
+        logical_source = source
+        if isinstance(scope, Scope):
+            # Always try to get the Scope from sources if it exists for this alias
+            logical_source = scope.sources.get(table) or source
+
+            if not logical_source:
+                # Nested 'rows_from' queries have their aliases in 'references'
+                logical_source = dict(scope.references).get(table)
+
+        self.source_scope = logical_source
+
+        # If it's in references but NOT in sources, it's a derived table (e.g. ROWS FROM)
+        if (
+            logical_source is not None
+            and isinstance(scope, Scope)
+            and table not in scope.sources
+            and table in dict(scope.references)
+        ):
+            self.parent_kind = TableType.DERIVED_TABLE
+            return
+
+        # Determine logical properties (CTE, UDTF, Derived Table)
+        if self._set_kind_of_derived_table(logical_source, table, gen_ctx):
+            return
+
+        # Determine physical properties (TABLE, VIEW, etc.)
+        self._set_kind_of_physical_table(catalog, schema, table, source or logical_source, gen_ctx)
+
+    def _set_kind_of_derived_table(
+        self,
+        source: TableOrScopeType | exp.Select | exp.Subquery | exp.Lateral | None,
+        table: str,
+        gen_ctx: GeneratorContext,
+    ) -> bool:
+        """
+        Identify if the source is a logical entity such as a CTE, UDTF, or derived table and set its properties.
+        """
+        if isinstance(source, Scope):
+            if source.scope_type == ScopeType.CTE:
+                self._set_cte_properties(source, table, gen_ctx)
+                return True
+
+            if source.scope_type == ScopeType.DERIVED_TABLE:
+                self.parent_kind = TableType.DERIVED_TABLE
+                return True
+
+            if source.scope_type == ScopeType.UDTF:
+                self.parent_kind = TableType.UDTF
+                return True
+
+        if isinstance(source, exp.Table) and "rows_from" in source.args:
+            self.parent_kind = TableType.DERIVED_TABLE
+            return True
+
+        if isinstance(source, exp.Table) and isinstance(source.this, exp.Anonymous):
+            # Table function, e.g. SELECT * FROM FUNCTION()
+            self.parent_kind = TableType.UDTF
+            return True
+
+        if isinstance(source, (exp.Select, exp.Subquery)):
+            self.parent_kind = TableType.DERIVED_TABLE
+            return True
+
+        return False
+
+    def _set_cte_properties(self, source: Scope, table: str, gen_ctx: GeneratorContext) -> None:
+        """
+        Resolve and set properties specific to Common Table Expressions, including recursive and materialized status.
+        """
+        self.parent_kind = TableType.CTE
+        logger.debug("Set node to be a CTE.")
+
+        # Check if it's recursive or materialized
+        if not source.parent:
+            return
+
+        scope = gen_ctx.scope
+        if not isinstance(scope, Scope):
+            return
+
+        selected_table, _ = scope.selected_sources.get(table, (None, None))
+        if not selected_table:
+            # Try the regular sources; may occur when column nested inside subquery references a CTE outside itself
+            # In this case, skip it.
+            selected_source = scope.sources.get(table)
+            if selected_source and selected_source.scope_type == ScopeType.CTE:
+                return
+            else:
+                message = f"Table '{table}' is referenced but there is no FROM containing it."
+                raise exception.MappingError(message=message)
+
+        for cte in source.parent.ctes:
+            if cte.alias_or_name == selected_table.name:
+                if cte.args.get("materialized"):
+                    self.parent_subkind = TableSubtype.MATERIALIZED
+                else:
+                    with_ = cte.parent
+                    if isinstance(with_, exp.With) and with_.recursive:
+                        logger.debug("Set node to be a recursive CTE.")
+                        self.parent_subkind = TableSubtype.RECURSIVE
+                break
+
+    def _set_kind_of_physical_table(
+        self,
+        catalog: str,
+        schema: str,
+        table: str,
+        source: TableOrScopeType | exp.Select | exp.Subquery | exp.Lateral | None,
+        gen_ctx: GeneratorContext,
+    ) -> None:
+        """
+        Determine the physical type of the table, such as a base table or view, by resolving its name against the object mapping.
+        """
+        if isinstance(source, exp.Table):
+            tokens = [str(s) for s in source.parts]
+        else:
+            tokens = [catalog, schema, table]
+
+        # Get the table type from the mapping
+        name = ".".join([tok for tok in tokens if tok])
+        if not name:
+            self.parent_kind = TableType.TABLE
+            return
+
+        tab = exp.to_table(name, dialect=gen_ctx.query.dialect)
+        query = gen_ctx.query.object_mapping.get_table_or_stage(table=tab, raise_on_missing=True)
+
+        if not query or query.kind == "ctas":
+            self.parent_kind = TableType.TABLE
+        else:
+            self.parent_kind = TableType(query.kind)
+
+        if query and query.property:
+            self.parent_subkind = TableSubtype(query.property)
+
+    @property
+    def friendly_name(self) -> str:
+        tokens = [self.catalog, self.schema, self.table, self.name]
+        name = ".".join([tok for tok in tokens if tok])
+        fields = self.friendly_fields()
+
+        parts = []
+        if name:
+            parts.append(name)
+
+        for k, v in fields.items():
+            if v:
+                parts.append(f"{k}={v}")
+
+        content = " ".join(parts)
+        return self.wrap(content)
+
+    def as_table(self) -> exp.Table:
+        return exp.table_(catalog=self.catalog, db=self.schema, table=self.table)
+
+    def fields(self) -> dict[str, str]:
+        f = {"kind": self.parent_kind}
+        if self.parent_subkind:
+            f["subkind"] = self.parent_subkind
+
+        f |= {
+            "table": self.table,
+            "schema": self.schema,
+            "type": self.data_type,
+        }
+
+        if self.parent_kind == TableType.CTE and self.parent_subkind == TableSubtype.RECURSIVE:
+            f["member"] = self.member
+
+        if self.parent_kind in [TableType.CTE, TableType.DERIVED_TABLE]:
+            f["statement"] = str(self.ctx.statement_index)
+
+        return f
+
+
+class FileColumnNode(NodeAttributes):
+    KIND = "column"
+
+    def __init__(
+        self,
+        column: str,
+        file_format: str,
+        file_path: str,
+        gen_ctx: GeneratorContext,
+        pos_ctx: PositionContext,
+    ):
+        super().__init__(gen_ctx, pos_ctx, name=column)
+        self.file_format = file_format
+        self.file_path = file_path
+
+    def to_dict(self):
+        d = super().to_dict()
+        d.update({
+            "format": self.file_format,
+            "path": self.file_path,
+        })
+        return d
+
+    def fields(self) -> dict[str, str]:
+        return {
+            "type": self.data_type,
+            "kind": "file",
+            "format": self.file_format,
+            "path": self.file_path,
+        }
+
+    def friendly_fields(self) -> dict[str, str]:
+        return {"path": self.file_path}
+
+
+class StageColumnNode(NodeAttributes):
+    KIND = "column"
+
+    def __init__(
+        self,
+        column: str,
+        stage: exp.Var,
+        gen_ctx: GeneratorContext,
+        pos_ctx: PositionContext,
+        path: str = "",
+    ):
+        super().__init__(gen_ctx, pos_ctx, name=column)
+        self.stage = stage.name.removeprefix("@").replace('"', "")
+        self.path = path
+
+    def to_dict(self):
+        d = super().to_dict()
+        d.update({
+            "stage": self.stage,
+            "path": self.path,
+        })
+        return d
+
+    def fields(self) -> dict[str, str]:
+        return {
+            "type": self.data_type,
+            "kind": "stage",
+            "stage": self.stage,
+            "path": self.path,
+        }
+
+    def friendly_fields(self) -> dict[str, str]:
+        return {
+            "stage": self.stage,
+            "path": self.path,
+        }

sqlleaf/models/node/dynamodb.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import logging
+
+from sqlleaf.models.context import GeneratorContext, PositionContext
+from sqlleaf.models.node import NodeAttributes
+
+logger = logging.getLogger("sqlleaf")
+
+
+class DynamoDbNode(NodeAttributes):
+    KIND = "column"
+
+    def __init__(
+        self,
+        gen_ctx: GeneratorContext,
+        pos_ctx: PositionContext,
+        column: str,
+    ):
+        super().__init__(gen_ctx, pos_ctx, name=column)
+        self.path = gen_ctx.expr.name
+        self.table_name = self.path.removeprefix("dynamodb://")
+
+    def to_dict(self):
+        d = super().to_dict()
+        d.update({"table": self.table_name})
+        return d
+
+    def fields(self) -> dict[str, str]:
+        return {
+            "type": self.data_type,
+            "kind": "dynamodb",
+            "table": self.table_name,
+        }
+
+    def friendly_fields(self) -> dict[str, str]:
+        return {
+            "path": self.path,
+        }

sqlleaf/models/node/function.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot import exp
+
+from sqlleaf import util
+from sqlleaf.models.context import GeneratorContext, PositionContext
+from sqlleaf.models.node import NodeAttributes
+
+
+class FunctionNode(NodeAttributes):
+    KIND = "function"
+    WITH_POSITIONS = True
+
+    def __init__(self, gen_ctx: GeneratorContext, pos_ctx: PositionContext):
+        expr = t.cast(t.Union[exp.Binary, exp.Func], gen_ctx.expr)
+        if isinstance(expr, exp.Binary):
+            name = expr.key
+        elif isinstance(expr, exp.Case):
+            name = "CASE"
+        elif isinstance(expr, exp.Anonymous):
+            _, name = util.get_udf_name(expr)
+        else:
+            name = util.calculate_function_name(expr, gen_ctx.query.dialect)
+
+        super().__init__(gen_ctx, pos_ctx, name=name)
+
+    def fields(self) -> dict[str, str]:
+        return {"type": self.data_type}
+
+    def get_name(self) -> str:
+        return f"{self.name}".upper()

sqlleaf/models/node/interval.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot import exp
+
+from sqlleaf.models.context import GeneratorContext, PositionContext
+from sqlleaf.models.node import NodeAttributes
+
+
+class IntervalNode(NodeAttributes):
+    KIND = "interval"
+    WITH_POSITIONS = True
+
+    def __init__(self, gen_ctx: GeneratorContext, pos_ctx: PositionContext):
+        expr = t.cast(exp.Interval, gen_ctx.expr)
+        name = f'"{str(expr.this.name)} {str(expr.unit)}"'
+        super().__init__(gen_ctx, pos_ctx, name=name)
+
+    def fields(self) -> dict[str, str]:
+        return {"type": self.data_type}

sqlleaf/models/node/json.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import logging
+import typing as t
+
+from sqlglot import exp
+
+from sqlleaf.models.context import GeneratorContext, PositionContext
+from sqlleaf.models.node import NodeAttributes
+
+logger = logging.getLogger("sqlleaf")
+
+
+class JsonPathNode(NodeAttributes):
+    KIND = "jsonpath"
+
+    def __init__(self, gen_ctx: GeneratorContext, pos_ctx: PositionContext):
+        expr = t.cast(exp.JSONExtract, gen_ctx.expr)
+
+        self.selectors = self.json_selectors(expr)
+        self.selector = "".join([str(s) for s in self.selectors])
+        self.selector_depth = len(self.selectors)
+
+        super().__init__(gen_ctx, pos_ctx, name=self.selector)
+
+    def json_selectors(self, expr: exp.JSONExtract):
+        """
+        Collect all the JSON path elements recursively.
+        e.g.
+            SELECT my_json -> 'a' -> 'b'
+        produces
+            ['a', 'b']
+        """
+        methods = {
+            # Mapping of expression -> attribute getter
+            exp.JSONPathKey: None,
+            exp.JSONPathSubscript: "this",
+        }
+        elements = list(expr.expression.find_all(tuple(methods)))
+
+        left = expr.left
+        while isinstance(left, (exp.JSONExtract, exp.JSONExtractScalar)):
+            elements.extend(list(left.expression.find_all(tuple(methods))))
+            left = left.left
+        elements.reverse()
+
+        result = []
+        for elem in elements:
+            getter = methods[type(elem)]
+            result.append(getattr(elem, getter) if getter else elem)
+        return result
+
+    def fields(self) -> dict[str, str]:
+        return {"depth": str(self.selector_depth)}
+
+    def friendly_fields(self) -> dict[str, str]:
+        return {}

sqlleaf/models/node/literal.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from sqlleaf.models.node import NodeAttributes
+
+
+class LiteralNode(NodeAttributes):
+    KIND = "literal"
+    WITH_POSITIONS = True
+
+    def fields(self) -> dict[str, str]:
+        return {"type": self.data_type}
+
+    def get_name(self) -> str:
+        if self.expr.is_string:
+            return f'"{self.name.replace("'", "")}"'
+        else:
+            return self.name

sqlleaf/models/node/null.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from sqlleaf.models.node import NodeAttributes
+
+
+class NullNode(NodeAttributes):
+    KIND = "null"
+    WITH_POSITIONS = True
+
+    def fields(self) -> dict[str, str]:
+        return {"type": "NULL"}
+
+    def get_name(self) -> str:
+        return "NULL"