metricflow-sql 0.1.1__py3-none-any.whl

metricflow_sql/compiler/rules/expand_metadata.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import logging
+
+from sqlglot import exp
+from sqlglot.optimizer.scope import traverse_scope
+
+from metricflow_sql.compiler.catalog.cluster import Cluster
+from metricflow_sql.compiler.catalog.metadata import MetadataTable
+from metricflow_sql.utils.sql import FINAL
+
+logger = logging.getLogger(__name__)
+
+# The key in a planned query's `meta` naming the engine that runs it, when it reads metadata.
+ENGINE = "engine"
+
+# A table's name and alias; anything else, e.g. TABLESAMPLE, can't carry over to its source.
+_NAME_ARGS = ("this", "db", "catalog", "alias")
+
+
+class ExpandMetadataRule:
+    """Replace each metadata table, e.g. `information_schema.tables`, with the query defining it.
+
+    A query that reads metadata reads only metadata, all from one engine: any other table,
+    including a table function, would run on that engine and reach past what the metadata
+    tables expose.
+    """
+
+    def __call__(self, expression: exp.Expr, cluster: Cluster) -> exp.Expr:
+        metadata: list[tuple[exp.Table, MetadataTable]] = []
+        reads_other_tables = False
+        for scope in traverse_scope(expression):
+            for source in scope.sources.values():
+                match source:
+                    # A warehouse table
+                    case exp.Table() if source.meta.get(FINAL):
+                        reads_other_tables = True
+                    case exp.Table():
+                        table = _metadata_table(cluster, source)
+                        if table is None:
+                            reads_other_tables = True
+                        else:
+                            metadata.append((source, table))
+
+        if not metadata:
+            return expression
+        if reads_other_tables:
+            raise ValueError("Queries mixing metadata and data tables are not supported")
+        _check(expression)
+
+        engines = {table.engine for _, table in metadata}
+        if len(engines) > 1:
+            raise ValueError("A query can only read metadata from one engine")
+
+        for source, table in metadata:
+            if any(value for key, value in source.args.items() if key not in _NAME_ARGS):
+                raise ValueError("Unsupported metadata table modifier")
+            source.replace(table.query().subquery(source.alias_or_name))
+
+        expression.meta[ENGINE] = engines.pop()
+        return expression
+
+
+def _check(expression: exp.Expr) -> None:
+    """Only tables (each a metadata table, as checked above), subqueries, UNNEST and VALUES may
+    supply rows, e.g. not a `LATERAL` table function, which could read the engine itself.
+
+    Only functions sqlglot knows may be called: it passes any other through as is, e.g.
+    `query_to_xml('SELECT * FROM objects', ...)`, which runs a query of its own on the engine."""
+    for node in expression.find_all(exp.From, exp.Join, exp.Lateral, exp.Anonymous):
+        match node, node.this:
+            case exp.Anonymous(), _:
+                raise ValueError(f"Unsupported function in a metadata query: {node.name}")
+            case _, exp.Subquery():
+                pass
+            # A LATERAL may only wrap a subquery.
+            case exp.From() | exp.Join(), exp.Table() | exp.Unnest() | exp.Values() | exp.Lateral():
+                pass
+            case _:
+                raise ValueError(f"Unsupported metadata source: {node.this.sql()}")
+
+
+def _metadata_table(cluster: Cluster, table: exp.Table) -> MetadataTable | None:
+    match cluster.find_table(table) if cluster.get_catalog(table.catalog) else None:
+        case MetadataTable() as metadata:
+            return metadata
+        case _:
+            return None

metricflow_sql/compiler/rules/isolate_fragments.py

@@ -0,0 +1,244 @@
+from __future__ import annotations
+
+from sqlglot import exp
+from sqlglot.helper import find_new_name, name_sequence
+from sqlglot.optimizer.scope import (
+    Scope,
+    find_all_in_scope,
+    find_in_scope,
+    traverse_scope,
+    walk_in_scope,
+)
+
+# Clauses that are not part of the fragment. They move to the outer select.
+_NON_FRAGMENT_ARGS = ("having", "qualify", "distinct", "order", "limit", "offset")
+
+# Nodes of a GROUP BY that wrap the grouped expressions instead of being one.
+_GROUPING_WRAPPERS = (exp.Cube, exp.GroupingSets, exp.Paren, exp.Rollup, exp.Tuple)
+
+
+def isolate_fragments(expression: exp.Expr) -> exp.Expr:
+    """
+    Isolate the fragment of every select in a derived table. The fragment is the
+    projections, joins, WHERE and GROUP BY of the select. Its HAVING, QUALIFY,
+    DISTINCT, ORDER BY, LIMIT and OFFSET move to an outer select:
+
+        SELECT day, AGG(a) FROM t GROUP BY 1 HAVING AGG(b) > 0 ORDER BY day
+        =>
+        SELECT day, a
+        FROM (SELECT day, AGG(a) AS a, AGG(b) AS b FROM t GROUP BY 1) AS t
+        WHERE b > 0
+        ORDER BY day
+
+    The derived table takes the alias of the FROM, so that the outer select reads
+    like the original. The DISTINCT of a select that does not aggregate becomes
+    its GROUP BY instead.
+    """
+    next_alias = name_sequence("_f_")
+
+    # The scopes are built once, before any select changes. An outer select finds
+    # the references of an isolated subquery through the columns its scope read before the
+    # change, so clearing the cache of a scope in this loop loses them. A fresh build_scope
+    # per select would avoid that, but each select would rebuild the scopes below it.
+    for scope in traverse_scope(expression):
+        select = scope.expression
+        if isinstance(select, exp.Select) and not select.is_star:
+            _group_by_distinct(select)
+            if any(select.args.get(arg) for arg in _NON_FRAGMENT_ARGS):
+                fragment_alias = _from_alias(select) or exp.to_identifier(next_alias())
+                outer = _isolate(select, scope, fragment_alias)
+                if select is expression:
+                    expression = outer
+
+    return expression
+
+
+def _from_alias(select: exp.Select) -> exp.Identifier | None:
+    from_ = select.args.get("from_")
+    table_alias = from_.this.args.get("alias") if from_ else None
+    return table_alias.this.copy() if table_alias else None
+
+
+def _group_by_distinct(select: exp.Select) -> None:
+    """
+    Replaces the DISTINCT of a select with a GROUP BY over its projections, which
+    returns the same rows. A select that aggregates or has a window function
+    keeps its DISTINCT, since it applies to the rows these produce.
+    """
+    distinct = select.args.get("distinct")
+    if (
+        not distinct
+        or distinct.args.get("on")
+        or select.args.get("group")
+        or select.args.get("having")
+        or any(
+            find_in_scope(projection, exp.AggFunc) or find_in_scope(projection, exp.Window)
+            for projection in select.expressions
+        )
+    ):
+        return
+
+    select.set("distinct", None)
+    select.group_by(*(projection.this.copy() for projection in select.expressions), copy=False)
+
+
+def _isolate(select: exp.Select, scope: Scope, fragment_alias: exp.Identifier) -> exp.Select:
+    taken: set[str] = set()
+    inner_columns: dict[exp.Expr, exp.Identifier] = {}
+    outer_projections = []
+
+    # Deduplicate projection names, so that they can be unambiguously referenced from outer scopes
+    for projection in select.expressions:
+        name = projection.args["alias"]
+        inner = _reserve_name(name, taken)
+        projection.set("alias", inner)
+        inner_columns.setdefault(projection.this, inner)
+        outer_projections.append(exp.alias_(exp.column(inner, fragment_alias), name, copy=False))
+    outer = exp.select(*outer_projections, copy=False)
+
+    clauses = {arg: clause for arg in _NON_FRAGMENT_ARGS if (clause := select.args.get(arg))}
+
+    _rewrite_for_outer(clauses, select, scope, fragment_alias, inner_columns, taken)
+
+    for arg, clause in clauses.items():
+        select.set(arg, None)
+        if arg == "having":
+            outer.where(clause.this, copy=False)
+        else:
+            outer.set(arg, clause)
+
+    with_ = select.args.get("with_")
+    if with_:
+        outer.set("with_", with_.pop())
+
+    select.replace(outer)
+
+    outer.from_(select.subquery(fragment_alias, copy=False), copy=False)
+    return outer
+
+
+def _rewrite_for_outer(
+    clauses: dict[str, exp.Expr],
+    select: exp.Select,
+    scope: Scope,
+    fragment_alias: exp.Identifier,
+    inner_columns: dict[exp.Expr, exp.Identifier],
+    taken: set[str],
+) -> None:
+    """
+    Rewrites the clauses in place to read what the select computes through its
+    derived table, projecting from the select whatever it does not project yet.
+    """
+    local_columns = {id(column) for column in scope.local_columns}
+    scope_columns = {id(column) for column in scope.columns}
+
+    grouped = _grouping_expressions(group) if (group := select.args.get("group")) else set()
+
+    # Subqueries that read a column of the select, at any depth. Nothing inside any
+    # other subquery can move, so the walk skips them. A climb stops at a query
+    # already added, since the queries above it are added too.
+    correlated_queries: set[int] = set()
+    for subquery_scope in scope.subquery_scopes:
+        for column in subquery_scope.external_columns:
+            if id(column) in local_columns:
+                query = column.find_ancestor(exp.Query)
+                while query and query is not select and id(query) not in correlated_queries:
+                    correlated_queries.add(id(query))
+                    query = query.find_ancestor(exp.Query)
+
+    def project(node: exp.Expr) -> None:
+        """
+        Replaces the node with the inner column that computes it, adding that
+        projection to the select unless it already has one.
+        """
+        identifier = inner_columns.get(node)
+        if identifier:
+            node.replace(exp.column(identifier, fragment_alias))
+            return
+
+        # A column, a path and a measure keep their identifier, anything nameless
+        # is numbered like qualify does.
+        identifier = _reserve_name(
+            _output_identifier(node)
+            or exp.to_identifier(node.output_name or f"_col_{len(select.expressions)}"),
+            taken,
+        )
+        inner_columns[node] = identifier
+        node.replace(exp.column(identifier, fragment_alias))
+        select.select(exp.alias_(node, identifier, copy=False), copy=False)
+
+    def move(node: exp.Expr, nested: bool) -> None:
+        root = node
+        while isinstance(root, exp.Dot):
+            root = root.this
+
+        # A local aggregate is evaluated by the select, wherever it is written.
+        # E.g. SUM(x.b) in: FROM x HAVING EXISTS (SELECT 1 FROM y WHERE y.b < SUM(x.b))
+        if _is_local_aggregate(node, nested, local_columns, scope_columns):
+            project(node)
+        elif id(root) in local_columns or (not nested and node in grouped):
+            project(node)
+        elif not isinstance(node, exp.Query) or id(node) in correlated_queries:
+            for child in node.iter_expressions():
+                move(child, nested or isinstance(node, exp.Query))
+
+    for clause in clauses.values():
+        move(clause, nested=False)
+
+
+def _grouping_expressions(group: exp.Group) -> set[exp.Expr]:
+    return {
+        node
+        for expression in group.iter_expressions()
+        for node in walk_in_scope(
+            expression, prune=lambda node: not isinstance(node, _GROUPING_WRAPPERS)
+        )
+        if not isinstance(node, _GROUPING_WRAPPERS)
+    }
+
+
+def _reserve_name(identifier: exp.Identifier, taken: set[str]) -> exp.Identifier:
+    name = find_new_name(taken, identifier.name)
+    taken.add(name)
+    if name == identifier.name:
+        return identifier
+
+    return exp.to_identifier(name, quoted=identifier.quoted)
+
+
+def _output_identifier(node: exp.Expr) -> exp.Identifier | None:
+    """
+    The identifier that names a column, the last field of a path or the measure
+    of an AGG, so that the projection keeps its quotes.
+    """
+    if isinstance(node, exp.Agg):
+        node = node.this
+
+    if isinstance(node, exp.Dot):
+        node = node.expression
+    elif isinstance(node, exp.Column):
+        node = node.this
+
+    return node if isinstance(node, exp.Identifier) else None
+
+
+def _is_local_aggregate(
+    node: exp.Expr, nested: bool, local_columns: set[int], scope_columns: set[int]
+) -> bool:
+    """
+    Whether the isolated select evaluates the aggregate. Only columns written
+    directly in the aggregate count.
+    """
+    # A FILTER moves with its aggregate, e.g. SUM(x) FILTER (WHERE y > 1), so the aggregate
+    # inside it isn't one on its own.
+    aggregate = node.this if isinstance(node, exp.Filter) else node
+    if not isinstance(aggregate, exp.AggFunc) or isinstance(node.parent, exp.Window | exp.Filter):
+        return False
+
+    columns = {id(column) for column in find_all_in_scope(node, exp.Column)}
+    if not columns:
+        # An aggregate without columns, like COUNT(*), belongs to the select it is written in
+        return not nested
+
+    # It reads at least one of the select's own columns and no column of a subquery
+    return not columns.isdisjoint(local_columns) and columns <= scope_columns

metricflow_sql/compiler/rules/passthrough_tables.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import logging
+
+from sqlglot import exp
+from sqlglot.optimizer.scope import traverse_scope
+
+from metricflow_sql.compiler.catalog.cluster import Cluster
+from metricflow_sql.compiler.catalog.passthrough import PassthroughCatalog
+from metricflow_sql.utils.sql import FINAL
+
+logger = logging.getLogger(__name__)
+
+
+def name_passthrough_tables(expression: exp.Expr, cluster: Cluster) -> exp.Expr:
+    """
+    Rename each passthrough table to the warehouse's name for it, e.g. `analytics.sales.orders`
+    to `"ANALYTICS"."SALES"."ORDERS"` on Snowflake, and mark it final.
+
+    Runs before the other planner rules, which skip final tables.
+    """
+    for scope in traverse_scope(expression):
+        for source in scope.sources.values():
+            match source:
+                case exp.Table() if isinstance(
+                    catalog := cluster.get_catalog(source.catalog), PassthroughCatalog
+                ):
+                    # Already looked up while qualifying
+                    found = catalog.get_table(source)
+                    assert found is not None, f"Unknown passthrough table: {source.sql()}"
+                    physical = found.physical.table
+                    for part in ("catalog", "db", "this"):
+                        name = physical.args.get(part)
+                        source.set(part, name.copy() if name is not None else None)
+                    source.meta[FINAL] = True
+    return expression

metricflow_sql/compiler/table_arguments.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import logging
+
+from sqlglot import exp
+
+logger = logging.getLogger(__name__)
+
+# Key into `exp.Table.meta` holding a table function's arguments, by parameter name.
+TABLE_ARGUMENTS = "arguments"
+
+
+def bind_table_arguments(table: exp.Table) -> None:
+    """
+    Turns a table function into a plain table that carries its arguments.
+
+    A dataset's parameters are arguments of the table it is read as, e.g.
+    `FROM bulk_store_sales(min_quantity => 90)` becomes `FROM bulk_store_sales`, with
+    `{"min_quantity": 90}` in the table's meta for dataset expansion to render with. Must
+    run before qualification, so it sees ordinary tables.
+    """
+    match table.this:
+        case exp.Anonymous(this=name, expressions=arguments):
+            bound: dict[str, exp.Expr] = {}
+            for argument in arguments:
+                match argument:
+                    case exp.Kwarg(this=parameter, expression=value):
+                        bound[parameter.name] = value
+                    case _:
+                        raise ValueError(
+                            f"Table function arguments must be named, e.g. "
+                            f"`{name}(parameter => value)`: {argument.sql()}"
+                        )
+            table.set("this", exp.to_identifier(name))
+            table.meta[TABLE_ARGUMENTS] = bound