From e5712ec87f7212d340f2e5c5022661ec83b70cc6 Mon Sep 17 00:00:00 2001 From: Krzysztof Macewicz Date: Fri, 25 Sep 2026 02:15:45 +0200 Subject: [PATCH] feat: the Weather COUNT runs while a staging maintains a copy A staging keeps the source authoritative, and the controller now carries saved queries through a staging, so the approved COUNT stays current on the staging's map. The starter's executor still refused any map with a copy ("Execute the count demo before staging or after completed cutover"), so a current COUNT could be neither exported nor run until the cutover. The controller's installed acceptance of queries surviving operations found it: its export after each staging was refused by this rule, not because the COUNT was out of date. The executor now names the source on such a map too; a query written for the copy stays refused until the cutover makes the copy the source. Relaxing the rule alone was not enough: the local run opened its session with the source's adapter only, and a session refuses a map naming an engine it has no adapter for. It now connects every engine the map names, as the application does during a staging, counts on the source's adapter, and closes every adapter even when one cannot close. Tests: a real staging through the starter's own operator on both engines, a second run writing through both copies, and the COUNT for the staging's map counting every row (red before the fix with the session's EngineError); the same test with a failing close; the metadata path on a map loaded from a signed staging packet, and a query for the copy refused. Co-Authored-By: Claude Opus 5.5 (1M context) --- docs/weather-starter.md | 10 +- python/src/sde_demo/query_count.py | 30 +++--- python/tests/test_demo_query_count.py | 148 +++++++++++++++++++++++++- 3 files changed, 170 insertions(+), 18 deletions(-) diff --git a/docs/weather-starter.md b/docs/weather-starter.md index 5b32ded..bc6b632 100644 --- a/docs/weather-starter.md +++ b/docs/weather-starter.md @@ -147,11 +147,15 @@ while trust in who supplied the file comes from the authenticated handoff. This is the Weather demo's bounded count oracle. It accepts only `SELECT COUNT(*) [AS alias] FROM exact_source_table`, with `FINAL` required on ClickHouse. It refuses filters, joins, arbitrary expressions/functions, settings, comments and additional -statements. Names and the query stamp must match the current signed source-only map; execute it -before staging or after controller completion. A general analyst SQL/code API remains separate. +statements. Names and the query stamp must match the current signed map and its source. While a +staging maintains a copy the source stays authoritative, so the count reads the source then too; a +query written for the copy is refused until the cutover makes the copy the source. A general analyst +SQL/code API remains separate. The command verifies all completed local runs, plans and executes the exact admitted SQL using -runtime credentials, and compares its count to the locally expected total. The query has a +runtime credentials, and compares its count to the locally expected total. During a staging it +connects to both engines, as the application does, because the local map names both; the count +itself runs on the source. The query has a server-side execution limit; PostgreSQL additionally uses a read-only transaction. A changed map, changed run catalog, unresolved run or unexpected data refuses a successful receipt. diff --git a/python/src/sde_demo/query_count.py b/python/src/sde_demo/query_count.py index 14efb0e..d696ecd 100644 --- a/python/src/sde_demo/query_count.py +++ b/python/src/sde_demo/query_count.py @@ -10,6 +10,7 @@ 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 @@ -61,10 +62,10 @@ def _source(placement: sde.PlacementMap, project_id: str) -> sde.Materialization or set(placement.groups) != {"WeatherReading"} ): raise DemoRefused("The count demo needs the loaded, signed Weather project map.") - group = placement.groups["WeatherReading"] - if len(group.all()) != 1: - raise DemoRefused("Execute the count demo before staging or after completed cutover.") - return group.source + # 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: @@ -257,18 +258,21 @@ def current() -> sde.PlacementMap: 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"]) - adapter = engine(request.dialect, dsns[request.engine]) - try: - adapter.connect() - with sde.Session( - logical, placement, {request.engine: adapter}, project_id=settings["project_id"] - ): - total = _count(adapter, request.sql) + 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.") - finally: - adapter.close() if ( current().fingerprint != placement.fingerprint or ids != sorted(path.name for path in (root / "runs").iterdir()) diff --git a/python/tests/test_demo_query_count.py b/python/tests/test_demo_query_count.py index 1bb931c..d92a2e7 100644 --- a/python/tests/test_demo_query_count.py +++ b/python/tests/test_demo_query_count.py @@ -8,6 +8,7 @@ from copy import deepcopy from pathlib import Path from typing import Any +from uuid import uuid4 import pytest from _weather_fixture import supplied @@ -19,11 +20,12 @@ def packet( - dialect: str = "postgres", bundle: dict[str, Any] | None = None + dialect: str = "postgres", bundle: dict[str, Any] | None = None, placement: Any = None ) -> tuple[dict[str, Any], Any, dict[str, Any]]: if bundle is None: bundle, _ = supplied(dialect) - _, placement = project.bootstrap(bundle) + if placement is None: + _, placement = project.bootstrap(bundle) source = placement.groups["WeatherReading"].source sql = "SELECT COUNT(*) AS total FROM " + QUOTE[dialect]( source.layout.table_for("WeatherReading") @@ -66,6 +68,82 @@ def test_request_binds_exact_query_to_source_and_map(dialect: str) -> None: assert loaded.map_fingerprint == placement.fingerprint +def staging(bundle: dict[str, Any], sign: Any) -> dict[str, Any]: + """A signed staging packet: the source kept, a maintained copy in the other engine.""" + group = sde.colocation_groups(model())[0] + source = bundle["current_map"]["groups"][group.name]["source"]["engine"] + target = "clickhouse" if source == "postgres" else "postgres" + layout = sde.default_layout(model(), group, dialect=target) + stage_id = uuid4().hex + prepared = deepcopy(bundle["current_map"]) + prepared["map_version"] = 2 + prepared["groups"][group.name].update( + derived=[ + { + "id": "stage-copy", + "engine": target, + "lag_budget_ms": 30000, + "layout": { + "tables": {name: sde.staging_table_name(stage_id, 1) for name in layout.tables}, + "columns": {name: dict(columns) for name, columns in layout.columns.items()}, + }, + } + ], + also_write=["stage-copy"], + ) + return sign( + { + "kind": "sde-stage", + "protocol": 1, + "stage_id": stage_id, + "project_id": bundle["project_id"], + "group": group.name, + "current": bundle["current_map"], + "prepared": sign(prepared), + } + ) + + +def staged(dialect: str = "postgres") -> tuple[dict[str, Any], Any, dict[str, Any]]: + """The COUNT's packet on the map a staging publishes, loaded and verified from its packet.""" + bundle, sign = supplied(dialect) + plan = sde.load_staging_plan( + staging(bundle, sign), + model=model(), + project_id=bundle["project_id"], + public_key=project.public_keys(bundle["public_keys"]), + ) + return packet(dialect, bundle, plan.prepared) + + +@pytest.mark.parametrize("dialect", ["postgres", "clickhouse"]) +def test_a_count_reads_the_source_while_a_staging_maintains_a_copy(dialect: str) -> None: + """Saved queries stay current through a staging now, so their COUNT runs there too.""" + document, placement, bundle = staged(dialect) + assert len(placement.groups["WeatherReading"].all()) == 2 + loaded = query_count.load_count_request( + document, project_id=bundle["project_id"], placement=placement, dialects=bundle["engines"] + ) + source = placement.groups["WeatherReading"].source + assert loaded.engine == source.engine + assert source.layout.table_for("WeatherReading") in loaded.sql + + +def test_a_query_written_for_the_maintained_copy_is_refused() -> None: + document, placement, bundle = staged() + value = deepcopy(document) + copy = placement.groups["WeatherReading"].derived[0] + table = copy.layout.table_for("WeatherReading") + value["query"].update(engine=copy.engine, materialization=copy.id, dialect="clickhouse", + sql=f"SELECT COUNT(*) FROM {table} FINAL") + value["query"]["stamp"]["materialization"] = copy.id + value["digest"] = query_count._digest(value) + with pytest.raises(project.DemoRefused, match="current source"): + query_count.load_count_request( + value, project_id=bundle["project_id"], placement=placement, dialects=bundle["engines"] + ) + + @pytest.mark.parametrize( "sql", [ @@ -290,6 +368,72 @@ def changed(adapter: Any, sql: str) -> int: resources.reset(root, admin) +@pytest.mark.parametrize("dialect", ["postgres", "clickhouse"]) +def test_the_count_runs_on_the_source_of_a_real_staging( + dialect: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The starter's operator stages a copy; the COUNT for the staging's map counts every row.""" + from sde_demo.__main__ import operator + + admin = { + name: os.environ.get(variable, "") + for name, variable in ( + ("postgres", "SDE_POSTGRES_DSN"), + ("clickhouse", "SDE_CLICKHOUSE_DSN"), + ) + } + if not all(admin.values()): + pytest.skip("both native engines are required for the Weather count demo") + root = tmp_path / "weather" + bundle, sign = supplied(dialect) + try: + project.setup(root, bundle, admin) + before = runtime.run(root, iterations=2, batch_size=3, interval_ms=0) + project.write(tmp_path / "stage.json", staging(bundle, sign)) + assert operator(root, "stage", tmp_path / "stage.json")["outcome"] == "prepared" + during = runtime.run(root, iterations=1, batch_size=2, interval_ms=0) + settings = project.config(root) + placement = sde.load_local_map( + root / "state", + model=model(), + project_id=settings["project_id"], + public_key=project.public_keys(settings["public_keys"]), + ) + assert placement.map_version == 2 + assert len(placement.groups["WeatherReading"].all()) == 2 + document, _, _ = packet(dialect, bundle, placement) + receipt = query_count.run_count_query(root, document) + assert receipt["verified"] is True + assert receipt["run_ids"] == sorted([before["run_id"], during["run_id"]]) + result_file = root / "query-results" / receipt["execution_id"] / "result.json" + assert project.read(result_file)["count"] == 8 + + # Both engines are opened for the session, so both are closed - even when one cannot be. + closed: list[str] = [] + opened = query_count.engine + + def recorded(kind: str, dsn: str) -> Any: + adapter = opened(kind, dsn) + close = adapter.close + + def closing() -> None: + close() + closed.append(kind) + if len(closed) == 1: + raise RuntimeError("controlled close failure") + + adapter.close = closing + return adapter + + monkeypatch.setattr(query_count, "engine", recorded) + with pytest.raises(RuntimeError, match="controlled close failure"): + query_count.run_count_query(root, document) + assert sorted(closed) == ["clickhouse", "postgres"] + finally: + if (root / "resources.json").exists(): + resources.reset(root, admin) + + @pytest.mark.parametrize( ("table", "spelling"), [