dbt-metricflow-sql 0.1.1__py3-none-any.whl

dbt_metricflow_sql/__init__.py

@@ -0,0 +1,11 @@
+"""A dbt plugin: dbt discovers it by its `dbt_` module prefix and `plugins` list."""
+
+from __future__ import annotations
+
+import logging
+
+from dbt_metricflow_sql.plugin import MetricFlowPlugin
+
+logger = logging.getLogger(__name__)
+
+plugins = [MetricFlowPlugin]

dbt_metricflow_sql/client.py

@@ -0,0 +1,173 @@
+"""Metricflow HTTP client"""
+
+from __future__ import annotations
+
+import json
+import logging
+import urllib.error
+import urllib.parse
+import urllib.request
+from collections.abc import Mapping
+from dataclasses import dataclass
+from email.message import Message
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+_TIMEOUT_SECONDS = 10
+# Bundles can be large, and the server reads one as it syncs it.
+_PUBLISH_TIMEOUT_SECONDS = 300
+# The header of a publish's response naming the catalog to query the bundle as, e.g. `env_220`.
+CATALOG_HEADER = "X-Dbt-Catalog"
+
+
+class MetricFlowError(Exception):
+    """A request to a MetricFlow server failed: it's unreachable, or it rejected the request."""
+
+
+class _NoRedirectHandler(urllib.request.HTTPRedirectHandler):
+    """
+    Follows no redirect, so a 3xx raises `HTTPError`: urllib would otherwise send the redirect's
+    host the request's headers, e.g. `Authorization`.
+    """
+
+    def redirect_request(self, *args: Any, **kwargs: Any) -> None:
+        return None
+
+
+# For the server's requests and the upload alike.
+_OPENER = urllib.request.build_opener(_NoRedirectHandler)
+
+
+@dataclass(frozen=True)
+class Compiled:
+    """Warehouse SQL compiled from semantic SQL."""
+
+    sql: str
+    # The physical tables `sql` reads, by their parts, e.g. `("analytics", "main", "orders")`.
+    tables: tuple[tuple[str, ...], ...]
+
+
+def compile_sql(url: str, sql: str, dialect: str) -> Compiled:
+    """Metricflow SQL, in the warehouse's `dialect`, to warehouse SQL."""
+    response = _request(
+        url,
+        "POST",
+        "/v1/compile",
+        {"sql": sql, "input_dialect": dialect, "output_dialect": dialect},
+    ).body
+    return Compiled(
+        sql=str(response["sql"]),
+        tables=tuple(tuple(str(part) for part in table) for table in response.get("tables", [])),
+    )
+
+
+@dataclass(frozen=True)
+class Published:
+    """A bundle published to a MetricFlow server."""
+
+    name: str
+    version: int
+    # False if the server already had this bundle, as its current version.
+    changed: bool
+    # The catalog to query it as, if the server names one, e.g. `env_220`.
+    catalog: str | None = None
+
+
+def publish_bundle(
+    url: str, name: str, archive: bytes, headers: Mapping[str, str] | None = None
+) -> Published:
+    """
+    Publishes the zip `archive` of a bundle to the server at `url`, as the next version of
+    the bundle `name`: uploads it through the URL the server gives, then asks the server to
+    publish it. `headers`, e.g. credentials, go with each request to the server, and never
+    with the upload.
+    """
+    timeout = _PUBLISH_TIMEOUT_SECONDS
+    upload = _request(
+        url, "POST", "/api/v1/bundles/uploads", {"size": len(archive)}, timeout, headers
+    ).body
+    upload_url = str(upload["url"])
+    # Only the headers the server gives: the URL is the store's, e.g. S3's, which mustn't see
+    # the server's credentials.
+    put = urllib.request.Request(
+        upload_url, data=archive, headers=dict(upload["headers"]), method=str(upload["method"])
+    )
+    # Named by its host: the URL itself carries a signature.
+    parts = urllib.parse.urlsplit(upload_url)
+    _send(put, f"{parts.scheme}://{parts.netloc}", timeout)
+    path = f"/api/v1/bundles/{urllib.parse.quote(name, safe='')}"
+    response = _request(url, "PUT", path, {"ref": upload["ref"]}, timeout, headers)
+    result = response.body
+    return Published(
+        name=str(result["name"]),
+        version=int(result["version"]),
+        changed=bool(result["changed"]),
+        catalog=response.headers.get(CATALOG_HEADER) or None,
+    )
+
+
+@dataclass(frozen=True)
+class _Response:
+    """A JSON response from the server."""
+
+    body: dict[str, Any]
+    headers: Message
+
+
+def _request(
+    url: str,
+    method: str,
+    path: str,
+    body: Mapping[str, Any],
+    timeout: float = _TIMEOUT_SECONDS,
+    headers: Mapping[str, str] | None = None,
+) -> _Response:
+    parts = urllib.parse.urlsplit(url)
+    if parts.scheme not in ("http", "https") or not parts.netloc:
+        raise MetricFlowError(f"invalid URL `{url}`: it must be like http://localhost:8080")
+    request = urllib.request.Request(
+        f"{url.rstrip('/')}{path}",
+        data=json.dumps(body).encode(),
+        headers={**(headers or {}), "Content-Type": "application/json"},
+        method=method,
+    )
+    content, response_headers = _send(request, url, timeout)
+    return _Response(body=json.loads(content), headers=response_headers)
+
+
+def _send(request: urllib.request.Request, url: str, timeout: float) -> tuple[bytes, Message]:
+    """
+    Sends `request`, returning the response's body and headers. A redirect is an error, naming its
+    host.
+    """
+    try:
+        with _OPENER.open(request, timeout=timeout) as response:
+            content: bytes = response.read()
+            return content, response.headers
+    except urllib.error.HTTPError as e:
+        location = e.headers.get("Location")
+        if 300 <= e.code < 400 and location:
+            # Named by its host, as the upload's URL is: a `Location` can carry a signature. A
+            # relative one, e.g. `/other`, is on the request's host.
+            parts = urllib.parse.urlsplit(urllib.parse.urljoin(request.full_url, location))
+            raise MetricFlowError(
+                f"redirected ({e.code}) to {parts.scheme}://{parts.netloc}, not followed"
+            ) from e
+        raise MetricFlowError(_error_message(e.read())) from e
+    # `URLError` (refused, DNS) and timeouts, which a read can raise unwrapped, are both `OSError`s.
+    except OSError as e:
+        raise MetricFlowError(f"unreachable at {url}: {e}") from e
+
+
+def _error_message(content: bytes) -> str:
+    """The message of an error response, e.g. `{"message": "..."}`, else its raw body."""
+    text = content.decode(errors="replace")
+    try:
+        body = json.loads(text)
+    except ValueError:
+        return text
+    message: Any | None = None
+    if type(body) is dict:
+        message = body.get("message", body.get("error"))
+    return text if message is None else str(message)

dbt_metricflow_sql/dependencies.py

@@ -0,0 +1,56 @@
+"""Maps the physical tables a compiled MetricFlow model reads to the dbt nodes that build them.
+
+dbt can't see what a model's semantic SQL reads, e.g. that `FROM sales` reads the `orders`
+model through the `sales` constellation; MetricFlow reports the tables, e.g.
+`("analytics", "main", "orders")`, and their nodes become the model's dependencies.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections import defaultdict
+from collections.abc import Iterable
+from typing import cast
+
+from dbt.artifacts.resources.types import NodeType
+from dbt.contracts.graph.manifest import Manifest
+from dbt.contracts.graph.nodes import ModelNode, SeedNode, SnapshotNode
+
+logger = logging.getLogger(__name__)
+
+# The node types that build tables a query can read.
+_RELATION_TYPES = (NodeType.Model, NodeType.Seed, NodeType.Snapshot)
+
+
+class RelationIndex:
+    """The dbt nodes and sources of a manifest, by the relation each builds or names."""
+
+    def __init__(self, manifest: Manifest) -> None:
+        # Unique ids by every suffix of the relation's parts, casefolded: warehouses differ in
+        # how they case unquoted names, and a query may leave off the database or schema.
+        self._ids: dict[tuple[str, ...], set[str]] = defaultdict(set)
+        for node in manifest.nodes.values():
+            if node.resource_type not in _RELATION_TYPES or node.config.materialized == "ephemeral":
+                continue
+            relation = cast("ModelNode | SeedNode | SnapshotNode", node)
+            self._add(relation.unique_id, (relation.database, relation.schema, relation.alias))
+        for source in manifest.sources.values():
+            self._add(source.unique_id, (source.database, source.schema, source.identifier))
+
+    def _add(self, unique_id: str, parts: tuple[str | None, ...]) -> None:
+        names = tuple(part.casefold() for part in parts if part)
+        for start in range(len(names)):
+            self._ids[names[start:]].add(unique_id)
+
+    def nodes(self, tables: Iterable[tuple[str, ...]]) -> list[str]:
+        """
+        The unique ids of the nodes that build `tables`, in order. A table that no node builds,
+        e.g. one outside the project, or that names several, e.g. `orders` in two schemas, has
+        none.
+        """
+        unique_ids: dict[str, None] = {}
+        for table in tables:
+            matches = self._ids.get(tuple(part.casefold() for part in table), set())
+            if len(matches) == 1:
+                unique_ids[next(iter(matches))] = None
+        return list(unique_ids)