smart-data-engine-sdk 0.1.0__py3-none-any.whl

sde_demo/project.py

@@ -0,0 +1,312 @@
+"""Validated customer configuration and local Weather bootstrap orchestration."""
+
+from __future__ import annotations
+
+import base64
+import binascii
+import importlib
+import json
+import re
+import stat
+from collections.abc import Iterable, Iterator, Mapping
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any
+
+import sde
+from sde._local_state import confirm_file, transaction, write_bytes
+from sde.generation import GENERATIONS_SINCE, check_map_project
+
+from .model import model
+
+
+class DemoRefused(ValueError):
+    """An actionable starter error whose message contains no driver or input values."""
+
+
+def payload(path: Path, *, private: bool = False) -> bytes:
+    if path.is_symlink() or not path.is_file():
+        raise DemoRefused("A required local file is missing or is a symlink.")
+    if private and stat.S_IMODE(path.stat().st_mode) != 0o600:
+        raise DemoRefused("Credential files must have mode 0600; fix their permissions.")
+    with path.open("rb") as stream:
+        data = stream.read(2 * 1024 * 1024 + 1)
+    if len(data) > 2 * 1024 * 1024:
+        raise DemoRefused("The local JSON input exceeds 2 MiB.")
+    return data
+
+
+def decode(data: bytes) -> dict[str, Any]:
+    def unique(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
+        result: dict[str, Any] = {}
+        for name, value in pairs:
+            if name in result:
+                raise DemoRefused("JSON contains duplicate keys; restore the original file.")
+            result[name] = value
+        return result
+
+    value = json.loads(data, object_pairs_hook=unique)
+    if not isinstance(value, dict):
+        raise DemoRefused("The local JSON input must be an object.")
+    return value
+
+
+def read(path: Path, *, private: bool = False) -> dict[str, Any]:
+    return decode(payload(path, private=private))
+
+
+def write(path: Path, value: Mapping[str, Any], *, once: bool = False) -> None:
+    payload = (
+        json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + "\n"
+    ).encode()
+    if once and path.exists():
+        if path.is_symlink() or path.read_bytes() != payload:
+            raise DemoRefused("This directory belongs to another setup; use a fresh directory.")
+        confirm_file(path)
+        return
+    write_bytes(path, payload, replace=not once)
+
+
+def public_keys(value: Any) -> dict[str, bytes]:
+    if not isinstance(value, dict) or not value:
+        raise DemoRefused("Bootstrap needs trusted public keys supplied by your operator.")
+    result = {}
+    for name, encoded in value.items():
+        if not isinstance(name, str) or not name or not isinstance(encoded, str):
+            raise DemoRefused("Invalid public key configuration.")
+        try:
+            decoded = base64.b64decode(encoded, validate=True)
+        except binascii.Error:
+            raise DemoRefused("Invalid public key configuration.") from None
+        if len(decoded) != 32:
+            raise DemoRefused("Ed25519 public keys must contain 32 bytes.")
+        result[name] = decoded
+    return result
+
+
+def bootstrap(value: dict[str, Any]) -> tuple[dict[str, Any], sde.PlacementMap]:
+    if (
+        set(value)
+        != {"kind", "protocol", "project_id", "model", "public_keys", "engines", "current_map"}
+        or value["kind"] != "sde-weather-bootstrap"
+        or type(value["protocol"]) is not int
+        or value["protocol"] != 1
+    ):
+        raise DemoRefused("Use the complete protocol-1 Weather bootstrap from your operator.")
+    from sde.testing.loader import model_from_neutral
+
+    logical = model()
+    if model_from_neutral(value["model"]).version != logical.version:
+        raise DemoRefused("This starter requires the Weather model; request a matching bootstrap.")
+    engines = value["engines"]
+    if (
+        not isinstance(engines, dict)
+        or not 1 <= len(engines) <= 2
+        or any(not re.fullmatch(r"[a-z][a-z0-9_-]{0,47}", name) for name in engines)
+        or any(dialect not in ("postgres", "clickhouse") for dialect in engines.values())
+        or len(set(engines.values())) != len(engines)
+    ):
+        raise DemoRefused(
+            "Weather demo needs one or two uniquely named PostgreSQL/ClickHouse bindings."
+        )
+    keys = public_keys(value["public_keys"])
+    try:
+        placement = sde.load_map(
+            value["current_map"], model=logical, public_key=keys, require_signature=True
+        )
+    except sde.MapError as exc:
+        # Refused before setup writes anything, and the reason is about the map - a missing
+        # 'signed' extra, a bad signature - never about a credential, so it can be said in full.
+        raise DemoRefused(f"The bootstrap's map does not load: {exc}") from None
+    check_map_project(placement, value["project_id"])
+    if placement.contract < GENERATIONS_SINCE or placement.map_version != 1:
+        raise DemoRefused(
+            "A new demo must start from a signed generation-bearing initial map (version 1)."
+        )
+    for group in placement.groups.values():
+        if len(group.all()) != 1 or group.source.engine not in engines:
+            raise DemoRefused("Bootstrap must contain only bound initial source materializations.")
+    config = {
+        "protocol": 1,
+        "project_id": value["project_id"],
+        "model": sde.neutral_declaration(logical),
+        "public_keys": value["public_keys"],
+        "engines": {
+            name: {
+                "dialect": dialect,
+                "operator_dsn_env": f"SDE_WEATHER_OPERATOR_{index}",
+                "runtime_dsn_envs": [f"SDE_WEATHER_RUNTIME_{index}"],
+            }
+            for index, (name, dialect) in enumerate(sorted(engines.items()))
+        },
+    }
+    return config, placement
+
+
+def config(root: Path, *, ready: bool = True) -> dict[str, Any]:
+    expected, _ = bootstrap(read(root / "bootstrap.json"))
+    if read(root / "config.json") != expected:
+        raise DemoRefused("Local config differs from the enrolled bootstrap; restore it.")
+    if ready:
+        if (root / "reset-request.json").exists():
+            raise DemoRefused("Reset was requested; finish reset and set up a fresh directory.")
+        marker = read(root / "setup-complete.json")
+        if marker != {"protocol": 1, "config_digest": sde.digest16(expected)}:
+            raise DemoRefused(
+                "Setup is incomplete or changed; rerun setup with its original input."
+            )
+    return expected
+
+
+def credentials(root: Path, purpose: str, bindings: Mapping[str, Any]) -> dict[str, str]:
+    import hashlib
+
+    filename = purpose + "-credentials.json"
+    data = payload(root / filename, private=True)
+    value = decode(data)
+    if set(value) != set(bindings) or any(
+        not isinstance(dsn, str) or not dsn for dsn in value.values()
+    ):
+        raise DemoRefused("The local credential bindings are incomplete.")
+    # Verify the file this caller needs, without opening the other privilege level's secrets.
+    metadata = read(root / "resources.json")
+    hashes = metadata.get("credential_hashes")
+    if (
+        not isinstance(hashes, dict)
+        or metadata.get("status") != "ready"
+        or hashes.get(purpose) != hashlib.sha256(data).hexdigest()
+    ):
+        raise DemoRefused("Local credentials differ from the resource allocation; restore them.")
+    return value
+
+
+DRIVERS = {"postgres": ("psycopg", "postgres"), "clickhouse": ("clickhouse_connect", "clickhouse")}
+"""The driver module each binding's adapter imports, and the SDK extra that installs it."""
+
+
+def require_drivers(dialects: Iterable[str]) -> None:
+    """Refuse before anything is written when a binding's driver cannot be imported.
+
+    Without this the first session open failed inside a run's retry loop - ten seconds of retries,
+    a message that could only say "incomplete", and a report left `incomplete` for later fleet runs
+    to stop at. The import is the adapter's own, so a module that is present but cannot load counts
+    as missing too.
+    """
+    missing = []
+    for dialect in sorted(set(dialects)):
+        module, extra = DRIVERS[dialect]
+        try:
+            importlib.import_module(module)
+        except ImportError:
+            missing.append(extra)
+    if missing:
+        raise DemoRefused(
+            "Install the drivers for this demo's engines: "
+            f"pip install 'smart-data-engine-sdk[{','.join(missing)}]'"
+        )
+
+
+def engine(dialect: str, dsn: str) -> Any:
+    if dialect == "postgres":
+        from sde.engines.postgres import PostgresEngine
+
+        return PostgresEngine(dsn)
+    if dialect == "clickhouse":
+        from sde.engines.clickhouse import ClickHouseEngine
+
+        return ClickHouseEngine(dsn)
+    raise DemoRefused("Unsupported engine dialect.")
+
+
+@contextmanager
+def connections(
+    root: Path, settings: dict[str, Any]
+) -> Iterator[tuple[dict[str, Any], dict[str, Any]]]:
+    opened: list[Any] = []
+    try:
+        operators: dict[str, Any] = {}
+        runtime: dict[str, Any] = {}
+        for purpose, target in (("operator", operators), ("runtime", runtime)):
+            dsns = credentials(root, purpose, settings["engines"])
+            for name, binding in sorted(settings["engines"].items()):
+                adapter = engine(binding["dialect"], dsns[name])
+                opened.append(adapter)
+                adapter.connect()
+                target[name] = adapter
+        yield operators, runtime
+    finally:
+        # Attempt every cleanup, even if an earlier adapter cannot close.
+        failures: list[Exception] = []
+        for adapter in reversed(opened):
+            try:
+                adapter.close()
+            except Exception as exc:
+                failures.append(exc)
+        if failures:
+            raise DemoRefused("An engine connection could not close; stop this process.")
+
+
+def setup(root: Path, supplied: dict[str, Any], admin_dsns: Mapping[str, str]) -> dict[str, Any]:
+    from .resources import allocate, grant_tables
+
+    settings, placement = bootstrap(supplied)
+    require_drivers(supplied["engines"].values())
+    logical = model()
+    with transaction(root):
+        if (root / "reset-request.json").exists():
+            raise DemoRefused(
+                "This allocation is being reset; use a new directory for the next demo."
+            )
+        write(root / "bootstrap.json", supplied, once=True)
+        write(root / "config.json", settings, once=True)
+        allocate(root, supplied["engines"], admin_dsns)
+        if (root / "setup-complete.json").exists():
+            config(root)
+            from sde._cutover_project import ProjectState
+
+            store = ProjectState(root / "state", settings["project_id"], logical.version)
+            with store.lock():
+                store.read()
+                current = sde.load_local_map(
+                    root / "state",
+                    model=logical,
+                    project_id=settings["project_id"],
+                    public_key=public_keys(settings["public_keys"]),
+                )
+                store.confirm()
+                confirm_file(root / "setup-complete.json")
+                return {"status": "ready", "map_version": current.map_version}
+        if (root / "state" / "active-map.json").exists():
+            current = sde.load_local_map(
+                root / "state",
+                model=logical,
+                project_id=settings["project_id"],
+                public_key=public_keys(settings["public_keys"]),
+            )
+            if current.fingerprint != placement.fingerprint:
+                raise DemoRefused(
+                    "An unfinished setup already advanced its map; inspect local state."
+                )
+        with connections(root, settings) as (operators, runtime):
+            sde.prepare_schema(logical, placement, operators, project_id=settings["project_id"])
+            tables: dict[str, list[str]] = {name: ["sde_map_state"] for name in operators}
+            for group in placement.groups.values():
+                tables[group.source.engine].extend(group.source.layout.tables.values())
+            grant_tables(root, tables, admin_dsns)
+            local = sde.LocalCutover(
+                root / "state",
+                model=logical,
+                project_id=settings["project_id"],
+                public_key=public_keys(settings["public_keys"]),
+                operators=operators,
+                runtime={name: [adapter] for name, adapter in runtime.items()},
+            )
+            local.enroll(supplied["current_map"])
+            with sde.Session(logical, placement, runtime, project_id=settings["project_id"]):
+                pass
+        write(
+            root / "setup-complete.json",
+            {"protocol": 1, "config_digest": sde.digest16(settings)},
+            once=True,
+        )
+        return {"status": "ready", "map_version": placement.map_version}

sde_demo/query_count.py

@@ -0,0 +1,301 @@
+"""Execute the Weather demo's approved COUNT locally; business values never enter receipts.
+
+This is deliberately a closed COUNT(*) grammar over the current source, not a general SQL runner.
+The metadata constructor is pure and can be used by the controller without importing drivers.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+from collections.abc import Mapping
+from contextlib import ExitStack
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+import sde
+from sde.generation import GENERATIONS_SINCE, check_map_project
+from sde.schema import QUOTE
+
+from .model import model
+from .project import DemoRefused, config, credentials, engine, public_keys, write
+
+KIND = "sde-weather-count-query"
+FIELDS = {
+    "kind",
+    "protocol",
+    "project_id",
+    "name",
+    "revision",
+    "map_fingerprint",
+    "query",
+    "digest",
+}
+
+
+def _bytes(value: Any) -> bytes:
+    return (
+        json.dumps(
+            value, sort_keys=True, ensure_ascii=False, separators=(",", ":"), allow_nan=False
+        )
+        + "\n"
+    ).encode()
+
+
+def _digest(value: Mapping[str, Any]) -> str:
+    return hashlib.sha256(
+        _bytes({key: item for key, item in value.items() if key != "digest"})
+    ).hexdigest()
+
+
+def _source(placement: sde.PlacementMap, project_id: str) -> sde.Materialization:
+    check_map_project(placement, project_id)
+    if (
+        placement.contract < GENERATIONS_SINCE
+        or not placement.signed
+        or placement.fingerprint is None
+        or placement.model_version != model().version
+        or set(placement.groups) != {"WeatherReading"}
+    ):
+        raise DemoRefused("The count demo needs the loaded, signed Weather project map.")
+    # The source is the one authoritative copy, also while a staging maintains another: the count
+    # reads it and a query must name it. A query written for the copy stays refused until the
+    # cutover makes the copy the source.
+    return placement.groups["WeatherReading"].source
+
+
+def count_sql(sql: Any, *, table: str, dialect: str) -> str:
+    """Permit exactly one count over the named source; no expressions, settings or comments."""
+    if (
+        not isinstance(sql, str)
+        or not 1 <= len(sql) <= 8192
+        or dialect not in ("postgres", "clickhouse")
+    ):
+        raise DemoRefused("The demo accepts one bounded COUNT(*) statement on its current source.")
+    identifier = re.escape(QUOTE[dialect](table))
+    if re.fullmatch(r"[a-z_][a-z0-9_]*", table):
+        # PostgreSQL folds unquoted names; quoted names and ClickHouse names remain exact.
+        bare = "(?i:" + re.escape(table) + ")" if dialect == "postgres" else re.escape(table)
+        quoted_alternative = "|" + re.escape('"' + table + '"') if dialect == "clickhouse" else ""
+        identifier = "(?:" + identifier + "|" + bare + quoted_alternative + ")"
+    alias_name = r'[A-Za-z_][A-Za-z0-9_]*|"[A-Za-z_][A-Za-z0-9_]*"'
+    if dialect == "clickhouse":
+        alias_name += r"|`[A-Za-z_][A-Za-z0-9_]*`"
+    alias = r"(?:\s+(?i:AS)\s+(?:" + alias_name + r"))?"
+    final = r"\s+(?i:FINAL)" if dialect == "clickhouse" else ""
+    pattern = (
+        r"\s*(?i:SELECT)\s+(?i:COUNT)\s*\(\s*\*\s*\)"
+        + alias
+        + r"\s+(?i:FROM)\s+"
+        + identifier
+        + final
+        + r"\s*;?\s*"
+    )
+    # SQL token folding is ASCII. Unicode IGNORECASE would also accept dotless i/long s/Kelvin
+    # as another identifier and could execute a count on a different accessible native table.
+    if re.fullmatch(pattern, sql, flags=re.ASCII) is None:
+        raise DemoRefused(
+            "Use SELECT COUNT(*) [AS alias] FROM the exact Weather source "
+            "(with FINAL on ClickHouse); other SQL is outside this demo."
+        )
+    return sql
+
+
+def _query(value: Any, placement: sde.PlacementMap, project_id: str, dialect: str) -> str:
+    source = _source(placement, project_id)
+    if not isinstance(value, dict) or not isinstance(value.get("stamp"), dict):
+        raise DemoRefused("Supply the approved query record, including its stamp.")
+    stamp = value["stamp"]
+    if (
+        value.get("group") != "WeatherReading"
+        or value.get("engine") != source.engine
+        or value.get("materialization") != source.id
+        or value.get("dialect") != dialect
+        or stamp.get("model_version") != placement.model_version
+        or type(stamp.get("map_version")) is not int
+        or stamp["map_version"] != placement.map_version
+        or stamp.get("materialization") != source.id
+    ):
+        raise DemoRefused("The query does not describe this project's current source and map.")
+    return count_sql(
+        value.get("sql"), table=source.layout.table_for("WeatherReading"), dialect=dialect
+    )
+
+
+def make_count_request(
+    *,
+    project_id: str,
+    name: str,
+    revision: int,
+    query: Mapping[str, Any],
+    placement: sde.PlacementMap,
+    dialect: str,
+) -> dict[str, Any]:
+    """Prepare metadata for an authenticated handoff; no connection or customer file is used."""
+    if (
+        not isinstance(name, str)
+        or not name
+        or len(name) > 128
+        or type(revision) is not int
+        or revision < 1
+    ):
+        raise DemoRefused("The count handoff needs a query name and positive integer revision.")
+    copied = json.loads(_bytes(dict(query)))
+    _query(copied, placement, project_id, dialect)
+    document = {
+        "kind": KIND,
+        "protocol": 1,
+        "project_id": project_id,
+        "name": name,
+        "revision": revision,
+        "map_fingerprint": placement.fingerprint,
+        "query": copied,
+    }
+    document["digest"] = _digest(document)
+    return document
+
+
+@dataclass(frozen=True)
+class CountRequest:
+    name: str
+    revision: int
+    sql: str
+    engine: str
+    dialect: str
+    digest: str
+    map_fingerprint: str
+
+
+def load_count_request(
+    value: Mapping[str, Any],
+    *,
+    project_id: str,
+    placement: sde.PlacementMap,
+    dialects: Mapping[str, str],
+) -> CountRequest:
+    copied = json.loads(_bytes(dict(value)))
+    if (
+        set(copied) != FIELDS
+        or copied["kind"] != KIND
+        or type(copied["protocol"]) is not int
+        or copied["protocol"] != 1
+        or copied["project_id"] != project_id
+        or copied["map_fingerprint"] != placement.fingerprint
+        or copied["digest"] != _digest(copied)
+    ):
+        raise DemoRefused(
+            "The count handoff is incomplete, changed or belongs to another map/project."
+        )
+    source = _source(placement, project_id)
+    if source.engine not in dialects or dialects[source.engine] not in ("postgres", "clickhouse"):
+        raise DemoRefused("The count source has no supported local runtime binding.")
+    rebuilt = make_count_request(
+        project_id=project_id,
+        name=copied["name"],
+        revision=copied["revision"],
+        query=copied["query"],
+        placement=placement,
+        dialect=dialects[source.engine],
+    )
+    if copied != rebuilt:
+        raise DemoRefused("The count handoff has inconsistent metadata.")
+    return CountRequest(
+        copied["name"],
+        copied["revision"],
+        copied["query"]["sql"],
+        source.engine,
+        dialects[source.engine],
+        copied["digest"],
+        str(placement.fingerprint),
+    )
+
+
+def _count(adapter: Any, query: str) -> int:
+    if adapter.dialect == "postgres":
+        with adapter._cx.transaction(), adapter._cx.cursor() as cursor:
+            cursor.execute("SET TRANSACTION READ ONLY")
+            cursor.execute("SET LOCAL statement_timeout = '10000ms'")
+            cursor.execute("SET LOCAL lock_timeout = '2000ms'")
+            sde.explain(adapter, query)
+            cursor.execute(query)
+            rows = cursor.fetchall()
+    else:
+        settings = {"readonly": 1, "max_execution_time": 10, "timeout_overflow_mode": "throw"}
+        # The SQL grammar cannot override these settings or call external table functions.
+        adapter._cx.query("EXPLAIN PLAN " + query, settings=settings)
+        rows = adapter._cx.query(query, settings=settings).result_rows
+    if len(rows) != 1 or len(rows[0]) != 1 or type(rows[0][0]) is not int or rows[0][0] < 0:
+        raise DemoRefused("The count result is not one exact nonnegative integer.")
+    return rows[0][0]
+
+
+def run_count_query(root: Path, document: Mapping[str, Any]) -> dict[str, Any]:
+    from .verification import load_completed_runs, verify_runs
+
+    settings, logical = config(root), model()
+    keys = public_keys(settings["public_keys"])
+
+    def current() -> sde.PlacementMap:
+        return sde.load_local_map(
+            root / "state", model=logical, project_id=settings["project_id"], public_key=keys
+        )
+
+    placement = current()
+    request = load_count_request(
+        document,
+        project_id=settings["project_id"],
+        placement=placement,
+        dialects={name: item["dialect"] for name, item in settings["engines"].items()},
+    )
+    ids = sorted(path.name for path in (root / "runs").iterdir())
+    runs = load_completed_runs(root, ids)
+    if not runs:
+        raise DemoRefused("Finish a demo workload before running its count oracle.")
+    verify_runs(root, ids)
+    if current().fingerprint != placement.fingerprint:
+        raise DemoRefused("The local map changed during data verification; retry after cutover.")
+    # A staging's map also names the engine of the copy it maintains, and the session checks the
+    # map against every engine it names; the count itself runs on the source's adapter. Each has a
+    # local binding: the run verification above refused this map otherwise.
+    named = sorted({item.engine for group in placement.groups.values() for item in group.all()})
+    dsns = credentials(root, "runtime", settings["engines"])
+    with ExitStack() as cleanup:  # every adapter is closed, even if an earlier one cannot close
+        adapters: dict[str, Any] = {}
+        for name in named:
+            adapters[name] = engine(settings["engines"][name]["dialect"], dsns[name])
+            cleanup.callback(adapters[name].close)
+            adapters[name].connect()
+        with sde.Session(logical, placement, adapters, project_id=settings["project_id"]):
+            total = _count(adapters[request.engine], request.sql)
+        if total != sum(run.expected_rows for run in runs):
+            raise DemoRefused("The local SQL result differs from the completed synthetic runs.")
+    if (
+        current().fingerprint != placement.fingerprint
+        or ids != sorted(path.name for path in (root / "runs").iterdir())
+        or load_completed_runs(root, ids) != runs
+    ):
+        raise DemoRefused(
+            "The map or run catalog changed during the query; retry after it settles."
+        )
+    identity = uuid4().hex
+    receipt = {
+        "protocol": 1,
+        "kind": "sde-weather-count-observation",
+        "execution_id": identity,
+        "project_id": settings["project_id"],
+        "map_fingerprint": request.map_fingerprint,
+        "query_name": request.name,
+        "query_revision": request.revision,
+        "query_digest": request.digest,
+        "verified": True,
+        "run_ids": ids,
+        "observed_at": datetime.now(UTC).isoformat(),
+    }
+    # Business values exist only in this private local result file, never in the exported receipt.
+    write(root / "query-results" / identity / "result.json", {"count": total, "receipt": receipt})
+    write(root / "query-results" / identity / "receipt.json", receipt)
+    return receipt