metricflow-sql 0.1.1__py3-none-any.whl

metricflow_sql/compiler/compiler.py

@@ -0,0 +1,297 @@
+from __future__ import annotations
+
+import inspect
+import logging
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+
+from sqlglot import Schema, parse
+from sqlglot import expressions as exp
+from sqlglot.dialects.dialect import DialectType, NormalizationStrategy
+from sqlglot.helper import find_new_name
+from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+from sqlglot.optimizer.optimizer import OptimizerFn
+from sqlglot.optimizer.qualify import qualify
+from sqlglot.optimizer.qualify_columns import (
+    qualify_columns,
+    quote_identifiers,
+    validate_qualify_columns,
+)
+from sqlglot.optimizer.scope import traverse_scope
+
+from metricflow_sql.compiler.catalog.cluster import Cluster
+from metricflow_sql.compiler.catalog.metadata import INFORMATION_SCHEMA
+from metricflow_sql.compiler.catalog.semantic import SemanticCatalog
+from metricflow_sql.compiler.rules.expand_constellations import ExpandConstellationsRule
+from metricflow_sql.compiler.rules.expand_datasets import ExpandDatasetsRule
+from metricflow_sql.compiler.rules.expand_metadata import ENGINE, ExpandMetadataRule
+from metricflow_sql.compiler.rules.isolate_fragments import isolate_fragments
+from metricflow_sql.compiler.rules.passthrough_tables import name_passthrough_tables
+from metricflow_sql.compiler.table_arguments import bind_table_arguments
+
+
+def qualify_tables(
+    expression: exp.Expr,
+    dialect: DialectType = None,
+    db: str | None = None,
+    catalog: str | None = None,
+    schema: dict[str, object] | Schema | None = None,
+) -> exp.Expr:
+    return qualify(
+        expression,
+        dialect=dialect,
+        db=db,
+        catalog=catalog,
+        schema=schema,
+        qualify_columns=False,
+        quote_identifiers=True,
+        validate_qualify_columns=False,
+        isolate_tables=False,
+    )
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class Plan:
+    out_sql: str
+
+    # The query's projection names, when `out_sql` names them by position: `_0`, `_1`, ...
+    projection_names: tuple[str, ...] | None = None
+
+    # The engine that runs `out_sql`, e.g. the metadata engine for information_schema queries;
+    # None for the warehouse
+    engine: str | None = None
+
+    # Whether the SQL was `EXPLAIN <query>`, whose result is `out_sql` itself rather than its rows
+    explain: bool = False
+
+    # The physical tables `out_sql` reads, by their parts as written, e.g. `("shop", "main","orders")`.
+    # None unless compiled `with_tables`
+    tables: tuple[tuple[str, ...], ...] | None = None
+
+
+# Names are quoted once normalized, so they compare exactly with the registry's.
+ANALYZER_RULES: Sequence[OptimizerFn] = (
+    qualify_tables,
+    qualify_columns,
+    validate_qualify_columns,
+    quote_identifiers,
+)
+
+PLANNER_RULES: Sequence[OptimizerFn] = (
+    isolate_fragments,
+    name_passthrough_tables,
+    ExpandMetadataRule(),
+    ExpandConstellationsRule(),
+    ExpandDatasetsRule(),
+)
+
+
+class Compiler:
+    def __init__(
+        self,
+        cluster: Cluster | None = None,
+        analyzer_rules: Sequence[OptimizerFn] = ANALYZER_RULES,
+        planner_rules: Sequence[OptimizerFn] = PLANNER_RULES,
+        catalog: str | None = None,
+        db: str | None = None,
+        input_dialect: str = "",
+        output_dialect: str = "",
+        engine_dialects: Mapping[str, str] | None = None,
+        normalization_strategy: NormalizationStrategy | None = None,
+    ) -> None:
+        """`engine_dialects` is the SQL dialect of each engine a plan can name, e.g.
+        `{"metadata": <its dialect>}`; the warehouse's is `output_dialect`.
+
+        `normalization_strategy` is the registry's, which queries read names with, whatever
+        their dialect's own, so their names match the registry's."""
+        self.cluster = cluster or Cluster()
+        self.analyzer_rules = analyzer_rules
+        self.planner_rules = planner_rules
+        self.catalog = catalog
+        self.db = db
+        self.input_dialect = input_dialect
+        self.output_dialect = output_dialect
+        self.engine_dialects = dict(engine_dialects or {})
+        self.normalization_strategy = normalization_strategy
+
+    def compile(
+        self,
+        sql: str,
+        *,
+        read_only: bool = False,
+        input_dialect: str | None = None,
+        output_dialect: str | None = None,
+        positional_columns: bool = False,
+        with_tables: bool = False,
+        db: str | None = None,
+    ) -> Plan:
+        """Compiles `sql`.
+
+        `db` names the db unqualified tables are in, as a client gave it, e.g. a Trino session
+        schema `Sales`, instead of the compiler's. It's normalized as a name in `sql` is, quoted
+        only if it must be, e.g. `my-schema`.
+
+        With `positional_columns`, a query's result columns are named `_0`, `_1`, ... and
+        `Plan.projection_names` holds its own names, which a warehouse may reject, e.g. Looker's
+        `accounts.segment` on BigQuery.
+
+        `EXPLAIN <query>` compiles the query, with `Plan.explain` set; EXPLAIN's options, e.g.
+        `EXPLAIN ANALYZE`, aren't supported.
+
+        With `with_tables`, `Plan.tables` holds the physical tables the compiled SQL reads, e.g.
+        for dbt to order a model after the models that build them.
+        """
+        input_dialect = self.input_dialect if input_dialect is None else input_dialect
+        output_dialect = self.output_dialect if output_dialect is None else output_dialect
+        # sqlglot's qualify takes the db as SQL, which it parses and normalizes with the query's
+        # dialect and normalization strategy.
+        default_db = self.db if db is None else exp.to_identifier(db).sql(dialect=input_dialect)
+        if self.normalization_strategy is not None:
+            input_dialect = (
+                f"{input_dialect}, normalization_strategy = "
+                f"{self.normalization_strategy.value.lower()}"
+            )
+        expressions = parse(sql, read=input_dialect)
+        explain = False
+        match expressions:
+            # sqlglot parses EXPLAIN as a command whose argument is the rest of the SQL, unparsed.
+            case [exp.Command(this=str(command), expression=exp.Literal(this=str(explained)))] if (
+                command.upper() == "EXPLAIN"
+            ):
+                explain = True
+                expressions = parse(explained, read=input_dialect)
+                match expressions:
+                    case [exp.Query() | exp.Values()]:
+                        pass
+                    case _:
+                        raise ValueError("EXPLAIN supports only one query statement")
+        if read_only:
+            match expressions:
+                case [exp.Query() | exp.Values()]:
+                    pass
+                case _:
+                    raise ValueError("Only one query statement is supported")
+
+        out = []
+        tables: dict[tuple[str, ...], None] = {}
+        engines: set[str | None] = set()
+        projection_names: tuple[str, ...] | None = None
+
+        for expression in expressions:
+            if not expression:
+                continue
+            self._prepare(expression, input_dialect, default_db, read_only=read_only)
+            planned = self._apply_rules(expression, input_dialect, default_db)
+            engine = planned.meta.get(ENGINE)
+            engines.add(engine)
+            if len(engines) > 1:
+                raise ValueError(
+                    "Statements that run on different engines must be compiled separately"
+                )
+            if positional_columns:
+                match planned:
+                    case exp.Query():
+                        projection_names, planned = _positional(planned)
+            if with_tables:
+                tables.update(dict.fromkeys(_physical_tables(planned)))
+            dialect = self._dialect(engine, output_dialect)
+            # The planned tree is discarded, so the generator needn't copy it first.
+            out.append(planned.sql(dialect=dialect, pretty=True, identify=True, copy=False))
+
+        if len(out) == 0:
+            raise ValueError("Failed to compile SQL")
+        return Plan(
+            out_sql=";\n".join(out),
+            projection_names=projection_names,
+            engine=engine,
+            explain=explain,
+            tables=tuple(tables) if with_tables else None,
+        )
+
+    def _apply_rules(self, expression: exp.Expr, dialect: str, db: str | None) -> exp.Expr:
+        arguments = {
+            "schema": self.cluster,
+            "cluster": self.cluster,
+            "dialect": dialect,
+            "catalog": self.catalog,
+            "db": db,
+        }
+        for rule in (*self.analyzer_rules, *self.planner_rules):
+            parameters = inspect.getfullargspec(rule).args
+            expression = rule(expression, **{p: arguments[p] for p in parameters if p in arguments})
+            if ENGINE in expression.meta:
+                break
+        return expression
+
+    def _dialect(self, engine: str | None, output_dialect: str) -> str:
+        if engine is None:
+            return output_dialect
+        if engine not in self.engine_dialects:
+            raise ValueError(f"No dialect for engine {engine}")
+        return self.engine_dialects[engine]
+
+    def _prepare(
+        self, expression: exp.Expr, dialect: str, db: str | None, *, read_only: bool
+    ) -> None:
+        for node in expression.walk():
+            match node:
+                case exp.Table(this=exp.Anonymous()) if self._is_dataset(node, dialect, db):
+                    bind_table_arguments(node)
+                case exp.Into() if read_only:
+                    raise ValueError("Only one query statement is supported")
+
+    def _is_dataset(self, table: exp.Table, dialect: str, db: str | None) -> bool:
+        """Nothing is loaded yet, so a constellation counts too."""
+        # Not normalized yet, so normalize as `qualify` will. `normalize_identifiers` changes its
+        # argument, so the query's own catalog is copied.
+        if table.args.get("catalog"):
+            catalog = table.args["catalog"].copy()
+        elif self.catalog:
+            catalog = exp.parse_identifier(self.catalog, dialect=dialect)
+        else:
+            return False
+        match self.cluster.get_catalog(normalize_identifiers(catalog, dialect=dialect).name):
+            case SemanticCatalog():
+                return (table.db or db or "").lower() != INFORMATION_SCHEMA
+            case _:
+                return False
+
+
+def _physical_tables(expression: exp.Expr) -> list[tuple[str, ...]]:
+    tables = [
+        source
+        for scope in traverse_scope(expression)
+        for _alias, (_node, source) in scope.selected_sources.items()
+        if isinstance(source, exp.Table)
+    ]
+    return [tuple(part.name for part in table.parts) for table in tables]
+
+
+def _positional(query: exp.Query) -> tuple[tuple[str, ...], exp.Query]:
+    """`query`'s column names, and a query that only renames its columns by position, e.g.
+    `SELECT 1 AS "a.b"` -> `SELECT q."a.b" AS _0 FROM (SELECT 1 AS "a.b") AS q`.
+
+    Qualification aliases every column. A repeated alias, which `q` couldn't tell apart, is first
+    renamed to one no column has, e.g. `_1` in `SELECT a, b AS a`.
+    """
+    names = tuple(select.alias for select in query.selects)
+    taken = {name.lower() for name in names}
+    seen: set[str] = set()
+    for index, (select, name) in enumerate(zip(query.selects, names, strict=True)):
+        if name.lower() in seen:
+            alias = find_new_name(taken, f"_{index}")
+            taken.add(alias)
+            select.replace(exp.alias_(select.unalias(), alias, copy=False))
+        seen.add(name.lower())
+
+    # The planned tree is discarded, so nothing here needs a copy, but the alias each column names.
+    renamed = exp.select(
+        *(
+            exp.column(select.args["alias"], table="q").as_(f"_{index}", copy=False)
+            for index, select in enumerate(query.selects)
+        )
+    ).from_(query.subquery("q", copy=False), copy=False)
+    return names, renamed