altimate-code 0.5.1 → 0.5.3

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/jsonpath.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+import typing as t
+
+import sqlglot.expressions as exp
+from sqlglot.errors import ParseError
+from sqlglot.tokens import Token, Tokenizer, TokenType
+
+if t.TYPE_CHECKING:
+    from sqlglot._typing import Lit
+    from sqlglot.dialects.dialect import DialectType
+
+
+class JSONPathTokenizer(Tokenizer):
+    SINGLE_TOKENS = {
+        "(": TokenType.L_PAREN,
+        ")": TokenType.R_PAREN,
+        "[": TokenType.L_BRACKET,
+        "]": TokenType.R_BRACKET,
+        ":": TokenType.COLON,
+        ",": TokenType.COMMA,
+        "-": TokenType.DASH,
+        ".": TokenType.DOT,
+        "?": TokenType.PLACEHOLDER,
+        "@": TokenType.PARAMETER,
+        "'": TokenType.QUOTE,
+        '"': TokenType.QUOTE,
+        "$": TokenType.DOLLAR,
+        "*": TokenType.STAR,
+    }
+
+    KEYWORDS = {
+        "..": TokenType.DOT,
+    }
+
+    IDENTIFIER_ESCAPES = ["\\"]
+    STRING_ESCAPES = ["\\"]
+
+
+def parse(path: str, dialect: DialectType = None) -> exp.JSONPath:
+    """Takes in a JSON path string and parses it into a JSONPath expression."""
+    from sqlglot.dialects import Dialect
+
+    jsonpath_tokenizer = Dialect.get_or_raise(dialect).jsonpath_tokenizer
+    tokens = jsonpath_tokenizer.tokenize(path)
+    size = len(tokens)
+
+    i = 0
+
+    def _curr() -> t.Optional[TokenType]:
+        return tokens[i].token_type if i < size else None
+
+    def _prev() -> Token:
+        return tokens[i - 1]
+
+    def _advance() -> Token:
+        nonlocal i
+        i += 1
+        return _prev()
+
+    def _error(msg: str) -> str:
+        return f"{msg} at index {i}: {path}"
+
+    @t.overload
+    def _match(token_type: TokenType, raise_unmatched: Lit[True] = True) -> Token:
+        pass
+
+    @t.overload
+    def _match(token_type: TokenType, raise_unmatched: Lit[False] = False) -> t.Optional[Token]:
+        pass
+
+    def _match(token_type, raise_unmatched=False):
+        if _curr() == token_type:
+            return _advance()
+        if raise_unmatched:
+            raise ParseError(_error(f"Expected {token_type}"))
+        return None
+
+    def _parse_literal() -> t.Any:
+        token = _match(TokenType.STRING) or _match(TokenType.IDENTIFIER)
+        if token:
+            return token.text
+        if _match(TokenType.STAR):
+            return exp.JSONPathWildcard()
+        if _match(TokenType.PLACEHOLDER) or _match(TokenType.L_PAREN):
+            script = _prev().text == "("
+            start = i
+
+            while True:
+                if _match(TokenType.L_BRACKET):
+                    _parse_bracket()  # nested call which we can throw away
+                if _curr() in (TokenType.R_BRACKET, None):
+                    break
+                _advance()
+
+            expr_type = exp.JSONPathScript if script else exp.JSONPathFilter
+            return expr_type(this=path[tokens[start].start : tokens[i].end])
+
+        number = "-" if _match(TokenType.DASH) else ""
+
+        token = _match(TokenType.NUMBER)
+        if token:
+            number += token.text
+
+        if number:
+            return int(number)
+
+        return False
+
+    def _parse_slice() -> t.Any:
+        start = _parse_literal()
+        end = _parse_literal() if _match(TokenType.COLON) else None
+        step = _parse_literal() if _match(TokenType.COLON) else None
+
+        if end is None and step is None:
+            return start
+
+        return exp.JSONPathSlice(start=start, end=end, step=step)
+
+    def _parse_bracket() -> exp.JSONPathPart:
+        literal = _parse_slice()
+
+        if isinstance(literal, str) or literal is not False:
+            indexes = [literal]
+            while _match(TokenType.COMMA):
+                literal = _parse_slice()
+
+                if literal:
+                    indexes.append(literal)
+
+            if len(indexes) == 1:
+                if isinstance(literal, str):
+                    node: exp.JSONPathPart = exp.JSONPathKey(this=indexes[0])
+                elif isinstance(literal, exp.JSONPathPart) and isinstance(
+                    literal, (exp.JSONPathScript, exp.JSONPathFilter)
+                ):
+                    node = exp.JSONPathSelector(this=indexes[0])
+                else:
+                    node = exp.JSONPathSubscript(this=indexes[0])
+            else:
+                node = exp.JSONPathUnion(expressions=indexes)
+        else:
+            raise ParseError(_error("Cannot have empty segment"))
+
+        _match(TokenType.R_BRACKET, raise_unmatched=True)
+
+        return node
+
+    def _parse_var_text() -> str:
+        """
+        Consumes & returns the text for a var. In BigQuery it's valid to have a key with spaces
+        in it, e.g JSON_QUERY(..., '$. a b c ') should produce a single JSONPathKey(' a b c ').
+        This is done by merging "consecutive" vars until a key separator is found (dot, colon etc)
+        or the path string is exhausted.
+        """
+        prev_index = i - 2
+
+        while _match(TokenType.VAR):
+            pass
+
+        start = 0 if prev_index < 0 else tokens[prev_index].end + 1
+
+        if i >= len(tokens):
+            # This key is the last token for the path, so it's text is the remaining path
+            text = path[start:]
+        else:
+            text = path[start : tokens[i].start]
+
+        return text
+
+    # We canonicalize the JSON path AST so that it always starts with a
+    # "root" element, so paths like "field" will be generated as "$.field"
+    _match(TokenType.DOLLAR)
+    expressions: t.List[exp.JSONPathPart] = [exp.JSONPathRoot()]
+
+    while _curr():
+        if _match(TokenType.DOT) or _match(TokenType.COLON):
+            recursive = _prev().text == ".."
+
+            if _match(TokenType.VAR):
+                value: t.Optional[str | exp.JSONPathWildcard] = _parse_var_text()
+            elif _match(TokenType.IDENTIFIER):
+                value = _prev().text
+            elif _match(TokenType.STAR):
+                value = exp.JSONPathWildcard()
+            else:
+                value = None
+
+            if recursive:
+                expressions.append(exp.JSONPathRecursive(this=value))
+            elif value:
+                expressions.append(exp.JSONPathKey(this=value))
+            else:
+                raise ParseError(_error("Expected key name or * after DOT"))
+        elif _match(TokenType.L_BRACKET):
+            expressions.append(_parse_bracket())
+        elif _match(TokenType.VAR):
+            expressions.append(exp.JSONPathKey(this=_parse_var_text()))
+        elif _match(TokenType.IDENTIFIER):
+            expressions.append(exp.JSONPathKey(this=_prev().text))
+        elif _match(TokenType.STAR):
+            expressions.append(exp.JSONPathWildcard())
+        else:
+            raise ParseError(_error(f"Unexpected {tokens[i].token_type}"))
+
+    return exp.JSONPath(expressions=expressions)
+
+
+JSON_PATH_PART_TRANSFORMS: t.Dict[t.Type[exp.Expression], t.Callable[..., str]] = {
+    exp.JSONPathFilter: lambda _, e: f"?{e.this}",
+    exp.JSONPathKey: lambda self, e: self._jsonpathkey_sql(e),
+    exp.JSONPathRecursive: lambda _, e: f"..{e.this or ''}",
+    exp.JSONPathRoot: lambda *_: "$",
+    exp.JSONPathScript: lambda _, e: f"({e.this}",
+    exp.JSONPathSelector: lambda self, e: f"[{self.json_path_part(e.this)}]",
+    exp.JSONPathSlice: lambda self, e: ":".join(
+        "" if p is False else self.json_path_part(p)
+        for p in [e.args.get("start"), e.args.get("end"), e.args.get("step")]
+        if p is not None
+    ),
+    exp.JSONPathSubscript: lambda self, e: self._jsonpathsubscript_sql(e),
+    exp.JSONPathUnion: lambda self,
+    e: f"[{','.join(self.json_path_part(p) for p in e.expressions)}]",
+    exp.JSONPathWildcard: lambda *_: "*",
+}
+
+ALL_JSON_PATH_PARTS = set(JSON_PATH_PART_TRANSFORMS)

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/lineage.py

@@ -0,0 +1,423 @@
+from __future__ import annotations
+
+import json
+import logging
+import typing as t
+from dataclasses import dataclass, field
+
+from sqlglot import Schema, exp, maybe_parse
+from sqlglot.errors import SqlglotError
+from sqlglot.optimizer import Scope, build_scope, find_all_in_scope, normalize_identifiers, qualify
+from sqlglot.optimizer.scope import ScopeType
+
+if t.TYPE_CHECKING:
+    from sqlglot.dialects.dialect import DialectType
+
+logger = logging.getLogger("sqlglot")
+
+
+@dataclass(frozen=True)
+class Node:
+    name: str
+    expression: exp.Expression
+    source: exp.Expression
+    downstream: t.List[Node] = field(default_factory=list)
+    source_name: str = ""
+    reference_node_name: str = ""
+
+    def walk(self) -> t.Iterator[Node]:
+        yield self
+
+        for d in self.downstream:
+            yield from d.walk()
+
+    def to_html(self, dialect: DialectType = None, **opts) -> GraphHTML:
+        nodes = {}
+        edges = []
+
+        for node in self.walk():
+            if isinstance(node.expression, exp.Table):
+                label = f"FROM {node.expression.this}"
+                title = f"<pre>SELECT {node.name} FROM {node.expression.this}</pre>"
+                group = 1
+            else:
+                label = node.expression.sql(pretty=True, dialect=dialect)
+                source = node.source.transform(
+                    lambda n: (
+                        exp.Tag(this=n, prefix="<b>", postfix="</b>") if n is node.expression else n
+                    ),
+                    copy=False,
+                ).sql(pretty=True, dialect=dialect)
+                title = f"<pre>{source}</pre>"
+                group = 0
+
+            node_id = id(node)
+
+            nodes[node_id] = {
+                "id": node_id,
+                "label": label,
+                "title": title,
+                "group": group,
+            }
+
+            for d in node.downstream:
+                edges.append({"from": node_id, "to": id(d)})
+        return GraphHTML(nodes, edges, **opts)
+
+
+def lineage(
+    column: str | exp.Column,
+    sql: str | exp.Expression,
+    schema: t.Optional[t.Dict | Schema] = None,
+    sources: t.Optional[t.Mapping[str, str | exp.Query]] = None,
+    dialect: DialectType = None,
+    scope: t.Optional[Scope] = None,
+    trim_selects: bool = True,
+    **kwargs,
+) -> Node:
+    """Build the lineage graph for a column of a SQL query.
+
+    Args:
+        column: The column to build the lineage for.
+        sql: The SQL string or expression.
+        schema: The schema of tables.
+        sources: A mapping of queries which will be used to continue building lineage.
+        dialect: The dialect of input SQL.
+        scope: A pre-created scope to use instead.
+        trim_selects: Whether or not to clean up selects by trimming to only relevant columns.
+        **kwargs: Qualification optimizer kwargs.
+
+    Returns:
+        A lineage node.
+    """
+
+    expression = maybe_parse(sql, dialect=dialect)
+    column = normalize_identifiers.normalize_identifiers(column, dialect=dialect).name
+
+    if sources:
+        expression = exp.expand(
+            expression,
+            {k: t.cast(exp.Query, maybe_parse(v, dialect=dialect)) for k, v in sources.items()},
+            dialect=dialect,
+        )
+
+    if not scope:
+        expression = qualify.qualify(
+            expression,
+            dialect=dialect,
+            schema=schema,
+            **{"validate_qualify_columns": False, "identify": False, **kwargs},  # type: ignore
+        )
+
+        scope = build_scope(expression)
+
+    if not scope:
+        raise SqlglotError("Cannot build lineage, sql must be SELECT")
+
+    if not any(select.alias_or_name == column for select in scope.expression.selects):
+        raise SqlglotError(f"Cannot find column '{column}' in query.")
+
+    return to_node(column, scope, dialect, trim_selects=trim_selects)
+
+
+def to_node(
+    column: str | int,
+    scope: Scope,
+    dialect: DialectType,
+    scope_name: t.Optional[str] = None,
+    upstream: t.Optional[Node] = None,
+    source_name: t.Optional[str] = None,
+    reference_node_name: t.Optional[str] = None,
+    trim_selects: bool = True,
+) -> Node:
+    # Find the specific select clause that is the source of the column we want.
+    # This can either be a specific, named select or a generic `*` clause.
+    select = (
+        scope.expression.selects[column]
+        if isinstance(column, int)
+        else next(
+            (select for select in scope.expression.selects if select.alias_or_name == column),
+            exp.Star() if scope.expression.is_star else scope.expression,
+        )
+    )
+
+    if isinstance(scope.expression, exp.Subquery):
+        for source in scope.subquery_scopes:
+            return to_node(
+                column,
+                scope=source,
+                dialect=dialect,
+                upstream=upstream,
+                source_name=source_name,
+                reference_node_name=reference_node_name,
+                trim_selects=trim_selects,
+            )
+    if isinstance(scope.expression, exp.SetOperation):
+        name = type(scope.expression).__name__.upper()
+        upstream = upstream or Node(name=name, source=scope.expression, expression=select)
+
+        index = (
+            column
+            if isinstance(column, int)
+            else next(
+                (
+                    i
+                    for i, select in enumerate(scope.expression.selects)
+                    if select.alias_or_name == column or select.is_star
+                ),
+                -1,  # mypy will not allow a None here, but a negative index should never be returned
+            )
+        )
+
+        if index == -1:
+            raise ValueError(f"Could not find {column} in {scope.expression}")
+
+        for s in scope.union_scopes:
+            to_node(
+                index,
+                scope=s,
+                dialect=dialect,
+                upstream=upstream,
+                source_name=source_name,
+                reference_node_name=reference_node_name,
+                trim_selects=trim_selects,
+            )
+
+        return upstream
+
+    if trim_selects and isinstance(scope.expression, exp.Select):
+        # For better ergonomics in our node labels, replace the full select with
+        # a version that has only the column we care about.
+        #   "x", SELECT x, y FROM foo
+        #     => "x", SELECT x FROM foo
+        source = t.cast(exp.Expression, scope.expression.select(select, append=False))
+    else:
+        source = scope.expression
+
+    # Create the node for this step in the lineage chain, and attach it to the previous one.
+    node = Node(
+        name=f"{scope_name}.{column}" if scope_name else str(column),
+        source=source,
+        expression=select,
+        source_name=source_name or "",
+        reference_node_name=reference_node_name or "",
+    )
+
+    if upstream:
+        upstream.downstream.append(node)
+
+    subquery_scopes = {
+        id(subquery_scope.expression): subquery_scope for subquery_scope in scope.subquery_scopes
+    }
+
+    for subquery in find_all_in_scope(select, exp.UNWRAPPED_QUERIES):
+        subquery_scope = subquery_scopes.get(id(subquery))
+        if not subquery_scope:
+            logger.warning(f"Unknown subquery scope: {subquery.sql(dialect=dialect)}")
+            continue
+
+        for name in subquery.named_selects:
+            to_node(
+                name,
+                scope=subquery_scope,
+                dialect=dialect,
+                upstream=node,
+                trim_selects=trim_selects,
+            )
+
+    # if the select is a star add all scope sources as downstreams
+    if select.is_star:
+        for source in scope.sources.values():
+            if isinstance(source, Scope):
+                source = source.expression
+            node.downstream.append(
+                Node(name=select.sql(comments=False), source=source, expression=source)
+            )
+
+    # Find all columns that went into creating this one to list their lineage nodes.
+    source_columns = set(find_all_in_scope(select, exp.Column))
+
+    # If the source is a UDTF find columns used in the UDTF to generate the table
+    if isinstance(source, exp.UDTF):
+        source_columns |= set(source.find_all(exp.Column))
+        derived_tables = [
+            source.expression.parent
+            for source in scope.sources.values()
+            if isinstance(source, Scope) and source.is_derived_table
+        ]
+    else:
+        derived_tables = scope.derived_tables
+
+    source_names = {
+        dt.alias: dt.comments[0].split()[1]
+        for dt in derived_tables
+        if dt.comments and dt.comments[0].startswith("source: ")
+    }
+
+    pivots = scope.pivots
+    pivot = pivots[0] if len(pivots) == 1 and not pivots[0].unpivot else None
+    if pivot:
+        # 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 = pivot.args["columns"]
+        pivot_aggs_count = len(pivot.expressions)
+
+        pivot_column_mapping = {}
+        for i, agg in enumerate(pivot.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] = agg_cols
+
+    for c in source_columns:
+        table = c.table
+        source = scope.sources.get(table)
+
+        if isinstance(source, Scope):
+            reference_node_name = None
+            if source.scope_type == ScopeType.DERIVED_TABLE and table not in source_names:
+                reference_node_name = table
+            elif source.scope_type == ScopeType.CTE:
+                selected_node, _ = scope.selected_sources.get(table, (None, None))
+                reference_node_name = selected_node.name if selected_node else None
+
+            # The table itself came from a more specific scope. Recurse into that one using the unaliased column name.
+            to_node(
+                c.name,
+                scope=source,
+                dialect=dialect,
+                scope_name=table,
+                upstream=node,
+                source_name=source_names.get(table) or source_name,
+                reference_node_name=reference_node_name,
+                trim_selects=trim_selects,
+            )
+        elif pivot and pivot.alias_or_name == c.table:
+            downstream_columns = []
+
+            column_name = c.name
+            if any(column_name == pivot_column.name for pivot_column in pivot_columns):
+                downstream_columns.extend(pivot_column_mapping[column_name])
+            else:
+                # The column is not in the pivot, so it must be an implicit column of the
+                # pivoted source -- adapt column to be from the implicit pivoted source.
+                downstream_columns.append(exp.column(c.this, table=pivot.parent.this))
+
+            for downstream_column in downstream_columns:
+                table = downstream_column.table
+                source = scope.sources.get(table)
+                if isinstance(source, Scope):
+                    to_node(
+                        downstream_column.name,
+                        scope=source,
+                        scope_name=table,
+                        dialect=dialect,
+                        upstream=node,
+                        source_name=source_names.get(table) or source_name,
+                        reference_node_name=reference_node_name,
+                        trim_selects=trim_selects,
+                    )
+                else:
+                    source = source or exp.Placeholder()
+                    node.downstream.append(
+                        Node(
+                            name=downstream_column.sql(comments=False),
+                            source=source,
+                            expression=source,
+                        )
+                    )
+        else:
+            # The source is not a scope and the column is not in any pivot - we've reached the end
+            # of the line. At this point, if a source is not found it means this column's lineage
+            # is unknown. This can happen if the definition of a source used in a query is not
+            # passed into the `sources` map.
+            source = source or exp.Placeholder()
+            node.downstream.append(
+                Node(name=c.sql(comments=False), source=source, expression=source)
+            )
+
+    return node
+
+
+class GraphHTML:
+    """Node to HTML generator using vis.js.
+
+    https://visjs.github.io/vis-network/docs/network/
+    """
+
+    def __init__(
+        self, nodes: t.Dict, edges: t.List, imports: bool = True, options: t.Optional[t.Dict] = None
+    ):
+        self.imports = imports
+
+        self.options = {
+            "height": "500px",
+            "width": "100%",
+            "layout": {
+                "hierarchical": {
+                    "enabled": True,
+                    "nodeSpacing": 200,
+                    "sortMethod": "directed",
+                },
+            },
+            "interaction": {
+                "dragNodes": False,
+                "selectable": False,
+            },
+            "physics": {
+                "enabled": False,
+            },
+            "edges": {
+                "arrows": "to",
+            },
+            "nodes": {
+                "font": "20px monaco",
+                "shape": "box",
+                "widthConstraint": {
+                    "maximum": 300,
+                },
+            },
+            **(options or {}),
+        }
+
+        self.nodes = nodes
+        self.edges = edges
+
+    def __str__(self):
+        nodes = json.dumps(list(self.nodes.values()))
+        edges = json.dumps(self.edges)
+        options = json.dumps(self.options)
+        imports = (
+            """<script type="text/javascript" src="https://unpkg.com/vis-data@latest/peer/umd/vis-data.min.js"></script>
+  <script type="text/javascript" src="https://unpkg.com/vis-network@latest/peer/umd/vis-network.min.js"></script>
+  <link rel="stylesheet" type="text/css" href="https://unpkg.com/vis-network/styles/vis-network.min.css" />"""
+            if self.imports
+            else ""
+        )
+
+        return f"""<div>
+  <div id="sqlglot-lineage"></div>
+  {imports}
+  <script type="text/javascript">
+    var nodes = new vis.DataSet({nodes})
+    nodes.forEach(row => row["title"] = new DOMParser().parseFromString(row["title"], "text/html").body.childNodes[0])
+
+    new vis.Network(
+        document.getElementById("sqlglot-lineage"),
+        {{
+            nodes: nodes,
+            edges: new vis.DataSet({edges})
+        }},
+        {options},
+    )
+  </script>
+</div>"""
+
+    def _repr_html_(self) -> str:
+        return self.__str__()

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/optimizer/__init__.py

@@ -0,0 +1,11 @@
+# ruff: noqa: F401
+
+from sqlglot.optimizer.optimizer import RULES as RULES, optimize as optimize
+from sqlglot.optimizer.scope import (
+    Scope as Scope,
+    build_scope as build_scope,
+    find_all_in_scope as find_all_in_scope,
+    find_in_scope as find_in_scope,
+    traverse_scope as traverse_scope,
+    walk_in_scope as walk_in_scope,
+)