metricflow-sql 0.1.1__py3-none-any.whl

metricflow_sql/registry/memory.py

@@ -0,0 +1,409 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Collection, Iterator, Mapping
+from dataclasses import dataclass
+from dataclasses import field as dcfield
+from functools import cached_property
+from typing import assert_never, cast
+
+from sqlglot import expressions as exp
+from sqlglot.dialects.dialect import NormalizationStrategy
+from sqlglot.errors import SchemaError
+
+from metricflow_sql.objects.base import Backlink, BaseObject, Id, Link, Prefix
+from metricflow_sql.objects.constellation import Assignment, Constellation, Import, Metric, Role
+from metricflow_sql.objects.dataset import (
+    Calculation,
+    Column,
+    Dataset,
+    DatasetKey,
+    Parameter,
+    Relationship,
+    Table,
+    View,
+)
+from metricflow_sql.objects.entity import Attribute, Entity, KeyColumn, Source
+from metricflow_sql.objects.wiki import Page
+from metricflow_sql.registry.operation import Batch, Delete, Insert
+from metricflow_sql.registry.registry import (
+    ConstellationProxy,
+    DatasetProxy,
+    DatasetTarget,
+    EntityProxy,
+    EntityTarget,
+    Registry,
+    SearchHit,
+    relationship_target,
+)
+from metricflow_sql.registry.type_resolver import TypeResolver
+from metricflow_sql.utils.composed_set import ComposedSet
+from metricflow_sql.utils.normalize import DEFAULT_STRATEGY
+from metricflow_sql.utils.ordered_set import OrderedSet
+from metricflow_sql.utils.sql import ref_type
+
+
+@dataclass(frozen=True)
+class InMemoryRegistry(Registry):
+    """
+    Graph defined entirely in-memory of the current process
+    """
+
+    dialect: str
+    normalization_strategy: NormalizationStrategy = DEFAULT_STRATEGY
+
+    # Datasets
+    datasets: dict[exp.Table, Table | View] = dcfield(default_factory=dict)
+    dataset_fields: dict[exp.Table, dict[exp.Identifier, Column | Calculation | Relationship]] = (
+        dcfield(default_factory=lambda: defaultdict(dict))
+    )
+    # reverse index
+    dataset_entity_sources: dict[exp.Table, dict[exp.Table, Source]] = dcfield(
+        default_factory=lambda: defaultdict(dict)
+    )
+
+    # Entities
+    entities: dict[exp.Table, Entity] = dcfield(default_factory=dict)
+    entity_fields: dict[exp.Table, dict[exp.Identifier, Attribute]] = dcfield(
+        default_factory=lambda: defaultdict(dict)
+    )
+    entity_sources: dict[exp.Table, dict[exp.Table, Source]] = dcfield(
+        default_factory=lambda: defaultdict(dict)
+    )
+
+    # Constellations
+    constellations: dict[exp.Table, Constellation] = dcfield(default_factory=dict)
+    constellation_fields: dict[exp.Table, dict[exp.Identifier, Metric | Role]] = dcfield(
+        default_factory=lambda: defaultdict(dict)
+    )
+    constellation_assignments: dict[
+        exp.Table, dict[exp.Table, dict[exp.Identifier, Assignment]]
+    ] = dcfield(default_factory=dict)
+
+    # Wiki
+    pages: dict[str, Page] = dcfield(default_factory=dict)
+
+    # Every object by id, and reverse index of links by target
+    objects: dict[Id, BaseObject] = dcfield(default_factory=dict)
+    backlinks: dict[Id, OrderedSet[Backlink]] = dcfield(
+        default_factory=lambda: defaultdict(OrderedSet)
+    )
+
+    @cached_property
+    def type_resolver(self) -> TypeResolver:
+        return TypeResolver(self)
+
+    def apply(self, batch: Batch) -> None:
+        for op in batch.ops:
+            match op:
+                case Insert():
+                    self.insert(op.object)
+                case Delete():
+                    self.delete(op.id)
+                case _:
+                    assert_never(op)
+
+    def insert(self, obj: BaseObject) -> BaseObject:
+        match obj:
+            case Table() | View():
+                self.datasets[obj.name] = obj
+            case Column() | Calculation() | Relationship():
+                self.dataset_fields[obj.dataset][obj.name] = obj
+            case Entity():
+                self.entities[obj.name] = obj
+            case Attribute():
+                self.entity_fields[obj.entity][obj.name] = obj
+            case Source():
+                self.entity_sources[obj.entity][obj.dataset] = obj
+                self.dataset_entity_sources[obj.dataset][obj.entity] = obj
+            case Constellation():
+                self.constellations[obj.name] = obj
+            case Metric() | Role():
+                self.constellation_fields[obj.constellation][obj.name] = obj
+            case Assignment():
+                assignments = self.constellation_assignments.setdefault(obj.constellation, {})
+                roles = assignments.setdefault(obj.dataset, {})
+                roles[obj.role] = obj
+            case Page():
+                self.pages[obj.name] = obj
+            case _:
+                raise ValueError(f"Unknown object type: {obj}")
+
+        replaced = self.objects.get(obj.id)
+        if replaced is not obj:
+            self.type_resolver.clear()
+
+        if replaced is not None:
+            for link in replaced.links():
+                self.backlinks[link.target].discard(Backlink(source=obj.id, label=link.label))
+
+        self.objects[obj.id] = obj
+        for link in obj.links():
+            self.backlinks[link.target].add(Backlink(source=obj.id, label=link.label))
+        return obj
+
+    def delete(self, id: Id) -> BaseObject:
+        raise NotImplementedError
+
+    def get_objects(self, ids: Collection[Id]) -> dict[Id, BaseObject | None]:
+        return {object_id: self.objects.get(object_id) for object_id in ids}
+
+    def get_links(self, ids: Collection[Id]) -> dict[Id, list[Link]]:
+        objects = self.get_objects(ids)
+        return {
+            object_id: list(obj.links()) if obj is not None else []
+            for object_id, obj in objects.items()
+        }
+
+    def get_backlinks(self, ids: Collection[Id]) -> dict[Id, list[Backlink]]:
+        return {object_id: list(self.backlinks.get(object_id, ())) for object_id in ids}
+
+    def get_dataset_field_ids(self, dataset: exp.Table) -> list[Id]:
+        return [field.id for field in self.dataset_fields.get(dataset, {}).values()]
+
+    def search(self, query: str, limit: int = 10) -> list[SearchHit]:
+        """
+        Naive text search.
+
+        Chunks containing every whitespace-separated term of `query`, case-insensitively,
+        scored by term occurrences.
+        """
+        terms = query.lower().split()
+        if not terms:
+            return []
+        hits: list[SearchHit] = []
+        for obj in self.objects.values():
+            for chunk in obj.chunks():
+                haystack = " ".join((*chunk.section, chunk.text)).lower()
+                if all(term in haystack for term in terms):
+                    score = sum(haystack.count(term) for term in terms)
+                    hits.append(SearchHit(obj.id, chunk.section, chunk.text, float(score)))
+        # Stable sort, so equal scores keep insertion order.
+        hits.sort(key=lambda hit: -hit.score)
+        return hits[:limit]
+
+    def get_dataset(self, name: exp.Table) -> InMemoryDatasetProxy | None:
+        table = self.datasets.get(name)
+        if not table:
+            return None
+        fields = self.dataset_fields.get(name) or {}
+        return InMemoryDatasetProxy(table, fields, self)
+
+    def get_entity(self, name: exp.Table) -> EntityProxy | None:
+        entity = self.entities.get(name)
+        if not entity:
+            return None
+        fields = self.entity_fields.get(name) or {}
+        sources = self.entity_sources.get(name) or {}
+        return InMemoryEntityProxy(entity, fields, sources, self)
+
+    def get_constellation(self, name: exp.Table) -> InMemoryConstellationProxy | None:
+        if name not in self.constellations:
+            return None
+        return self._constellation_proxy(name, frozenset())
+
+    def _constellation_proxy(
+        self, name: exp.Table, importing: frozenset[exp.Table]
+    ) -> InMemoryConstellationProxy:
+        """`importing` are the constellations importing `name`, to catch a cycle."""
+        if name in importing:
+            raise SchemaError(f"Constellation imports itself: {name.sql(self.dialect)}")
+        constellation = self.constellations.get(name)
+        if not constellation:
+            raise SchemaError(f"Unknown constellation: {name.sql(self.dialect)}")
+        return InMemoryConstellationProxy(
+            constellation,
+            self.constellation_fields.get(name) or {},
+            self.constellation_assignments.get(name) or {},
+            tuple(
+                (import_, self._constellation_proxy(import_.constellation, importing | {name}))
+                for import_ in constellation.imports
+            ),
+            self,
+        )
+
+    def entity_struct_type(self, entity: exp.Table, *, max_depth: int = 1) -> exp.DataType:
+        """Describe an entity as a bounded STRUCT, with only keys at the depth limit."""
+        if entity not in self.entities:
+            raise SchemaError(f"Unknown entity: {entity.sql(self.dialect)}")
+        return self.type_resolver.expand_type(ref_type(entity), max_depth=max_depth)
+
+    def dataset_struct_type(self, dataset: exp.Table, *, max_depth: int = 1) -> exp.DataType:
+        """Describe a dataset as a bounded STRUCT, with only key fields at the depth limit."""
+        if dataset not in self.datasets:
+            raise SchemaError(f"Unknown dataset: {dataset.sql(self.dialect)}")
+        return self.type_resolver.expand_type(ref_type(dataset), max_depth=max_depth)
+
+
+@dataclass(frozen=True)
+class InMemoryDatasetProxy(DatasetProxy):
+    dataset: Dataset
+    fields: Mapping[exp.Identifier, Column | Calculation | Relationship]
+    registry: InMemoryRegistry
+
+    @property
+    def parameters(self) -> dict[exp.Identifier, Parameter]:
+        return self.dataset.parameters
+
+    @property
+    def key(self) -> DatasetKey | None:
+        return self.dataset.key
+
+    def from_(self, arguments: dict[str, exp.Expr]) -> exp.Expr:
+        return self.dataset.from_(arguments)
+
+    @property
+    def name(self) -> exp.Table:
+        return self.dataset.name
+
+    def column_names(self) -> Collection[exp.Identifier]:
+        return sorted(self.fields, key=lambda name: name.name)
+
+    def get_column(self, name: exp.Identifier) -> Column | Calculation | Relationship | None:
+        return self.fields.get(name)
+
+    def get_column_type(self, name: exp.Identifier) -> exp.DataType:
+        match self.get_column(name):
+            case Column() | Calculation() as field:
+                return self.registry.type_resolver.expand_type(
+                    self.registry.type_resolver.infer(field.id)
+                )
+            case Relationship() as relationship:
+                match relationship_target(self.registry, relationship):
+                    case EntityTarget(entity=entity):
+                        return self.registry.entity_struct_type(entity.name)
+                    case DatasetTarget(dataset=dataset):
+                        return self.registry.dataset_struct_type(dataset.name)
+            case _:
+                dialect = self.registry.dialect
+                raise SchemaError(f"Unknown column: {self.name.sql(dialect)}.{name.sql(dialect)}")
+
+    def get_entity_sources(self) -> Mapping[exp.Table, Source]:
+        return self.registry.dataset_entity_sources.get(self.name) or {}
+
+
+@dataclass(frozen=True)
+class InMemoryEntityProxy(EntityProxy):
+    entity: Entity
+    fields: Mapping[exp.Identifier, Attribute]
+    sources: Mapping[exp.Table, Source]
+    registry: InMemoryRegistry
+
+    @property
+    def key(self) -> dict[exp.Identifier, KeyColumn]:
+        return self.entity.key
+
+    @property
+    def name(self) -> exp.Table:
+        return self.entity.name
+
+    def column_names(self) -> Collection[exp.Identifier]:
+        return self.fields
+
+    def get_column(self, name: exp.Identifier) -> Attribute | None:
+        return self.fields.get(name)
+
+    def get_column_type(self, name: exp.Identifier) -> exp.DataType:
+        field = self.get_column(name)
+        if not field:
+            dialect = self.registry.dialect
+            raise SchemaError(f"Unknown column: {self.name.sql(dialect)}.{name.sql(dialect)}")
+
+        return self.registry.type_resolver.expand_type(self.registry.type_resolver.infer(field.id))
+
+    def get_source(self, name: exp.Table) -> Source | None:
+        return self.sources.get(name)
+
+
+@dataclass(frozen=True)
+class ImportedNames(Collection[exp.Identifier]):
+    """The names an import lets through, filtered when read instead of copied."""
+
+    import_: Import
+    names: Collection[exp.Identifier]
+
+    def __contains__(self, name: object) -> bool:
+        return name in self.names and self.import_.includes(cast(exp.Identifier, name))
+
+    def __iter__(self) -> Iterator[exp.Identifier]:
+        return filter(self.import_.includes, self.names)
+
+    def __len__(self) -> int:
+        return sum(1 for _ in self)
+
+
+@dataclass(frozen=True)
+class InMemoryConstellationProxy(ConstellationProxy):
+    """
+    A constellation with the fields it imports. Its own fields override imported ones, and a
+    later import overrides an earlier one.
+    """
+
+    constellation: Constellation
+    fields: Mapping[exp.Identifier, Metric | Role]
+    # Assignments by dataset, then role.
+    assignments: Mapping[exp.Table, Mapping[exp.Identifier, Assignment]]
+    imports: tuple[tuple[Import, InMemoryConstellationProxy], ...]
+    registry: InMemoryRegistry
+
+    @property
+    def name(self) -> exp.Table:
+        return self.constellation.name
+
+    def column_names(self) -> Collection[exp.Identifier]:
+        return ComposedSet(
+            *(
+                ImportedNames(import_, imported.column_names())
+                for import_, imported in self.imports
+            ),
+            self.fields,
+        )
+
+    def get_column(self, name: exp.Identifier) -> Metric | Role | None:
+        if (field := self.fields.get(name)) is not None:
+            return field
+        for import_, imported in reversed(self.imports):
+            if import_.includes(name) and (field := imported.get_column(name)) is not None:
+                return field
+        return None
+
+    def get_column_type(self, name: exp.Identifier) -> exp.DataType:
+        field = self.get_column(name)
+        if not field:
+            raise SchemaError(
+                f"Unknown column: {self.name.sql(self.registry.dialect)}.{name.sql(self.registry.dialect)}"
+            )
+        return self.registry.type_resolver.expand_type(
+            self.registry.type_resolver.infer(Id(Prefix.FIELD, (self.name, name)))
+        )
+
+    def get_assignment(self, dataset: exp.Table, role: exp.Identifier) -> Assignment | None:
+        if (assignment := self.assignments.get(dataset, {}).get(role)) is not None:
+            return assignment
+        for import_, imported in reversed(self.imports):
+            if (
+                import_.includes(role)
+                and (assignment := imported.get_assignment(dataset, role)) is not None
+            ):
+                return assignment
+        return None
+
+    def assigned_datasets(self) -> Collection[exp.Table]:
+        return ComposedSet(
+            self.assignments,
+            *(
+                # Skips datasets that only provide roles the import leaves out.
+                tuple(
+                    dataset
+                    for dataset in imported.assigned_datasets()
+                    if any(map(import_.includes, imported._assigned_roles(dataset)))
+                )
+                for import_, imported in self.imports
+            ),
+        )
+
+    def _assigned_roles(self, dataset: exp.Table) -> Iterator[exp.Identifier]:
+        """Yields the roles `dataset` provides, including imported ones."""
+        yield from self.assignments.get(dataset, {})
+        for import_, imported in self.imports:
+            yield from filter(import_.includes, imported._assigned_roles(dataset))

metricflow_sql/registry/operation.py

@@ -0,0 +1,25 @@
+"""
+Operations that can be applied on a `Registry`.
+"""
+
+from dataclasses import dataclass
+
+from metricflow_sql.objects.base import BaseObject, Id
+
+
+@dataclass(frozen=True)
+class Insert:
+    object: BaseObject
+
+
+@dataclass(frozen=True)
+class Delete:
+    id: Id
+
+
+Operation = Insert | Delete
+
+
+@dataclass(frozen=True)
+class Batch:
+    ops: list[Operation]

metricflow_sql/registry/reference_resolution.py

@@ -0,0 +1,68 @@
+"""Reference lookup shared by compilation and type inference."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from dataclasses import dataclass
+
+from sqlglot import exp
+
+from metricflow_sql.objects.base import resolve_dataset_path
+from metricflow_sql.objects.constellation import Metric
+from metricflow_sql.objects.entity import Source
+from metricflow_sql.registry.registry import ConstellationProxy, EntityProxy, Registry
+from metricflow_sql.utils.sql import column_references, is_shadowed, path_names
+
+
+@dataclass(frozen=True)
+class AttributeReference:
+    """A qualified attribute reference, with its entity source when one exists."""
+
+    column: exp.Column | exp.Dot
+    dataset: exp.Table
+
+    # None can mean a sibling struct attribute (profile.city) or an unknown name.
+    source: Source | None
+
+    # The dataset field and any remaining access, preserving normalized identifiers.
+    path: tuple[exp.Identifier, ...]
+
+
+def resolve_attribute_reference(
+    entity: EntityProxy, column: exp.Column | exp.Dot
+) -> AttributeReference:
+    """Use the compiler's namespace precedence to resolve a qualified reference."""
+    dataset, path = resolve_dataset_path(
+        path_names(column),
+        entity.name.args["db"],
+        lambda name: entity.get_source(name) is not None,
+    )
+    return AttributeReference(
+        column=column, dataset=dataset, source=entity.get_source(dataset), path=path
+    )
+
+
+def resolve_attribute_references(
+    entity: EntityProxy, expression: exp.Expr
+) -> Iterator[AttributeReference]:
+    """Resolve source references, excluding local SQL aliases and retaining unresolved paths."""
+    for column in column_references(expression):
+        if not is_shadowed(column, column.parts[0].name):
+            yield resolve_attribute_reference(entity, column)
+
+
+def resolve_metric(
+    registry: Registry,
+    constellation: ConstellationProxy,
+    name: exp.Identifier,
+    source: exp.Table | None = None,
+) -> Metric | None:
+    """Prefer the queried constellation's metric, then its caller's defining constellation."""
+    match constellation.get_column(name):
+        case Metric() as metric:
+            return metric
+    if source is not None and (defined := registry.get_constellation(source)) is not None:
+        match defined.get_column(name):
+            case Metric() as metric:
+                return metric
+    return None