sqlleaf 0.0.1__py3-none-any.whl

sqlleaf/models/node/pivot.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from sqlleaf.models.context import GeneratorContext, PositionContext
+from sqlleaf.models.node import NodeAttributes
+
+
+class _PivotNode(NodeAttributes):
+    def __init__(self, gen_ctx: GeneratorContext, pos_ctx: PositionContext):
+        super().__init__(gen_ctx, pos_ctx)
+        self.source: str = ""
+        self.target: str = ""
+
+    def set(self, source: str, target: str):
+        self.source = source
+        self.target = target
+
+    def fields(self) -> dict[str, str]:
+        f = {}
+        # Keep empty strings as values to match expected test output
+        f["source"] = self.source
+        f["target"] = self.target
+        f["statement"] = str(self.ctx.statement_index)
+        return f
+
+    def get_name(self) -> str:
+        return ""
+
+
+class PivotNode(_PivotNode):
+    KIND = "pivot"
+
+
+class UnpivotNode(_PivotNode):
+    KIND = "unpivot"

sqlleaf/models/node/program.py

@@ -0,0 +1,27 @@
+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 ProgramNode(NodeAttributes):
+    KIND = "program"
+    WITH_POSITIONS = True
+
+    def __init__(self, gen_ctx: GeneratorContext, pos_ctx: PositionContext):
+        pos_ctx = pos_ctx.replace(select_index=0)  # Prevent duplicate nodes
+        copy_stmt = gen_ctx.query.get_original_self().statement
+        expr = t.cast(exp.Copy, copy_stmt)
+
+        # TODO: put these inside COPY to prevent lookups to the original query here
+        program = expr.args["params"][0].name
+        name, args = (program + " ").split(" ", maxsplit=1)
+        super().__init__(gen_ctx, pos_ctx, name=name)
+        self.program_args = args.strip()
+
+    def fields(self) -> dict[str, str]:
+        return {"args": f"'{self.program_args}'"}

sqlleaf/models/node/sequence.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from sqlleaf.models.context import GeneratorContext, PositionContext
+from sqlleaf.models.node import NodeAttributes
+
+
+class SequenceNode(NodeAttributes):
+    KIND = "sequence"
+
+    def __init__(self, name: str, gen_ctx: GeneratorContext, pos_ctx: PositionContext, subkind: str = ""):
+        super().__init__(gen_ctx, pos_ctx, name=name)
+        self.subkind = subkind
+
+    def fields(self) -> dict[str, str]:
+        f = {"type": self.data_type}
+        if self.subkind:
+            f["subkind"] = self.subkind
+        return f

sqlleaf/models/node/star.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from sqlleaf.models.node import NodeAttributes
+
+
+class StarNode(NodeAttributes):
+    KIND = "star"
+    WITH_POSITIONS = True
+
+    def fields(self) -> dict[str, str]:
+        return {}

sqlleaf/models/node/stream.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from sqlleaf.models.context import GeneratorContext, PositionContext
+from sqlleaf.models.node import NodeAttributes
+
+
+class StreamNode(NodeAttributes):
+    KIND = "stream"
+
+    def __init__(self, name: str, gen_ctx: GeneratorContext, pos_ctx: PositionContext):
+        super().__init__(gen_ctx, pos_ctx, name=name)
+
+    def fields(self) -> dict[str, str]:
+        return {}

sqlleaf/models/node/user_defined_function.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from sqlleaf.models.context import GeneratorContext, PositionContext
+from sqlleaf.models.node import NodeAttributes
+
+
+class UserDefinedFunctionNode(NodeAttributes):
+    KIND = "udf"
+    WITH_POSITIONS = True
+
+    def __init__(
+        self,
+        schema: str,
+        gen_ctx: GeneratorContext,
+        pos_ctx: PositionContext,
+    ):
+        super().__init__(gen_ctx, pos_ctx, name=gen_ctx.expr.this)
+        self.schema = schema
+
+    def get_name(self) -> str:
+        return str(self.name).upper()
+
+    @property
+    def friendly_name(self) -> str:
+        tokens = [self.schema, self.name]
+        name = ".".join([tok for tok in tokens if tok]).upper()
+        return self.wrap(name)
+
+    def fields(self) -> dict[str, str]:
+        return {
+            "type": self.data_type,
+            "schema": self.schema.upper() if self.schema else "",
+        }

sqlleaf/models/node/var.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+from sqlleaf.models.node import NodeAttributes
+
+
+class VarNode(NodeAttributes):
+    KIND = "var"

sqlleaf/models/node/variable.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+from sqlleaf.models.node import NodeAttributes
+
+
+class VariableNode(NodeAttributes):
+    KIND = "variable"
+    WITH_POSITIONS = True

sqlleaf/models/node/window.py

@@ -0,0 +1,22 @@
+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 WindowNode(NodeAttributes):
+    KIND = "window"
+
+    def __init__(self, gen_ctx: GeneratorContext, pos_ctx: PositionContext):
+        expr = t.cast(exp.Window, gen_ctx.expr.this)
+        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}

sqlleaf/models/query/__init__.py

@@ -0,0 +1,33 @@
+import typing as t
+
+from sqlleaf.models.query.base import Query as Query
+from sqlleaf.models.query.call import CallQuery as CallQuery
+from sqlleaf.models.query.copy import CopyQuery as CopyQuery
+from sqlleaf.models.query.ctas import CTASQuery as CTASQuery
+from sqlleaf.models.query.database import DatabaseQuery as DatabaseQuery
+from sqlleaf.models.query.delete import DeleteQuery as DeleteQuery
+from sqlleaf.models.query.execute import ExecuteQuery as ExecuteQuery
+from sqlleaf.models.query.holder import QueryHolder as QueryHolder
+from sqlleaf.models.query.insert import InsertQuery as InsertQuery
+from sqlleaf.models.query.merge import MergeQuery as MergeQuery
+from sqlleaf.models.query.multitable_insert import MultitableInsertQuery as MultitableInsertQuery
+from sqlleaf.models.query.prepare import PrepareQuery as PrepareQuery
+from sqlleaf.models.query.procedure import ProcedureQuery as ProcedureQuery
+from sqlleaf.models.query.put import PutQuery as PutQuery
+from sqlleaf.models.query.replace import ReplaceQuery as ReplaceQuery
+from sqlleaf.models.query.schema import SchemaQuery as SchemaQuery
+from sqlleaf.models.query.select import SelectQuery as SelectQuery
+from sqlleaf.models.query.sequence import SequenceQuery as SequenceQuery
+from sqlleaf.models.query.set import SetQuery as SetQuery
+from sqlleaf.models.query.stage import StageQuery as StageQuery
+from sqlleaf.models.query.table import TableQuery as TableQuery
+from sqlleaf.models.query.trigger import TriggerQuery as TriggerQuery
+from sqlleaf.models.query.type import TypeQuery as TypeQuery
+from sqlleaf.models.query.unload import UnloadQuery as UnloadQuery
+from sqlleaf.models.query.update import UpdateQuery as UpdateQuery
+from sqlleaf.models.query.user_defined_function import FunctionParam as FunctionParam
+from sqlleaf.models.query.user_defined_function import UserDefinedFunctionQuery as UserDefinedFunctionQuery
+from sqlleaf.models.query.values import ValuesQuery as ValuesQuery
+from sqlleaf.models.query.view import ViewQuery as ViewQuery
+
+Q = t.TypeVar("Q", bound=Query)

sqlleaf/models/query/base.py

@@ -0,0 +1,248 @@
+from __future__ import annotations
+
+import logging
+import typing as t
+
+if t.TYPE_CHECKING:
+    from sqlleaf.models.query import QueryHolder
+
+from sqlglot import exp
+from sqlglot.optimizer.annotate_types import annotate_types
+
+from sqlleaf import exception, mappings, util
+from sqlleaf.typing import SourceExprType, SourceInfo, SqlObjectType, TargetExprType, TargetInfo
+
+logger = logging.getLogger("sqlleaf")
+
+
+class Query:
+    KIND: str = ""
+
+    def __init__(
+        self,
+        dialect: str,
+        statement: exp.Expr,
+        statement_index: int,
+        object_mapping: mappings.ObjectMapping,
+        source_info: SourceInfo | None,
+        target_info: TargetInfo,
+        skip_type_annotation: bool = False,
+    ):
+        self.kind = self.KIND
+        self.dialect = dialect
+        self.object_mapping = object_mapping
+        self.parent_query = None
+        self.child_queries = []
+        self.column_defs: t.List[exp.ColumnDef] = []
+        self.property = ""
+
+        # Remove comments at initialization
+        for expr in statement.walk():
+            expr.pop_comments()
+
+        if skip_type_annotation:
+            # Prevent overwriting the types with UNKNOWN if they've already been annotated
+            annotate_types(statement, dialect=dialect, schema=object_mapping)
+
+        self.statement: exp.Expr = statement
+        self.statement_index = statement_index  # The position of this query within a list of queries
+
+        self.source_info = source_info
+        self.target_info = target_info
+
+        logger.debug(
+            f"Created new query. Query => {self.__class__.__name__} | SourceType => {self.source_info and self.source_info.type.name} | TargetType => {self.target_info and self.target_info.type.name}"
+        )
+
+    def set_holder(self, holder: QueryHolder):
+        self.holder = holder
+
+    def _determine_expression_type(self, expr: exp.Expr | t.List[exp.Expr], dialect: str) -> SqlObjectType:
+        """
+        Determine the type of object that an expression represents.
+        For example, the literal "/tmp/data.csv" is of type FILE; identifier "STDIN" is of type STREAM.
+        """
+        if isinstance(expr, exp.Literal):
+            _type = SqlObjectType.DYNAMODB if expr.this.lower().startswith("dynamodb://") else SqlObjectType.FILE
+
+        elif isinstance(expr, exp.Identifier):
+            if expr.name.lower() in ["stdin", "stdout"]:
+                _type = SqlObjectType.STREAM
+            elif expr.name.lower() in ["program"]:
+                _type = SqlObjectType.PROGRAM
+            else:
+                raise exception.InvalidQueryError(f"Unknown object type identifier: {expr.name}")
+
+        elif isinstance(expr, exp.Var):
+            if dialect == "snowflake":
+                _type = SqlObjectType.STAGE
+            else:
+                _type = SqlObjectType.PREPARED_STATEMENT
+
+        elif isinstance(expr, exp.Table):
+            if isinstance(expr.this, exp.Var) and dialect == "snowflake":
+                _type = SqlObjectType.STAGE
+            else:
+                _type = SqlObjectType.TABLE
+
+        elif isinstance(expr, exp.Anonymous):
+            if expr.this == "PROCEDURE":
+                # CALL() statement
+                _type = SqlObjectType.PROCEDURE
+            else:
+                # CTAS EXECUTE <statement(args)>
+                _type = SqlObjectType.PREPARED_STATEMENT
+
+        elif dialect == "postgres" and isinstance(expr, exp.Alias):
+            if expr.this.name.upper() == "TABLE":
+                # PREPARE X AS TABLE Y;
+                _type = SqlObjectType.SELECT
+
+        elif isinstance(expr, exp.Select):
+            _type = SqlObjectType.SELECT
+
+        elif isinstance(expr, exp.Values):
+            _type = SqlObjectType.VALUES
+
+        elif isinstance(expr, exp.OnConflict):
+            # temporary; figure out how to handle this case
+            _type = SqlObjectType.SET
+
+        elif isinstance(expr, list) and isinstance(expr[0], exp.EQ):
+            # UPDATE .. SET
+            _type = SqlObjectType.SET
+
+        elif isinstance(expr, (exp.Insert, exp.Update, exp.Delete, exp.Merge)):
+            _type = SqlObjectType.DML
+
+        else:
+            raise exception.InvalidQueryError(f"Unknown source/target object type in query: {type(expr)}")
+
+        return _type
+
+    def get_original_self(self) -> Query:
+        return self.holder.original
+
+    def get_transformed_self(self) -> Query | None:
+        return self.holder.transformed
+
+    def get_columns_from_target(self) -> t.List[exp.ColumnDef]:
+        """
+        Given a query, figure out its target object, including its columns.
+
+        The structure is: INSERT INTO <target> .. <source>
+
+        This is straightforward if 'source' isn't a JOIN: we just use the source object's columns.
+        But if it is a JOIN, we use the selected columns rather than the source's columns.
+        """
+        if self.source_info is None:
+            # TODO: temp. This occurs when CALL() invokes an SP with only a SELECT
+            #  Remove this after CopyQuery refactor is complete
+            return []
+
+        source_expr = self.source_info.expression
+        target_expr = self.target_info.expression
+        target_type = self.target_info.type
+
+        if SqlObjectType.type_has_no_column_defs(target_type):
+            object_with_columns = source_expr
+        elif target_type == SqlObjectType.TABLE:
+            object_with_columns = target_expr
+        else:
+            raise exception.InvalidQueryError(f"Unhandled target object type: {target_type}")
+
+        column_defs = self._get_column_defs(object_with_columns)
+        if not column_defs:
+            raise exception.InvalidQueryError(f"Could not find any columns for expression: {object_with_columns}")
+
+        return column_defs
+
+    def _get_column_defs(
+        self,
+        expr: SourceExprType | TargetExprType,
+    ) -> t.List[exp.ColumnDef]:
+        """
+        Most of the time, the sources and target are tables.
+        However, with COPY/UNLOAD, they can be files or streams.
+
+        If the target is not a table and the source is a SELECT,
+        there may be a JOIN with many tables as the source.
+        """
+        if not isinstance(expr, exp.Table):
+            # Fall back to the source
+            source = self.source_info.expression
+            if isinstance(source, exp.Select):
+                # TODO: this can't handle functions
+                return [
+                    exp.ColumnDef(this=exp.to_identifier(col.alias_or_name), kind=col.unalias().type)
+                    for col in source.expressions
+                ]
+
+        table_query = self.object_mapping.get_table_or_stage(expr)
+        if not table_query:
+            return []
+
+        return table_query.get_column_defs()
+
+    def get_target_as_table(self) -> exp.Table:
+        """
+        For functions that only accept tables.
+        """
+        if not isinstance(self.target_info.expression, exp.Table):
+            raise exception.InvalidQueryError(
+                message=f"Expected the target object to be a table but it is a {repr(self.target_info.type)}"
+            )
+        return self.target_info.expression
+
+    def get_statement_index(self) -> str:
+        """
+        Get the statement index for this query (including its parents).
+        """
+        if self.parent_query:
+            index = self.parent_query.get_statement_index()
+            return index + ":" + str(self.statement_index)
+        else:
+            return str(self.statement_index)
+
+    def get_ctes(self) -> t.List:
+        return []
+
+    def get_column_defs(self, include_system: bool = False) -> t.List[exp.ColumnDef]:
+        return self.column_defs
+
+    def get_column_names_with_types(self, include_system: bool = False) -> t.Dict[str, str]:
+        """
+        Used by sqlglot's MappingSchema
+        """
+        columns = {col.name: str(col.kind) for col in self.get_column_defs(include_system=include_system)}
+        return columns
+
+    @property
+    def id(self) -> str:
+        return "query:" + util.short_sha256_hash(self.statement.sql() + ":" + str(self.statement_index))
+
+    def add_child_query(self, child_query):
+        child_query.parent_query = self
+        self.child_queries.append(child_query)
+
+    def get_all_queries(self, types: t.Tuple | None = None):
+        """
+        Collect all queries (self + children recursively), optionally filtered by type.
+        """
+        queries = [self]
+
+        for child in self.child_queries:
+            queries.extend(child.get_all_queries(types))
+
+        if types:
+            queries = [q for q in queries if isinstance(q, types)]
+
+        return queries
+
+    def to_dict(self):
+        result = {
+            "id": self.id,
+            "kind": self.kind,
+            "statement_index": self.statement_index,
+        }
+        return result

sqlleaf/models/query/call.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import typing as t
+from dataclasses import dataclass
+
+import sqlglot
+from sqlglot import exp
+
+from sqlleaf import util
+from sqlleaf.models.query.base import Query
+from sqlleaf.typing import SqlObjectType, TargetInfo
+
+if t.TYPE_CHECKING:
+    from sqlleaf import mappings
+
+
+@dataclass(frozen=True)
+class CallQueryProperties:
+    schema: t.Optional[str]
+    procedure: str
+    args: t.List[exp.Expr]
+
+    @classmethod
+    def from_expression(cls, statement: exp.Command, dialect: str) -> CallQueryProperties:
+        # The 'expression' part of the command contains the procedure call
+        # e.g., CALL hello('world') -> expression is "hello('world')"
+        call_str = statement.args.get("expression").this
+
+        # Parse the call string as a function call
+        inner_expr = sqlglot.parse_one(call_str, read=dialect)
+
+        schema = None
+        if isinstance(inner_expr, exp.Anonymous):
+            procedure_name = inner_expr.this
+            args = inner_expr.expressions
+        elif isinstance(inner_expr, exp.Column):
+            procedure_name = inner_expr.name
+            args = []
+        elif isinstance(inner_expr, exp.Dot):
+            schema = inner_expr.left.name
+            if isinstance(inner_expr.right, exp.Anonymous):
+                procedure_name = inner_expr.right.this
+                args = inner_expr.right.expressions
+            else:
+                procedure_name = inner_expr.right.name
+                args = []
+        else:
+            procedure_name = call_str
+            args = []
+
+        return cls(schema=schema, procedure=procedure_name, args=args)
+
+
+class CallQuery(Query):
+    """
+    Holds metadata related to procedure calls (CALL statement).
+    """
+
+    KIND = "call"
+
+    def __init__(
+        self,
+        expr: exp.Command,
+        dialect: str,
+        object_mapping: mappings.ObjectMapping,
+        statement_index: int,
+    ):
+        self.properties = CallQueryProperties.from_expression(expr, dialect)
+
+        # We can represent the procedure being called as the target
+        target_info = TargetInfo(
+            expression=exp.Table(
+                this=exp.to_identifier(self.properties.procedure),
+                db=exp.to_identifier(self.properties.schema) if self.properties.schema else None,
+            ),
+            type=SqlObjectType.PROCEDURE,
+        )
+
+        super().__init__(
+            dialect=dialect,
+            statement=expr,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=None,
+            target_info=target_info,
+        )
+
+    @property
+    def schema(self) -> t.Optional[str]:
+        return self.properties.schema
+
+    @property
+    def procedure(self) -> str:
+        return self.properties.procedure
+
+    @property
+    def args(self) -> t.List[exp.Expr]:
+        return self.properties.args
+
+    @property
+    def name(self):
+        return ".".join([var for var in [self.schema, self.procedure] if var])
+
+    def to_dict(self):
+        return {
+            "id": self.id,
+            "name": self.name,
+            "procedure": self.procedure,
+            "schema": self.schema,
+            "args": [str(arg) for arg in self.args],
+        }