altimate-code 0.5.2 → 0.5.3

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

@@ -0,0 +1,222 @@
+from __future__ import annotations
+
+import itertools
+import typing as t
+
+from sqlglot import exp
+from sqlglot.dialects.dialect import Dialect, DialectType
+from sqlglot.helper import is_date_unit, is_iso_date, is_iso_datetime
+from sqlglot.optimizer.annotate_types import TypeAnnotator
+
+
+def canonicalize(expression: exp.Expression, dialect: DialectType = None) -> exp.Expression:
+    """Converts a sql expression into a standard form.
+
+    This method relies on annotate_types because many of the
+    conversions rely on type inference.
+
+    Args:
+        expression: The expression to canonicalize.
+    """
+
+    dialect = Dialect.get_or_raise(dialect)
+
+    def _canonicalize(expression: exp.Expression) -> exp.Expression:
+        expression = add_text_to_concat(expression)
+        expression = replace_date_funcs(expression, dialect=dialect)
+        expression = coerce_type(expression, dialect.PROMOTE_TO_INFERRED_DATETIME_TYPE)
+        expression = remove_redundant_casts(expression)
+        expression = ensure_bools(expression, _replace_int_predicate)
+        expression = remove_ascending_order(expression)
+        return expression
+
+    return exp.replace_tree(expression, _canonicalize)
+
+
+def add_text_to_concat(node: exp.Expression) -> exp.Expression:
+    if isinstance(node, exp.Add) and node.type and node.type.this in exp.DataType.TEXT_TYPES:
+        node = exp.Concat(expressions=[node.left, node.right])
+    return node
+
+
+def replace_date_funcs(node: exp.Expression, dialect: DialectType) -> exp.Expression:
+    if (
+        isinstance(node, (exp.Date, exp.TsOrDsToDate))
+        and not node.expressions
+        and not node.args.get("zone")
+        and node.this.is_string
+        and is_iso_date(node.this.name)
+    ):
+        return exp.cast(node.this, to=exp.DataType.Type.DATE)
+    if isinstance(node, exp.Timestamp) and not node.args.get("zone"):
+        if not node.type:
+            from sqlglot.optimizer.annotate_types import annotate_types
+
+            node = annotate_types(node, dialect=dialect)
+        return exp.cast(node.this, to=node.type or exp.DataType.Type.TIMESTAMP)
+
+    return node
+
+
+COERCIBLE_DATE_OPS = (
+    exp.Add,
+    exp.Sub,
+    exp.EQ,
+    exp.NEQ,
+    exp.GT,
+    exp.GTE,
+    exp.LT,
+    exp.LTE,
+    exp.NullSafeEQ,
+    exp.NullSafeNEQ,
+)
+
+
+def coerce_type(node: exp.Expression, promote_to_inferred_datetime_type: bool) -> exp.Expression:
+    if isinstance(node, COERCIBLE_DATE_OPS):
+        _coerce_date(node.left, node.right, promote_to_inferred_datetime_type)
+    elif isinstance(node, exp.Between):
+        _coerce_date(node.this, node.args["low"], promote_to_inferred_datetime_type)
+    elif isinstance(node, exp.Extract) and not node.expression.type.is_type(
+        *exp.DataType.TEMPORAL_TYPES
+    ):
+        _replace_cast(node.expression, exp.DataType.Type.DATETIME)
+    elif isinstance(node, (exp.DateAdd, exp.DateSub, exp.DateTrunc)):
+        _coerce_timeunit_arg(node.this, node.unit)
+    elif isinstance(node, exp.DateDiff):
+        _coerce_datediff_args(node)
+
+    return node
+
+
+def remove_redundant_casts(expression: exp.Expression) -> exp.Expression:
+    if (
+        isinstance(expression, exp.Cast)
+        and expression.this.type
+        and expression.to == expression.this.type
+    ):
+        return expression.this
+
+    if (
+        isinstance(expression, (exp.Date, exp.TsOrDsToDate))
+        and expression.this.type
+        and expression.this.type.this == exp.DataType.Type.DATE
+        and not expression.this.type.expressions
+    ):
+        return expression.this
+
+    return expression
+
+
+def ensure_bools(
+    expression: exp.Expression, replace_func: t.Callable[[exp.Expression], None]
+) -> exp.Expression:
+    if isinstance(expression, exp.Connector):
+        replace_func(expression.left)
+        replace_func(expression.right)
+    elif isinstance(expression, exp.Not):
+        replace_func(expression.this)
+        # We can't replace num in CASE x WHEN num ..., because it's not the full predicate
+    elif isinstance(expression, exp.If) and not (
+        isinstance(expression.parent, exp.Case) and expression.parent.this
+    ):
+        replace_func(expression.this)
+    elif isinstance(expression, (exp.Where, exp.Having)):
+        replace_func(expression.this)
+
+    return expression
+
+
+def remove_ascending_order(expression: exp.Expression) -> exp.Expression:
+    if isinstance(expression, exp.Ordered) and expression.args.get("desc") is False:
+        # Convert ORDER BY a ASC to ORDER BY a
+        expression.set("desc", None)
+
+    return expression
+
+
+def _coerce_date(
+    a: exp.Expression,
+    b: exp.Expression,
+    promote_to_inferred_datetime_type: bool,
+) -> None:
+    for a, b in itertools.permutations([a, b]):
+        if isinstance(b, exp.Interval):
+            a = _coerce_timeunit_arg(a, b.unit)
+
+        a_type = a.type
+        if (
+            not a_type
+            or a_type.this not in exp.DataType.TEMPORAL_TYPES
+            or not b.type
+            or b.type.this not in exp.DataType.TEXT_TYPES
+        ):
+            continue
+
+        if promote_to_inferred_datetime_type:
+            if b.is_string:
+                date_text = b.name
+                if is_iso_date(date_text):
+                    b_type = exp.DataType.Type.DATE
+                elif is_iso_datetime(date_text):
+                    b_type = exp.DataType.Type.DATETIME
+                else:
+                    b_type = a_type.this
+            else:
+                # If b is not a datetime string, we conservatively promote it to a DATETIME,
+                # in order to ensure there are no surprising truncations due to downcasting
+                b_type = exp.DataType.Type.DATETIME
+
+            target_type = (
+                b_type if b_type in TypeAnnotator.COERCES_TO.get(a_type.this, {}) else a_type
+            )
+        else:
+            target_type = a_type
+
+        if target_type != a_type:
+            _replace_cast(a, target_type)
+
+        _replace_cast(b, target_type)
+
+
+def _coerce_timeunit_arg(arg: exp.Expression, unit: t.Optional[exp.Expression]) -> exp.Expression:
+    if not arg.type:
+        return arg
+
+    if arg.type.this in exp.DataType.TEXT_TYPES:
+        date_text = arg.name
+        is_iso_date_ = is_iso_date(date_text)
+
+        if is_iso_date_ and is_date_unit(unit):
+            return arg.replace(exp.cast(arg.copy(), to=exp.DataType.Type.DATE))
+
+        # An ISO date is also an ISO datetime, but not vice versa
+        if is_iso_date_ or is_iso_datetime(date_text):
+            return arg.replace(exp.cast(arg.copy(), to=exp.DataType.Type.DATETIME))
+
+    elif arg.type.this == exp.DataType.Type.DATE and not is_date_unit(unit):
+        return arg.replace(exp.cast(arg.copy(), to=exp.DataType.Type.DATETIME))
+
+    return arg
+
+
+def _coerce_datediff_args(node: exp.DateDiff) -> None:
+    for e in (node.this, node.expression):
+        if e.type.this not in exp.DataType.TEMPORAL_TYPES:
+            e.replace(exp.cast(e.copy(), to=exp.DataType.Type.DATETIME))
+
+
+def _replace_cast(node: exp.Expression, to: exp.DATA_TYPE) -> None:
+    node.replace(exp.cast(node.copy(), to=to))
+
+
+# this was originally designed for presto, there is a similar transform for tsql
+# this is different in that it only operates on int types, this is because
+# presto has a boolean type whereas tsql doesn't (people use bits)
+# with y as (select true as x) select x = 0 FROM y -- illegal presto query
+def _replace_int_predicate(expression: exp.Expression) -> None:
+    if isinstance(expression, exp.Coalesce):
+        for child in expression.iter_expressions():
+            _replace_int_predicate(child)
+    elif expression.type and expression.type.this in exp.DataType.INTEGER_TYPES:
+        expression.replace(expression.neq(0))

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

@@ -0,0 +1,43 @@
+from sqlglot.optimizer.scope import Scope, build_scope
+
+
+def eliminate_ctes(expression):
+    """
+    Remove unused CTEs from an expression.
+
+    Example:
+        >>> import sqlglot
+        >>> sql = "WITH y AS (SELECT a FROM x) SELECT a FROM z"
+        >>> expression = sqlglot.parse_one(sql)
+        >>> eliminate_ctes(expression).sql()
+        'SELECT a FROM z'
+
+    Args:
+        expression (sqlglot.Expression): expression to optimize
+    Returns:
+        sqlglot.Expression: optimized expression
+    """
+    root = build_scope(expression)
+
+    if root:
+        ref_count = root.ref_count()
+
+        # Traverse the scope tree in reverse so we can remove chains of unused CTEs
+        for scope in reversed(list(root.traverse())):
+            if scope.is_cte:
+                count = ref_count[id(scope)]
+                if count <= 0:
+                    cte_node = scope.expression.parent
+                    with_node = cte_node.parent
+                    cte_node.pop()
+
+                    # Pop the entire WITH clause if this is the last CTE
+                    if with_node and len(with_node.expressions) <= 0:
+                        with_node.pop()
+
+                    # Decrement the ref count for all sources this CTE selects from
+                    for _, source in scope.selected_sources.values():
+                        if isinstance(source, Scope):
+                            ref_count[id(source)] -= 1
+
+    return expression

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

@@ -0,0 +1,181 @@
+from sqlglot import expressions as exp
+from sqlglot.optimizer.normalize import normalized
+from sqlglot.optimizer.scope import Scope, traverse_scope
+
+
+def eliminate_joins(expression):
+    """
+    Remove unused joins from an expression.
+
+    This only removes joins when we know that the join condition doesn't produce duplicate rows.
+
+    Example:
+        >>> import sqlglot
+        >>> sql = "SELECT x.a FROM x LEFT JOIN (SELECT DISTINCT y.b FROM y) AS y ON x.b = y.b"
+        >>> expression = sqlglot.parse_one(sql)
+        >>> eliminate_joins(expression).sql()
+        'SELECT x.a FROM x'
+
+    Args:
+        expression (sqlglot.Expression): expression to optimize
+    Returns:
+        sqlglot.Expression: optimized expression
+    """
+    for scope in traverse_scope(expression):
+        # If any columns in this scope aren't qualified, it's hard to determine if a join isn't used.
+        # It's probably possible to infer this from the outputs of derived tables.
+        # But for now, let's just skip this rule.
+        if scope.unqualified_columns:
+            continue
+
+        joins = scope.expression.args.get("joins", [])
+
+        # Reverse the joins so we can remove chains of unused joins
+        for join in reversed(joins):
+            alias = join.alias_or_name
+            if _should_eliminate_join(scope, join, alias):
+                join.pop()
+                scope.remove_source(alias)
+    return expression
+
+
+def _should_eliminate_join(scope, join, alias):
+    inner_source = scope.sources.get(alias)
+    return (
+        isinstance(inner_source, Scope)
+        and not _join_is_used(scope, join, alias)
+        and (
+            (join.side == "LEFT" and _is_joined_on_all_unique_outputs(inner_source, join))
+            or (not join.args.get("on") and _has_single_output_row(inner_source))
+        )
+    )
+
+
+def _join_is_used(scope, join, alias):
+    # We need to find all columns that reference this join.
+    # But columns in the ON clause shouldn't count.
+    on = join.args.get("on")
+    if on:
+        on_clause_columns = {id(column) for column in on.find_all(exp.Column)}
+    else:
+        on_clause_columns = set()
+    return any(
+        column for column in scope.source_columns(alias) if id(column) not in on_clause_columns
+    )
+
+
+def _is_joined_on_all_unique_outputs(scope, join):
+    unique_outputs = _unique_outputs(scope)
+    if not unique_outputs:
+        return False
+
+    _, join_keys, _ = join_condition(join)
+    remaining_unique_outputs = unique_outputs - {c.name for c in join_keys}
+    return not remaining_unique_outputs
+
+
+def _unique_outputs(scope):
+    """Determine output columns of `scope` that must have a unique combination per row"""
+    if scope.expression.args.get("distinct"):
+        return set(scope.expression.named_selects)
+
+    group = scope.expression.args.get("group")
+    if group:
+        grouped_expressions = set(group.expressions)
+        grouped_outputs = set()
+
+        unique_outputs = set()
+        for select in scope.expression.selects:
+            output = select.unalias()
+            if output in grouped_expressions:
+                grouped_outputs.add(output)
+                unique_outputs.add(select.alias_or_name)
+
+        # All the grouped expressions must be in the output
+        if not grouped_expressions.difference(grouped_outputs):
+            return unique_outputs
+        else:
+            return set()
+
+    if _has_single_output_row(scope):
+        return set(scope.expression.named_selects)
+
+    return set()
+
+
+def _has_single_output_row(scope):
+    return isinstance(scope.expression, exp.Select) and (
+        all(isinstance(e.unalias(), exp.AggFunc) for e in scope.expression.selects)
+        or _is_limit_1(scope)
+        or not scope.expression.args.get("from")
+    )
+
+
+def _is_limit_1(scope):
+    limit = scope.expression.args.get("limit")
+    return limit and limit.expression.this == "1"
+
+
+def join_condition(join):
+    """
+    Extract the join condition from a join expression.
+
+    Args:
+        join (exp.Join)
+    Returns:
+        tuple[list[str], list[str], exp.Expression]:
+            Tuple of (source key, join key, remaining predicate)
+    """
+    name = join.alias_or_name
+    on = (join.args.get("on") or exp.true()).copy()
+    source_key = []
+    join_key = []
+
+    def extract_condition(condition):
+        left, right = condition.unnest_operands()
+        left_tables = exp.column_table_names(left)
+        right_tables = exp.column_table_names(right)
+
+        if name in left_tables and name not in right_tables:
+            join_key.append(left)
+            source_key.append(right)
+            condition.replace(exp.true())
+        elif name in right_tables and name not in left_tables:
+            join_key.append(right)
+            source_key.append(left)
+            condition.replace(exp.true())
+
+    # find the join keys
+    # SELECT
+    # FROM x
+    # JOIN y
+    #   ON x.a = y.b AND y.b > 1
+    #
+    # should pull y.b as the join key and x.a as the source key
+    if normalized(on):
+        on = on if isinstance(on, exp.And) else exp.and_(on, exp.true(), copy=False)
+
+        for condition in on.flatten():
+            if isinstance(condition, exp.EQ):
+                extract_condition(condition)
+    elif normalized(on, dnf=True):
+        conditions = None
+
+        for condition in on.flatten():
+            parts = [part for part in condition.flatten() if isinstance(part, exp.EQ)]
+            if conditions is None:
+                conditions = parts
+            else:
+                temp = []
+                for p in parts:
+                    cs = [c for c in conditions if p == c]
+
+                    if cs:
+                        temp.append(p)
+                        temp.extend(cs)
+                conditions = temp
+
+        for condition in conditions:
+            extract_condition(condition)
+
+    return source_key, join_key, on

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

@@ -0,0 +1,189 @@
+from __future__ import annotations
+
+import itertools
+import typing as t
+
+from sqlglot import expressions as exp
+from sqlglot.helper import find_new_name
+from sqlglot.optimizer.scope import Scope, build_scope
+
+if t.TYPE_CHECKING:
+    ExistingCTEsMapping = t.Dict[exp.Expression, str]
+    TakenNameMapping = t.Dict[str, t.Union[Scope, exp.Expression]]
+
+
+def eliminate_subqueries(expression: exp.Expression) -> exp.Expression:
+    """
+    Rewrite derived tables as CTES, deduplicating if possible.
+
+    Example:
+        >>> import sqlglot
+        >>> expression = sqlglot.parse_one("SELECT a FROM (SELECT * FROM x) AS y")
+        >>> eliminate_subqueries(expression).sql()
+        'WITH y AS (SELECT * FROM x) SELECT a FROM y AS y'
+
+    This also deduplicates common subqueries:
+        >>> expression = sqlglot.parse_one("SELECT a FROM (SELECT * FROM x) AS y CROSS JOIN (SELECT * FROM x) AS z")
+        >>> eliminate_subqueries(expression).sql()
+        'WITH y AS (SELECT * FROM x) SELECT a FROM y AS y CROSS JOIN y AS z'
+
+    Args:
+        expression (sqlglot.Expression): expression
+    Returns:
+        sqlglot.Expression: expression
+    """
+    if isinstance(expression, exp.Subquery):
+        # It's possible to have subqueries at the root, e.g. (SELECT * FROM x) LIMIT 1
+        eliminate_subqueries(expression.this)
+        return expression
+
+    root = build_scope(expression)
+
+    if not root:
+        return expression
+
+    # Map of alias->Scope|Table
+    # These are all aliases that are already used in the expression.
+    # We don't want to create new CTEs that conflict with these names.
+    taken: TakenNameMapping = {}
+
+    # All CTE aliases in the root scope are taken
+    for scope in root.cte_scopes:
+        taken[scope.expression.parent.alias] = scope
+
+    # All table names are taken
+    for scope in root.traverse():
+        taken.update(
+            {
+                source.name: source
+                for _, source in scope.sources.items()
+                if isinstance(source, exp.Table)
+            }
+        )
+
+    # Map of Expression->alias
+    # Existing CTES in the root expression. We'll use this for deduplication.
+    existing_ctes: ExistingCTEsMapping = {}
+
+    with_ = root.expression.args.get("with")
+    recursive = False
+    if with_:
+        recursive = with_.args.get("recursive")
+        for cte in with_.expressions:
+            existing_ctes[cte.this] = cte.alias
+    new_ctes = []
+
+    # We're adding more CTEs, but we want to maintain the DAG order.
+    # Derived tables within an existing CTE need to come before the existing CTE.
+    for cte_scope in root.cte_scopes:
+        # Append all the new CTEs from this existing CTE
+        for scope in cte_scope.traverse():
+            if scope is cte_scope:
+                # Don't try to eliminate this CTE itself
+                continue
+            new_cte = _eliminate(scope, existing_ctes, taken)
+            if new_cte:
+                new_ctes.append(new_cte)
+
+        # Append the existing CTE itself
+        new_ctes.append(cte_scope.expression.parent)
+
+    # Now append the rest
+    for scope in itertools.chain(root.union_scopes, root.subquery_scopes, root.table_scopes):
+        for child_scope in scope.traverse():
+            new_cte = _eliminate(child_scope, existing_ctes, taken)
+            if new_cte:
+                new_ctes.append(new_cte)
+
+    if new_ctes:
+        query = expression.expression if isinstance(expression, exp.DDL) else expression
+        query.set("with", exp.With(expressions=new_ctes, recursive=recursive))
+
+    return expression
+
+
+def _eliminate(
+    scope: Scope, existing_ctes: ExistingCTEsMapping, taken: TakenNameMapping
+) -> t.Optional[exp.Expression]:
+    if scope.is_derived_table:
+        return _eliminate_derived_table(scope, existing_ctes, taken)
+
+    if scope.is_cte:
+        return _eliminate_cte(scope, existing_ctes, taken)
+
+    return None
+
+
+def _eliminate_derived_table(
+    scope: Scope, existing_ctes: ExistingCTEsMapping, taken: TakenNameMapping
+) -> t.Optional[exp.Expression]:
+    # This makes sure that we don't:
+    # - drop the "pivot" arg from a pivoted subquery
+    # - eliminate a lateral correlated subquery
+    if scope.parent.pivots or isinstance(scope.parent.expression, exp.Lateral):
+        return None
+
+    # Get rid of redundant exp.Subquery expressions, i.e. those that are just used as wrappers
+    to_replace = scope.expression.parent.unwrap()
+    name, cte = _new_cte(scope, existing_ctes, taken)
+    table = exp.alias_(exp.table_(name), alias=to_replace.alias or name)
+    table.set("joins", to_replace.args.get("joins"))
+
+    to_replace.replace(table)
+
+    return cte
+
+
+def _eliminate_cte(
+    scope: Scope, existing_ctes: ExistingCTEsMapping, taken: TakenNameMapping
+) -> t.Optional[exp.Expression]:
+    parent = scope.expression.parent
+    name, cte = _new_cte(scope, existing_ctes, taken)
+
+    with_ = parent.parent
+    parent.pop()
+    if not with_.expressions:
+        with_.pop()
+
+    # Rename references to this CTE
+    for child_scope in scope.parent.traverse():
+        for table, source in child_scope.selected_sources.values():
+            if source is scope:
+                new_table = exp.alias_(exp.table_(name), alias=table.alias_or_name, copy=False)
+                table.replace(new_table)
+
+    return cte
+
+
+def _new_cte(
+    scope: Scope, existing_ctes: ExistingCTEsMapping, taken: TakenNameMapping
+) -> t.Tuple[str, t.Optional[exp.Expression]]:
+    """
+    Returns:
+        tuple of (name, cte)
+        where `name` is a new name for this CTE in the root scope and `cte` is a new CTE instance.
+        If this CTE duplicates an existing CTE, `cte` will be None.
+    """
+    duplicate_cte_alias = existing_ctes.get(scope.expression)
+    parent = scope.expression.parent
+    name = parent.alias
+
+    if not name:
+        name = find_new_name(taken=taken, base="cte")
+
+    if duplicate_cte_alias:
+        name = duplicate_cte_alias
+    elif taken.get(name):
+        name = find_new_name(taken=taken, base=name)
+
+    taken[name] = scope
+
+    if not duplicate_cte_alias:
+        existing_ctes[scope.expression] = name
+        cte = exp.CTE(
+            this=scope.expression,
+            alias=exp.TableAlias(this=exp.to_identifier(name)),
+        )
+    else:
+        cte = None
+    return name, cte

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

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot import alias, exp
+from sqlglot.errors import OptimizeError
+from sqlglot.optimizer.scope import traverse_scope
+from sqlglot.schema import ensure_schema
+
+if t.TYPE_CHECKING:
+    from sqlglot._typing import E
+    from sqlglot.schema import Schema
+    from sqlglot.dialects.dialect import DialectType
+
+
+def isolate_table_selects(
+    expression: E,
+    schema: t.Optional[t.Dict | Schema] = None,
+    dialect: DialectType = None,
+) -> E:
+    schema = ensure_schema(schema, dialect=dialect)
+
+    for scope in traverse_scope(expression):
+        if len(scope.selected_sources) == 1:
+            continue
+
+        for _, source in scope.selected_sources.values():
+            assert source.parent
+
+            if (
+                not isinstance(source, exp.Table)
+                or not schema.column_names(source)
+                or isinstance(source.parent, exp.Subquery)
+                or isinstance(source.parent.parent, exp.Table)
+            ):
+                continue
+
+            if not source.alias:
+                raise OptimizeError("Tables require an alias. Run qualify_tables optimization.")
+
+            source.replace(
+                exp.select("*")
+                .from_(
+                    alias(source, source.alias_or_name, table=True),
+                    copy=False,
+                )
+                .subquery(source.alias, copy=False)
+            )
+
+    return expression