sqlleaf 0.0.1__py3-none-any.whl

sqlleaf/processors/generator/dialects/postgres.py

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+import logging
+import typing as t
+from dataclasses import replace
+
+from sqlglot import exp
+
+from sqlleaf import exception, util
+from sqlleaf.models.context import GeneratorContext, PositionContext
+from sqlleaf.models.node import (
+    ColumnNode,
+    SequenceNode,
+)
+from sqlleaf.processors.generator.dialects.base import BaseGenerator, EdgeToCreate, singledispatchmethodlogger
+from sqlleaf.typing import SourceInfo, SqlObjectType
+
+logger = logging.getLogger("sqlleaf")
+
+
+class PostgresGenerator(BaseGenerator):
+    dialect = "postgres"
+
+    @singledispatchmethodlogger
+    def process(self, expr: exp.Expr, gen_ctx: GeneratorContext, pos_ctx: PositionContext) -> t.Iterator[EdgeToCreate]:
+        yield from super().process(expr, gen_ctx, pos_ctx)
+
+    @process.register
+    def process_table(
+        self, expr: exp.Table, gen_ctx: GeneratorContext, pos_ctx: PositionContext
+    ) -> t.Iterator[EdgeToCreate]:
+        """
+        Process a table or a table function.
+        This is a bit awkward as we have the sequence: Table -> ColumnDef -> Table
+        for table functions.
+        """
+        if "rows_from" in expr.args:
+            yield from self._process_rows_from(expr, gen_ctx, pos_ctx)
+        elif isinstance(expr.this, exp.Anonymous):
+            yield from super().process(expr.this, gen_ctx.replace(expr=expr.this), pos_ctx.replace(select_index=0))
+        else:
+            yield from super().process(expr, gen_ctx, pos_ctx)
+
+    def _process_rows_from(
+        self, expr: exp.TableColumn, gen_ctx: GeneratorContext, pos_ctx: PositionContext
+    ) -> t.Iterator[EdgeToCreate]:
+        """
+        SELECT * FROM ROWS FROM
+        (
+            unnest(ARRAY['x', 'y']),
+            json_to_recordset('[{"a":40,"b":"foo"}]')
+                AS y(a INTEGER, b TEXT),
+            generate_series(1, 3)
+        ) AS x (name, age, kind, amount)
+        """
+        downstream_exprs = []
+        for table_function in expr.args["rows_from"]:
+            # Determine the immediate children of the expression.
+            # These are either table functions or aliases to table functions (ColumnDefs)
+            cols = list(table_function.find_all(exp.ColumnDef))
+            downstream_exprs.extend(cols if cols else [table_function])
+
+        # Get the expression associated with the column name
+        child_column_name = gen_ctx.get_child_node().expr.name
+        for i, col in enumerate(expr.alias_column_names):
+            if col == child_column_name:
+                # Returns ColumnDef | Function | Table
+                down_expr = downstream_exprs[i]
+                if isinstance(down_expr, exp.Table) and down_expr.arg_key == "rows_from":
+                    # A table function inside a 'ROWS FROM'
+                    down_expr = down_expr.this
+
+                gen_ctx = gen_ctx.replace(expr=down_expr)
+                yield from self.process(down_expr, gen_ctx, pos_ctx.replace(select_index=i))
+                return
+
+    @process.register
+    def process_lateral(
+        self, expr: exp.Lateral, gen_ctx: GeneratorContext, pos_ctx: PositionContext
+    ) -> t.Iterator[EdgeToCreate]:
+        """
+        Example query: "SELECT .. LATERAL function() AS t(name)"
+        This only handles cases where a function follows LATERAL. If a subquery follows instead,
+        its expressions are walked inside walk_query_scope().
+        """
+        gen_ctx = gen_ctx.replace(expr=expr.this)
+        yield from self.process(expr.this, gen_ctx, pos_ctx)
+
+    @process.register
+    def process_anonymous(
+        self, expr: exp.Anonymous, gen_ctx: GeneratorContext, pos_ctx: PositionContext
+    ) -> t.Iterator[EdgeToCreate]:
+        """
+        Either user-defined functions or sequence functions.
+
+        SELECT my.func() or SELECT nextval('my_sequence')
+        """
+        schema, function = util.get_udf_name(expr)
+
+        # Process a sequence
+        if not schema and function in [
+            "nextval",
+            "currval",
+            "setval",
+            # 'lastval()' is not supported since it requires tracking state
+        ]:
+            yield from self._process_sequence(expr, gen_ctx, pos_ctx)
+        else:
+            yield from super().process(expr, gen_ctx, pos_ctx)
+
+    def _process_sequence(
+        self, expr: exp.Anonymous, gen_ctx: GeneratorContext, pos_ctx: PositionContext
+    ) -> t.Iterator[EdgeToCreate]:
+        """
+        Either user-defined functions or sequence functions.
+
+        SELECT my.func() or SELECT nextval('my_sequence')
+        """
+        seq_name_expr: exp.Literal = expr.args["expressions"][0]
+        schema, function = util.get_udf_name(expr)
+        full_name = ".".join([schema, function])
+
+        # Ensure the sequence exists
+        seq_table = exp.table_(table=seq_name_expr.name, db=schema)
+        seq_query = gen_ctx.query.object_mapping.lookup_sequence_query(table=seq_table)
+        if not seq_query:
+            raise exception.MappingError(f"Sequence '{full_name}' not found.")
+
+        subkind = seq_query.property if seq_query else ""
+        gen_ctx = gen_ctx.replace(new_data_type=exp.DataType.build("INT"))
+        parent = SequenceNode(name=seq_name_expr.name, gen_ctx=gen_ctx, pos_ctx=pos_ctx, subkind=subkind)
+        yield EdgeToCreate(parent, gen_ctx.child_node)
+
+    @process.register
+    def process_column_def(
+        self, expr: exp.ColumnDef, gen_ctx: GeneratorContext, pos_ctx: PositionContext
+    ) -> t.Iterator[EdgeToCreate]:
+        gen_ctx = gen_ctx.replace(new_data_type=expr.kind)
+
+        if isinstance(expr.parent, exp.TableAlias):
+            # A table alias inside a 'ROWS FROM' cannot have a name
+            parent = ColumnNode(
+                catalog="",
+                schema="",
+                table="",
+                column=expr.name,
+                gen_ctx=gen_ctx,
+                pos_ctx=pos_ctx,
+            )
+            yield EdgeToCreate(parent, gen_ctx.child_node)
+
+            # Process the inner table function of the table alias
+            table_function: exp.Table = t.cast(exp.Table, expr.parent.parent)
+            yield from self.do_grandparents([table_function.this], parent, gen_ctx, pos_ctx.replace(select_index=0))
+
+    @process.register
+    def process_column(
+        self, expr: exp.Column, gen_ctx: GeneratorContext, pos_ctx: PositionContext
+    ) -> t.Iterator[EdgeToCreate]:
+        """
+        COPY x FROM/TO y
+        """
+        # TODO: create a new Node called column[kind=table subkind=function] for table functions
+        #  (JSON_EACH_TEXT, UNNEST, UDFs that return sets)
+        source_info: SourceInfo = gen_ctx.query.source_info
+
+        if SqlObjectType.type_has_no_column_defs(source_info.type):
+            if source_info.type == SqlObjectType.FILE:
+                gen_ctx = gen_ctx.replace(
+                    expr=source_info.expression, new_data_type=gen_ctx.get_child_node().get_data_type()
+                )
+
+            node = self.create_node_from_type(
+                object_type=source_info.type,
+                expression=source_info.expression,
+                column_name=gen_ctx.get_child_node().name,
+                gen_ctx=gen_ctx,
+                pos_ctx=pos_ctx,
+            )
+            yield EdgeToCreate(node, gen_ctx.child_node)
+        else:
+            yield from super().process(expr, gen_ctx, pos_ctx)

sqlleaf/processors/generator/dialects/redshift.py

@@ -0,0 +1,192 @@
+from __future__ import annotations
+
+import logging
+import typing as t
+from dataclasses import replace
+
+from sqlglot import exp
+from sqlglot.optimizer import Scope
+
+from sqlleaf import exception
+from sqlleaf.models.context import GeneratorContext, PositionContext
+from sqlleaf.models.node import (
+    FileColumnNode,
+    PivotNode,
+    UnpivotNode,
+)
+from sqlleaf.processors.generator.dialects.base import BaseGenerator, EdgeToCreate, singledispatchmethodlogger
+from sqlleaf.typing import SourceInfo, SqlObjectType
+
+logger = logging.getLogger("sqlleaf")
+
+
+class RedshiftGenerator(BaseGenerator):
+    dialect = "redshift"
+
+    @singledispatchmethodlogger
+    def process(self, expr: exp.Expr, gen_ctx: GeneratorContext, pos_ctx: PositionContext) -> t.Iterator[EdgeToCreate]:
+        yield from super().process(expr, gen_ctx, pos_ctx)
+
+    @process.register
+    def process_unpivot(
+        self, expr: exp.Pivot, gen_ctx: GeneratorContext, pos_ctx: PositionContext
+    ) -> t.Iterator[EdgeToCreate]:
+        """
+        SELECT * FROM ... UNPIVOT ( ... )
+        """
+        # Given expression:
+        #   UNPIVOT ( <expression> FOR <field> IN (<column> AS <value>) )
+        # We have lineage:
+        #   <column> -> UNPIVOT -> <expression>
+        #   <value> -> UNPIVOT -> <field>
+        scope = t.cast(Scope, gen_ctx.scope)
+        selected_column = scope.columns[pos_ctx.select_index]
+        pivot_expression = expr.expressions[0]
+        pivot_field = expr.fields[0]
+
+        if selected_column.name == pivot_expression.name:
+            arg = "this"
+        elif selected_column.name == pivot_field.this.name:
+            arg = "alias"
+        else:
+            message = f"Could not find column '{selected_column.name}' in UNPIVOT expression"
+            raise exception.SqlLeafException(message=message)
+
+        # Create an unpivot node for each downstream column/expression
+        for pivot_alias in pivot_field.expressions:
+            pivot_value = pivot_alias.args[arg]
+
+            unpivot_node = self.create_node(UnpivotNode(
+                gen_ctx=gen_ctx,
+                pos_ctx=pos_ctx,
+            ))
+            if unpivot_node:
+                source = pivot_value.name if arg == "this" else ""  # Only columns are sources for now
+                unpivot_node.set(source=source, target=selected_column.name)
+                yield EdgeToCreate(unpivot_node, gen_ctx.child_node)
+
+            yield from self.do_grandparents([pivot_value], unpivot_node, gen_ctx, pos_ctx)
+
+    @process.register
+    def process_pivot(
+        self, expr: exp.Pivot, gen_ctx: GeneratorContext, pos_ctx: PositionContext
+    ) -> t.Iterator[EdgeToCreate]:
+        """
+        SELECT * FROM (SELECT  ...) PIVOT ( ... )
+        """
+        # Find the associated expression for the column, and process it
+        scope = t.cast(Scope, gen_ctx.scope)
+        selected_column = scope.columns[pos_ctx.select_index]
+        pivot_column_mapping = _get_pivot_mapping(expr)
+
+        # The associated column and expression
+        column_and_expr = pivot_column_mapping[selected_column.name]
+        pivot_expr = column_and_expr["expression"]
+
+        pivot_node = self.create_node(PivotNode(
+            gen_ctx=gen_ctx,
+            pos_ctx=pos_ctx,
+        ))
+        if pivot_node:
+            pivot_node.set(source=pivot_expr.alias_or_name, target=selected_column.alias_or_name)
+            yield EdgeToCreate(pivot_node, gen_ctx.child_node)
+
+        grandparents = [pivot_expr]
+        yield from self.do_grandparents(grandparents, pivot_node, gen_ctx, pos_ctx)
+
+    @process.register
+    def process_column(
+        self, expr: exp.Column, gen_ctx: GeneratorContext, pos_ctx: PositionContext
+    ) -> t.Iterator[EdgeToCreate]:
+        scope = t.cast(Scope, gen_ctx.scope)
+        pivot = _get_pivot_expr(scope)
+        if (
+            pivot
+            and pivot.alias_or_name == expr.table
+            and not isinstance(gen_ctx.child_node, UnpivotNode)  # Prevent infinite recursion
+        ):
+            gen_ctx = gen_ctx.replace(expr=pivot)
+            if pivot.unpivot:
+                yield from self.process_unpivot(pivot, gen_ctx, pos_ctx)
+            else:
+                yield from self.process_pivot(pivot, gen_ctx, pos_ctx)
+        else:
+            yield from self._process_column_other(expr, gen_ctx, pos_ctx)
+
+    def _process_column_other(
+        self, expr: exp.Column, gen_ctx: GeneratorContext, pos_ctx: PositionContext
+    ) -> t.Iterator[EdgeToCreate]:
+        source_info: SourceInfo = gen_ctx.query.source_info
+
+        if SqlObjectType.type_has_no_column_defs(source_info.type):
+            # A filename from COPY/UNLOAD
+            gen_ctx = gen_ctx.replace(
+                expr=source_info.expression, new_data_type=gen_ctx.get_child_node().get_data_type()
+            )
+
+            node = self.create_node_from_type(
+                object_type=source_info.type,
+                expression=source_info.expression,
+                column_name=gen_ctx.get_child_node().name,
+                gen_ctx=gen_ctx,
+                pos_ctx=pos_ctx,
+            )
+            yield EdgeToCreate(node, gen_ctx.child_node)
+        else:
+            yield from super().process(expr, gen_ctx, pos_ctx)
+
+    @process.register
+    def process_location(
+        self, expr: exp.LocationProperty, gen_ctx: GeneratorContext, pos_ctx: PositionContext
+    ) -> t.Iterator[EdgeToCreate]:
+        """
+        CREATE EXTERNAL TABLE ... LOCATION
+        """
+        location = expr.this
+        child_node = gen_ctx.child_node
+
+        # TODO: change to source/target info
+
+        file_format_property = gen_ctx.query.statement.args["properties"].find(exp.FileFormatProperty)
+        file_format = file_format_property.this.name if file_format_property else ""
+
+        column_node = self.create_node(FileColumnNode(
+            column=child_node.name,
+            file_format=file_format,
+            file_path=location.name,
+            gen_ctx=gen_ctx,
+            pos_ctx=pos_ctx,
+        ))
+
+        yield EdgeToCreate(column_node, child_node)
+
+
+def _get_pivot_expr(scope: Scope) -> exp.Pivot | None:
+    pivots = scope.pivots if scope else []
+    pivot = pivots[0] if len(pivots) == 1 else None
+    return pivot
+
+
+def _get_pivot_mapping(expr: exp.Pivot) -> dict:
+    """
+    Get information related to PIVOT statements.
+    """
+    pivot_column_mapping = {}
+    # For each aggregation function, the pivot creates a new column for each field in category
+    # combined with the aggfunc. So the columns parsed have this order: cat_a_value_sum, cat_a,
+    # b_value_sum, b. Because of this step wise manner the aggfunc 'sum(value) as value_sum'
+    # belongs to the column indices 0, 2, and the aggfunc 'max(price)' without an alias belongs
+    # to the column indices 1, 3. Here, only the columns used in the aggregations are of interest
+    # in the lineage, so lookup the pivot column name by index and map that with the columns used
+    # in the aggregation.
+    #
+    # Example: PIVOT (SUM(value) AS value_sum, MAX(price)) FOR category IN ('a' AS cat_a, 'b')
+    pivot_columns = expr.args["columns"]
+    pivot_aggs_count = len(expr.expressions)
+
+    for i, agg in enumerate(expr.expressions):
+        agg_cols = list(agg.find_all(exp.Column))
+        for col_index in range(i, len(pivot_columns), pivot_aggs_count):
+            pivot_column_mapping[pivot_columns[col_index].name] = {"column": agg_cols, "expression": agg}
+
+    return pivot_column_mapping

sqlleaf/processors/generator/dialects/snowflake.py

@@ -0,0 +1,90 @@
+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 (
+    FileColumnNode,
+    StageColumnNode,
+)
+from sqlleaf.processors.generator.dialects.base import BaseGenerator, EdgeToCreate, singledispatchmethodlogger
+from sqlleaf.typing import SqlObjectType
+
+logger = logging.getLogger("sqlleaf")
+
+
+class SnowflakeGenerator(BaseGenerator):
+    dialect = "snowflake"
+
+    @singledispatchmethodlogger
+    def process(self, expr: exp.Expr, gen_ctx: GeneratorContext, pos_ctx: PositionContext) -> t.Iterator[EdgeToCreate]:
+        yield from super().process(expr, gen_ctx, pos_ctx)
+
+    @process.register
+    def process_put(
+        self, expr: exp.Put, gen_ctx: GeneratorContext, pos_ctx: PositionContext
+    ) -> t.Iterator[EdgeToCreate]:
+        """
+        PUT 'file:///tmp/data/mydata.csv' @my_int_stage;
+        - Creates two node: FileColumnNode and StageColumnNode
+        """
+        source = gen_ctx.query.source_info.expression
+        target = gen_ctx.query.target_info.expression
+
+        file_format = gen_ctx.query.get_original_self().parameters.file_format
+
+        file_ctx = gen_ctx.replace(expr=source)
+        stage_ctx = gen_ctx.replace(expr=target)
+
+        stage_query = gen_ctx.query.object_mapping.get_table_or_stage(table=target, raise_on_missing=False)
+
+        file_node = self.create_node(FileColumnNode(
+            column="?",
+            file_format=file_format,
+            file_path=source.name,
+            gen_ctx=file_ctx,
+            pos_ctx=pos_ctx,
+        ))
+        stage_node = self.create_node(StageColumnNode(
+            column="?",
+            stage=target,
+            gen_ctx=stage_ctx,
+            pos_ctx=pos_ctx,
+            path=stage_query.path,
+        ))
+
+        yield EdgeToCreate(file_node, stage_node)
+
+    @process.register
+    def process_copy(
+        self, expr: exp.Copy, gen_ctx: GeneratorContext, pos_ctx: PositionContext
+    ) -> t.Iterator[EdgeToCreate]:
+        yield from self.process_column(expr, gen_ctx, pos_ctx)
+
+    @process.register
+    def process_column(
+        self, expr: exp.Column, gen_ctx: GeneratorContext, pos_ctx: PositionContext
+    ) -> t.Iterator[EdgeToCreate]:
+        """
+        If the source is a Stage, create a StageColumnNode.
+        """
+        query = gen_ctx.query
+        if hasattr(query, "source_info") and query.source_info.type == SqlObjectType.STAGE:
+            stage_expr = query.source_info.expression
+            stage_name: exp.Var = stage_expr.this
+            stage_query = query.object_mapping.get_table_or_stage(table=stage_expr, raise_on_missing=False)
+            stage_path = stage_query.path
+
+            parent = self.create_node(StageColumnNode(
+                column=expr.name,
+                stage=stage_name,
+                gen_ctx=gen_ctx,
+                pos_ctx=pos_ctx,
+                path=stage_path,
+            ))
+            yield EdgeToCreate(parent, gen_ctx.child_node)
+        else:
+            yield from super().process_column(expr, gen_ctx, pos_ctx)