sqlleaf 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

sqlleaf/__init__.py

@@ -2,3 +2,4 @@ from sqlleaf.lineage import Lineage as Lineage
 from sqlleaf.processors.collector.collector import collect_queries as collect_queries
 from sqlleaf.processors.generator.generator import generate_lineage_for_query as generate_lineage_for_query
 from sqlleaf.processors.transformer.transformer import transform_query as transform_query
+from sqlleaf.dialects import PlPgSQL

sqlleaf/dialects/__init__.py

@@ -0,0 +1 @@
+from sqlleaf.dialects.plpgsql.p_dialect import PlPgSQL

sqlleaf/dialects/plpgsql/__init__.py

@@ -0,0 +1 @@
+from sqlleaf.dialects.plpgsql import p_classes as pgexp

sqlleaf/dialects/plpgsql/p_classes.py

@@ -146,7 +146,7 @@ class PGGetDiagnostics(exp.Expression):
     arg_types = {"current": False, "stacked": False, "expressions": True}
 
 
-class PGOpenCursor(exp.Expression):
+class PGOpen(exp.Expression):
     """An OPEN cursor statement."""
     arg_types = {"this": True, "expression": False, "scroll": False, "expressions": False}
 

sqlleaf/dialects/plpgsql/p_generator.py

@@ -25,7 +25,7 @@ class PlPgSQLGenerator(PostgresGenerator):
         PGContinue: lambda self, e: self.pgcontinue_sql(e),
         PGRaise: lambda self, e: self.pgraise_sql(e),
         PGSqlState: lambda self, e: self.pgsqlstate_sql(e),
-        PGOpenCursor: lambda self, e: self.pgopencursor_sql(e),
+        PGOpen: lambda self, e: self.pgopen_sql(e),
         PGFetch: lambda self, e: self.pgfetch_sql(e),
         PGMove: lambda self, e: self.pgmove_sql(e),
         PGClose: lambda self, e: self.pgclose_sql(e),
@@ -419,7 +419,7 @@ class PlPgSQLGenerator(PostgresGenerator):
             return f"{self.sql(dir_expr)} {prep} {self.sql(expression.this)}"
         return self.sql(expression.this)
 
-    def pgopencursor_sql(self, expression: PGOpenCursor) -> str:
+    def pgopen_sql(self, expression: PGOpen) -> str:
         # Unbound form: OPEN c [NO|SCROLL] FOR <query>
         if expression.args.get("expression") is not None:
             parts: list[str] = ["OPEN", self.sql(expression.this)]

sqlleaf/dialects/plpgsql/p_parser.py

@@ -1053,7 +1053,7 @@ class PlPgSQLParser(PostgresParser):
 
         return args
 
-    def _parse_pgopen(self) -> PGOpenCursor:
+    def _parse_pgopen(self) -> PGOpen:
         # Cursor variable identifier
         cursor = self._parse_id_var()
 
@@ -1062,7 +1062,7 @@ class PlPgSQLParser(PostgresParser):
         if self._match(TokenType.L_PAREN, advance=False):
             # Parse cursor call arguments directly (positional, name := value, name => value)
             args = self._parse_cursor_call_args()
-            return self.expression(PGOpenCursor(this=cursor, expressions=args))
+            return self.expression(PGOpen(this=cursor, expressions=args))
 
         # Optional [[NO] SCROLL]
         scroll: bool | None = None
@@ -1075,7 +1075,7 @@ class PlPgSQLParser(PostgresParser):
         # If there's no FOR and no SCROLL/NO SCROLL, treat as bound cursor with zero args: OPEN c;
         if not self._match(TokenType.FOR):
             if scroll is None:
-                return self.expression(PGOpenCursor(this=cursor))
+                return self.expression(PGOpen(this=cursor))
             # If SCROLL/NO SCROLL was provided, FOR is required
             self._match_expect(TokenType.FOR)
 
@@ -1088,7 +1088,7 @@ class PlPgSQLParser(PostgresParser):
             if query is None:
                 self.raise_error("Expected query after OPEN ... FOR")
 
-        return self.expression(PGOpenCursor(this=cursor, expression=query, scroll=scroll))
+        return self.expression(PGOpen(this=cursor, expression=query, scroll=scroll))
 
     def _parse_pgfetch(self) -> PGFetch:
         # Parse optional direction, optional preposition (required when direction given), and cursor

sqlleaf/exception.py

@@ -6,12 +6,12 @@ ERROR_LEVEL = "stop"
 logger = logging.getLogger("sqlleaf")
 
 
-def raise_error(clazz: type, message: str):
+def raise_error(cls: type, message: str):
     """
     Raise a custom exception based on the error level.
     """
     if ERROR_LEVEL == "stop":
-        raise clazz(message)
+        raise cls(message)
     elif ERROR_LEVEL == "continue":
         logger.error(message)
 

sqlleaf/lineage.py

@@ -8,21 +8,11 @@ from sqlleaf import exception, mappings, path, util
 from sqlleaf.models.hook import Hook
 from sqlleaf.models.node import EdgeAttributes, GraphAttributes, N
 from sqlleaf.models.query import (
-    CopyQuery,
-    CTASQuery,
-    InsertQuery,
-    PutQuery,
     Q,
-    TableQuery,
-    UnloadQuery,
-    UpdateQuery,
-    ViewQuery,
 )
 from sqlleaf.path import LineagePath
 from sqlleaf.processors.collector import collector
 from sqlleaf.processors.generator import generator
-from sqlleaf.processors.transformer import transformer
-from sqlleaf.typing import SqlObjectType
 
 logging.getLogger("sqlglot").setLevel(logging.WARNING)
 logger = logging.getLogger("sqlleaf")
@@ -35,8 +25,8 @@ class Lineage:
 
     def __init__(self):
         self.graph = new_graph()  # The graph that contains all lineage
-        self.subgraphs: t.List[nx.MultiDiGraph] = []  # The subgraphs that make up the main graph
-        self.paths: t.Dict[str, t.List[LineagePath]] = {}  # The paths throughout the graph
+        self.subgraphs: list[nx.MultiDiGraph] = []  # The subgraphs that make up the main graph
+        self.paths: dict[str, list[LineagePath]] = {}  # The paths throughout the graph
         self.collected_queries: collector.CollectQueryResult | None = None
         self.user_defined_hooks: list[Hook] = []
         self.object_mapping: mappings.ObjectMapping | None = None
@@ -60,11 +50,8 @@ class Lineage:
             query_holders = parent_holder.get_all_holders()
 
             for holder in query_holders:
-                transformer.transform_query(holder)
-
-            for holder in query_holders:
-                # Transform and produce lineage only for certain queries
-                if query_has_lineage(holder.original):
+                # Produce lineage only for certain queries
+                if generator.query_has_lineage(holder.original):
                     generator.generate_lineage_for_query(holder, graph, self.user_defined_hooks)
 
             # Associate the query with the graph even if it has no lineage
@@ -96,14 +83,14 @@ class Lineage:
 
         self.graph.add_edges_from(subgraph.edges(data=True))
 
-    def get_edges(self) -> t.List[EdgeAttributes]:
+    def get_edges(self) -> list[EdgeAttributes]:
         edges = [data["attrs"] for par, chi, data in self.graph.edges(data=True)]
         edges = sorted(edges, key=lambda e: (e.parent.full_name, e.child.full_name))
         for edge in edges:
             logger.debug(f"Found edge: {edge.parent.friendly_name} -> {edge.child.friendly_name}")
         return edges
 
-    def get_nodes(self) -> t.List[N]:
+    def get_nodes(self) -> list[N]:
         nodes = [data["attrs"] for n, data in self.graph.nodes(data=True)]
         # TODO: sort on selected index?
         nodes = sorted(
@@ -111,13 +98,13 @@ class Lineage:
         )
         return nodes
 
-    def get_original_queries(self) -> t.List[Q]:
+    def get_original_queries(self) -> list[Q]:
         """
         Get the original queries from the graph.
         """
         return [q.original for q in self.graph.graph["attrs"].queries]
 
-    def get_transformed_queries(self) -> t.List[Q]:
+    def get_transformed_queries(self) -> list[Q]:
         """
         Get the transformed queries from the graph.
         """
@@ -280,47 +267,3 @@ def new_graph() -> nx.MultiDiGraph:
     A graph has attributes along with its node and edges.
     """
     return nx.MultiDiGraph(attrs=GraphAttributes())
-
-
-QUERIES_WITH_LINEAGE = (
-    CTASQuery,
-    CopyQuery,
-    InsertQuery,
-    PutQuery,
-    TableQuery,
-    UnloadQuery,
-    UpdateQuery,
-    ViewQuery,
-)
-
-
-def query_has_lineage(query: Q) -> bool:
-    """
-    Check if a query has lineage within its expressions.
-
-    We distinguish between a query's definition and a query that is called.
-    That is, in order for a `CREATE FUNCTION` or `CREATE PROCEDURE` to have
-    lineage, it must be executed by an invoking statement, e.g. `CALL()` or `SELECT UDF()`
-    Simply having statements inside its definition is not sufficient to produce lineage.
-    """
-    has_lineage = True
-    if not isinstance(query, QUERIES_WITH_LINEAGE):
-        has_lineage = False
-    elif isinstance(query, CopyQuery) and query.source_info.type == SqlObjectType.VALUES:
-        # COPY TO STDOUT VALUES (..) # TODO: this should have lineage
-        has_lineage = False
-    elif isinstance(query, CopyQuery) and not query.is_query_active():
-        has_lineage = False
-    elif isinstance(query, CTASQuery) and not query.load_data:
-        # CREATE TABLE WITH NO DATA
-        has_lineage = False
-    elif isinstance(query, CTASQuery) and query.source_info.type == SqlObjectType.PREPARED_STATEMENT:
-        # CREATE TABLE AS EXECUTE
-        has_lineage = False
-    elif isinstance(query, TableQuery) and query.property != "external":
-        # CREATE EXTERNAL TABLE
-        has_lineage = False
-
-    if not has_lineage:
-        logger.debug(f"Query type '{query.__class__.__name__}' does NOT have lineage. Skipping.")
-    return has_lineage

sqlleaf/mappings.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 import typing as t
 
 from sqlglot import MappingSchema, exp
@@ -24,68 +25,119 @@ from sqlleaf.models.query import (
     ViewQuery,
 )
 
-ColumnMapping = t.Union[t.Dict, str, t.List]
+ColumnMapping = t.Union[dict, str, list]
+
+logger = logging.getLogger("sqlleaf")
 
 
 class ObjectMapping(MappingSchema):
     """
-    Extends sqlglot.MappingSchema to provide additional functionality related to tracking exp.Table
+    Contains all the objects that get created, such as tables, procedures, functions, and sequences.
 
-    Specifically, we need to track the exp.Table inside the exp.Create statements, as they contain more information
-    than the exp.Table that we encounter later when parsing INSERT statements.
+    We extends sqlglot.MappingSchema to provide additional functionality related to tracking these objects,
+    while still allowing optimizer inside sqlglot to work (as it needs a mapping of objects to perform functions
+    like qualification and simplification).
     """
 
     def __init__(self, dialect: str):
-        """
-        Initialize a mapping of tables parts to exp.Table
-        """
         super().__init__(
             dialect=dialect, normalize=False
         )  # Set `normalize=False` to prevent an unnecessary second parse.
         self.kind_mapping = {}
         self.kind_mapping_trie = {}
+        # Session/global variables (e.g., SET var = ...)
         self.session_variables: dict[str, exp.Expr] = {}
+        # Block-scoped variables: one dict per active block, innermost scope at the end
+        self.variable_scopes: list[dict[str, exp.Expr]] = []
+
+    # ---- Variable scope API ----
+    def push_variable_scope(self) -> None:
+        """Enter a new block scope."""
+        self.variable_scopes.append({})
+
+    def pop_variable_scope(self) -> None:
+        """Leave the current block scope."""
+        if self.variable_scopes:
+            self.variable_scopes.pop()
+
+    def set_variable(self, name: str, value: exp.Expr) -> None:
+        """Assign or declare a variable in the current scope."""
+        self.variable_scopes[-1][name] = value
+        logger.debug(f"Set variable value: {name} := {value.sql()}")
+
+    def get_variable(self, name: str, raise_on_missing: bool = True) -> exp.Expr | None:
+        """Resolve a variable by searching from innermost to outermost scopes, then session vars."""
+        for scope in reversed(self.variable_scopes):
+            if name in scope:
+                return scope[name]
+        if raise_on_missing:
+            raise exception.MappingError(f"Could not find value for variable: {name}")
+        return None
+
+    def get_variables_for(self, expressions: list[exp.Expr]) -> list[exp.Expr]:
+        """
+        Given a list of expressions as keys, fetch their values from the variable scope if they exist.
+
+        Example:
+            expressions = [ Literal[42], Column[name] ]
+            variable scope = { Column[name] -> Literal['john'] }
+            ->
+            [ Literal['john'] ]
+        """
+        resolved = []
+        for expr in expressions:
+            if isinstance(expr, exp.Column) and not expr.table:
+                variable = self.get_variable(expr.name)
+                if variable is not None:
+                    resolved.append(variable.copy())
+                    continue
+            resolved.append(expr)
+        return resolved
+
+    def get_scope(self) -> dict[str, exp.Expr]:
+        """Return the current scope dict (innermost)."""
+        return self.variable_scopes[-1]
 
     def add_database_query(self, query: DatabaseQuery) -> None:
-        self._add_query(kind="database", query=query, dialect=query.dialect)
+        self._add_query(query=query, dialect=query.dialect)
 
     def add_prepare_query(self, query: PrepareQuery) -> None:
-        self._add_query(kind="prepare", query=query, dialect=query.dialect)
+        self._add_query(query=query, dialect=query.dialect)
 
     def add_procedure_query(self, query: ProcedureQuery) -> None:
-        self._add_query(kind="procedure", query=query, dialect=query.dialect)
+        self._add_query(query=query, dialect=query.dialect)
 
     def add_schema_query(self, query: SchemaQuery) -> None:
-        self._add_query(kind="schema", query=query, dialect=query.dialect)
+        self._add_query(query=query, dialect=query.dialect)
 
     def add_sequence_query(self, query: SequenceQuery) -> None:
-        self._add_query(kind="sequence", query=query, dialect=query.dialect)
+        self._add_query(query=query, dialect=query.dialect)
 
     def add_stage_query(self, query: StageQuery) -> None:
-        self._add_query(kind="stage", query=query, dialect=query.dialect)
+        self._add_query(query=query, dialect=query.dialect)
 
     def add_table_query(
-        self, query: TableQuery | ViewQuery | CTASQuery, column_mapping: t.Optional[ColumnMapping] = None
+        self, query: TableQuery | ViewQuery | CTASQuery, column_mapping: ColumnMapping | None = None
     ) -> None:
-        self._add_query(kind="table", query=query, column_mapping=column_mapping, dialect=query.dialect)
+        self._add_query(query=query, column_mapping=column_mapping, dialect=query.dialect, kind="table")
 
     def add_trigger_query(self, query: TriggerQuery) -> None:
-        self._add_query(kind="trigger", query=query, dialect=query.dialect)
+        self._add_query(query=query, dialect=query.dialect)
 
     def add_type_query(self, query: TypeQuery) -> None:
-        self._add_query(kind="type", query=query, dialect=query.dialect)
+        self._add_query( query=query, dialect=query.dialect)
 
-    def add_udf_query(self, query: UserDefinedFunctionQuery, column_mapping: t.Optional[ColumnMapping] = None) -> None:
-        self._add_query(kind="udf", query=query, column_mapping=column_mapping, dialect=query.dialect)
+    def add_udf_query(self, query: UserDefinedFunctionQuery, column_mapping: ColumnMapping | None = None) -> None:
+        self._add_query(query=query, column_mapping=column_mapping, dialect=query.dialect)
 
     def _add_query(
         self,
-        kind: str,
         query: Q,
-        column_mapping: t.Optional[ColumnMapping] = None,
+        column_mapping: ColumnMapping | None = None,
         dialect: DialectType = None,
-        normalize: t.Optional[bool] = None,
+        normalize: bool | None = None,
         match_depth: bool = False,
+        kind: str = ""
     ) -> None:
         """
         Register or update a table. Updates are only performed if a new column mapping is provided.
@@ -100,6 +152,9 @@ class ObjectMapping(MappingSchema):
             match_depth: whether to enforce that the table must match the schema's depth or not.
         """
         # Initialize the structures. These are required by sqlglot.
+        if not kind:
+            kind = query.kind
+
         if kind not in self.kind_mapping:
             self.kind_mapping[kind] = {}
             self.kind_mapping_trie[kind] = new_trie({})
@@ -133,11 +188,11 @@ class ObjectMapping(MappingSchema):
     def _add_columns_for_table(
         self,
         table: exp.Table,
-        column_mapping: t.Optional[ColumnMapping] = None,
+        column_mapping: ColumnMapping | None = None,
         dialect: DialectType = None,
-        normalize: t.Optional[bool] = None,
+        normalize: bool | None = None,
         match_depth: bool = False,
-    ):
+    ) -> None:
         super().add_table(
             table=table,
             column_mapping=column_mapping,
@@ -151,7 +206,7 @@ class ObjectMapping(MappingSchema):
         table: exp.Table,
         raise_on_missing: bool = True,
         ensure_data_types: bool = False,
-    ):
+    ) -> None:
         """
         A nicer name for the parent's function.
         """
@@ -184,7 +239,7 @@ class ObjectMapping(MappingSchema):
 
     def lookup_udf_query(
         self, table: exp.Table, raise_on_missing: bool = True
-    ) -> t.List[UserDefinedFunctionQuery] | None:
+    ) -> list[UserDefinedFunctionQuery] | None:
         # A UDF can have multiple names but different types, thus we return a list
         return self._lookup_query(kind="udf", object=table, raise_on_missing=raise_on_missing)
 
@@ -210,7 +265,12 @@ class ObjectMapping(MappingSchema):
         """
         if kind not in self.kind_mapping:
             if raise_on_missing:
-                raise exception.MappingError(f"Could not find '{object.this}' of type '{kind}' in mapping.")
+                if kind == "prepare":
+                    knd = "PREPARE"
+                else:
+                    knd = "CREATE " + {"udf": "FUNCTION"}.get(kind, kind).upper()
+
+                raise exception.MappingError(f"Could not find '{object.this}' of type '{kind}' in mapping. Did you forget to define a '{knd}' statement?")
             return None
 
         parts = self.table_parts(object)[0: len(self.supported_table_args)]
@@ -234,7 +294,7 @@ class ObjectMapping(MappingSchema):
 
     # Override sqlglot's property. It seems to be buggy when using different dict sizes (catalog, schema, etc)
     @property
-    def supported_table_args(self) -> t.Tuple[str, ...]:
+    def supported_table_args(self) -> tuple[str, ...]:
         return exp.TABLE_PARTS
 
     def get_table_or_stage(
@@ -248,7 +308,7 @@ class ObjectMapping(MappingSchema):
             table_expr = exp.Table(this=table)
 
         if str(table_expr).startswith("@"):
-            child_object_query = self.lookup_stage_query(table=table_expr, raise_on_missing=raise_on_missing)
+            child_object_query = self.lookup_stage_query(table=table_expr, raise_on_missing=False)
         else:
             child_object_query = self.lookup_table_query(table=table_expr, raise_on_missing=raise_on_missing)
 
@@ -260,7 +320,7 @@ class ObjectMapping(MappingSchema):
     def lookup_udf_call(
         self,
         node: exp.Anonymous,
-    ) -> t.Optional[UserDefinedFunctionQuery]:
+    ) -> UserDefinedFunctionQuery | None:
         """
         Looks up the UDF definition for a single exp.Anonymous node.
         Returns the matched UDF definition, or None if not found.
@@ -275,8 +335,8 @@ class ObjectMapping(MappingSchema):
 
 
 def resolve_overloaded_function(
-    node: exp.Anonymous, candidates: t.List[UserDefinedFunctionQuery]
-) -> t.Optional[UserDefinedFunctionQuery]:
+    node: exp.Anonymous, candidates: list[UserDefinedFunctionQuery]
+) -> UserDefinedFunctionQuery | None:
     """
     Resolves the best function candidate for an overloaded function call.
     Applies some precedence rules (e.g., type matching, or non-variadic > variadic) but many are currently
@@ -323,7 +383,7 @@ def resolve_overloaded_function(
     return matches[0]
 
 
-def match_function_arguments(args: t.List[exp.Expr], candidate: UserDefinedFunctionQuery) -> bool:
+def match_function_arguments(args: list[exp.Expr], candidate: UserDefinedFunctionQuery) -> bool:
     """
     Checks if the provided arguments match the function candidate's parameters.
     Handles exact types, polymorphic types (anyelement, anyarray), and VARIADIC parameters.

sqlleaf/models/context.py

@@ -17,7 +17,7 @@ logger = logging.getLogger("sqlleaf")
 
 class ScopePositions:
     def __init__(self):
-        self.positions: t.Dict[int, t.Dict[int, int]] = {}
+        self.positions: dict[int, dict[int, int]] = {}
 
     def get_position_for_expr(self, expr: exp.Expr) -> dict[int, int]:
         return self.positions[id(expr)]

sqlleaf/models/node/base.py

@@ -142,7 +142,7 @@ class EdgeAttributes:
 
 class GraphAttributes:
     def __init__(self):
-        self.queries: t.List[QueryHolder] = []
+        self.queries: list[QueryHolder] = []
 
     def add_query_to_graph(self, holder: QueryHolder):
         self.queries.append(holder)

sqlleaf/models/query/__init__.py

@@ -1,22 +1,30 @@
 import typing as t
 
 from sqlleaf.models.query.base import Query as Query
+from sqlleaf.models.query.block import BlockQuery as BlockQuery
 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.declare_item import AssignmentQuery as AssignmentQuery
 from sqlleaf.models.query.delete import DeleteQuery as DeleteQuery
 from sqlleaf.models.query.execute import ExecuteQuery as ExecuteQuery
+from sqlleaf.models.query.execute_dynamic import ExecuteDynamicQuery as ExecuteDynamicQuery
+from sqlleaf.models.query.fetch import FetchQuery as FetchQuery
+from sqlleaf.models.query.for_in import ForInQuery as ForInQuery
 from sqlleaf.models.query.holder import QueryHolder as QueryHolder
 from sqlleaf.models.query.insert import InsertQuery as InsertQuery
+from sqlleaf.models.query.loop import LoopQuery as LoopQuery
 from sqlleaf.models.query.merge import MergeQuery as MergeQuery
 from sqlleaf.models.query.multitable_insert import MultitableInsertQuery as MultitableInsertQuery
+from sqlleaf.models.query.open import OpenQuery as OpenQuery
 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.perform import PerformQuery as PerformQuery
 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
@@ -29,5 +37,48 @@ from sqlleaf.models.query.user_defined_function import FunctionParam as Function
 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
+from sqlleaf.models.query.while_ import WhileQuery as WhileQuery
 
 Q = t.TypeVar("Q", bound=Query)
+
+__all__ = [
+    "AssignmentQuery",
+    "BlockQuery",
+    "CallQuery",
+    "CopyQuery",
+    "CTASQuery",
+    "DatabaseQuery",
+    "DeleteQuery",
+    "ExecuteDynamicQuery",
+    "ExecuteQuery",
+    "FetchQuery",
+    "ForInQuery",
+    "FunctionParam",
+    "InsertQuery",
+    "LoopQuery",
+    "MergeQuery",
+    "MultitableInsertQuery",
+    "OpenQuery",
+    "PerformQuery",
+    "PrepareQuery",
+    "ProcedureQuery",
+    "PutQuery",
+    "Q",
+    "Query",
+    "QueryHolder",
+    "ReplaceQuery",
+    "SchemaQuery",
+    "SelectQuery",
+    "SequenceQuery",
+    "SetQuery",
+    "StageQuery",
+    "TableQuery",
+    "TriggerQuery",
+    "TypeQuery",
+    "UnloadQuery",
+    "UpdateQuery",
+    "UserDefinedFunctionQuery",
+    "ValuesQuery",
+    "ViewQuery",
+    "WhileQuery",
+]

sqlleaf/models/query/base.py

@@ -25,15 +25,14 @@ class Query:
         statement_index: int,
         object_mapping: mappings.ObjectMapping,
         source_info: SourceInfo | None,
-        target_info: TargetInfo,
+        target_info: TargetInfo | None,
         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.column_defs: list[exp.ColumnDef] = []
         self.property = ""
 
         # Remove comments at initialization
@@ -54,10 +53,10 @@ class Query:
             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):
+    def set_holder(self, holder: QueryHolder) -> None:
         self.holder = holder
 
-    def _determine_expression_type(self, expr: exp.Expr | t.List[exp.Expr], dialect: str) -> SqlObjectType:
+    def _determine_expression_type(self, expr: exp.Expr | 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.
@@ -126,7 +125,7 @@ class Query:
     def get_transformed_self(self) -> Query | None:
         return self.holder.transformed
 
-    def get_columns_from_target(self) -> t.List[exp.ColumnDef]:
+    def get_columns_from_target(self) -> list[exp.ColumnDef]:
         """
         Given a query, figure out its target object, including its columns.
 
@@ -160,7 +159,7 @@ class Query:
     def _get_column_defs(
         self,
         expr: SourceExprType | TargetExprType,
-    ) -> t.List[exp.ColumnDef]:
+    ) -> list[exp.ColumnDef]:
         """
         Most of the time, the sources and target are tables.
         However, with COPY/UNLOAD, they can be files or streams.
@@ -198,19 +197,15 @@ class Query:
         """
         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)
+        return str(self.statement_index)
 
-    def get_ctes(self) -> t.List:
+    def get_ctes(self) -> list:
         return []
 
-    def get_column_defs(self, include_system: bool = False) -> t.List[exp.ColumnDef]:
+    def get_column_defs(self, include_system: bool = False) -> list[exp.ColumnDef]:
         return self.column_defs
 
-    def get_column_names_with_types(self, include_system: bool = False) -> t.Dict[str, str]:
+    def get_column_names_with_types(self, include_system: bool = False) -> dict[str, str]:
         """
         Used by sqlglot's MappingSchema
         """
@@ -219,13 +214,9 @@ class Query:
 
     @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)
+        return "query:" + util.short_sha256_hash(self.statement.sql(dialect=self.dialect) + ":" + str(self.statement_index))
 
-    def get_all_queries(self, types: t.Tuple | None = None):
+    def get_all_queries(self, types: tuple | None = None):
         """
         Collect all queries (self + children recursively), optionally filtered by type.
         """