metricflow-sql 0.1.1__py3-none-any.whl

metricflow_sql/__init__.py

@@ -0,0 +1,8 @@
+"""Core SQL library for MetricFlow."""
+
+from importlib import metadata
+
+__all__ = ["__version__"]
+
+# The installed package's, which a release sets from its tag.
+__version__ = metadata.version("metricflow-sql")

metricflow_sql/compiler/catalog/cluster.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, override
+
+from sqlglot import Dialect, Schema
+from sqlglot import expressions as exp
+from sqlglot.dialects.dialect import DialectType, NormalizationStrategy
+from sqlglot.errors import OptimizeError, SchemaError
+from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+
+from metricflow_sql.compiler.catalog.protocols import CatalogLike, ColumnLike, TableLike
+from metricflow_sql.registry.registry import ConstellationProxy, DatasetProxy, Registry
+from metricflow_sql.utils.normalize import exact_identifier, normalizing_dialect
+
+if TYPE_CHECKING:
+    from sqlglot.schema import ColumnMapping
+
+
+class Cluster(Schema):
+    """
+    Collection of catalogs.
+
+    This serves as the database schema for the compiler and sqlglot optimizer.
+    """
+
+    def __init__(
+        self,
+        catalogs: dict[str, CatalogLike] | None = None,
+        dialect: DialectType = None,
+        fallback: Callable[[str], CatalogLike | None] | None = None,
+        normalization_strategy: NormalizationStrategy | None = None,
+    ):
+        """
+        `catalogs` maps names to catalogs. A name matches a query's once both are normalized with
+        `normalization_strategy`, e.g. `env_1` matches `ENV_1` if names are uppercase. Give it
+        the registry's strategy, as the compiler's.
+
+        `fallback` builds a catalog for any other name, e.g. a warehouse catalog.
+        """
+        self.catalogs = catalogs or {}
+        self._dialect = dialect
+        self._normalizing = (
+            normalizing_dialect(None, normalization_strategy) if normalization_strategy else None
+        )
+        self._fallback = fallback
+        self._fallen_back: dict[str, CatalogLike | None] = {}
+        self._normalized_names: dict[str, str] = {}
+
+    def get_catalog(self, name: str) -> CatalogLike | None:
+        """The catalog named `name`, which is normalized: one of `catalogs`, else the fallback's."""
+        catalog = next(
+            (own for key, own in self.catalogs.items() if self._normalized(key) == name), None
+        )
+        if catalog is not None or self._fallback is None or not name:
+            return catalog
+        # Built once, so it looks up each table once.
+        if name not in self._fallen_back:
+            self._fallen_back[name] = self._fallback(name)
+        return self._fallen_back[name]
+
+    def _normalized(self, name: str) -> str:
+        if name not in self._normalized_names:
+            identifier = exp.to_identifier(name)
+            self._normalized_names[name] = normalize_identifiers(
+                identifier, dialect=self._normalizing
+            ).name
+        return self._normalized_names[name]
+
+    def find_table(self, table: exp.Table | str) -> TableLike[ColumnLike] | None:
+        catalog, name = self._resolve(table)
+        return catalog.get_table(name)
+
+    def get_table(self, table: exp.Table | str) -> TableLike[ColumnLike]:
+        catalog, name = self._resolve(table)
+        table_obj = catalog.get_table(name)
+        if not table_obj:
+            catalog_name = exp.to_table(table, dialect=self._dialect, copy=False).args["catalog"]
+            name.set("catalog", catalog_name.copy())
+            raise SchemaError(f"Unknown table: {name.sql(self._dialect)}")
+        return table_obj
+
+    def get_dataset(self, table: exp.Table | str) -> tuple[Registry, DatasetProxy] | None:
+        catalog, name = self._resolve(table)
+        return catalog.get_dataset(name)
+
+    def get_constellation(
+        self, table: exp.Table | str
+    ) -> tuple[Registry, ConstellationProxy] | None:
+        catalog, name = self._resolve(table)
+        return catalog.get_constellation(name)
+
+    def _resolve(self, table: exp.Table | str) -> tuple[CatalogLike, exp.Table]:
+        """The table's catalog, and its name in it."""
+        table = exp.to_table(table, dialect=self._dialect, copy=False)
+        if not table.catalog:
+            raise OptimizeError(f"Expected table to have been qualified with catalog: {table}")
+        if not table.db:
+            raise OptimizeError(f"Expected table to have been qualified with db: {table}")
+        catalog = self.get_catalog(table.catalog)
+        if not catalog:
+            raise SchemaError(f"Unknown catalog: {table.sql(self._dialect)}")
+        # Not `exp.table_`, which a table function's name, e.g. `f('x')`, can't be passed to.
+        return catalog, exp.Table(this=table.this.copy(), db=table.args["db"].copy())
+
+    @property
+    @override
+    def empty(self) -> bool:
+        # Tables load lazily, so there's no mapping to show; to sqlglot, an empty schema knows no
+        # columns, and an alias shadows one, e.g. `b` in `SELECT CONCAT(a, b) AS b ... WHERE b = 1`.
+        return not self.catalogs
+
+    @property
+    @override
+    def dialect(self) -> Dialect | None:
+        return Dialect.get_or_raise(self._dialect)
+
+    @override
+    def add_table(
+        self,
+        table: exp.Table | str,
+        column_mapping: ColumnMapping | None = None,
+        dialect: DialectType = None,
+        normalize: bool | None = None,
+        match_depth: bool = True,
+    ) -> None:
+        raise NotImplementedError
+
+    @override
+    def column_names(
+        self,
+        table: exp.Table | str,
+        only_visible: bool = False,
+        dialect: DialectType = None,
+        normalize: bool | None = None,
+    ) -> Sequence[str]:
+        return [name.name for name in self.get_table(table).column_names()]
+
+    @override
+    def get_column_type(
+        self,
+        table: exp.Table | str,
+        column: exp.Column | str,
+        dialect: DialectType = None,
+        normalize: bool | None = None,
+    ) -> exp.DataType:
+        return self.get_table(table).get_column_type(exact_identifier(column))
+
+    @property
+    @override
+    def supported_table_args(self) -> tuple[str, ...]:
+        return "this", "db", "catalog"

metricflow_sql/compiler/catalog/information_schema.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import logging
+
+from sqlglot import exp
+
+from metricflow_sql.compiler.catalog.metadata import MetadataColumn, MetadataTable
+
+logger = logging.getLogger(__name__)
+
+# Trino's public information_schema.tables columns, all with SQL type VARCHAR.
+TABLE_COLUMNS = (
+    "table_catalog",
+    "table_schema",
+    "table_name",
+    "table_type",
+)
+
+# The information_schema.columns columns clients read, e.g. Looker, all VARCHAR.
+COLUMN_COLUMNS = (
+    "table_catalog",
+    "table_schema",
+    "table_name",
+    "column_name",
+    "data_type",
+)
+
+
+def table_metadata(source: exp.Query, *, engine: str) -> MetadataTable:
+    """information_schema.tables, whose rows `source` produces on `engine`."""
+    return _varchar_table("tables", TABLE_COLUMNS, source, engine)
+
+
+def column_metadata(source: exp.Query, *, engine: str) -> MetadataTable:
+    """information_schema.columns, whose rows `source` produces on `engine`."""
+    return _varchar_table("columns", COLUMN_COLUMNS, source, engine)
+
+
+def _varchar_table(
+    name: str, columns: tuple[str, ...], source: exp.Query, engine: str
+) -> MetadataTable:
+    return MetadataTable(
+        name=name,
+        columns={column: MetadataColumn(exp.DType.VARCHAR.into_expr()) for column in columns},
+        source=source,
+        engine=engine,
+    )

metricflow_sql/compiler/catalog/metadata.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable, Collection, Mapping, Sequence
+from dataclasses import dataclass
+
+from sqlglot import exp
+from sqlglot.errors import SchemaError
+
+from metricflow_sql.compiler.catalog.protocols import CatalogLike, ColumnLike, TableLike
+
+logger = logging.getLogger(__name__)
+
+# The db a semantic catalog serves its metadata tables under.
+INFORMATION_SCHEMA = "information_schema"
+
+
+@dataclass(frozen=True)
+class MetadataColumn:
+    datatype: exp.DataType
+    description: str | None = None
+
+
+@dataclass(frozen=True)
+class MetadataTable:
+    name: str
+    columns: dict[str, MetadataColumn]
+
+    # The query that replaces the table, or a function that builds it when it's read
+    source: exp.Query | Callable[[], exp.Query]
+
+    # The engine that runs queries over the table, e.g. "metadata"
+    engine: str
+
+    def query(self) -> exp.Query:
+        """A copy of the query that replaces the table."""
+        return self.source() if callable(self.source) else self.source.copy()
+
+    def column_names(self) -> Collection[exp.Identifier]:
+        return [exp.to_identifier(name) for name in self.columns]
+
+    # information_schema's names are case-insensitive.
+    def get_column(self, name: exp.Identifier) -> MetadataColumn | None:
+        return self.columns.get(name.name.lower())
+
+    def get_column_type(self, name: exp.Identifier) -> exp.DataType:
+        column = self.get_column(name)
+        if column is None:
+            raise SchemaError(f"Unknown column: {self.name}.{name.name}")
+        return column.datatype
+
+
+class MetadataCatalog(CatalogLike):
+    """A catalog of only metadata tables, e.g. a Trino client's `system.runtime.nodes`."""
+
+    def __init__(self, tables: Mapping[str, Sequence[MetadataTable]]) -> None:
+        """`tables` are the catalog's metadata tables, by db."""
+        self._tables = {
+            (db.lower(), table.name.lower()): table
+            for db, db_tables in tables.items()
+            for table in db_tables
+        }
+
+    def get_table(self, table: exp.Table) -> TableLike[ColumnLike] | None:
+        return self._tables.get((table.db.lower(), table.name.lower()))

metricflow_sql/compiler/catalog/passthrough.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable, Collection, Iterable, Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any, Protocol, override
+
+from sqlglot import exp
+from sqlglot.dialects.dialect import DialectType
+from sqlglot.errors import ParseError, SchemaError
+from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+
+from metricflow_sql.compiler.catalog.metadata import MetadataColumn
+from metricflow_sql.compiler.catalog.protocols import CatalogLike, TableLike
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class PhysicalTable:
+    """A warehouse table, named as the warehouse names it."""
+
+    # Quoted, e.g. `"ANALYTICS"."SALES"."ORDERS"`
+    table: exp.Table
+    # In the table's order
+    columns: Mapping[str, MetadataColumn]
+
+
+class TableLoader(Protocol):
+    """Looks up warehouse tables for a `PassthroughCatalog`."""
+
+    def find_tables(self, tables: Sequence[exp.Table]) -> Sequence[PhysicalTable | None]:
+        """
+        The warehouse tables named `tables`, in order, with None for each not found. Names are
+        normalized, so they match the warehouse's, e.g. `ANALYTICS.SALES.ORDERS` on Snowflake.
+
+        Raises `SchemaError` if more than one table matches a name.
+        """
+        ...
+
+    def list_catalogs(self) -> Sequence[str]:
+        """The warehouse's catalog names."""
+        ...
+
+
+def choose_table(
+    table: exp.Table, found: Iterable[PhysicalTable], dialect: DialectType = None
+) -> PhysicalTable | None:
+    """
+    The table in `found` named `table`, or None. Names match exactly, or ignoring case if
+    `dialect` ignores case. Raises `SchemaError` if more than one matches.
+    """
+    candidates = [
+        physical
+        for physical in found
+        if all(
+            _normalized(getattr(physical.table, part), dialect) == getattr(table, part)
+            for part in ("catalog", "db", "name")
+        )
+    ]
+    if len(candidates) > 1:
+        names = ", ".join(sorted(physical.table.sql() for physical in candidates))
+        raise SchemaError(f"Ambiguous table: {table.sql()} could be any of {names}")
+    return candidates[0] if candidates else None
+
+
+def find_in_information_schema(
+    table: exp.Table, run: Callable[[str], Iterable[Sequence[Any]]], dialect: DialectType
+) -> PhysicalTable | None:
+    """
+    Look up `table` in the warehouse's `information_schema.columns`, as `TableLoader.find_tables`
+    does. `run` runs a query in `dialect` and returns its rows.
+    """
+    query = (
+        exp.select("table_catalog", "table_schema", "table_name", "column_name", "data_type")
+        .from_(exp.to_table("information_schema.columns"))
+        .where(
+            # Match ignoring case; `choose_table` picks the right one.
+            exp.Lower(this=exp.column("table_catalog")).eq(
+                exp.Literal.string(table.catalog.lower())
+            ),
+            exp.Lower(this=exp.column("table_schema")).eq(exp.Literal.string(table.db.lower())),
+            exp.Lower(this=exp.column("table_name")).eq(exp.Literal.string(table.name.lower())),
+        )
+        .order_by("table_catalog", "table_schema", "table_name", "ordinal_position")
+    )
+    tables: dict[tuple[str, str, str], dict[str, MetadataColumn]] = {}
+    for catalog_name, schema, name, column, datatype in run(query.sql(dialect=dialect)):
+        tables.setdefault((catalog_name, schema, name), {})[column] = MetadataColumn(
+            _datatype(datatype, dialect)
+        )
+    return choose_table(
+        table,
+        (
+            PhysicalTable(
+                exp.table_(
+                    exp.to_identifier(name, quoted=True),
+                    db=exp.to_identifier(schema, quoted=True),
+                    catalog=exp.to_identifier(catalog_name, quoted=True),
+                ),
+                columns,
+            )
+            for (catalog_name, schema, name), columns in tables.items()
+        ),
+        dialect,
+    )
+
+
+def _datatype(datatype: str, dialect: DialectType) -> exp.DataType:
+    """A column's type, or `UNKNOWN` if sqlglot can't parse it."""
+    try:
+        return exp.DataType.build(datatype, dialect=dialect)
+    except ParseError:
+        return exp.DType.UNKNOWN.into_expr()
+
+
+class PassthroughTable(TableLike[MetadataColumn]):
+    """A warehouse table, with its column names normalized as a query's are."""
+
+    def __init__(self, physical: PhysicalTable, dialect: DialectType) -> None:
+        self.physical = physical
+        self._columns = {
+            _normalized(name, dialect): column for name, column in physical.columns.items()
+        }
+
+    def column_names(self) -> Collection[exp.Identifier]:
+        return [exp.to_identifier(name, quoted=True) for name in self._columns]
+
+    def get_column(self, name: exp.Identifier) -> MetadataColumn | None:
+        return self._columns.get(name.name)
+
+    def get_column_type(self, name: exp.Identifier) -> exp.DataType:
+        column = self.get_column(name)
+        if column is None:
+            raise SchemaError(f"Unknown column: {self.physical.table.sql()}.{name.name}")
+        return column.datatype
+
+
+class PassthroughCatalog(CatalogLike):
+    """
+    A warehouse catalog the semantic layer doesn't model, e.g. `analytics` in
+    `analytics.sales.orders`. `loader` looks up its tables, in a batch with `load`, or one by
+    one as queries read them.
+
+    It caches tables, so build one per compile.
+    """
+
+    def __init__(self, loader: TableLoader, dialect: DialectType, catalog: str) -> None:
+        """`dialect` normalizes column names as the compiler does: give it the registry's
+        normalization strategy."""
+        self.loader = loader
+        self.dialect = dialect
+        self.catalog = catalog
+        self._tables: dict[tuple[str, str], PassthroughTable | None] = {}
+
+    @override
+    def get_table(self, table: exp.Table) -> PassthroughTable | None:
+        # More than three parts, which sqlglot reads as a `Dot`
+        if not isinstance(table.this, exp.Identifier):
+            return None
+        self.load([table])
+        return self._tables[(table.db, table.name)]
+
+    def load(self, tables: Iterable[exp.Table]) -> None:
+        """Look up those of `tables` not looked up yet, in one batch."""
+        missing = {
+            (table.db, table.name): exp.table_(
+                table.this, db=table.args["db"], catalog=exp.to_identifier(self.catalog)
+            )
+            for table in tables
+            if isinstance(table.this, exp.Identifier) and (table.db, table.name) not in self._tables
+        }
+        if not missing:
+            return
+        found = self.loader.find_tables(list(missing.values()))
+        for key, physical in zip(missing, found, strict=True):
+            self._tables[key] = (
+                PassthroughTable(physical, self.dialect) if physical is not None else None
+            )
+
+
+def _normalized(name: str, dialect: DialectType) -> str:
+    """`name` as the compiler normalizes it, quoted."""
+    return normalize_identifiers(exp.to_identifier(name, quoted=True), dialect=dialect).name

metricflow_sql/compiler/catalog/protocols.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from collections.abc import Collection
+from typing import TYPE_CHECKING, Protocol
+
+from sqlglot import expressions as exp
+
+if TYPE_CHECKING:
+    from metricflow_sql.registry.registry import ConstellationProxy, DatasetProxy, Registry
+
+
+class ColumnLike(Protocol):
+    """
+    Protocol for objects that can behave like a table column.
+
+    Multiple semantic object types act like columns: `Calculation`, `Relationship`, `Metric`,
+    `Role`, etc.
+    But this also works for non-semantic objects, e.g. `MetadataColumn`.
+    """
+
+    @property
+    def description(self) -> str | None: ...
+
+    @property
+    def datatype(self) -> exp.DataType | None: ...
+
+
+class TableLike[C: ColumnLike](Protocol):
+    """
+    Protocol for objects that can behave like a table.
+
+    Multiple semantic object types act like tables: `Table`, `Registry`, etc.
+    But this also works for non-semantic objects, e.g. `MetadataTable`.
+    """
+
+    def column_names(self) -> Collection[exp.Identifier]: ...
+
+    def get_column(self, name: exp.Identifier) -> C | None: ...
+
+    def get_column_type(self, name: exp.Identifier) -> exp.DataType: ...
+
+
+class CatalogLike(Protocol):
+    """
+    Protocol for objects that can behave like a catalog.
+
+    `Registry` acts like a catalog.
+    But this also works for non-semantic objects, e.g. `MetadataCatalog`.
+    """
+
+    def get_table(self, table: exp.Table) -> TableLike[ColumnLike] | None: ...
+
+    def get_dataset(self, table: exp.Table) -> tuple[Registry, DatasetProxy] | None:
+        """
+        Compiler rules need to map tables to semantic datasets.
+
+        This defaults to `None` for non-semantic catalogs.
+        """
+        return None
+
+    def get_constellation(self, table: exp.Table) -> tuple[Registry, ConstellationProxy] | None:
+        """
+        Compiler rules need to map tables to semantic constellations.
+
+        This defaults to `None` for non-semantic catalogs.
+        """
+        return None

metricflow_sql/compiler/catalog/semantic.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import override
+
+from sqlglot import exp
+from sqlglot.errors import SchemaError
+from sqlglot.optimizer.qualify_columns import quote_identifiers
+
+from metricflow_sql.compiler.catalog.metadata import INFORMATION_SCHEMA, MetadataTable
+from metricflow_sql.compiler.catalog.protocols import CatalogLike, ColumnLike, TableLike
+from metricflow_sql.registry.registry import ConstellationProxy, DatasetProxy, Registry
+
+logger = logging.getLogger(__name__)
+
+
+class SemanticCatalog(CatalogLike):
+    """
+    Adapts a `Registry` so it can serve as a catalog in a sqlglot `Schema`.
+    """
+
+    def __init__(
+        self,
+        registry: Registry,
+        dialect: str | None = None,
+        *,
+        metadata_tables: Sequence[MetadataTable] = (),
+    ) -> None:
+        """`metadata_tables` are this catalog's information_schema, e.g. `tables`."""
+        self._metadata_tables = {table.name: table for table in metadata_tables}
+
+        self.registry = registry
+        self.dialect = dialect
+
+    @override
+    def get_table(self, table: exp.Table) -> TableLike[ColumnLike] | None:
+        if table.db.lower() == INFORMATION_SCHEMA:
+            return self._metadata_tables.get(table.name.lower())
+
+        name = self._name(table)
+        found = self.registry.get_constellation(name) or self.registry.get_dataset(name)
+        if found is None:
+            raise SchemaError(f"Unknown table: {name.sql(self.registry.dialect)}")
+        return found
+
+    @override
+    def get_dataset(self, table: exp.Table) -> tuple[Registry, DatasetProxy] | None:
+        dataset = self.registry.get_dataset(self._name(table))
+        return (self.registry, dataset) if dataset else None
+
+    @override
+    def get_constellation(self, table: exp.Table) -> tuple[Registry, ConstellationProxy] | None:
+        constellation = self.registry.get_constellation(self._name(table))
+        return (self.registry, constellation) if constellation else None
+
+    def _name(self, table: exp.Table) -> exp.Table:
+        """A query's table, which sqlglot has normalized, quoted as the registry's names are."""
+        return quote_identifiers(table.copy(), dialect=self.registry.dialect)