diff --git a/.github/workflows/db-migrations.yml b/.github/workflows/db-migrations.yml index 1f54f0340..f37859bea 100644 --- a/.github/workflows/db-migrations.yml +++ b/.github/workflows/db-migrations.yml @@ -132,7 +132,7 @@ jobs: env: PGDATABASE: rapid run: | - python -m pytest tests/db -q --ignore=tests/db/test_schema.py --ignore=tests/db/test_publish.py --ignore=tests/db/test_change.py + python -m pytest tests/db -q --ignore=tests/db/test_schema.py --ignore=tests/db/test_publish.py --ignore=tests/db/test_change.py --ignore=tests/db/proof - name: drift check — a modified already-applied file must be refused env: diff --git a/.github/workflows/foundation.yml b/.github/workflows/foundation.yml index 5e7626210..b5ec41cb2 100644 --- a/.github/workflows/foundation.yml +++ b/.github/workflows/foundation.yml @@ -87,4 +87,4 @@ jobs: env: PGDATABASE: rapid run: | - python -m pytest tests/db/test_schema.py tests/db/test_publish.py tests/db/test_change.py -q -rs + python -m pytest tests/db/test_schema.py tests/db/test_publish.py tests/db/test_change.py tests/db/proof -q -rs diff --git a/tests/db/proof/__init__.py b/tests/db/proof/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/db/proof/fixture.py b/tests/db/proof/fixture.py new file mode 100644 index 000000000..a0a11a084 --- /dev/null +++ b/tests/db/proof/fixture.py @@ -0,0 +1,276 @@ +"""A synthetic world for the procedure proof (foundation plan v5, Part 3). + +Not a test module and no science: one dataset, one survey, one field, two +epochs of one exposure each on two SCAs in one filter. Every product is +written the way a stage writes it, as Parquet plus a manifest under a +``file://`` prefix holding the job id and attempt, and loaded by +:func:`rapidpipe.foundation.publish`; every run and job is created through +:func:`create_run` and :func:`create_job` with the inputs, expected +bindings and exclusions Part 1 statement 4 names. + +* :meth:`World.import_image` stands for the import stage: one difference + image and its source set of :data:`NSOURCES` sources at fixed positions, + one job per (epoch, SCA), so each source set is its own set. +* :meth:`World.associate` writes an object set: the predecessor's objects + with this epoch's sources added to their sufficient sums, the excluded + source sets' members taken out of them, and this epoch's memberships. + A source matches the object whose id is its position hash + (``radec_index``), which is exact here because positions repeat. +* :meth:`World.packets` writes a packet set for one image key: one index + row and one forced row per source of the epoch's source set. + +Expected bindings default to :data:`CURRENT`: the key's open binding at +launch, what the loop would choose. +""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path + +from rapidpipe.foundation import change, create_job, create_run, publish +from rapidpipe.science.spatial import radec_index +from tests.db.foundation_helpers import write_attempt + +SURVEY, DATASET, FIELD, FID = "S", "D", 1, 1 +SCAS = (1, 2) +EPOCHS = (1, 2) +NSOURCES = 24 +FIELD_KEY = {"dataset": DATASET, "survey": SURVEY, "field": FIELD} + +#: ``expects`` default: the key's open binding at launch (or none). +CURRENT = object() + + +def image_key(epoch: int, sca: int) -> dict: + """The scientific key of an image product, source set or packet set.""" + return {"dataset": DATASET, "expid": epoch, "sca": sca, "fid": FID} + + +def position(sca: int, i: int) -> tuple[float, float]: + """Source ``i`` of SCA ``sca``: the same sky position in every epoch.""" + return 10.0 + 0.5 * (sca - 1) + 0.01 * i, -5.0 + 0.002 * i + + +def aid_of(ra: float, dec: float) -> int: + return int(radec_index(ra, dec)) + + +class World: + """The synthetic world on one connection, with a root directory for attempts.""" + + def __init__(self, conn, root: Path): + self.conn = conn + self.root = root + self._pids = 0 + + # -- bookkeeping ------------------------------------------------------- + + def run(self, why: str, settings: dict | None = None, live: bool = False) -> int: + run =create_run(self.conn, "proof", why, "synthetic", + settings or {"recipe": why}, SURVEY, DATASET) + if live: + change(self.conn, [{"open": {"survey": SURVEY, "dataset": DATASET, "run": run}}], + "proof", f"{why} goes live") + return run + + def binding(self, product: str, key: dict) -> int | None: + """The open binding on (product, key), or None.""" + return self._one("SELECT id FROM official WHERE closed_at IS NULL AND product = %s" + " AND key = %s::jsonb", (product, json.dumps(key))) + + def live(self) -> tuple: + """The open live binding for the survey and dataset and its run: (binding, run).""" + rows = self.rows("SELECT id, run FROM official WHERE closed_at IS NULL" + " AND run IS NOT NULL AND survey = %s AND dataset = %s", (SURVEY, DATASET)) + return rows[0] if rows else (None, None) + + def _one(self, query, params): + rows = self.rows(query, params) + return rows[0][0] if rows else None + + def rows(self, query, params=()) -> list[tuple]: + """Run one read in its own short transaction.""" + with self.conn.cursor() as cur: + cur.execute(query, params) + rows = cur.fetchall() + self.conn.commit() + return rows + + def job(self, run: int, stage: str, scope: dict, outputs: dict[str, list[dict]], + inputs=(), expects=(), excludes=(), attempt: int = 1, + published: bool = True) -> int: + """Create a job, write its attempt, and (unless told not to) publish it; return its id.""" + expects = [{**e, "binding": self.binding(e["product"], e["key"])} + if e["binding"] is CURRENT else e for e in expects] + job = create_job(self.conn, run, stage, scope, inputs, expects, excludes) + self.write(job, attempt, outputs) + if published: + publish(self.conn, job, attempt, self.prefix(job, attempt)) + return job + + def prefix(self, job: int, attempt: int) -> str: + return (self.root / f"job-{job}" / f"attempt-{attempt}").as_uri() + + def write(self, job: int, attempt: int, outputs: dict[str, list[dict]]) -> str: + """Write an attempt's declared outputs under its job-and-attempt prefix.""" + return write_attempt(self.root / f"job-{job}" / f"attempt-{attempt}", outputs) + + # -- the stages, standing in ------------------------------------------- + + def import_image(self, run: int, epoch: int, sca: int, flux_scale: float = 1.0) -> int: + """One difference image and its source set; the job is the source set.""" + key = image_key(epoch, sca) + self._pids += 1 + pid = self._pids + diffimage = {"pid": pid, **key, "field": FIELD, "jd": 2460000.5 + epoch, + "infobitssci": 0, "infobitsref": 0, "filename": f"diff-{epoch}-{sca}.fits", + **{f"{c}{n}": v for n in range(5) + for c, v in (("ra", position(sca, 0)[0]), ("dec", position(sca, 0)[1]))}} + sources = [] + for i in range(NSOURCES): + ra, dec = position(sca, i) + sources.append({ + "sid": i, **key, "pid": pid, "id": i, "field": FIELD, "hp6": 0, "hp9": 0, + "mjdobs": 60000.0 + epoch, "ra": ra, "dec": dec, + "xfit": float(i), "yfit": float(i), "xerr": 0.1, "yerr": 0.1, + "fluxfit": flux_scale * (100.0 + i + 10.0 * epoch), "fluxerr": 1.0, + "flags": 0, "qfit": 0.0, "cfit": 0.0, "npixfit": 9, "sharpness": 0.5, + "roundness1": 0.0, "roundness2": 0.0, "peak": 10.0}) + expected = [{"product": product, "key": key, "binding": CURRENT} + for product in ("diffimages", "sources")] + return self.job(run, "import", {"epoch": epoch, "sca": sca}, + {"diffimages": [diffimage], "sources": sources}, expects=expected) + + def import_epochs(self, run: int) -> dict[tuple[int, int], int]: + """Import every (epoch, SCA); return the source-set job per (epoch, SCA).""" + return {(e, s): self.import_image(run, e, s) for e in EPOCHS for s in SCAS} + + def associate(self, run: int, epoch: int, source_sets: list[int], predecessor: int | None, + excludes: list[int] = (), expects=CURRENT, published: bool = True) -> int: + """An object set: the predecessor's objects updated by ``source_sets``, minus ``excludes``.""" + objects = {} + if predecessor is not None: + for row in self.rows( + "SELECT aid, ra0, dec0, nsources, sumra, sumra2, sumdec, sumdec2, sumflux," + " sumflux2 FROM objects WHERE job = %s", (predecessor,)): + objects[row[0]] = dict(zip(_OBJECT_COLUMNS, row)) + for aid, ra, dec, flux in self._excluded_members(predecessor, excludes): + _accumulate(objects[aid], ra, dec, flux, -1) + memberships = [] + for source_job, sid, ra, dec, flux in self.rows( + "SELECT job, sid, ra, dec, fluxfit FROM sources WHERE job = ANY(%s)" + " AND flags = 0 ORDER BY job, sid", (list(source_sets),)): + aid = aid_of(ra, dec) + obj = objects.setdefault(aid, {"aid": aid, "ra0": ra, "dec0": dec, **_ZERO_SUMS}) + _accumulate(obj, ra, dec, flux, +1) + memberships.append({"aid": aid, "source_job": source_job, "sid": sid}) + rows = [{**obj, "dataset": DATASET, "survey": SURVEY, "field": FIELD} + for obj in objects.values()] + inputs = [{"product": "sources", "key": self._key_of(s), "source_job": s} + for s in source_sets] + if predecessor is not None: + inputs.append({"product": "objects", "key": FIELD_KEY, "source_job": predecessor}) + return self.job(run, "associate", {"field": FIELD, "epoch": epoch}, + {"objects": rows, "memberships": memberships}, inputs=inputs, + expects=[{"product": "objects", "key": FIELD_KEY, "binding": expects}], + excludes=[{"key": self._key_of(s), "source_job": s} for s in excludes], + published=published) + + def packets(self, run: int, epoch: int, sca: int, source_set: int, object_set: int, + expects=CURRENT, attempt: int = 1, published: bool = True) -> int: + """A packet set for one image key: the index and forced rows for its sources.""" + key = image_key(epoch, sca) + index, forced = [], [] + for n, (sid, pid, ra, dec, mjd, flux) in enumerate(self.rows( + "SELECT sid, pid, ra, dec, mjdobs, fluxfit FROM sources WHERE job = %s" + " ORDER BY sid", (source_set,))): + aid = aid_of(ra, dec) + index.append({**key, "object_job": object_set, "aid": aid, "source_job": source_set, + "sid": sid, "trigger_mjd": mjd, "band": "F184", + "schema_version": "00.06", "filename": f"packets-{epoch}-{sca}.avro", + "byte_offset": 1000 * n, "byte_length": 1000}) + forced.append({"forced_id": _forced_id(pid, aid, ra, dec), "aid": aid, "rid": 0, + "pid": pid, "expid": epoch, "sca": sca, "fid": FID, "band": "F184", + "mjdobs": mjd, "ra": ra, "dec": dec, "x": 0.0, "y": 0.0, + "psf_flux": flux, "psf_fluxerr": 1.0, "science_flux": flux, + "science_fluxerr": 1.0, "flags": 0, "time_proc": 0.0}) + inputs = [{"product": "objects", "key": FIELD_KEY, "source_job": object_set}, + {"product": "sources", "key": key, "source_job": source_set}] + return self.job(run, "packet", {"epoch": epoch, "sca": sca}, + {"packets": index, "forced": forced}, inputs=inputs, + expects=[{"product": "packets", "key": key, "binding": expects}], + attempt=attempt, published=published) + + def recipe(self, run: int, source_sets: dict[tuple[int, int], int], epochs=EPOCHS, + predecessor: int | None = None) -> dict: + """Associate and packets for ``epochs`` under ``run``, each epoch's set the next's predecessor. + + Returns ``{"objects": {epoch: job}, "packets": {(epoch, sca): job}}``. + """ + made = {"objects": {}, "packets": {}} + for epoch in epochs: + sets = [source_sets[(epoch, s)] for s in SCAS] + predecessor = self.associate(run, epoch, sets, predecessor) + made["objects"][epoch] = predecessor + for sca in SCAS: + made["packets"][(epoch, sca)] = self.packets( + run, epoch, sca, source_sets[(epoch, sca)], predecessor) + return made + + def cutover(self, to_run: int, made: dict) -> list[dict]: + """The cutover group: every key ``made`` names rebound to its job, and ``to_run`` live.""" + group = [{"close": self.binding("objects", FIELD_KEY), + "open": {"product": "objects", "key": FIELD_KEY, + "job": made["objects"][max(made["objects"])]}}] + for (epoch, sca), job in sorted(made["packets"].items()): + key = image_key(epoch, sca) + group.append({"close": self.binding("packets", key), + "open": {"product": "packets", "key": key, "job": job}}) + group.append({"close": self.live()[0], "open": {"survey": SURVEY, "dataset": DATASET, + "run": to_run}}) + return group + + # -- internals --------------------------------------------------------- + + def _key_of(self, source_set: int) -> dict: + (key,) = self.rows("SELECT DISTINCT jsonb_build_object('dataset', dataset, 'expid', expid," + " 'sca', sca, 'fid', fid) FROM sources WHERE job = %s", (source_set,))[0] + return key + + def _excluded_members(self, predecessor: int | None, excludes) -> list[tuple]: + """The members of the predecessor's chain whose source set is excluded, with their values.""" + if predecessor is None or not excludes: + return [] + return self.rows(""" + WITH RECURSIVE chain (job) AS ( + SELECT %s::bigint + UNION ALL + SELECT i.source_job FROM chain c + JOIN job_inputs i ON i.job = c.job AND i.kind = 'input' AND i.product = 'objects' + ) + SELECT m.aid, s.ra, s.dec, s.fluxfit + FROM chain c + JOIN memberships m ON m.job = c.job + JOIN sources s ON s.job = m.source_job AND s.sid = m.sid + WHERE m.source_job = ANY(%s)""", (predecessor, list(excludes))) + + +_OBJECT_COLUMNS = ("aid", "ra0", "dec0", "nsources", "sumra", "sumra2", "sumdec", "sumdec2", + "sumflux", "sumflux2") +_ZERO_SUMS = {"nsources": 0, "sumra": 0.0, "sumra2": 0.0, "sumdec": 0.0, "sumdec2": 0.0, + "sumflux": 0.0, "sumflux2": 0.0} + + +def _accumulate(obj: dict, ra: float, dec: float, flux: float, sign: int) -> None: + obj["nsources"] += sign + for name, value in (("ra", ra), ("dec", dec), ("flux", flux)): + obj[f"sum{name}"] += sign * value + obj[f"sum{name}2"] += sign * value * value + + +def _forced_id(pid: int, aid: int, ra: float, dec: float) -> int: + """A stable 63-bit id from image, object and position.""" + digest = hashlib.sha256(f"{pid}:{aid}:{ra:.9f}:{dec:.9f}".encode()).digest() + return int.from_bytes(digest[:8], "big") >> 1 diff --git a/tests/db/proof/test_steps.py b/tests/db/proof/test_steps.py new file mode 100644 index 000000000..9c758a637 --- /dev/null +++ b/tests/db/proof/test_steps.py @@ -0,0 +1,271 @@ +"""The procedure proof, milestone 0 (foundation plan v5, Part 3, steps 2, 4, 6, 8, 10, 11, 12, 13). + +One test per step, each asserting the step's observation as stated, read +back from the database through the official views and the ``packets`` +table. Steps 1, 3, 5, 7 and 9 need the real stages and are milestone 3. +Run by foundation.yml against database/schema.sql; db-migrations.yml +ignores this directory. +""" + +from __future__ import annotations + +import importlib +import shutil +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from rapidpipe.foundation import Refused, apply, change, job_show, publish, undo +from tests.db.foundation_helpers import wait_until_blocked +from tests.db.proof.fixture import ( + DATASET, EPOCHS, FIELD_KEY, NSOURCES, SCAS, SURVEY, World, image_key, +) + + +@pytest.fixture() +def world(connect, tmp_path): + return World(connect(), tmp_path) + + +def _official_jobs(world, view): + return {r[0] for r in world.rows(f"SELECT DISTINCT job FROM {view}")} + + +def _official_table(world): + return world.rows("SELECT * FROM official ORDER BY id") + + +def _two_recipes(world): + """Run 1 live through both epochs; run 2, not live, over the same imports.""" + first = world.run("first recipe", live=True) + sets = world.import_epochs(first) + made_first = world.recipe(first, sets) + second = world.run("second recipe", {"recipe": "second", "radius": 2.0}) + made_second = world.recipe(second, sets) + return sets, first, made_first, second, made_second + + +def test_step_2_two_candidates_from_one_predecessor(world): + run = world.run("recipe", live=True) + sets = world.import_epochs(run) + predecessor = world.recipe(run, sets, epochs=(1,))["objects"][1] + expected = world.binding("objects", FIELD_KEY) + epoch2 = [sets[(2, s)] for s in SCAS] + + first, second = (world.associate(run, 2, epoch2, predecessor, published=False) + for _ in range(2)) + for job in (first, second): + publish(world.conn, job, 1, world.prefix(job, 1)) + + assert _official_jobs(world, "objects_official") == {first} + assert world.rows("SELECT count(*) FROM objects WHERE job = %s", (second,))[0][0] > 0 + shown = job_show(world.conn, second) + assert (shown["state"], shown["refusal"]) == ( + "published", f"binding {expected} is no longer open") + assert {"kind": "input", "product": "objects", "key": FIELD_KEY, + "source_job": predecessor, "binding": None} in shown["inputs"] + assert {"kind": "expects", "product": "objects", "key": FIELD_KEY, + "source_job": None, "binding": expected} in shown["inputs"] + assert shown["bindings"] == [] + + +def test_step_4_cutover_is_one_group(world): + sets, first, made_first, second, made_second = _two_recipes(world) + assert _official_jobs(world, "packets_official") == set(made_first["packets"].values()) + + group = change(world.conn, world.cutover(second, made_second), "proof", "cutover") + + opened = world.rows("SELECT product, job, run FROM official WHERE opened_group = %s", (group,)) + closed = world.rows("SELECT product, job, run FROM official WHERE closed_group = %s", (group,)) + assert len(opened) == len(closed) == 1 + len(made_second["packets"]) + 1 + assert world.live()[1] == second + assert _official_jobs(world, "packets_official") == set(made_second["packets"].values()) + assert _official_jobs(world, "objects_official") == {made_second["objects"][2]} + for packet_job, object_job, source_job in world.rows( + "SELECT DISTINCT job, object_job, source_job FROM packets_official"): + inputs = {(i["product"], i["source_job"]) + for i in job_show(world.conn, packet_job)["inputs"] if i["kind"] == "input"} + assert inputs == {("objects", object_job), ("sources", source_job)} + assert object_job in made_second["objects"].values() + + +def test_step_6_undo_restores_the_first_recipe(world): + sets, first, made_first, second, made_second = _two_recipes(world) + cutover = change(world.conn, world.cutover(second, made_second), "proof", "cutover") + replaced = {r[0] for r in world.rows( + "SELECT id FROM official WHERE closed_group = %s", (cutover,))} + + undo(world.conn, cutover, "proof", "undo the cutover") + + assert world.live()[1] == first + assert _official_jobs(world, "packets_official") == set(made_first["packets"].values()) + assert _official_jobs(world, "objects_official") == {made_first["objects"][2]} + open_ids = {r[0] for r in world.rows("SELECT id FROM official WHERE closed_at IS NULL")} + assert not replaced & open_ids + second_jobs = list(made_second["packets"].values()) + assert world.rows("SELECT count(*) FROM packets WHERE job = ANY(%s)", (second_jobs,))[0][0] \ + == NSOURCES * len(second_jobs) + + +class Killed(BaseException): + """The execution was killed.""" + + +def test_step_8_an_interrupted_publication(world, monkeypatch): + run = world.run("recipe", live=True) + sets = world.import_epochs(run) + object_set = world.recipe(run, sets, epochs=(1,))["objects"][1] + job = world.packets(run, 2, 1, sets[(2, 1)], object_set, published=False) + publisher = importlib.import_module("rapidpipe.foundation.publish") + loaded_when_killed = [] + + def killed(cur, job_id, run_id): + cur.execute("SELECT count(*) FROM packets WHERE job = %s", (job_id,)) + loaded_when_killed.append(cur.fetchone()[0]) + raise Killed + + def count(): + return world.rows("SELECT (SELECT count(*) FROM packets WHERE job = %(j)s)," + " (SELECT count(*) FROM forced WHERE job = %(j)s)", {"j": job})[0] + + with monkeypatch.context() as patch: + patch.setattr(publisher, "_bind", killed) + with pytest.raises(Killed): + publish(world.conn, job, 1, world.prefix(job, 1)) + assert loaded_when_killed == [NSOURCES] + assert count() == (0, 0) + assert job_show(world.conn, job)["state"] == "launched" + + shutil.copytree(world.root / f"job-{job}" / "attempt-1", world.root / f"job-{job}" / "attempt-2") + retry = publish(world.conn, job, 2, world.prefix(job, 2)) + again = publish(world.conn, job, 3, world.prefix(job, 3)) + + assert retry["published"] and not again["published"] + assert count() == (NSOURCES, NSOURCES) + assert job_show(world.conn, job)["published_attempt"] == 2 + assert world.rows("SELECT job FROM official WHERE id = %s", + (world.binding("packets", image_key(2, 1)),)) == [(job,)] + + +def test_step_10_two_change_groups_at_once(world, connect): + run = world.run("recipe", live=True) + sets = world.import_epochs(run) + world.recipe(run, sets, epochs=(1,)) + contested = world.binding("objects", FIELD_KEY) + bystander = world.binding("diffimages", image_key(1, 1)) + assert bystander < contested + campaign = world.run("campaign") + one, other = (world.associate(campaign, 1, [sets[(1, s)] for s in SCAS], None) + for _ in range(2)) + a, b = world.conn, connect() + + with a.cursor() as cur: + apply(cur, [{"close": contested, + "open": {"product": "objects", "key": FIELD_KEY, "job": one}}], "a", "first") + with ThreadPoolExecutor(1) as pool: + pending = pool.submit(change, b, [ + {"close": bystander}, + {"close": contested, "open": {"product": "objects", "key": FIELD_KEY, "job": other}}, + ], "b", "second") + wait_until_blocked(connect(), b.get_backend_pid()) + a.commit() + with pytest.raises(Refused) as refused: + pending.result(timeout=10) + + assert refused.value.binding == contested + assert f"binding {contested}" in str(refused.value) + assert _official_jobs(world, "objects_official") == {one} + assert world.binding("diffimages", image_key(1, 1)) == bystander + assert world.rows("SELECT count(*) FROM official WHERE job = %s OR closed_by = 'b'" + " OR opened_by = 'b'", (other,))[0][0] == 0 + + +def test_step_11_publication_racing_a_cutover(world, connect): + sets, first, made_first, second, made_second = _two_recipes(world) + key = image_key(2, 1) + expected = world.binding("packets", key) + late = world.packets(first, 2, 1, sets[(2, 1)], made_first["objects"][2], published=False) + a, b = world.conn, connect() + + with a.cursor() as cur: + group = apply(cur, world.cutover(second, made_second), "a", "cutover") + with ThreadPoolExecutor(1) as pool: + pending = pool.submit(publish, b, late, 1, world.prefix(late, 1)) + wait_until_blocked(connect(), b.get_backend_pid()) + a.commit() + result = pending.result(timeout=10) + + assert result["published"] and result["group"] is None + assert result["refusal"] == f"binding {expected} is no longer open" + assert world.rows("SELECT count(*) FROM packets WHERE job = %s", (late,))[0][0] == NSOURCES + assert late not in _official_jobs(world, "packets_official") + packet_bindings = world.rows("SELECT key, job, opened_group FROM official" + " WHERE closed_at IS NULL AND product = 'packets'") + assert sorted((k["expid"], k["sca"]) for k, _, _ in packet_bindings) == \ + sorted(made_second["packets"]) + assert {(j, g) for _, j, g in packet_bindings} == \ + {(j, group) for j in made_second["packets"].values()} + assert world.live()[1] == second + + +def test_step_12_undo_after_progress(world): + first = world.run("first recipe", live=True) + sets = world.import_epochs(first) + made_first = world.recipe(first, sets, epochs=(1,)) + second = world.run("second recipe", {"recipe": "second"}) + made_second = world.recipe(second, sets, epochs=(1,)) + cutover = change(world.conn, world.cutover(second, made_second), "proof", "cutover") + moved = world.binding("objects", FIELD_KEY) + progress = world.recipe(second, sets, epochs=(2,), predecessor=made_second["objects"][1]) + assert _official_jobs(world, "objects_official") == {progress["objects"][2]} + before = _official_table(world) + + with pytest.raises(Refused) as refused: + undo(world.conn, cutover, "proof", "undo the cutover") + + assert refused.value.binding == moved + assert _official_table(world) == before + + still = [{"close": world.binding("packets", image_key(1, s)), + "open": {"product": "packets", "key": image_key(1, s), + "job": made_first["packets"][(1, s)]}} for s in SCAS] + still.append({"close": world.live()[0], + "open": {"survey": SURVEY, "dataset": DATASET, "run": first}}) + change(world.conn, still, "proof", "restore the keys still as expected") + + assert {(e, j) for e, j in world.rows("SELECT DISTINCT expid, job FROM packets_official")} == \ + {(1, made_first["packets"][(1, s)]) for s in SCAS} | \ + {(2, progress["packets"][(2, s)]) for s in SCAS} + assert _official_jobs(world, "objects_official") == {progress["objects"][2]} + assert world.live()[1] == first + + +def test_step_13_a_superseded_source_set(world): + run = world.run("recipe", live=True) + sets = world.import_epochs(run) + set_1 = world.recipe(run, sets, epochs=(1,))["objects"][1] + superseded = sets[(1, 1)] + replacement = world.import_image(run, 1, 1, flux_scale=2.0) + assert _official_jobs(world, "sources_official") >= {replacement} + assert superseded not in _official_jobs(world, "sources_official") + + set_2 = world.associate(run, 2, [sets[(2, s)] for s in SCAS] + [replacement], set_1, + excludes=[superseded]) + + assert _official_jobs(world, "objects_official") == {set_2} + members = world.rows("SELECT DISTINCT source_job FROM members_official") + assert superseded not in {r[0] for r in members} + sums = world.rows(""" + SELECT o.aid, o.nsources, o.sumflux, count(s.sid), coalesce(sum(s.fluxfit), 0) + FROM objects_official o + LEFT JOIN members_official m ON m.set_job = o.job AND m.aid = o.aid + LEFT JOIN sources s ON s.job = m.source_job AND s.sid = m.sid + GROUP BY o.aid, o.nsources, o.sumflux""") + assert len(sums) == NSOURCES * len(SCAS) + for _, nsources, sumflux, members_count, members_flux in sums: + assert nsources == members_count == len(EPOCHS) + assert sumflux == pytest.approx(members_flux) + assert world.rows("SELECT count(*) FROM sources WHERE job = %s", (superseded,))[0][0] \ + == NSOURCES + assert world.rows("SELECT count(*) FROM memberships WHERE job = %s AND source_job = %s", + (set_1, superseded))[0][0] == NSOURCES diff --git a/tests/db/proof/timing.py b/tests/db/proof/timing.py new file mode 100644 index 000000000..c753a2559 --- /dev/null +++ b/tests/db/proof/timing.py @@ -0,0 +1,60 @@ +"""The wall time of one change group over N keys (foundation plan v5, Part 4, milestone 0). + +A script, not a test; CI does not run it. Against a database that +``database/schema.sql`` was applied to (it adds rows and deletes none, so +point it at a throwaway database), using the ``PG*`` environment: + + python tests/db/proof/timing.py 10000 + +builds, by SQL, two jobs, N ``catalogs`` rows for each and N open bindings +to the first job's rows, analyses them, then applies one change group that closes every +binding and opens the second job's row on the same key, through +:func:`rapidpipe.foundation.change`, and prints N and the group's wall time. +""" + +from __future__ import annotations + +import os +import sys +import time + +import psycopg2 + +from rapidpipe.foundation import change, create_job, create_run + + +def main(n: int) -> None: + conn = psycopg2.connect(host=os.environ["PGHOST"], port=os.environ.get("PGPORT", "5432"), + dbname=os.environ["PGDATABASE"], user=os.environ["PGUSER"], + password=os.environ.get("PGPASSWORD", "")) + run = create_run(conn, "timing", f"change group over {n} keys", "synthetic", {}, "T", "T") + old, new = (create_job(conn, run, "timing", {"n": n}) for _ in range(2)) + tag = f"timing-{old}" + with conn.cursor() as cur: + cur.execute("INSERT INTO catalogs (job, name, version, location)" + " SELECT j, %s || g, 'v1', 'file:///dev/null'" + " FROM generate_series(1, %s) g, unnest(ARRAY[%s, %s]::bigint[]) j", + (tag, n, old, new)) + cur.execute("INSERT INTO official (product, key, job, opened_group, opened_by, opened_why)" + " SELECT 'catalogs', jsonb_build_object('name', %s || g, 'version', 'v1')," + " %s, nextval('change_groups'), 'timing', 'setup'" + " FROM generate_series(1, %s) g RETURNING id, key", (tag, old, n)) + bindings = cur.fetchall() + # Bulk-loaded tables have no statistics until autovacuum reaches them; without + # them the planner answers the open-binding check by (product, job) and the + # group is quadratic in N. Measure the analysed state a live table is in. + cur.execute("ANALYZE official") + cur.execute("ANALYZE catalogs") + conn.commit() + group = [{"close": binding, "open": {"product": "catalogs", "key": key, "job": new}} + for binding, key in bindings] + + start = time.monotonic() + change(conn, group, "timing", f"one group over {n} keys") + elapsed = time.monotonic() - start + print(f"N={n} change_group_seconds={elapsed:.2f}") + conn.close() + + +if __name__ == "__main__": + main(int(sys.argv[1]))