sqlleaf 0.0.1__py3-none-any.whl

sqlleaf/models/query/copy.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+import logging
+import typing as t
+from dataclasses import dataclass
+
+from sqlglot import exp
+
+from sqlleaf import exception, mappings, util
+from sqlleaf.models.query.base import Query
+from sqlleaf.typing import SourceExprType, SourceInfo, TargetExprType, TargetInfo
+
+logger = logging.getLogger("sqlleaf")
+
+
+@dataclass(frozen=True)
+class CopyQueryProperties:
+    file_format: str = "TEXT"
+    load_data: bool = True
+    is_a_job: bool = False
+    job_action: str = ""
+    job_name: str = ""
+    job_auto_run: bool = True
+
+    @classmethod
+    def from_expression(cls, expr: exp.Copy, source_info: SourceInfo, target_info: TargetInfo) -> CopyQueryProperties:
+        """
+        Extract the parameters of the COPY statement.
+
+        For example,
+            Input: "COPY FROM ... FORMAT AS CSV NOLOAD"
+            Params: ["FORMAT AS CSV", "NOLOAD"]
+        """
+        params_dict: t.Dict[str, t.Any] = {}
+
+        params_list = []
+        # Filter for the parameters we need
+        for param in expr.args.get("params", []):
+            if isinstance(param, exp.CopyParameter) and isinstance(param.this, exp.Var):
+                params_list.append(param)
+
+        for i, param in enumerate(params_list):
+            param_name = param.this.name.upper()
+            if param_name == "FORMAT":
+                params_dict["file_format"] = str(param.expression)
+            elif param_name == "NOLOAD":
+                params_dict["load_data"] = False
+            elif param_name == "JOB":
+                params_dict["is_a_job"] = True
+                try:
+                    # COPY <command> JOB CREATE <name>
+                    job_command = params_list[i + 1].this.name.upper()
+                    job_name = params_list[i + 1].expression.name
+                    params_dict["job_action"] = job_command
+                    params_dict["job_name"] = job_name
+                except IndexError:
+                    message = "Missing one or more parameters to follow expression 'COPY .. JOB'"
+                    raise exception.SqlLeafException(message=message)
+            elif param_name == "AUTO":
+                # COPY <command> JOB CREATE <name> AUTO ON | OFF
+                params_dict["job_auto_run"] = str(param.expression) == "ON"
+        return cls(**params_dict)
+
+    @property
+    def is_active(self) -> bool:
+        """
+        An active COPY query is one that has lineage.
+        It is either:
+        - a regular COPY query
+        - a JOB that is either 'RUN' or 'CREATE and AUTO=ON'
+        Otherwise (e.g. `COPY JOB LIST`) it is not.
+        """
+        if not self.is_a_job and self.load_data:
+            return True
+        if self.job_action == "RUN":
+            return True
+        if self.job_auto_run and self.job_action == "CREATE":
+            return True
+        return False
+
+
+class CopyQuery(Query):
+    KIND = "copy"
+
+    def __init__(self, expr: exp.Copy, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int):
+        source, target = self.get_source_and_target_expressions(expr, dialect)
+
+        if dialect == "snowflake":
+            util.rename_if_stage(source, target)
+
+        source_type = self._determine_expression_type(source, dialect)
+        target_type = self._determine_expression_type(target, dialect)
+
+        super().__init__(
+            dialect=dialect,
+            statement=expr,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=SourceInfo(expression=source, type=source_type),
+            target_info=TargetInfo(expression=target, type=target_type),
+        )
+        self.parameters = self.get_params()
+        util.qualify_and_annotate(self.source_info.expression, self.dialect, self.object_mapping)
+
+    def get_params(
+        self,
+    ) -> CopyQueryProperties:
+        return CopyQueryProperties.from_expression(self.statement, self.source_info, self.target_info)
+
+    def is_query_active(self) -> bool:
+        logger.debug(self.parameters)
+        return self.parameters.is_active
+
+    def get_source_and_target_expressions(
+        self, expr: exp.Copy, dialect: str
+    ) -> t.Tuple[SourceExprType, TargetExprType]:
+        """
+        Determine the source and target expressions of the query.
+        """
+        if dialect in ["postgres", "redshift"]:
+            # Postgres treats STDOUT and STDIN the same
+
+            if expr.args["kind"]:
+                # COPY X FROM STDOUT/STDIN
+                source = expr.args["files"][0]
+                target = expr.args["this"]
+                if isinstance(target, exp.Schema):
+                    target = target.this
+            else:
+                # COPY X TO STDOUT/STDIN
+                source = expr.args["this"]
+                target = expr.args["files"][0]
+                if isinstance(source, exp.Schema):
+                    source = source.this
+
+        elif dialect == "snowflake":
+            source = expr.args["files"][0]
+            target = expr.args["this"]
+
+        # It may be a subquery
+        source = source.unnest()
+        target = target.unnest()
+
+        return source, target

sqlleaf/models/query/ctas.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot import exp
+
+from sqlleaf import mappings, util
+from sqlleaf.models.query.base import Query
+from sqlleaf.models.query.table import TableQuery
+from sqlleaf.typing import SourceInfo, TargetInfo
+
+
+class CTASQuery(Query):
+    KIND = "ctas"
+
+    def __init__(
+        self,
+        expr: exp.Create,
+        dialect: str,
+        object_mapping: mappings.ObjectMapping,
+        columns: t.List[exp.ColumnDef],
+        statement_index: int,
+    ):
+        source = expr.expression
+        if source is None:
+            if exec_prop := expr.find(exp.ExecuteAsProperty):
+                source = exec_prop.this
+
+        target = util.get_table(expr)
+
+        source_type = self._determine_expression_type(source, dialect)
+        target_type = self._determine_expression_type(target, dialect)
+
+        super().__init__(
+            dialect=dialect,
+            statement=expr,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=SourceInfo(expression=source, type=source_type),
+            target_info=TargetInfo(expression=target, type=target_type),
+        )
+        self.column_defs: t.List[exp.ColumnDef] = columns
+        self.system_column_defs: t.List[exp.ColumnDef] = []
+        self.inherited_by: t.List[TableQuery] = []
+
+        self.load_data: bool = True
+        if props := expr.args["properties"]:
+            if load_data := props.find(exp.WithDataProperty):
+                self.load_data: bool = not load_data.args["no"]
+
+        self.property: str = util.find_property(expr, self.target_info.expression, dialect)
+
+    def get_column_defs(self, include_system: bool = False) -> t.List[exp.ColumnDef]:
+        return self.column_defs + self.system_column_defs if include_system else 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

sqlleaf/models/query/database.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from sqlglot import exp
+
+from sqlleaf import mappings, util
+from sqlleaf.models.query.base import Query
+from sqlleaf.typing import SqlObjectType, TargetInfo
+
+
+@dataclass(frozen=True)
+class DatabaseQueryProperties:
+    location: str | None
+
+    @classmethod
+    def from_expression(cls, expr: exp.Create, dialect: str):
+        location = util.get_location_property(expr, dialect)
+        return cls(location=location)
+
+
+class DatabaseQuery(Query):
+    KIND = "database"
+
+    def __init__(
+        self,
+        expr: exp.Create,
+        dialect: str,
+        object_mapping: mappings.ObjectMapping,
+        statement_index: int,
+    ):
+        source = None
+        target = expr.this
+
+        super().__init__(
+            dialect=dialect,
+            statement=expr,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=source,
+            target_info=TargetInfo(expression=target, type=SqlObjectType.DATABASE),
+        )
+
+        self.properties = DatabaseQueryProperties.from_expression(expr, dialect)
+
+    @property
+    def location(self) -> str | None:
+        return self.properties.location

sqlleaf/models/query/delete.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from sqlglot import exp
+
+from sqlleaf import mappings
+from sqlleaf.models.query.base import Query
+from sqlleaf.typing import SqlObjectType, TargetInfo
+
+
+class DeleteQuery(Query):
+    KIND = "delete"
+
+    def __init__(self, expr: exp.Delete, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int):
+        super().__init__(
+            dialect=dialect,
+            statement=expr,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=None,
+            target_info=TargetInfo(expression=expr.this, type=SqlObjectType.TABLE),
+        )
+
+    def get_ctes(self):
+        if "with_" in self.statement.args:
+            return self.statement.args["with_"].expressions
+        return []

sqlleaf/models/query/execute.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import typing as t
+from dataclasses import dataclass
+
+import sqlglot
+from sqlglot import exp
+
+from sqlleaf import mappings
+from sqlleaf.models.query.base import Query
+from sqlleaf.typing import SqlObjectType, TargetInfo
+
+
+@dataclass(frozen=True)
+class ExecuteQueryProperties:
+    name: str
+    arguments: t.List[exp.Literal]
+
+    @classmethod
+    def from_expression(cls, expr: exp.Command, dialect: str):
+        """
+        EXECUTE <name>(<arguments>);
+        """
+        expression_name = expr.expression.name
+        parsed = sqlglot.parse_one(expression_name, read=dialect)
+
+        if isinstance(parsed, exp.Anonymous):
+            name = parsed.this
+            arguments = parsed.expressions
+        else:
+            name = parsed.name
+            arguments = []
+
+        return cls(name=name, arguments=arguments)
+
+
+class ExecuteQuery(Query):
+    KIND = "execute"
+
+    def __init__(
+        self,
+        expr: exp.Command,
+        dialect: str,
+        object_mapping: mappings.ObjectMapping,
+        statement_index: int,
+    ):
+        self.parameters = ExecuteQueryProperties.from_expression(expr, dialect)
+
+        super().__init__(
+            dialect=dialect,
+            statement=expr,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=None,
+            target_info=TargetInfo(
+                expression=exp.to_table(self.parameters.name), type=SqlObjectType.PREPARED_STATEMENT
+            ),
+        )
+
+    @property
+    def name(self) -> str:
+        return self.parameters.name
+
+    @property
+    def arguments(self) -> t.List[exp.Literal]:
+        return self.parameters.arguments

sqlleaf/models/query/holder.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import typing as t
+
+if t.TYPE_CHECKING:
+    from sqlleaf.models.query import Q
+    from sqlleaf.models.query.base import Query
+
+
+class QueryHolder:
+    """
+    A container that holds up to three versions of the same SQL statement
+    as three distinct Query instances.
+
+    - original:           The parsed, unmodified Query.
+    - transformed:        The Query after transformer.py has processed it.
+    - downstream_holders: The containers for downstream queries. Used both for structural
+                          children (MERGE branches, INSERT ON CONFLICT, UDF bodies) and for
+                          call-site substituted queries (CallQuery, ExecuteQuery, UDF call sites).
+                          Populated entirely during the Collector phase.
+    """
+
+    def __init__(self, original: Query):
+        self.original: Query = original
+        self.transformed: Query | None = None
+        self.downstream_holders: t.List[QueryHolder] = []
+        self.parent_holder: QueryHolder | None = None
+
+        self.original.set_holder(self)
+
+    def add_downstream_holder(self, holder: QueryHolder) -> None:
+        holder.parent_holder = self
+        self.downstream_holders.append(holder)
+
+    def get_all_holders(self, types: t.Tuple | None = None) -> t.List[QueryHolder]:
+        """
+        Recursively collect all holders (self + children).
+        Optionally filter by the type of their `original` query.
+        """
+        holders = [self]
+        for child in self.downstream_holders:
+            holders.extend(child.get_all_holders(types))
+        if types:
+            holders = [h for h in holders if isinstance(h.original, types)]
+        return holders
+
+    def set_transformed_query(self, query: Q) -> None:
+        self.transformed = query
+        self.transformed.set_holder(self)

sqlleaf/models/query/insert.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import sqlglot
+from sqlglot import exp
+
+from sqlleaf import mappings, util
+from sqlleaf.models.query.base import Query
+from sqlleaf.typing import SourceInfo, SqlObjectType, TargetInfo
+
+
+class InsertQuery(Query):
+    KIND = "insert"
+
+    def __init__(
+        self,
+        expr: exp.Insert,
+        dialect: str,
+        object_mapping: mappings.ObjectMapping,
+        statement_index: int,
+        table: exp.Table | None = None,
+        # Whether to skip column annotation of types.
+        # Used to preserve column types after conversion from COPY to INSERT.
+        skip_annotate: bool = False,
+    ):
+        if not table:
+            table = util.get_table(expr)
+
+        target = table
+        source = expr.expression
+        if source:
+            # Subquery(Select) -> Select
+            source = source.unnest()
+        elif src := expr.args.get("source"):
+            # INSERT .. DEFAULT VALUES
+            source = src
+        else:
+            # INSERT .. (VALUES())
+            exprs = expr.this.expressions
+            if len(exprs) == 1 and (ex := exprs[0]):
+                if isinstance(ex, exp.Anonymous) and ex.name.upper() == "VALUES":
+                    # sqlglot parses this incorrectly; fix it by re-parsing it to an exp.Values
+                    source = sqlglot.parse_one(ex.sql(dialect=dialect), dialect=dialect)
+                    expr.set("expression", source)
+                    [e.pop() for e in (expr.this.expressions or [])]
+
+        # Set the correct source type. Hacky, but works for now
+        is_default_values = expr.args.get("default", False)
+        if is_default_values:
+            source_type = SqlObjectType.VALUES
+        elif isinstance(source, exp.Tuple):
+            source_type = SqlObjectType.TUPLE
+        elif isinstance(source, exp.SetOperation):
+            source_type = SqlObjectType.SELECT
+        else:
+            source_type = self._determine_expression_type(source, dialect)
+
+        target_type = self._determine_expression_type(target, dialect)
+
+        super().__init__(
+            dialect=dialect,
+            statement=expr,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=SourceInfo(expression=source, type=source_type),
+            target_info=TargetInfo(expression=target, type=target_type),
+        )
+
+    def get_ctes(self):
+        return getattr(self.statement, "ctes", [])

sqlleaf/models/query/merge.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from sqlglot import exp
+
+from sqlleaf import mappings
+from sqlleaf.models.query.base import Query
+from sqlleaf.typing import TargetInfo
+
+
+class MergeQuery(Query):
+    KIND = "merge"
+
+    def __init__(self, expr: exp.Merge, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int):
+        source = None
+        target = expr.this
+
+        target_type = self._determine_expression_type(target, dialect)
+
+        super().__init__(
+            dialect=dialect,
+            statement=expr,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=source,
+            target_info=TargetInfo(expression=target, type=target_type),
+        )
+
+    def get_ctes(self):
+        return getattr(self.statement, "ctes", [])

sqlleaf/models/query/multitable_insert.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot import exp
+
+from sqlleaf import mappings
+from sqlleaf.models.query.base import Query
+from sqlleaf.typing import SourceInfo, SqlObjectType, TargetInfo
+
+
+class MultitableInsertQuery(Query):
+    KIND = "multitableinserts"
+
+    def __init__(
+        self,
+        expr: exp.MultitableInserts,
+        dialect: str,
+        object_mapping: mappings.ObjectMapping,
+        statement_index: int,
+    ):
+        source = expr.args["source"]
+        source_type = SqlObjectType.SELECT
+
+        # For multi-table inserts, the target is actually multiple tables,
+        # but the base Query model expects a single TargetInfo.
+        # We'll use the first one or a dummy for the parent query.
+        # The children will have the actual target info.
+        # TODO: fix this.
+        target = expr.expressions[0].this.this
+
+        super().__init__(
+            dialect=dialect,
+            statement=expr,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=SourceInfo(expression=source, type=source_type),
+            target_info=TargetInfo(expression=target, type=SqlObjectType.TABLE),
+        )
+
+    def get_ctes(self) -> t.List[exp.CTE]:
+        return getattr(self.statement, "ctes", [])

sqlleaf/models/query/prepare.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+import logging
+import re
+import typing as t
+
+import sqlglot
+from sqlglot import exp
+
+from sqlleaf import exception, mappings, util
+from sqlleaf.models.query.base import Query
+from sqlleaf.typing import SourceInfo, SqlObjectType, TargetInfo
+
+logger = logging.getLogger("sqlleaf")
+
+
+class PrepareQuery(Query):
+    KIND = "prepare"
+
+    def __init__(self, expr: exp.Command, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int):
+        source, target = self.get_source_and_target_expressions(expr)
+
+        source_type = self._determine_expression_type(source, dialect)
+        target_type = SqlObjectType.PREPARED_STATEMENT
+
+        super().__init__(
+            dialect=dialect,
+            statement=source,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=SourceInfo(expression=source, type=source_type),
+            target_info=TargetInfo(expression=target, type=target_type),
+        )
+        util.qualify_and_annotate(self.source_info.expression, self.dialect, self.object_mapping)
+        self.parameter_count = self._count_parameters(source)
+
+    def _count_parameters(self, source: exp.Expr) -> int:
+        """
+        Count the number of expected arguments for this prepared statement.
+        Positional parameters are $1, $2, etc.
+        """
+        params = [int(p.name) for p in source.find_all(exp.Parameter) if p.name.isdigit()]
+        return max(params) if params else 0
+
+    def get_source_and_target_expressions(self, statement: exp.Command) -> t.Tuple[exp.Expr, exp.Table]:
+        """
+        Parse a PREPARE statement for Postgres.
+        Syntax: PREPARE name AS statement
+        """
+        expression_name = statement.expression.name
+
+        # Find the 'AS' keyword to split name and statement
+        match = re.search(r"\s+AS\s+", expression_name, re.IGNORECASE)
+        if not match:
+            raise exception.SqlLeafException(message=f"Could not find 'AS' in PREPARE expression: {expression_name}")
+
+        name = expression_name[: match.start()].strip()
+        statement_text = expression_name[match.end() :].strip()
+
+        prepared_expr = sqlglot.parse_one(statement_text, dialect="postgres")
+        return prepared_expr, exp.to_table(name)

sqlleaf/models/query/procedure.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot import exp
+
+from sqlleaf import mappings, util
+from sqlleaf.models.query.base import Query
+from sqlleaf.models.query.user_defined_function import _extract_function_info
+from sqlleaf.typing import TargetInfo
+
+
+class ProcedureQuery(Query):
+    """
+    Holds metadata related to stored procedures.
+    """
+
+    KIND = "procedure"
+
+    def __init__(
+        self,
+        expr: exp.Create,
+        dialect: str,
+        object_mapping: mappings.ObjectMapping,
+        statement_index: int,
+    ):
+        table = util.get_table(expr)
+        target_type = self._determine_expression_type(table, dialect)
+
+        super().__init__(
+            dialect=dialect,
+            statement=expr,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=None,
+            target_info=TargetInfo(expression=table, type=target_type),
+        )
+        self.schema = table.db
+        self.procedure = table.name
+        self.signature = str(expr.this)  # e.g. etl.my_proc(v_session_id VARCHAR)
+
+        # TODO: support 'default'
+        self.column_defs: t.List[exp.ColumnDef] = expr.this.expressions
+        _, _, self.parameters = _extract_function_info(expr)
+        self.args = [  # e.g. {'name': 'v_session_id', 'type': 'VARCHAR'}
+            {"name": p.name, "type": str(p.type)} for p in self.parameters
+        ]
+        self.inner_statements = self._extract_inner_statements(expr)
+
+    def _extract_inner_statements(self, expr: exp.Create) -> t.List[exp.Expr]:
+        body_expr = expr.args.get("expression")
+        if not body_expr:
+            return []
+
+        inner_statements = util.iter_inner_statements(body_expr, self.dialect, wrap=True)
+
+        # Filter out statements that do not contain lineage (e.g. END;)
+        return [
+            stmt for stmt in inner_statements if not isinstance(stmt, (exp.EndStatement, exp.Column, exp.Identifier))
+        ]
+
+    def to_dict(self):
+        return {
+            "id": self.id,
+            "name": self.name,
+            "signature": self.signature,
+            "args": self.args,
+        }
+
+    @property
+    def name(self):
+        return ".".join([var for var in [self.schema, self.procedure] if var])