metricflow-sql 0.1.1__tar.gz → 0.2.0__tar.gz

{metricflow_sql-0.1.1 → metricflow_sql-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: metricflow-sql
-Version: 0.1.1
+Version: 0.2.0
 Summary: MetricFlow's core semantic layer engine
 Project-URL: Documentation, https://congenial-bassoon-7zkjv8w.pages.github.io/
 License-Expression: LicenseRef-Fivetran

{metricflow_sql-0.1.1 → metricflow_sql-0.2.0}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "metricflow-sql"
 # A placeholder: a release sets its tag's version, with scripts/set_version.py.
-version = "0.1.1"
+version = "0.2.0"
 description = "MetricFlow's core semantic layer engine"
 # Its PyPI page: README.md is for this repository.
 readme = "PYPI.md"

{metricflow_sql-0.1.1 → metricflow_sql-0.2.0}/src/metricflow_sql/server/__init__.py

@@ -1,4 +1,5 @@
-"""An HTTP server that compiles semantic SQL, e.g. for dbt-metricflow-sql's models.
+"""An HTTP server that compiles semantic SQL, e.g. for dbt-metricflow-sql's models, and serves
+the JSON API reading its objects.
 
-It requires the `server` extra (`metricflow-sql[server]`), except `errors`.
+It requires the `server` extra (`metricflow-sql[server]`), except `errors` and `serde`.
 """

{metricflow_sql-0.1.1 → metricflow_sql-0.2.0}/src/metricflow_sql/server/app.py

@@ -6,14 +6,18 @@ from starlette.applications import Starlette
 from starlette.concurrency import run_in_threadpool
 
 from metricflow_sql.compiler.compiler import Compiler, Plan
+from metricflow_sql.registry.registry import Registry
 from metricflow_sql.server.compile import compile_routes
+from metricflow_sql.server.registry import registry_routes
 
 logger = logging.getLogger(__name__)
 
 
-def create_compile_app(compiler: Compiler) -> Starlette:
+def create_compile_app(compiler: Compiler, registry: Registry | None = None) -> Starlette:
     """
-    A server of `/v1/compile` alone (see `metricflow_sql.server.compile`), with no warehouse.
+    A server of `/v1/compile` (see `metricflow_sql.server.compile`), with no warehouse, and of
+    `/v1/search` and `/v1/objects` over `registry`, the objects `compiler` compiles against (see
+    `metricflow_sql.server.registry`).
 
     It compiles with `compiler` on Starlette's threadpool, so concurrent compiles share it, and
     its registry must be safe to read from several threads, e.g. an `InMemoryRegistry` that's
@@ -30,4 +34,6 @@ def create_compile_app(compiler: Compiler) -> Starlette:
             with_tables=True,
         )
 
-    return Starlette(routes=compile_routes(lambda request: compile_sql))
+    return Starlette(
+        routes=[*compile_routes(lambda request: compile_sql), *registry_routes(registry)]
+    )

{metricflow_sql-0.1.1 → metricflow_sql-0.2.0}/src/metricflow_sql/server/compile.py

@@ -5,12 +5,18 @@
         "input_dialect": "bigquery",
         "output_dialect": "bigquery"
     }
-    -> 200 {"sql": "SELECT SUM(amount) ...", "tables": [["warehouse", "main", "sales"]]}
+    -> 200 {
+        "sql": "SELECT SUM(amount) ...",
+        "tables": [["warehouse", "main", "sales"]],
+        "explain": false
+    }
        or 400 {"error": "..."}
 
 `input_dialect` is the SQL's, and `output_dialect` the compiled SQL's, e.g. both a warehouse's
 for dbt: no Trino in the path. `tables` are the physical tables the compiled SQL reads, by
 their parts as written, e.g. for dbt to order the model after the models that build them.
+`explain` is true for `EXPLAIN <query>`, whose `sql` is the query's compiled SQL: its result,
+rather than SQL to run.
 """
 
 from __future__ import annotations
@@ -64,6 +70,12 @@ def compile_routes(compile_for: CompileFor) -> list[BaseRoute]:
             return JSONResponse({"error": str(e)}, status_code=400)
         logger.info("Compiled:\n%s\n=>\n%s", body["sql"].strip(), plan.out_sql)
         assert plan.tables is not None, "compiled `with_tables`"
-        return JSONResponse({"sql": plan.out_sql, "tables": [list(table) for table in plan.tables]})
+        return JSONResponse(
+            {
+                "sql": plan.out_sql,
+                "tables": [list(table) for table in plan.tables],
+                "explain": plan.explain,
+            }
+        )
 
     return [Route("/v1/compile", post_compile, methods=["POST"])]

metricflow_sql-0.2.0/src/metricflow_sql/server/http.py

@@ -0,0 +1,44 @@
+"""Helpers shared by the JSON API's routes."""
+
+from __future__ import annotations
+
+import functools
+import logging
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+import msgspec
+from starlette.requests import Request
+from starlette.responses import Response
+
+logger = logging.getLogger(__name__)
+
+Handler = Callable[[Request], Awaitable[Response]]
+
+
+class ApiError(Exception):
+    """A request error, rendered as a JSON error response by `endpoint`."""
+
+    def __init__(self, status_code: int, message: str):
+        super().__init__(message)
+        self.status_code = status_code
+        self.message = message
+
+
+def endpoint(handler: Handler) -> Handler:
+    """Renders `ApiError`s raised by `handler` as JSON error responses."""
+
+    @functools.wraps(handler)
+    async def wrapped(request: Request) -> Response:
+        try:
+            return await handler(request)
+        except ApiError as error:
+            return json_response({"message": error.message}, status_code=error.status_code)
+
+    return wrapped
+
+
+def json_response(payload: Any, *, status_code: int = 200) -> Response:
+    return Response(
+        content=msgspec.json.encode(payload), status_code=status_code, media_type="application/json"
+    )

metricflow_sql-0.2.0/src/metricflow_sql/server/registry.py

@@ -0,0 +1,173 @@
+"""The JSON API reading a registry's objects: searching them, and looking them up by id.
+
+    GET /v1/search?q=return+rate&limit=10
+    -> 200 {"results": [{"id": "f:...", "section": [...], "text": "...", "score": 1.0}]}
+
+    GET /v1/objects?id=metric:sales:revenue&id=p:homepage
+    -> 200 {"objects": {"metric:sales:revenue": {...}, "p:homepage": null}}
+
+Each handler opens the registry a `RequestRegistryFactory` selects for the request, e.g. the
+environment its caller is admitted to, off the event loop, as registries block. Errors are
+`{"message": "..."}`, with a 400 for a bad request.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable, Mapping, Sequence
+from contextlib import AbstractContextManager, nullcontext
+from typing import Any
+
+from starlette.concurrency import run_in_threadpool
+from starlette.requests import Request
+from starlette.responses import Response
+from starlette.routing import BaseRoute, Route
+
+from metricflow_sql.objects.base import Id
+from metricflow_sql.registry.registry import Registry
+from metricflow_sql.server.http import ApiError, Handler, endpoint, json_response
+from metricflow_sql.server.serde import hit_to_json, object_to_json
+
+logger = logging.getLogger(__name__)
+
+RegistryFactory = Callable[[], AbstractContextManager[Registry]]
+RequestRegistryFactory = Callable[[Request], RegistryFactory]
+
+DEFAULT_SEARCH_LIMIT = 10
+MAX_SEARCH_LIMIT = 100
+MAX_BATCH_IDS = 100
+
+type BatchLookup[T] = Callable[[Registry, Sequence[Id]], Mapping[Id, T | None]]
+
+
+async def read_registry[T](registry_factory: RegistryFactory, read: Callable[[Registry], T]) -> T:
+    """Runs `read` against a freshly opened registry, off the event loop - registries block."""
+
+    def run() -> T:
+        with registry_factory() as registry:
+            return read(registry)
+
+    return await run_in_threadpool(run)
+
+
+async def read_batch[T](
+    registry_factory: RegistryFactory,
+    keys: list[str],
+    lookup: BatchLookup[T],
+    to_json: Callable[[T, str], Any],
+) -> dict[str, Any]:
+    """
+    Looks up the ids `keys` name in one batch, as JSON keyed like the input.
+
+    Keys are what the caller passed, e.g. an alias like `metric:a:b`, which can differ from the
+    id it names, and whose names are read as the registry normalizes names. Values are None
+    where `lookup` has none. Raises `ValueError` for a key that isn't an id.
+    """
+
+    def read(registry: Registry) -> dict[str, Any]:
+        dialect = registry.dialect
+        ids_by_key = {
+            key: Id.resolve(key, dialect, registry.normalization_strategy) for key in keys
+        }
+        values = lookup(registry, list(dict.fromkeys(ids_by_key.values())))
+        return {
+            key: None if (value := values[object_id]) is None else to_json(value, dialect)
+            for key, object_id in ids_by_key.items()
+        }
+
+    return await read_registry(registry_factory, read)
+
+
+def batch_endpoint[T](
+    result_key: str,
+    lookup: BatchLookup[T],
+    to_json: Callable[[T, str], Any],
+    registry_factory: RequestRegistryFactory,
+) -> Handler:
+    """
+    A handler looking up the ids given as repeated `id` query parameters, e.g.
+    `?id=metric:sales:revenue&id=p:homepage`, returning `{result_key: {id: value or null}}`.
+    """
+
+    @endpoint
+    async def handler(request: Request) -> Response:
+        keys = request.query_params.getlist("id")
+        if not 1 <= len(keys) <= MAX_BATCH_IDS:
+            raise ApiError(400, f"Pass between 1 and {MAX_BATCH_IDS} `id` query parameters")
+        try:
+            values = await read_batch(registry_factory(request), keys, lookup, to_json)
+        except ValueError as error:
+            raise ApiError(400, str(error)) from error
+        return json_response({result_key: values})
+
+    return handler
+
+
+def objects_endpoint(registry_factory: RequestRegistryFactory) -> Handler:
+    """Object lookup with the registry selected per request (see `batch_endpoint`)."""
+    return batch_endpoint(
+        "objects", lambda registry, ids: registry.get_objects(ids), object_to_json, registry_factory
+    )
+
+
+def search_endpoint(registry_factory: RequestRegistryFactory) -> Handler:
+    """Search, `?q=<query>&limit=<n>`, with the registry selected per request."""
+
+    @endpoint
+    async def handler(request: Request) -> Response:
+        query = request.query_params.get("q", "").strip()
+        if not query:
+            raise ApiError(400, "Missing query parameter `q`")
+        try:
+            limit = int(request.query_params.get("limit", DEFAULT_SEARCH_LIMIT))
+        except ValueError as error:
+            raise ApiError(400, "Query parameter `limit` must be an integer") from error
+        if not 1 <= limit <= MAX_SEARCH_LIMIT:
+            raise ApiError(400, f"Query parameter `limit` must be between 1 and {MAX_SEARCH_LIMIT}")
+
+        results = await read_registry(
+            registry_factory(request),
+            lambda registry: [
+                hit_to_json(hit, registry.dialect) for hit in registry.search(query, limit)
+            ],
+        )
+        return json_response({"results": results})
+
+    return handler
+
+
+def registry_routes(registry: Registry | None) -> list[BaseRoute]:
+    """
+    `/v1/search` and `/v1/objects` over one `registry`, which must be safe to read from several
+    threads, e.g. an `InMemoryRegistry` that's no longer written. Without one, e.g. for a dbt
+    project that declares no semantics, search finds nothing and every id names nothing.
+    """
+    if registry is None:
+        return [
+            Route("/v1/search", _empty_search, methods=["GET"]),
+            Route("/v1/objects", _no_objects, methods=["GET"]),
+        ]
+    served = registry
+
+    def factory(request: Request) -> RegistryFactory:
+        return lambda: nullcontext(served)
+
+    return [
+        Route("/v1/search", search_endpoint(factory), methods=["GET"]),
+        Route("/v1/objects", objects_endpoint(factory), methods=["GET"]),
+    ]
+
+
+@endpoint
+async def _empty_search(request: Request) -> Response:
+    if not request.query_params.get("q", "").strip():
+        raise ApiError(400, "Missing query parameter `q`")
+    return json_response({"results": []})
+
+
+@endpoint
+async def _no_objects(request: Request) -> Response:
+    keys = request.query_params.getlist("id")
+    if not 1 <= len(keys) <= MAX_BATCH_IDS:
+        raise ApiError(400, f"Pass between 1 and {MAX_BATCH_IDS} `id` query parameters")
+    return json_response({"objects": dict.fromkeys(keys)})

metricflow_sql-0.2.0/src/metricflow_sql/server/serde.py

@@ -0,0 +1,247 @@
+"""
+JSON representation of registry objects for the API.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Any, TypedDict, assert_never
+
+from sqlglot import exp
+
+from metricflow_sql.objects.base import Backlink, BaseObject, Link
+from metricflow_sql.objects.constellation import (
+    Assignment,
+    Constellation,
+    Exclude,
+    ImportRule,
+    Include,
+    IncludeAll,
+    Metric,
+    Role,
+)
+from metricflow_sql.objects.dataset import (
+    Calculation,
+    Column,
+    DatasetKey,
+    KeyCondition,
+    Parameter,
+    Relationship,
+    Table,
+    View,
+)
+from metricflow_sql.objects.entity import Attribute, Entity, Source
+from metricflow_sql.objects.wiki import Page
+from metricflow_sql.registry.registry import SearchHit
+
+logger = logging.getLogger(__name__)
+
+
+class ObjectDoc(TypedDict):
+    id: str
+    type: str
+    label: str | None
+    description: str | None
+    meta: dict[str, Any] | None
+    spec: dict[str, Any]
+
+
+class LinkDoc(TypedDict):
+    target: str
+    label: str | None
+
+
+class BackLinkDoc(TypedDict):
+    source: str
+    label: str | None
+
+
+class SearchHitDoc(TypedDict):
+    id: str
+    section: list[str]
+    text: str
+    score: float
+
+
+def object_to_json(obj: BaseObject, dialect: str) -> ObjectDoc:
+    """`obj` with its type-specific `spec`, its names written in the registry's `dialect`."""
+    return {
+        "id": obj.id.text(),
+        "type": type(obj).__name__,
+        "label": obj.label,
+        "description": obj.description,
+        "meta": obj.meta,
+        "spec": _Spec(dialect).to_json(obj),
+    }
+
+
+def links_to_json(links: list[Link], dialect: str) -> list[LinkDoc]:
+    return [{"target": link.target.text(), "label": link.label} for link in links]
+
+
+def backlinks_to_json(backlinks: list[Backlink], dialect: str) -> list[BackLinkDoc]:
+    return [{"source": backlink.source.text(), "label": backlink.label} for backlink in backlinks]
+
+
+def hit_to_json(hit: SearchHit, dialect: str) -> SearchHitDoc:
+    return {
+        "id": hit.id.text(),
+        "section": list(hit.section),
+        "text": hit.text,
+        "score": hit.score,
+    }
+
+
+@dataclass(frozen=True)
+class _Spec:
+    dialect: str
+
+    def name(self, name: exp.Table | exp.Identifier) -> str:
+        return name.sql(self.dialect)
+
+    def sql(self, expression: exp.Expr | None) -> str | None:
+        """The expression's source text when it was parsed from one, else generated SQL."""
+        if expression is None:
+            return None
+        match expression.meta:
+            case {"sql": str(sql)}:
+                return sql
+            case _:
+                return expression.sql(self.dialect)
+
+    def import_rule(self, rule: ImportRule) -> str:
+        """The fragment spelling of `rule` - the inverse of `import_rule_from_str`."""
+        match rule:
+            case IncludeAll():
+                return "*"
+            case Include(field=field):
+                return self.name(field)
+            case Exclude(field=field):
+                return f"!{self.name(field)}"
+            case _:
+                assert_never(rule)
+
+    def parameters(self, parameters: dict[exp.Identifier, Parameter]) -> dict[str, Any]:
+        return {
+            self.name(name): {
+                "datatype": self.sql(parameter.datatype),
+                "default": self.sql(parameter.default),
+            }
+            for name, parameter in parameters.items()
+        }
+
+    def key(self, key: DatasetKey | None) -> list[Any]:
+        """The fragment spelling of `key`, e.g. `['user_id', {'params': [...], 'on': '...'}]`."""
+        parts: list[Any] = []
+        for part in key.parts if key is not None else ():
+            match part:
+                case exp.Identifier():
+                    parts.append(self.name(part))
+                case KeyCondition():
+                    params = [
+                        {"name": self.name(name), "datatype": self.sql(parameter.datatype)}
+                        for name, parameter in part.parameters.items()
+                    ]
+                    parts.append({"params": params, "on": self.sql(part.on)})
+        return parts
+
+    def to_json(self, obj: BaseObject) -> dict[str, Any]:
+        name, sql = self.name, self.sql
+        match obj:
+            case Table():
+                return {
+                    "name": name(obj.name),
+                    "location": sql(obj.location),
+                    "parameters": self.parameters(obj.parameters),
+                    "key": self.key(obj.key),
+                }
+            case View():
+                return {
+                    "name": name(obj.name),
+                    "sql": obj.sql,
+                    "dialect": obj.dialect,
+                    "parameters": self.parameters(obj.parameters),
+                    "key": self.key(obj.key),
+                }
+            case Column():
+                return {
+                    "dataset": name(obj.dataset),
+                    "name": name(obj.name),
+                    "datatype": sql(obj.datatype),
+                }
+            case Calculation():
+                return {
+                    "dataset": name(obj.dataset),
+                    "name": name(obj.name),
+                    "expr": sql(obj.expr),
+                    "datatype": sql(obj.datatype),
+                }
+            case Relationship():
+                return {
+                    "dataset": name(obj.dataset),
+                    "name": name(obj.name),
+                    "target": name(obj.target),
+                    "key": {name(key): sql(expr) for key, expr in obj.key.items()},
+                    "cardinality": obj.cardinality,
+                }
+            case Entity():
+                return {
+                    "name": name(obj.name),
+                    "key": {
+                        name(key): {"datatype": sql(column.datatype), "optional": column.optional}
+                        for key, column in obj.key.items()
+                    },
+                }
+            case Source():
+                return {
+                    "dataset": name(obj.dataset),
+                    "entity": name(obj.entity),
+                    "on": sql(obj.on),
+                }
+            case Attribute():
+                return {
+                    "entity": name(obj.entity),
+                    "name": name(obj.name),
+                    "expr": sql(obj.expr),
+                    "datatype": sql(obj.datatype),
+                }
+            case Constellation():
+                return {
+                    "name": name(obj.name),
+                    "imports": [
+                        {
+                            "constellation": name(import_.constellation),
+                            "fields": [self.import_rule(rule) for rule in import_.fields],
+                        }
+                        for import_ in obj.imports
+                    ],
+                }
+            case Role():
+                return {
+                    "constellation": name(obj.constellation),
+                    "name": name(obj.name),
+                    "datatype": sql(obj.datatype),
+                }
+            case Metric():
+                return {
+                    "constellation": name(obj.constellation),
+                    "name": name(obj.name),
+                    "expr": sql(obj.expr),
+                    "datatype": sql(obj.datatype),
+                }
+            case Assignment():
+                return {
+                    "constellation": name(obj.constellation),
+                    "dataset": name(obj.dataset),
+                    "role": name(obj.role),
+                    "expr": sql(obj.expr),
+                }
+            case Page():
+                return {
+                    "name": obj.name,
+                    "content": obj.content,
+                    "content_type": obj.content_type,
+                }
+            case _:
+                raise ValueError(f"Unknown object type: {type(obj).__name__}")

{metricflow_sql-0.1.1 → metricflow_sql-0.2.0}/src/metricflow_sql/server/serve.py

@@ -1,4 +1,5 @@
-"""Compiles semantic SQL over bundle directories, run by a process manager, e.g. dbt's.
+"""Compiles semantic SQL over bundle directories, and searches and reads their objects, run by a
+process manager, e.g. dbt's.
 
     metricflow-serve --bundle target/metricflow --catalog jaffle_shop
 
@@ -8,7 +9,8 @@ closes, e.g. when its manager exits, however it does.
 
 The bundles' objects are in one semantic catalog, `--catalog`, e.g. a dbt project's name, so
 unqualified tables resolve to it. It only compiles, and each compile names its dialects: see
-`metricflow_sql.server.compile`.
+`metricflow_sql.server.compile`. It also serves `/v1/search` and `/v1/objects` over them: see
+`metricflow_sql.server.registry`.
 
 It requires the `server` extra (`metricflow-sql[server]`).
 """
@@ -145,7 +147,7 @@ def main() -> None:
         sys.exit(f"Invalid bundle: {error}")
 
     logger.info("Serving them as catalog `%s`", args.catalog)
-    app = create_compile_app(_compiler(registry, args.catalog))
+    app = create_compile_app(_compiler(registry, args.catalog), registry)
     # No logging config of uvicorn's own, which logs requests to stdout: its loggers propagate to
     # the root's, on stderr, as stdout is for the address.
     server = _Server(

{metricflow_sql-0.1.1 → metricflow_sql-0.2.0}/src/metricflow_sql/utils/sql.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 
 import itertools
 import logging
+import re
 from collections.abc import Iterator
 from typing import overload
 
@@ -24,6 +25,11 @@ logger = logging.getLogger(__name__)
 UNKNOWN = exp.DType.UNKNOWN.into_expr()
 
 
+# Reference datatypes stored before `REF(name)` are spelled `TABLE('name')`, or earlier
+# `entity('name')`: bundles and registry rows hold them for good, e.g. `TABLE('item')` -> `item`.
+_LEGACY_REF_DATATYPE = re.compile(r"^(?:table|entity)\('(?P<target>[^']+)'\)$", re.IGNORECASE)
+
+
 def type_or_unknown(datatype: exp.DataType | None) -> exp.DataType:
     """Return a copy of the datatype, defaulting to UNKNOWN."""
     return (UNKNOWN if datatype is None else datatype).copy()
@@ -62,7 +68,13 @@ def parse_sql_with_source(
     datatype's. The text and dialect kept are the original, so it can be parsed again with
     another strategy.
     """
-    expression = parse_one(text, dialect=dialect, into=into)
+    legacy = _LEGACY_REF_DATATYPE.match(text) if into is exp.DataType else None
+    # The original text is kept below, so the stored form round-trips unchanged.
+    expression = (
+        parse_one(text, dialect=dialect, into=into)
+        if legacy is None
+        else ref_type(exp.to_table(legacy["target"], dialect=dialect))
+    )
     if normalization_strategy is not None:
         if into is exp.DataType:
             # Normalize REF targets, including ARRAY<REF(Customer)>, but preserve physical