From 218146988e83e0aca6fce546006c11d72419d3a6 Mon Sep 17 00:00:00 2001 From: Zoheb Shaikh <26975142+ZohebShaikh@users.noreply.github.com> Date: Fri, 16 Oct 2026 15:42:38 +0100 Subject: [PATCH 1/8] Let the policy read beamline and proposal tags as well as session ones Beamline and proposal nodes are created with only some of the access blob fields, and existing nodes carry composite tags like "beamline:i22,proposal:12345" rather than a JSON blob. Both now become OPA input instead of failing validation. Co-Authored-By: Claude Opus 5.5 --- src/dls_tiled/policy.py | 49 ++++++++++++++++++++++++++++---- tests/test_policy.py | 63 +++++++++++++++++++++++++++++++++++++---- 2 files changed, 101 insertions(+), 11 deletions(-) diff --git a/src/dls_tiled/policy.py b/src/dls_tiled/policy.py index b27fc3e..9cdc6be 100644 --- a/src/dls_tiled/policy.py +++ b/src/dls_tiled/policy.py @@ -8,7 +8,7 @@ import logging from fastapi import HTTPException -from pydantic import BaseModel, HttpUrl, TypeAdapter +from pydantic import BaseModel, HttpUrl, TypeAdapter, ValidationError from starlette.status import ( HTTP_401_UNAUTHORIZED, ) @@ -27,9 +27,40 @@ class DiamondAccessBlob(BaseModel): - proposal: int - visit: int + """What blueapi asks for when creating a node at each level of the tree. + + beamline alone for a beamline node, plus proposal for a proposal node, + plus visit for a session node. + """ + beamline: str + # The full proposal code, e.g. "cm12345" - tiled.rego strips the leading + # letters itself where it needs the bare number. + proposal: str | None = None + visit: int | None = None + + +# Maps a composite access tag's "key" (as produced by tiled.rego's +# beamline_tag/proposal_tag/session_tag, e.g. +# "beamline:i22,proposal:12345,session:12345-1") to the corresponding OPA +# input field. All of these are strings on an existing node's tag - "session" +# here is the "12345-1" proposal-visit pair, not the internal numeric session +# id used by modify_session. +_TAG_KEY_TO_INPUT_FIELD = { + "beamline": "beamline", + "proposal": "proposal", + "session": "session", +} + + +def _parse_composite_tag(tag: str) -> dict[str, str]: + fields: dict[str, str] = {} + for part in tag.split(","): + key, _, value = part.partition(":") + field = _TAG_KEY_TO_INPUT_FIELD.get(key) + if field is not None: + fields[field] = value + return fields def _check_principal(principal: Principal | None): @@ -137,9 +168,17 @@ def build_input( and "tags" in access_blob and len(access_blob["tags"]) > 0 ): - blob = self._type_adapter.validate_json(access_blob["tags"][0]) + tag = access_blob["tags"][0] + try: + blob = self._type_adapter.validate_json(tag) + except ValidationError: + # Not a create-request JSON blob - it's a composite tag + # already assigned to an existing node (e.g. tiled checking + # scopes on a node before reading it). + blob = None + _input.update(_parse_composite_tag(tag)) if isinstance(blob, DiamondAccessBlob): - _input.update(blob.model_dump()) + _input.update(blob.model_dump(exclude_none=True)) elif isinstance(blob, int): _input["session"] = str(blob) diff --git a/tests/test_policy.py b/tests/test_policy.py index d50ff90..4ae480a 100644 --- a/tests/test_policy.py +++ b/tests/test_policy.py @@ -6,7 +6,7 @@ import httpx import pytest from fastapi import HTTPException -from pydantic import HttpUrl, SecretStr, ValidationError +from pydantic import HttpUrl, SecretStr from tiled.access_control import access_policies from tiled.access_control.access_policies import ALL_ACCESS, NO_ACCESS from tiled.config import parse_configs @@ -20,8 +20,8 @@ OPA = "http://opa.example:8181/v1/data/" AUDIENCE = "tiled-audience" -BLOB_TAG = '{"proposal": 1234, "visit": 5, "beamline": "i22"}' -BLOB_FIELDS = {"proposal": 1234, "visit": 5, "beamline": "i22"} +BLOB_TAG = '{"proposal": "cm1234", "visit": 5, "beamline": "i22"}' +BLOB_FIELDS = {"proposal": "cm1234", "visit": 5, "beamline": "i22"} # async tests run on anyio's pytest plugin, which ships with tiled's httpx pytestmark = pytest.mark.anyio @@ -89,6 +89,23 @@ def test_build_input_with_diamond_access_blob(policy, user): assert body["input"] == {"audience": AUDIENCE, "token": "the-token", **BLOB_FIELDS} +@pytest.mark.parametrize( + "tag, fields", + [ + ('{"beamline": "i22"}', {"beamline": "i22"}), + ( + '{"beamline": "i22", "proposal": "cm1234"}', + {"beamline": "i22", "proposal": "cm1234"}, + ), + ], +) +def test_build_input_with_container_blob(policy, user, tag, fields): + # Beamline and proposal containers have no visit; absent fields aren't sent. + body = json.loads(policy.build_input(user, None, set(), {"tags": [tag]})) + + assert body["input"] == {"audience": AUDIENCE, "token": "the-token", **fields} + + def test_build_input_with_session_int(policy, user): body = json.loads(policy.build_input(user, None, set(), {"tags": ["98765"]})) @@ -115,9 +132,43 @@ def test_build_input_ignores_blob_without_tags(policy, user, blob): assert body["input"] == {"audience": AUDIENCE, "token": "the-token"} -def test_build_input_rejects_unparseable_tag(policy, user): - with pytest.raises(ValidationError): - policy.build_input(user, None, set(), {"tags": ['{"proposal": 1}']}) +@pytest.mark.parametrize( + "tag, fields", + [ + ("beamline:i22", {"beamline": "i22"}), + ("beamline:i22,proposal:12345", {"beamline": "i22", "proposal": "12345"}), + ( + "beamline:i22,proposal:12345,session:12345-1", + {"beamline": "i22", "proposal": "12345", "session": "12345-1"}, + ), + ("beamline:i22,unknown:x", {"beamline": "i22"}), + ], +) +def test_build_input_with_composite_tag(policy, user, tag, fields): + body = json.loads(policy.build_input(user, None, set(), {"tags": [tag]})) + + assert body["input"] == {"audience": AUDIENCE, "token": "the-token", **fields} + + +def test_build_input_ignores_unparseable_tag(policy, user): + blob = {"tags": ['{"proposal": 1}']} + + body = json.loads(policy.build_input(user, None, set(), blob)) + + assert body["input"] == {"audience": AUDIENCE, "token": "the-token"} + + +async def test_allowed_scopes_on_migrated_node(policy, user, opa): + # tiled sends a node's own tag when reading it; this used to raise. + opa.responses[f"{OPA}tiled/scopes"] = {"result": ["read:metadata"]} + node = SimpleNamespace( + access_blob={"tags": ["beamline:i22,proposal:12345,session:12345-1"]} + ) + + scopes = await policy.allowed_scopes(node, user, None, set()) + + assert scopes == {"read:metadata"} + assert opa.requests[0][1]["input"]["session"] == "12345-1" def test_build_input_omits_token_when_absent(policy): From 388b739e2c099b484ae94333e39b0ba531cc69ae Mon Sep 17 00:00:00 2001 From: Zoheb Shaikh <26975142+ZohebShaikh@users.noreply.github.com> Date: Fri, 16 Oct 2026 15:42:39 +0100 Subject: [PATCH 2/8] Add a migration that moves runs into the beamline/proposal/session tree Runs at the catalog root move under new beamline, proposal and session containers, and are retagged with session tags so the new policy can see them. Includes postgres backup, restore and read-only role tooling, and an integration test that restores a production dump and reads it back through dls-tiled. Co-Authored-By: Claude Opus 5.5 --- database_migration/README.md | 171 +++++++++ database_migration/integration/.gitignore | 2 + database_migration/integration/compose.yaml | 60 ++++ .../integration/opa/config.yaml | 8 + .../integration/opa/entrypoint.sh | 4 + .../integration/tiled_config/config.yml | 23 ++ database_migration/migrate_node_structure.py | 251 +++++++++++++ database_migration/postgres_tools.py | 141 ++++++++ database_migration/test_integration.py | 332 ++++++++++++++++++ .../test_migrate_node_structure.py | 25 ++ 10 files changed, 1017 insertions(+) create mode 100644 database_migration/README.md create mode 100644 database_migration/integration/.gitignore create mode 100644 database_migration/integration/compose.yaml create mode 100644 database_migration/integration/opa/config.yaml create mode 100644 database_migration/integration/opa/entrypoint.sh create mode 100644 database_migration/integration/tiled_config/config.yml create mode 100644 database_migration/migrate_node_structure.py create mode 100644 database_migration/postgres_tools.py create mode 100644 database_migration/test_integration.py create mode 100644 database_migration/test_migrate_node_structure.py diff --git a/database_migration/README.md b/database_migration/README.md new file mode 100644 index 0000000..49613b1 --- /dev/null +++ b/database_migration/README.md @@ -0,0 +1,171 @@ +# Database migration + +Commands to add a read-only role to the tiled Postgres database running in +Kubernetes, back up every database as that role, and move runs into the +`beamline/proposal/session` layout. + +Set these once in your shell: + +```bash +NS= +POD= # find it with: kubectl get pods -n $NS +``` + +## Passwords + +Postgres tools have no password option on the command line. `-W` only forces +a prompt, and a prompt can't be answered through `kubectl exec` unless it has +`-it`. `-t` must never be used when writing a dump to a file: it turns +newlines into `\r\n` and corrupts the dump. + +So the commands below pass the password in the `PGPASSWORD` environment +variable. Read it once without it showing on screen: + +```bash +read -rs -p "Password: " PGPASSWORD; echo +``` + +It ends up on the `kubectl exec` command line, so it is briefly visible in the +process list on your machine and in the pod. Run `unset PGPASSWORD` when +you're done. + +## 1. Create the read-only role (as postgres) + +This step needs the `postgres` superuser, once. Open `psql` in the pod; it +asks for the postgres password itself: + +```bash +kubectl exec -it -n $NS $POD -- psql -U postgres +``` + +Inside `psql`, create the role and set its password (`\password` asks without +showing it): + +```sql +CREATE ROLE tiled_readonly LOGIN NOSUPERUSER NOCREATEDB NOCREATEROLE NOREPLICATION; +\password tiled_readonly +ALTER ROLE tiled_readonly SET default_transaction_read_only = on; +``` + +Find which role owns tiled's tables, needed below: + +```sql +\c tiled +SELECT DISTINCT tableowner FROM pg_tables WHERE schemaname = 'public'; +``` + +Then, still connected to `tiled`, grant read access. Replace `` with +the owner from the query above: + +```sql +GRANT CONNECT ON DATABASE tiled TO tiled_readonly; +GRANT USAGE ON SCHEMA public TO tiled_readonly; +GRANT SELECT ON ALL TABLES IN SCHEMA public TO tiled_readonly; +GRANT SELECT ON ALL SEQUENCES IN SCHEMA public TO tiled_readonly; +-- tables created later (e.g. by a tiled schema upgrade) are readable too +ALTER DEFAULT PRIVILEGES FOR ROLE IN SCHEMA public GRANT SELECT ON TABLES TO tiled_readonly; +ALTER DEFAULT PRIVILEGES FOR ROLE IN SCHEMA public GRANT SELECT ON SEQUENCES TO tiled_readonly; +\q +``` + +Repeat the `\c ` and grant block for every other database you want +backed up. + +Why not `GRANT pg_read_all_data`: it also lets the role read `pg_authid`, the +stored password hash of every role including `postgres`. The grants above +don't. + +### Check it can read but not write + +Read the `tiled_readonly` password into `PGPASSWORD` (see +[Passwords](#passwords)), then: + +```bash +kubectl exec -n $NS $POD -- env PGPASSWORD=$PGPASSWORD psql -U tiled_readonly -d tiled -c "SELECT count(*) FROM nodes" +kubectl exec -n $NS $POD -- env PGPASSWORD=$PGPASSWORD psql -U tiled_readonly -d tiled -c "DELETE FROM nodes WHERE false" +kubectl exec -n $NS $POD -- env PGPASSWORD=$PGPASSWORD psql -U tiled_readonly -d postgres -c "SELECT count(*) FROM pg_authid" +``` + +The first prints a count. The second must fail with +`cannot execute DELETE in a read-only transaction`. The third must fail with +`permission denied for table pg_authid`. + +## 2. Back up every database (as tiled_readonly) + +With the `tiled_readonly` password in `PGPASSWORD`: + +```bash +mkdir -p backup + +# Roles and tablespaces (pg_dump doesn't include them). --no-role-passwords +# because tiled_readonly can't read password hashes; restored roles get no +# password and need one set again. +kubectl exec -n $NS $POD -- env PGPASSWORD=$PGPASSWORD \ + pg_dumpall -U tiled_readonly -l postgres --globals-only --no-role-passwords > backup/globals.sql + +# List the databases +kubectl exec -n $NS $POD -- env PGPASSWORD=$PGPASSWORD psql -U tiled_readonly -d postgres -AtXc \ + "SELECT datname FROM pg_database WHERE datallowconn AND NOT datistemplate" + +# Dump each database listed above, one command per database +kubectl exec -n $NS $POD -- env PGPASSWORD=$PGPASSWORD \ + pg_dump -U tiled_readonly --format=custom -d tiled > backup/tiled.dump +``` + +A database the role wasn't granted in step 1 fails with +`permission denied`; grant it or skip it. + +Check the files aren't empty before going further: + +```bash +ls -lh backup +``` + +## 3. Move runs into the new layout + +`migrate_node_structure.py` writes to the database, so it runs as the tiled +owner or `postgres`, not `tiled_readonly`. It connects from your machine, so +forward the Postgres port first. Leave this running in its own terminal: + +```bash +kubectl port-forward -n $NS pod/$POD 5432:5432 +``` + +In another terminal, from a Python environment with `tiled[server]` installed: + +```bash +URI=postgresql://postgres:@localhost:5432/tiled + +# Preview: prints what would move, changes nothing +python migrate_node_structure.py $URI --dry-run + +# Upgrade the tiled schema (to --revision, default head), then move the runs +python migrate_node_structure.py $URI --revision head +``` + +The dry run only works if the database is already at tiled schema revision +`e05e918092c3` or later; otherwise it says so, and the real run upgrades it +first. Running the migration twice is safe: the second run moves nothing. + +## Restore from the backup (as postgres) + +With the `postgres` password in `PGPASSWORD`: + +```bash +# Roles, only when restoring onto a fresh server +kubectl exec -i -n $NS $POD -- env PGPASSWORD=$PGPASSWORD psql -U postgres -d postgres < backup/globals.sql + +# Stop tiled first: on startup it creates the schema in an empty database, +# and the restore then fails with "already exists" errors +kubectl scale -n $NS deploy/ --replicas=0 + +# Drop the broken database first if it still exists: +# kubectl exec -n $NS $POD -- env PGPASSWORD=$PGPASSWORD dropdb -U postgres --force tiled +kubectl exec -n $NS $POD -- env PGPASSWORD=$PGPASSWORD createdb -U postgres tiled +kubectl exec -i -n $NS $POD -- env PGPASSWORD=$PGPASSWORD pg_restore -U postgres -d tiled --exit-on-error < backup/tiled.dump + +kubectl scale -n $NS deploy/ --replicas=1 +``` + +If the restore fails partway, drop the database and start again; don't try +to patch a partial restore. diff --git a/database_migration/integration/.gitignore b/database_migration/integration/.gitignore new file mode 100644 index 0000000..3323785 --- /dev/null +++ b/database_migration/integration/.gitignore @@ -0,0 +1,2 @@ +# Written by test_integration.py from the restored dump +opa/opa_data/ diff --git a/database_migration/integration/compose.yaml b/database_migration/integration/compose.yaml new file mode 100644 index 0000000..5c63b15 --- /dev/null +++ b/database_migration/integration/compose.yaml @@ -0,0 +1,60 @@ +# Stack for test_integration.py: a production dump is restored into postgres, +# migrated, then read back through dls-tiled with real Keycloak tokens and the +# Diamond OPA policy. Ports: postgres 5433, keycloak 8081, opa 8181, tiled 8407. +# Modelled on blueapi's tests/system_tests/compose.yaml. + +services: + postgres: + image: docker.io/postgres:16 + environment: + - POSTGRES_PASSWORD=postgres + ports: + - "5433:5432" + healthcheck: + test: pg_isready -U postgres + interval: 2s + retries: 30 + security_opt: + - label=disable + + keycloak: + image: keycloak/keycloak:26.4 + environment: + - KC_BOOTSTRAP_ADMIN_PASSWORD=admin + - KC_BOOTSTRAP_ADMIN_USERNAME=admin + - KC_HOSTNAME=http://localhost:8081 + command: ["start-dev"] + ports: + - "8081:8080" + healthcheck: + test: /opt/keycloak/bin/kcadm.sh config credentials --server http://keycloak:8080 --realm master --user admin --password admin + interval: 5s + timeout: 5s + retries: 20 + start_period: 30s + security_opt: + - label=disable + + # Started by the test once it has written opa_data/data.json from the dump. + opa: + image: openpolicyagent/opa:edge-static-debug + network_mode: host # Port 8181 + volumes: + - ./opa:/mnt + environment: + - ISSUER=http://localhost:8081/realms/master + entrypoint: "sh /mnt/entrypoint.sh" + security_opt: + - label=disable + + # dls-tiled's own image, so its access policy is what gets tested. + tiled: + build: + context: ../.. + target: runtime + network_mode: host # Port 8407 + volumes: + - ./tiled_config:/deploy/config:ro + command: ["tiled", "serve", "config", "/deploy/config", "--host", "0.0.0.0", "--port", "8407"] + security_opt: + - label=disable diff --git a/database_migration/integration/opa/config.yaml b/database_migration/integration/opa/config.yaml new file mode 100644 index 0000000..9c2668b --- /dev/null +++ b/database_migration/integration/opa/config.yaml @@ -0,0 +1,8 @@ +services: + ghcr: + url: https://ghcr.io + type: oci +bundles: + diamond-policies: + service: ghcr + resource: ghcr.io/zohebshaikh/authz-policy:0.0.25-alpha diff --git a/database_migration/integration/opa/entrypoint.sh b/database_migration/integration/opa/entrypoint.sh new file mode 100644 index 0000000..7b6f476 --- /dev/null +++ b/database_migration/integration/opa/entrypoint.sh @@ -0,0 +1,4 @@ +#!/bin/sh +# Same as blueapi's system tests: local data bundle + the published policy bundle. +/opa build -b /mnt/opa_data -o /tmp/bundle.tar.gz +/opa run --server --addr localhost:8181 -b /tmp/bundle.tar.gz --config-file /mnt/config.yaml diff --git a/database_migration/integration/tiled_config/config.yml b/database_migration/integration/tiled_config/config.yml new file mode 100644 index 0000000..d20e099 --- /dev/null +++ b/database_migration/integration/tiled_config/config.yml @@ -0,0 +1,23 @@ +authentication: + providers: + - provider: keycloak_oidc + authenticator: tiled.authenticators:ProxiedOIDCAuthenticator + args: + audience: account + client_id: tiled + device_flow_client_id: tiled-cli + # Keycloak tokens have no "scp" claim, so this grants these to all. + scopes_map: + read: ["read:metadata", "read:data"] + well_known_uri: "http://localhost:8081/realms/master/.well-known/openid-configuration" +trees: + - path: / + tree: catalog + args: + uri: "postgresql://postgres:postgres@localhost:5433/tiled" + writable_storage: "/tmp/data" +access_control: + access_policy: "dls_tiled.policy:DiamondOpenPolicyAgentAuthorizationPolicy" + args: + authorization_provider: "http://localhost:8181/v1/data/diamond/policy/" + token_audience: account diff --git a/database_migration/migrate_node_structure.py b/database_migration/migrate_node_structure.py new file mode 100644 index 0000000..de475c4 --- /dev/null +++ b/database_migration/migrate_node_structure.py @@ -0,0 +1,251 @@ +"""Move runs from the flat catalog layout into the beamline/proposal/session tree. + +Before: root/ +After: root//// + e.g. root/i22/cm12345/cm12345-1/ + +This is the layout blueapi's begin_task now writes into. The tiled schema is +upgraded first (to --revision, default head) so the move runs against the +nodes/nodes_closure layout this script expects. + +The new containers get the same composite access tags tiled.rego stamps on +nodes it creates, and every tagged node under a moved run is retagged with +its session tag. Runs written under the old policy carry a bare numeric +session id, which the new policy's user_sessions never returns, so they +would become invisible to users otherwise. + +Runs at the root without start.instrument / start.instrument_session, or whose +instrument_session doesn't parse, are left where they are and listed. Safe to +rerun: anything already moved is no longer a child of the root. + +Postgres only. + +Example: + python database_migration/migrate_node_structure.py \\ + postgresql://user:pass@host:5432/tiled --revision head --dry-run +""" + +import argparse +import asyncio +import json +import re + +from sqlalchemy import text +from sqlalchemy.ext.asyncio import create_async_engine +from tiled.alembic_utils import get_current_revision, upgrade +from tiled.catalog.alembic_constants import ALEMBIC_DIR, ALEMBIC_INI_TEMPLATE_PATH +from tiled.catalog.core import ALL_REVISIONS +from tiled.utils import ensure_specified_sql_driver + +ROOT_ID = 0 +# The closure table migration adds the explicit root node (id 0) and +# nodes_closure, both of which the move below relies on. +CLOSURE_REVISION = "e05e918092c3" +# Same as blueapi's INSTRUMENT_SESSION_RE, e.g. "cm12345-1" +INSTRUMENT_SESSION_RE = re.compile( + r"^(?P[a-z]{2}(?P\d+))-(?P\d+)$" +) + + +def tags(beamline: str, proposal_number: int | None = None, visit=None) -> dict: + """The access blob tiled.rego's beamline_tag/proposal_tag/session_tag give.""" + tag = f"beamline:{beamline}" + if proposal_number is not None: + tag += f",proposal:{proposal_number}" + if visit is not None: + tag += f",session:{proposal_number}-{visit}" + return {"tags": [tag]} + + +def plan_moves(runs): + """Split root-level runs into (moves, skipped). + + moves: (node_id, [(key, access_blob) for beamline, proposal, session]) + """ + moves, skipped = [], [] + for node_id, key, beamline, session in runs: + match = INSTRUMENT_SESSION_RE.match(session or "") + if not beamline or match is None: + skipped.append((key, beamline, session)) + continue + number, visit = int(match["proposal_number"]), int(match["visit"]) + path = [ + (beamline, tags(beamline)), + (match["proposal"], tags(beamline, number)), + (session, tags(beamline, number, visit)), + ] + moves.append((node_id, path)) + return moves, skipped + + +async def current_revision(uri: str) -> str | None: + engine = create_async_engine(uri) + try: + return await get_current_revision(engine, ALL_REVISIONS) + finally: + await engine.dispose() + + +async def get_or_create(conn, parent: int, key: str, access_blob: dict) -> int: + # Containers blueapi already made are kept as they are. + node_id = await conn.scalar( + text( + "INSERT INTO nodes (parent, key, structure_family, metadata, specs, " + "access_blob) VALUES (:parent, :key, 'container', '{}', '[]', " + "CAST(:blob AS jsonb)) ON CONFLICT ON CONSTRAINT " + "key_parent_unique_constraint DO NOTHING RETURNING id" + ), + {"parent": parent, "key": key, "blob": json.dumps(access_blob)}, + ) + if node_id is None: + node_id = await conn.scalar( + text("SELECT id FROM nodes WHERE parent = :parent AND key = :key"), + {"parent": parent, "key": key}, + ) + return node_id + + +async def migrate(uri: str, dry_run: bool) -> None: + engine = create_async_engine(uri) + try: + async with engine.begin() as conn: + runs = ( + await conn.execute( + text( + "SELECT id, key, metadata #>> '{start,instrument}', " + "metadata #>> '{start,instrument_session}' " + "FROM nodes WHERE parent = :root" + ), + {"root": ROOT_ID}, + ) + ).all() + moves, skipped = plan_moves(runs) + + # Root children with no start doc are the new beamline containers + # (or anything else that isn't a run) - not worth reporting. + for key, beamline, session in skipped: + if beamline or session: + print(f"Skipping {key}: {beamline=} {session=}") + counts: dict[str, int] = {} + for _, path in moves: + where = "/".join(key for key, _ in path) + counts[where] = counts.get(where, 0) + 1 + for where, n in sorted(counts.items()): + print(f" root/{where}: {n} runs") + print(f"{len(moves)} runs to move into {len(counts)} sessions.") + if dry_run or not moves: + print("Nothing changed.") + return + + containers: dict[tuple, int] = {} + rows = [] + for node_id, path in moves: + parent = ROOT_ID + for depth in range(1, len(path) + 1): + keys = tuple(key for key, _ in path[:depth]) + if keys not in containers: + key, blob = path[depth - 1] + containers[keys] = await get_or_create(conn, parent, key, blob) + parent = containers[keys] + rows.append( + { + "node_id": node_id, + "new_parent": parent, + "blob": json.dumps(path[-1][1]), + } + ) + + await conn.execute( + text( + "CREATE TEMP TABLE moves (node_id integer PRIMARY KEY, " + "new_parent integer NOT NULL, blob jsonb NOT NULL) " + "ON COMMIT DROP" + ) + ) + await conn.execute( + text( + "INSERT INTO moves VALUES " + "(:node_id, :new_parent, CAST(:blob AS jsonb))" + ), + rows, + ) + # nodes_closure is only kept up to date by an insert trigger, so a + # reparent has to fix it by hand: drop the links from each run's + # subtree to its old ancestors (just the root), then link the + # subtree to every ancestor of its new parent. + await conn.execute( + text( + "DELETE FROM nodes_closure nc USING moves m, nodes_closure sub " + "WHERE sub.ancestor = m.node_id AND nc.descendant = sub.descendant " + "AND nc.ancestor IN (SELECT ancestor FROM nodes_closure " + "WHERE descendant = m.node_id AND ancestor != m.node_id)" + ) + ) + await conn.execute( + text( + "INSERT INTO nodes_closure (ancestor, descendant, depth) " + "SELECT p.ancestor, c.descendant, p.depth + c.depth + 1 " + "FROM moves m " + "JOIN nodes_closure c ON c.ancestor = m.node_id " + "JOIN nodes_closure p ON p.descendant = m.new_parent" + ) + ) + await conn.execute( + text( + "UPDATE nodes n SET parent = m.new_parent " + "FROM moves m WHERE n.id = m.node_id" + ) + ) + # Untagged ({}) nodes are public under the DLS policy; leave them. + retagged = await conn.execute( + text( + "UPDATE nodes n SET access_blob = m.blob " + "FROM moves m JOIN nodes_closure c ON c.ancestor = m.node_id " + "WHERE n.id = c.descendant AND n.access_blob -> 'tags' IS NOT NULL" + ) + ) + print( + f"Moved {len(rows)} runs, {len(containers)} containers in place, " + f"retagged {retagged.rowcount} nodes." + ) + finally: + await engine.dispose() + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("database_uri", help="e.g. postgresql://user:pass@host/db") + parser.add_argument( + "--revision", + default="head", + help="tiled catalog schema revision to upgrade to first (default: head)", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Skip the upgrade and only print what would move", + ) + args = parser.parse_args() + uri = ensure_specified_sql_driver(args.database_uri) + + if (revision := asyncio.run(current_revision(uri))) is None: + raise SystemExit("Database is not a tiled catalog; run `tiled catalog init`.") + if not args.dry_run: + # Alembic runs its own event loop, so this has to sit outside asyncio.run. + upgrade(ALEMBIC_INI_TEMPLATE_PATH, ALEMBIC_DIR, uri, args.revision) + revision = asyncio.run(current_revision(uri)) + print(f"Catalog schema at revision {revision}") + if revision not in ALL_REVISIONS or ALL_REVISIONS.index( + revision + ) > ALL_REVISIONS.index(CLOSURE_REVISION): + raise SystemExit( + f"Revision {revision} predates the closure table ({CLOSURE_REVISION}); " + "upgrade to it or later before migrating." + ) + asyncio.run(migrate(uri, args.dry_run)) + + +if __name__ == "__main__": + main() diff --git a/database_migration/postgres_tools.py b/database_migration/postgres_tools.py new file mode 100644 index 0000000..097f9c8 --- /dev/null +++ b/database_migration/postgres_tools.py @@ -0,0 +1,141 @@ +"""Back up a Postgres pod in Kubernetes, or add a read-only role to it. + +Every command runs inside the pod with `kubectl exec`, so the pod's own +Postgres tools (same version as the server) are used and nothing needs +installing locally. Dumps are written to this machine. + +Examples: + # Write globals.sql plus one .dump per database into ./backup + python postgres_tools.py --pod tiled-postgres-0 --namespace tiled dump backup + + # Create a role that can read every database but change nothing. + # Asks for the new role's password. + python postgres_tools.py --pod tiled-postgres-0 --namespace tiled readonly-role + +Restore (roles first on a fresh server, then each database). Scale tiled to +zero first, or it creates the schema in the empty database and the restore +fails with "already exists" errors: + kubectl scale -n NS deploy/TILED --replicas=0 + kubectl exec -i -n NS POD -- psql -U postgres -d postgres < backup/globals.sql + kubectl exec -n NS POD -- createdb -U postgres tiled + kubectl exec -i -n NS POD -- pg_restore -U postgres -d tiled --exit-on-error < backup/tiled.dump + kubectl scale -n NS deploy/TILED --replicas=1 +""" + +import argparse +import getpass +import subprocess +from datetime import datetime +from pathlib import Path + +LIST_DATABASES = ( + "SELECT datname FROM pg_database WHERE datallowconn AND NOT datistemplate" +) + + +def kubectl_exec(args, command: list[str]) -> list[str]: + """The kubectl command that runs `command` inside the postgres pod.""" + cmd = ["kubectl", "exec", "-i"] + if args.namespace: + cmd += ["-n", args.namespace] + if args.container: + cmd += ["-c", args.container] + return cmd + [args.pod, "--", *command, "-U", args.user] + + +def dump(args) -> None: + out_dir = Path(args.out_dir or datetime.now().strftime("pg_dump_%Y%m%d_%H%M%S")) + out_dir.mkdir(parents=True, exist_ok=True) + + # Roles and tablespaces live outside any one database, pg_dump skips them. + with open(out_dir / "globals.sql", "wb") as f: + subprocess.run( + kubectl_exec(args, ["pg_dumpall", "--globals-only"]), stdout=f, check=True + ) + + databases = subprocess.run( + kubectl_exec(args, ["psql", "-d", "postgres", "-AtX", "-c", LIST_DATABASES]), + capture_output=True, + text=True, + check=True, + ).stdout.split() + + for database in databases: + print(f"Dumping {database}") + # Custom format: compressed, and pg_restore can pick out single tables. + with open(out_dir / f"{database}.dump", "wb") as f: + subprocess.run( + kubectl_exec(args, ["pg_dump", "--format=custom", "-d", database]), + stdout=f, + check=True, + ) + print(f"Dumps written to {out_dir}") + + +def readonly_role(args) -> None: + password = getpass.getpass(f"Password for new role {args.role}: ") + if password != getpass.getpass("Again: "): + raise SystemExit("Passwords don't match.") + + # psql's \set reads backslashes as escapes and needs quotes doubled. + quoted = password.replace("\\", "\\\\").replace("'", "''") + # pg_read_all_data (Postgres 14+) is SELECT on every table, view and + # sequence in every database, including ones created later. It grants no + # INSERT/UPDATE/DELETE/DDL, so the role cannot change anything. + # default_transaction_read_only is a second guard on top of that. + # On Postgres < 15 every role may CREATE in the public schema unless that + # was revoked (REVOKE CREATE ON SCHEMA public FROM PUBLIC), so new tables + # are only blocked there by the read-only default. + sql = f""" +\\set pw '{quoted}' +CREATE ROLE :"role" LOGIN PASSWORD :'pw' + NOSUPERUSER NOCREATEDB NOCREATEROLE NOREPLICATION; +GRANT pg_read_all_data TO :"role"; +ALTER ROLE :"role" SET default_transaction_read_only = on; +""" + # The password goes in on stdin, not the command line, so it never shows + # up in the pod's process list. + subprocess.run( + kubectl_exec( + args, + [ + "psql", + "-d", + "postgres", + "-v", + "ON_ERROR_STOP=1", + "-v", + f"role={args.role}", + ], + ), + input=sql, + text=True, + check=True, + ) + print(f"Created read-only role {args.role}") + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--pod", required=True, help="postgres pod name") + parser.add_argument("--namespace", help="default: current kubectl namespace") + parser.add_argument("--container", help="default: the pod's default container") + parser.add_argument("--user", default="postgres", help="superuser to run as") + commands = parser.add_subparsers(dest="command", required=True) + + dump_parser = commands.add_parser("dump", help="dump every database") + dump_parser.add_argument("out_dir", nargs="?", help="default: pg_dump_") + dump_parser.set_defaults(run=dump) + + role_parser = commands.add_parser("readonly-role", help="create a read-only role") + role_parser.add_argument("role", nargs="?", default="tiled_readonly") + role_parser.set_defaults(run=readonly_role) + + args = parser.parse_args() + args.run(args) + + +if __name__ == "__main__": + main() diff --git a/database_migration/test_integration.py b/database_migration/test_integration.py new file mode 100644 index 0000000..041060f --- /dev/null +++ b/database_migration/test_integration.py @@ -0,0 +1,332 @@ +"""Migrate a real tiled dump and read it back through dls-tiled, Keycloak and OPA. + +Needs docker compose and a custom-format dump of the tiled database, e.g. +backup/tiled.dump from README step 2: + + PROD_DUMP=backup/tiled.dump uv run pytest database_migration/test_integration.py + +Skipped without PROD_DUMP. The stack in integration/compose.yaml is torn down +afterwards; set KEEP_STACK=1 to leave it up for poking at. + +What it checks, for every run the migration moves: + - it is at /// with its metadata unchanged + - its whole subtree came with it and is tagged with its session tag + - an admin can read it through tiled, a user on one session sees only that + session's runs, and runs the migration skips stay at the root +Array data isn't read back: the dump holds the catalog, not the files. +""" + +import json +import os +import subprocess +import sys +from pathlib import Path + +import httpx +import pytest +from migrate_node_structure import INSTRUMENT_SESSION_RE, ROOT_ID, plan_moves + +HERE = Path(__file__).parent +STACK = HERE / "integration" +COMPOSE = ["docker", "compose", "-f", str(STACK / "compose.yaml")] +URI = "postgresql://postgres:postgres@localhost:5433/tiled" +KEYCLOAK = "http://localhost:8081" +TILED = "http://localhost:8407/api/v1" +DUMP = os.environ.get("PROD_DUMP") + +pytestmark = pytest.mark.skipif(not DUMP, reason="set PROD_DUMP to a tiled pg_dump") + + +def compose(*args: str, **kwargs) -> subprocess.CompletedProcess: + return subprocess.run([*COMPOSE, *args], check=True, **kwargs) + + +def sql(query: str): + """Run a query that returns one JSON value and decode it.""" + out = compose( + "exec", + "-T", + "postgres", + "psql", + "-U", + "postgres", + "-d", + "tiled", + "-AtXc", + query, + capture_output=True, + text=True, + ).stdout + return json.loads(out) + + +def snapshot() -> list[dict]: + """Every root child with its start doc fields, metadata and subtree size.""" + return sql( + "SELECT coalesce(json_agg(json_build_object(" + "'id', n.id, 'key', n.key, 'metadata', n.metadata, " + "'beamline', n.metadata #>> '{start,instrument}', " + "'session', n.metadata #>> '{start,instrument_session}', " + "'subtree', (SELECT count(*) FROM nodes_closure c WHERE c.ancestor = n.id)" + ")), '[]') FROM nodes n WHERE n.parent = 0" + ) + + +def make_client(admin: httpx.Client, fedid: str) -> None: + """A service account whose tokens carry fedid and the audience tiled wants.""" + client_id = f"migration-test-{fedid}" + if admin.get("/clients", params={"clientId": client_id}).json(): + return + mapper = {"protocol": "openid-connect", "consentRequired": False} + admin.post( + "/clients", + json={ + "clientId": client_id, + "secret": "secret", + "standardFlowEnabled": False, + "serviceAccountsEnabled": True, + "protocolMappers": [ + { + **mapper, + "name": "fedid", + "protocolMapper": "oidc-hardcoded-claim-mapper", + "config": { + "claim.name": "fedid", + "claim.value": fedid, + "jsonType.label": "String", + "access.token.claim": "true", + }, + }, + { + **mapper, + "name": "audience", + "protocolMapper": "oidc-audience-mapper", + "config": { + "included.custom.audience": "account", + "access.token.claim": "true", + }, + }, + ], + }, + ).raise_for_status() + + +def tiled_client(fedid: str) -> httpx.Client: + token = httpx.post( + f"{KEYCLOAK}/realms/master/protocol/openid-connect/token", + data={ + "client_id": f"migration-test-{fedid}", + "client_secret": "secret", + "grant_type": "client_credentials", + }, + ) + token.raise_for_status() + return httpx.Client( + base_url=TILED, + headers={"Authorization": f"Bearer {token.json()['access_token']}"}, + timeout=30, + ) + + +def opa_data(sessions: list[tuple[str, int, int]]) -> dict: + """OPA data giving admin everything and alice the first session only.""" + ids = {s: i for i, s in enumerate(sessions, start=1)} + beamlines: dict = {} + proposals: dict = {} + for (beamline, number, visit), i in ids.items(): + beamlines.setdefault(beamline, {"sessions": []})["sessions"].append(i) + proposals.setdefault(str(number), {"sessions": {}})["sessions"][str(visit)] = i + subject = {"permissions": [], "proposals": [], "sessions": []} + return { + "diamond": { + "data": { + "beamlines": beamlines, + "proposals": proposals, + "sessions": { + str(i): { + "beamline": beamline, + "proposal_number": number, + "visit_number": visit, + } + for (beamline, number, visit), i in ids.items() + }, + "subjects": { + "admin": {**subject, "permissions": ["super_admin"]}, + "alice": {**subject, "sessions": [1]}, + }, + } + } + } + + +def wait_for_tiled() -> None: + for _ in range(60): + try: + if httpx.get(f"{TILED}/").status_code < 500: + return + except httpx.TransportError: + pass + subprocess.run(["sleep", "2"]) + compose("logs", "tiled") + raise TimeoutError("tiled didn't come up") + + +@pytest.fixture(scope="module") +def migrated(): + """Restore the dump, migrate it, bring up tiled; yield before and plan.""" + compose("up", "-d", "--wait", "postgres", "keycloak") + try: + compose( + "exec", + "-T", + "postgres", + "dropdb", + "-U", + "postgres", + "--if-exists", + "--force", + "tiled", + ) + compose("exec", "-T", "postgres", "createdb", "-U", "postgres", "tiled") + with open(DUMP, "rb") as f: + # Production roles don't exist here, so drop ownership and grants. + compose( + "exec", + "-T", + "postgres", + "pg_restore", + "-U", + "postgres", + "-d", + "tiled", + "--no-owner", + "--no-acl", + stdin=f, + ) + # Upgrade first so the snapshot can count subtrees with nodes_closure. + subprocess.run( + [sys.executable, "-m", "tiled", "catalog", "upgrade-database", URI], + check=True, + ) + before = snapshot() + moves, skipped = plan_moves( + [(r["id"], r["key"], r["beamline"], r["session"]) for r in before] + ) + assert moves, "the dump has no runs to move" + + subprocess.run( + [sys.executable, str(HERE / "migrate_node_structure.py"), URI], + check=True, + ) + + sessions = sorted( + { + (path[0][0], int(m["proposal_number"]), int(m["visit"])) + for _, path in moves + if (m := INSTRUMENT_SESSION_RE.match(path[2][0])) + } + ) + data_dir = STACK / "opa" / "opa_data" + data_dir.mkdir(exist_ok=True) + (data_dir / "data.json").write_text(json.dumps(opa_data(sessions))) + # Without a root the data bundle claims everything and OPA refuses to + # activate the policy bundle alongside it. + (data_dir / ".manifest").write_text(json.dumps({"roots": ["diamond/data"]})) + + with httpx.Client(base_url=f"{KEYCLOAK}/admin/realms/master") as admin: + token = httpx.post( + f"{KEYCLOAK}/realms/master/protocol/openid-connect/token", + data={ + "client_id": "admin-cli", + "username": "admin", + "password": "admin", + "grant_type": "password", + }, + ).json()["access_token"] + admin.headers["Authorization"] = f"Bearer {token}" + for fedid in ("admin", "alice"): + make_client(admin, fedid) + + compose("up", "-d", "--build", "opa", "tiled") + wait_for_tiled() + yield {r["id"]: r for r in before}, moves, skipped, sessions + finally: + if not os.environ.get("KEEP_STACK"): + compose("down", "-v") + + +def path_of(path: list) -> str: + return "/".join(key for key, _ in path) + + +def test_runs_moved_with_subtrees_and_tags(migrated): + before, moves, _, _ = migrated + after = { + r["id"]: r + for r in sql( + "SELECT json_agg(json_build_object('id', n.id, 'parent', n.parent, " + "'metadata', n.metadata, " + "'subtree', (SELECT count(*) FROM nodes_closure c " + "WHERE c.ancestor = n.id), " + "'ancestors', (SELECT count(*) FROM nodes_closure c " + "WHERE c.descendant = n.id), " + "'tags', (SELECT json_agg(DISTINCT d.access_blob -> 'tags') " + "FROM nodes_closure c JOIN nodes d ON d.id = c.descendant " + "WHERE c.ancestor = n.id AND d.access_blob -> 'tags' IS NOT NULL))) " + "FROM nodes n WHERE n.id = ANY(ARRAY[" + + ",".join(str(node_id) for node_id, _ in moves) + + "])" + ) + } + for node_id, path in moves: + run = after[node_id] + assert run["parent"] != ROOT_ID + assert run["metadata"] == before[node_id]["metadata"] + assert run["subtree"] == before[node_id]["subtree"] + # itself, session, proposal, beamline, root + assert run["ancestors"] == 5 + assert run["tags"] in (None, [path[-1][1]["tags"]]) + + +def test_admin_reads_every_moved_run(migrated): + before, moves, _, _ = migrated + with tiled_client("admin") as client: + for node_id, path in moves: + response = client.get(f"/metadata/{path_of(path)}/{before[node_id]['key']}") + assert response.status_code == 200, response.text + assert ( + response.json()["data"]["attributes"]["metadata"] + == (before[node_id]["metadata"]) + ) + expected: dict[str, int] = {} + for _, path in moves: + expected[path_of(path)] = expected.get(path_of(path), 0) + 1 + for where, count in expected.items(): + response = client.get(f"/search/{where}", params={"page[limit]": 0}) + assert response.status_code == 200, response.text + assert response.json()["meta"]["count"] == count, where + + +def test_skipped_runs_stay_at_root(migrated): + before, _, skipped, _ = migrated + root = set(sql("SELECT coalesce(json_agg(key), '[]') FROM nodes WHERE parent = 0")) + moved_from_root = {r["key"] for r in before.values()} - {k for k, *_ in skipped} + assert {k for k, *_ in skipped} <= root + assert not moved_from_root & root + + +def test_user_sees_only_their_session(migrated): + before, moves, _, sessions = migrated + beamline, number, visit = sessions[0] + mine = f"beamline:{beamline},proposal:{number},session:{number}-{visit}" + with tiled_client("alice") as client: + top = client.get("/search/", params={"page[limit]": 300}) + assert top.status_code == 200, top.text + beamlines = {path[0][0] for _, path in moves} + assert {n["id"] for n in top.json()["data"]} & beamlines == {beamline} + for node_id, path in moves: + response = client.get(f"/metadata/{path_of(path)}/{before[node_id]['key']}") + if path[-1][1]["tags"] == [mine]: + assert response.status_code == 200, response.text + else: + assert response.status_code == 404, path_of(path) diff --git a/database_migration/test_migrate_node_structure.py b/database_migration/test_migrate_node_structure.py new file mode 100644 index 0000000..082623a --- /dev/null +++ b/database_migration/test_migrate_node_structure.py @@ -0,0 +1,25 @@ +from migrate_node_structure import plan_moves + + +def test_plan_moves(): + moves, skipped = plan_moves( + [ + (1, "run-a", "i22", "cm12345-1"), + (2, "run-b", "i22", "garbage"), + (3, "i22", None, None), + ] + ) + assert moves == [ + ( + 1, + [ + ("i22", {"tags": ["beamline:i22"]}), + ("cm12345", {"tags": ["beamline:i22,proposal:12345"]}), + ( + "cm12345-1", + {"tags": ["beamline:i22,proposal:12345,session:12345-1"]}, + ), + ], + ) + ] + assert [key for key, *_ in skipped] == ["run-b", "i22"] From 5fd0531e0ed0efdefb8c536d3ff31d53497b5b0d Mon Sep 17 00:00:00 2001 From: Zoheb Shaikh <26975142+ZohebShaikh@users.noreply.github.com> Date: Fri, 16 Oct 2026 15:51:34 +0100 Subject: [PATCH 3/8] Run the migration integration test in CI against a seeded catalog Without PROD_DUMP the test now seeds a small catalog in the old flat layout, so it needs nothing but docker compose and uv. The migration scripts carry their dependencies inline and run with `uv run --script`, with tiled pinned to the runtime image's version. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/ci.yml | 14 +++ database_migration/README.md | 7 +- database_migration/migrate_node_structure.py | 7 +- database_migration/postgres_tools.py | 14 ++- database_migration/test_integration.py | 95 ++++++++++++++------ 5 files changed, 104 insertions(+), 33 deletions(-) mode change 100644 => 100755 database_migration/migrate_node_structure.py mode change 100644 => 100755 database_migration/postgres_tools.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9d99496..a85871b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -27,6 +27,20 @@ jobs: runs-on: ${{ matrix.runs-on }} python-version: ${{ matrix.python-version }} + migration: + # Migrates a seeded catalog, then reads it back through dls-tiled, Keycloak + # and OPA in docker compose + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v7 + + - name: Install uv + uses: astral-sh/setup-uv@v10.0.0 + + - name: Run integration test + run: uv run --locked pytest database_migration/test_integration.py + container: needs: test if: always() diff --git a/database_migration/README.md b/database_migration/README.md index 49613b1..4431af8 100644 --- a/database_migration/README.md +++ b/database_migration/README.md @@ -131,16 +131,17 @@ forward the Postgres port first. Leave this running in its own terminal: kubectl port-forward -n $NS pod/$POD 5432:5432 ``` -In another terminal, from a Python environment with `tiled[server]` installed: +In another terminal (uv installs the script's dependencies, pinned to the +tiled version in the runtime image): ```bash URI=postgresql://postgres:@localhost:5432/tiled # Preview: prints what would move, changes nothing -python migrate_node_structure.py $URI --dry-run +uv run --script migrate_node_structure.py $URI --dry-run # Upgrade the tiled schema (to --revision, default head), then move the runs -python migrate_node_structure.py $URI --revision head +uv run --script migrate_node_structure.py $URI --revision head ``` The dry run only works if the database is already at tiled schema revision diff --git a/database_migration/migrate_node_structure.py b/database_migration/migrate_node_structure.py old mode 100644 new mode 100755 index de475c4..1159b1a --- a/database_migration/migrate_node_structure.py +++ b/database_migration/migrate_node_structure.py @@ -1,3 +1,8 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = ["tiled[server]==0.2.18"] +# /// """Move runs from the flat catalog layout into the beamline/proposal/session tree. Before: root/ @@ -21,7 +26,7 @@ Postgres only. Example: - python database_migration/migrate_node_structure.py \\ + uv run --script database_migration/migrate_node_structure.py \\ postgresql://user:pass@host:5432/tiled --revision head --dry-run """ diff --git a/database_migration/postgres_tools.py b/database_migration/postgres_tools.py old mode 100644 new mode 100755 index 097f9c8..5e6c482 --- a/database_migration/postgres_tools.py +++ b/database_migration/postgres_tools.py @@ -1,3 +1,8 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.11" +# dependencies = [] +# /// """Back up a Postgres pod in Kubernetes, or add a read-only role to it. Every command runs inside the pod with `kubectl exec`, so the pod's own @@ -6,11 +11,13 @@ Examples: # Write globals.sql plus one .dump per database into ./backup - python postgres_tools.py --pod tiled-postgres-0 --namespace tiled dump backup + uv run --script postgres_tools.py --pod tiled-postgres-0 --namespace tiled \\ + dump backup # Create a role that can read every database but change nothing. # Asks for the new role's password. - python postgres_tools.py --pod tiled-postgres-0 --namespace tiled readonly-role + uv run --script postgres_tools.py --pod tiled-postgres-0 --namespace tiled \\ + readonly-role Restore (roles first on a fresh server, then each database). Scale tiled to zero first, or it creates the schema in the empty database and the restore @@ -18,7 +25,8 @@ kubectl scale -n NS deploy/TILED --replicas=0 kubectl exec -i -n NS POD -- psql -U postgres -d postgres < backup/globals.sql kubectl exec -n NS POD -- createdb -U postgres tiled - kubectl exec -i -n NS POD -- pg_restore -U postgres -d tiled --exit-on-error < backup/tiled.dump + kubectl exec -i -n NS POD -- pg_restore -U postgres -d tiled --exit-on-error \\ + < backup/tiled.dump kubectl scale -n NS deploy/TILED --replicas=1 """ diff --git a/database_migration/test_integration.py b/database_migration/test_integration.py index 041060f..7789404 100644 --- a/database_migration/test_integration.py +++ b/database_migration/test_integration.py @@ -1,12 +1,17 @@ """Migrate a real tiled dump and read it back through dls-tiled, Keycloak and OPA. -Needs docker compose and a custom-format dump of the tiled database, e.g. -backup/tiled.dump from README step 2: +Needs docker compose and uv. Runs against a small seeded catalog in the old +flat layout (what CI does): + + uv run pytest database_migration/test_integration.py + +or against a custom-format dump of the real database, e.g. backup/tiled.dump +from README step 2: PROD_DUMP=backup/tiled.dump uv run pytest database_migration/test_integration.py -Skipped without PROD_DUMP. The stack in integration/compose.yaml is torn down -afterwards; set KEEP_STACK=1 to leave it up for poking at. +The stack in integration/compose.yaml is torn down afterwards; set KEEP_STACK=1 +to leave it up for poking at. What it checks, for every run the migration moves: - it is at /// with its metadata unchanged @@ -33,17 +38,24 @@ KEYCLOAK = "http://localhost:8081" TILED = "http://localhost:8407/api/v1" DUMP = os.environ.get("PROD_DUMP") - -pytestmark = pytest.mark.skipif(not DUMP, reason="set PROD_DUMP to a tiled pg_dump") +# Runs at the root of a catalog in the old layout: (key, beamline, session). +# The last two have nothing to place them by, so the migration skips them. +SEED_RUNS = [ + ("run-1", "i22", "cm12345-1"), + ("run-2", "i22", "cm12345-1"), + ("run-3", "i22", "cm12345-2"), + ("run-4", "p45", "mx67890-3"), + ("no-session", "i22", None), + ("bad-session", "i22", "not-a-session"), +] def compose(*args: str, **kwargs) -> subprocess.CompletedProcess: return subprocess.run([*COMPOSE, *args], check=True, **kwargs) -def sql(query: str): - """Run a query that returns one JSON value and decode it.""" - out = compose( +def psql(query: str) -> str: + return compose( "exec", "-T", "postgres", @@ -57,7 +69,33 @@ def sql(query: str): capture_output=True, text=True, ).stdout - return json.loads(out) + + +def sql(query: str): + """Run a query that returns one JSON value and decode it.""" + return json.loads(psql(query)) + + +def seed() -> None: + """A fresh catalog with SEED_RUNS at the root, tagged with the bare numeric + session id the old policy used, each with a tagged and an untagged child.""" + subprocess.run([sys.executable, "-m", "tiled", "catalog", "init", URI], check=True) + runs = [] + for i, (key, beamline, session) in enumerate(SEED_RUNS, start=1): + start = {"uid": key, "instrument": beamline} + if session: + start["instrument_session"] = session + metadata = json.dumps({"start": start}) + blob = json.dumps({"tags": [str(i)]}) + runs.append(f"(0, '{key}', 'container', '{metadata}', '[]', '{blob}')") + columns = "(parent, key, structure_family, metadata, specs, access_blob)" + psql(f"INSERT INTO nodes {columns} VALUES {', '.join(runs)}") + # Separate statements, so the closure trigger has the runs' rows to join. + for child, blob in (("primary", "access_blob"), ("public", "'{}'")): + psql( + f"INSERT INTO nodes {columns} SELECT id, '{child}', 'container', " + f"'{{}}', '[]', {blob} FROM nodes WHERE parent = 0" + ) def snapshot() -> list[dict]: @@ -188,21 +226,25 @@ def migrated(): "tiled", ) compose("exec", "-T", "postgres", "createdb", "-U", "postgres", "tiled") - with open(DUMP, "rb") as f: - # Production roles don't exist here, so drop ownership and grants. - compose( - "exec", - "-T", - "postgres", - "pg_restore", - "-U", - "postgres", - "-d", - "tiled", - "--no-owner", - "--no-acl", - stdin=f, - ) + if DUMP: + with open(DUMP, "rb") as f: + # Production roles don't exist here, so drop ownership and grants. + compose( + "exec", + "-T", + "postgres", + "pg_restore", + "-U", + "postgres", + "-d", + "tiled", + "--no-owner", + "--no-acl", + "--exit-on-error", + stdin=f, + ) + else: + seed() # Upgrade first so the snapshot can count subtrees with nodes_closure. subprocess.run( [sys.executable, "-m", "tiled", "catalog", "upgrade-database", URI], @@ -214,8 +256,9 @@ def migrated(): ) assert moves, "the dump has no runs to move" + # The way the README runs it, with the script's own pinned dependencies. subprocess.run( - [sys.executable, str(HERE / "migrate_node_structure.py"), URI], + ["uv", "run", "--script", str(HERE / "migrate_node_structure.py"), URI], check=True, ) From 24816c525e1559dc3d90db45094d59069006e516 Mon Sep 17 00:00:00 2001 From: Zoheb Shaikh <26975142+ZohebShaikh@users.noreply.github.com> Date: Sat, 17 Oct 2026 11:40:49 +0100 Subject: [PATCH 4/8] add more docs --- database_migration/README.md | 47 ++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/database_migration/README.md b/database_migration/README.md index 4431af8..4f6ad92 100644 --- a/database_migration/README.md +++ b/database_migration/README.md @@ -170,3 +170,50 @@ kubectl scale -n $NS deploy/ --replicas=1 If the restore fails partway, drop the database and start again; don't try to patch a partial restore. + +## Troubleshooting + +### `KeyError: 'postgresql://:5432/data'` when reading data + +After you restore a backup into a different deployment (for example, a prod +dump restored into tiled-testing), reads of tabular data fail in +`tiled/storage.py`, in `get_storage`. + +Each row in the catalog's `assets` table has a `data_uri` that names the +storage it was written to. Tiled keys its storage by URI without credentials, +built from the `writable_storage` and `readable_storage` entries in its config. +If the restored rows still name the old host, the lookup fails. The migration +script doesn't touch `data_uri`. + +To fix it, point the assets at the new host. Run this against the catalog +database (`tiled`), not `data`. The `data` database must already exist in +the new deployment with the same tables, or the `KeyError` becomes a +missing-table error. Tiled reads `data_uri` on every request, so it doesn't +need a restart. + +With the `postgres` password in `PGPASSWORD`, set the host that's in the +error (`OLD`) and the host in this deployment's `writable_storage` (`NEW`): + +```bash +OLD=tiled-postgres +NEW=tiled-testing-postgres + +# Preview: the rows that will change +kubectl exec -i -n $NS $POD -- env PGPASSWORD=$PGPASSWORD psql -U postgres -d tiled < Date: Sat, 17 Oct 2026 12:24:45 +0100 Subject: [PATCH 5/8] Stamp every tag the policy returns and cache user_sessions user_session now returns a node's full tag list (its own tag, then its ancestors' grant tags), so init_node stamps all of them rather than wrapping a single tag. filters caches each access token's user_sessions answer, for 1 minute for users and 7.5 minutes for service accounts (tokens without a fedid, as tiled.rego tells them apart), so a change in authz takes up to that long to reach tiled. A token that can't be read gets the user TTL, and a missing decision isn't cached. The tests for both move here from blueapi's copy of this policy. Co-Authored-By: Claude Opus 5.5 --- src/dls_tiled/policy.py | 58 ++++++++++++++++++++++++++++++----- tests/test_policy.py | 67 +++++++++++++++++++++++++++++++++++++++-- 2 files changed, 115 insertions(+), 10 deletions(-) diff --git a/src/dls_tiled/policy.py b/src/dls_tiled/policy.py index 9cdc6be..8682a9d 100644 --- a/src/dls_tiled/policy.py +++ b/src/dls_tiled/policy.py @@ -4,8 +4,10 @@ (https://github.com/DiamondLightSource/diamond-tiled). """ +import base64 import json import logging +import time from fastapi import HTTPException from pydantic import BaseModel, HttpUrl, TypeAdapter, ValidationError @@ -25,6 +27,12 @@ logger = logging.getLogger(__name__) +# How long a caller's user_sessions answer is reused before OPA is asked again, +# so a change in authz takes up to this long to reach tiled. Service accounts +# (tokens without a fedid, as tiled.rego tells them apart) change rarely. +USER_TAGS_TTL = 60 # seconds +SERVICE_ACCOUNT_TAGS_TTL = 450 # seconds, 7.5 minutes + class DiamondAccessBlob(BaseModel): """What blueapi asks for when creating a node at each level of the tree. @@ -63,6 +71,21 @@ def _parse_composite_tag(tag: str) -> dict[str, str]: return fields +def _tags_ttl(token: str) -> int: + """How long to cache this token's tags. Reads the JWT payload unverified, + since tiled has already verified the token; one that can't be read gets the + shorter user TTL rather than the service-account one.""" + try: + payload = token.split(".")[1] + claims = json.loads( + base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4)) + ) + is_user = "fedid" in claims + except (IndexError, ValueError, TypeError): + is_user = True + return USER_TAGS_TTL if is_user else SERVICE_ACCOUNT_TAGS_TTL + + def _check_principal(principal: Principal | None): if not isinstance(principal, Principal): raise HTTPException( @@ -93,6 +116,8 @@ def __init__( ): self._token_audience = token_audience self._type_adapter = TypeAdapter(DiamondAccessBlob | int) + # access token -> (monotonic expiry, filters) + self._filters_cache: dict[str, tuple[float, Filters]] = {} super().__init__( authorization_provider=authorization_provider, @@ -117,10 +142,12 @@ async def init_node( decision = await self._get_external_decision( self._create_node, self.build_input(principal, authn_access_tags, authn_scopes, access_blob), - ResultHolder[str], + ResultHolder[list[str]], ) if decision and decision.result is not None: - return (True, {"tags": [decision.result]}) + # The node's own tag first (build_input reads tags[0] back), then + # its ancestors' grant tags; see Node Tags in tiled.rego. + return (True, {"tags": decision.result}) raise ValueError("Permission denied not able to add the node") async def modify_node( @@ -193,14 +220,31 @@ async def filters( scopes: Scopes, ) -> Filters: _check_principal(principal) + token = ( + principal.access_token.get_secret_value() + if principal.access_token + else None + ) + now = time.monotonic() + if token and (cached := self._filters_cache.get(token)) and cached[0] > now: + return cached[1] tags = await self._get_external_decision( self._user_tags, self.build_input(principal, authn_access_tags, authn_scopes), ResultHolder[list[str]], ) - if tags is not None: - if tags.result == ["*"]: - return ALL_ACCESS - return [AccessBlobFilter(tags=tags.result, user_id=None)] # type: ignore - else: + if tags is None: + # Not cached, so OPA recovering takes effect at once. return NO_ACCESS # type: ignore + if tags.result == ["*"]: + filters = ALL_ACCESS + else: + filters = [AccessBlobFilter(tags=tags.result, user_id=None)] # type: ignore + if token: + # ponytail: drops expired entries on every miss, O(cached tokens); fine + # for hundreds of callers, use a TTL cache (e.g. cachetools) if that grows. + self._filters_cache = { + k: v for k, v in self._filters_cache.items() if v[0] > now + } + self._filters_cache[token] = (now + _tags_ttl(token), filters) + return filters diff --git a/tests/test_policy.py b/tests/test_policy.py index 4ae480a..4bc5cba 100644 --- a/tests/test_policy.py +++ b/tests/test_policy.py @@ -1,3 +1,4 @@ +import base64 import json import uuid from types import SimpleNamespace @@ -13,7 +14,10 @@ from tiled.queries import AccessBlobFilter from tiled.server.schemas import Principal, PrincipalType +from dls_tiled import policy as policy_module from dls_tiled.policy import ( + SERVICE_ACCOUNT_TAGS_TTL, + USER_TAGS_TTL, DiamondOpenPolicyAgentAuthorizationPolicy, _check_principal, ) @@ -255,6 +259,57 @@ async def test_filters_without_result_is_no_access(policy, user, opa): assert await policy.filters(None, user, None, set(), set()) is NO_ACCESS +def jwt(claims: dict) -> str: + payload = base64.urlsafe_b64encode(json.dumps(claims).encode()).rstrip(b"=") + return f"header.{payload.decode()}.signature" + + +@pytest.fixture +def clock(monkeypatch) -> list[float]: + now = [1000.0] + monkeypatch.setattr(policy_module.time, "monotonic", lambda: now[0]) + return now + + +@pytest.mark.parametrize( + "token, ttl", + [ + (jwt({"fedid": "abc12345"}), USER_TAGS_TTL), + (jwt({"fedid": "abc12345", "beamline": "i22"}), USER_TAGS_TTL), + (jwt({"beamline": "i22"}), SERVICE_ACCOUNT_TAGS_TTL), + ("not-a-jwt", USER_TAGS_TTL), # unreadable: the shorter TTL + ], +) +async def test_filters_cached_until_ttl(policy, opa, clock, token, ttl): + opa.responses[OPA + "tiled/user_sessions"] = {"result": ["beamline:i22"]} + user = make_principal(token=token) + + first = await policy.filters(None, user, None, set(), set()) + clock[0] += ttl - 1 + assert await policy.filters(None, user, None, set(), set()) is first + assert len(opa.requests) == 1 + clock[0] += 2 + await policy.filters(None, user, None, set(), set()) + assert len(opa.requests) == 2 + + +async def test_filters_cache_is_per_token(policy, opa, clock): + opa.responses[OPA + "tiled/user_sessions"] = {"result": ["beamline:i22"]} + + for token in ("token-a", "token-b"): + await policy.filters(None, make_principal(token=token), None, set(), set()) + + assert len(opa.requests) == 2 + + +async def test_filters_no_result_not_cached(policy, user, opa, clock): + opa.responses[OPA + "tiled/user_sessions"] = {} + assert await policy.filters(None, user, None, set(), set()) is NO_ACCESS + + opa.responses[OPA + "tiled/user_sessions"] = {"result": ["*"]} + assert await policy.filters(None, user, None, set(), set()) is ALL_ACCESS + + # --- init_node ----------------------------------------------------------------- @@ -274,12 +329,18 @@ async def test_init_node_empty_blob_not_public(user, opa): assert opa.requests == [] -async def test_init_node_returns_tag_from_opa(policy, user, opa): - opa.responses[OPA + "tiled/user_session"] = {"result": "12345"} +async def test_init_node_returns_tags_from_opa(policy, user, opa): + # the session's own tag first, then its ancestors' grant tags + tags = [ + "beamline:i22,proposal:1234,session:1234-5", + "beamline:i22,proposal:1234", + "beamline:i22", + ] + opa.responses[OPA + "tiled/user_session"] = {"result": tags} result = await policy.init_node(user, None, set(), {"tags": [BLOB_TAG]}) - assert result == (True, {"tags": ["12345"]}) + assert result == (True, {"tags": tags}) assert opa.requests == [ ( OPA + "tiled/user_session", From 5c6ac5ffb8d492ed8f6b05b9866646cb12dff7fa Mon Sep 17 00:00:00 2001 From: Zoheb Shaikh <26975142+ZohebShaikh@users.noreply.github.com> Date: Sat, 17 Oct 2026 12:24:45 +0100 Subject: [PATCH 6/8] Give migrated nodes their full tags and retag bare session ids The migration now stamps the tag lists the current policy does: each node's own tag plus its ancestors' grant tags, and visible: tags on the beamline and proposal containers. Nodes already in the tree with only their own tag (from blueapi under an earlier policy, or an earlier run of this script) get the full list too. With --authz-data (the authz bundle's data.json), nodes still tagged with a bare numeric session id, such as runs without start.instrument_session that can't be moved, are retagged in place from that session; otherwise only admins could see them. Unknown ids are counted. The retag is one join through a temp table: the GIN index on access_blob can't serve "access_blob -> 'tags' = ...", so an UPDATE per tag scanned nodes once per tag. The OPA in the integration stack moves to authz-policy 0.0.27-alpha, which has the matching policy. Co-Authored-By: Claude Opus 5.5 --- database_migration/README.md | 23 ++- .../integration/opa/config.yaml | 2 +- database_migration/migrate_node_structure.py | 135 +++++++++++++++--- database_migration/test_integration.py | 2 +- .../test_migrate_node_structure.py | 60 +++++++- 5 files changed, 197 insertions(+), 25 deletions(-) diff --git a/database_migration/README.md b/database_migration/README.md index 4f6ad92..b8650dd 100644 --- a/database_migration/README.md +++ b/database_migration/README.md @@ -137,16 +137,33 @@ tiled version in the runtime image): ```bash URI=postgresql://postgres:@localhost:5432/tiled -# Preview: prints what would move, changes nothing +# Preview: prints what would move and retag, changes nothing uv run --script migrate_node_structure.py $URI --dry-run # Upgrade the tiled schema (to --revision, default head), then move the runs -uv run --script migrate_node_structure.py $URI --revision head +uv run --script migrate_node_structure.py $URI --revision head --authz-data data.json ``` +`--authz-data` is the authz bundle's `data.json` (in `bundle.tar.gz` from +authz.diamond.ac.uk, or `GET /v1/data` from an OPA that has it loaded). The +script uses it for nodes still tagged with a bare numeric session id that it +can't move, such as runs without `start.instrument_session`: without the data +they stay tagged with the id, which the current policy never grants, so only +admins see them. With it they are retagged in place from that session's +beamline, proposal and visit. They stay at the root, since the data has the +proposal number but not the code (`cm`) the tree's keys use. The script prints +how many ids it couldn't find. + +Besides moving runs, it gives nodes already in the tree that carry only their +own tag (from blueapi under an earlier policy, or an earlier run of this +script) the full tag list the current policy stamps, so beamline and proposal +grants reach them. Deploy the matching authz policy with it: under an older +policy, a session member's beamline tag would match every node on the beamline. + The dry run only works if the database is already at tiled schema revision `e05e918092c3` or later; otherwise it says so, and the real run upgrades it -first. Running the migration twice is safe: the second run moves nothing. +first. Running the migration twice is safe: the second run moves and retags +nothing. ## Restore from the backup (as postgres) diff --git a/database_migration/integration/opa/config.yaml b/database_migration/integration/opa/config.yaml index 9c2668b..d751f94 100644 --- a/database_migration/integration/opa/config.yaml +++ b/database_migration/integration/opa/config.yaml @@ -5,4 +5,4 @@ services: bundles: diamond-policies: service: ghcr - resource: ghcr.io/zohebshaikh/authz-policy:0.0.25-alpha + resource: ghcr.io/zohebshaikh/authz-policy:0.0.27-alpha diff --git a/database_migration/migrate_node_structure.py b/database_migration/migrate_node_structure.py index 1159b1a..eb78f1a 100755 --- a/database_migration/migrate_node_structure.py +++ b/database_migration/migrate_node_structure.py @@ -13,14 +13,25 @@ upgraded first (to --revision, default head) so the move runs against the nodes/nodes_closure layout this script expects. -The new containers get the same composite access tags tiled.rego stamps on -nodes it creates, and every tagged node under a moved run is retagged with -its session tag. Runs written under the old policy carry a bare numeric -session id, which the new policy's user_sessions never returns, so they -would become invisible to users otherwise. +The new containers get the same access tags tiled.rego stamps on nodes it +creates (its Node Tags: the node's own tag, then each ancestor's grant tag), +and every tagged node under a moved run is retagged with its session's tags. +Runs written under the old policy carry a bare numeric session id, which the +new policy's user_sessions never returns, so they would become invisible to +users otherwise. + +Nodes already in the tree tagged with only their own composite tag (made by +blueapi under an earlier version of the policy, or by an earlier run of this +script) get the full tag list too; without the ancestor tags a beamline or +proposal grant would not reach them. Runs at the root without start.instrument / start.instrument_session, or whose -instrument_session doesn't parse, are left where they are and listed. Safe to +instrument_session doesn't parse, are left where they are and listed. Their +bare numeric session id would make them invisible to everyone but admins, so +with --authz-data (the authz bundle's data.json) any node still tagged with one +is retagged in place from that session's beamline, proposal and visit. They +can't be moved: the data has the proposal number but not its code ("cm"), which +the tree's keys use. Ids the data doesn't know are listed. Safe to rerun: anything already moved is no longer a child of the root. Postgres only. @@ -52,14 +63,46 @@ ) +# A node's own tag as tiled.rego's beamline_tag/proposal_tag/session_tag build it +COMPOSITE_TAG_RE = re.compile( + r"^beamline:(?P[^,]+)" + r"(?:,proposal:(?P\d+)(?:,session:(?P=proposal)-(?P\d+))?)?$" +) + + def tags(beamline: str, proposal_number: int | None = None, visit=None) -> dict: - """The access blob tiled.rego's beamline_tag/proposal_tag/session_tag give.""" - tag = f"beamline:{beamline}" - if proposal_number is not None: - tag += f",proposal:{proposal_number}" - if visit is not None: - tag += f",session:{proposal_number}-{visit}" - return {"tags": [tag]} + """The access blob tiled.rego stamps on a node at this level (its Node Tags): + the node's own tag first, then each ancestor's grant tag so a beamline or + proposal grant covers the subtree, and on the two container levels a + visible: tag so callers granted something below can still list them.""" + beamline_tag = f"beamline:{beamline}" + if proposal_number is None: + return {"tags": [beamline_tag, f"visible:{beamline_tag}"]} + proposal_tag = f"{beamline_tag},proposal:{proposal_number}" + if visit is None: + return {"tags": [proposal_tag, beamline_tag, f"visible:{proposal_tag}"]} + session_tag = f"{proposal_tag},session:{proposal_number}-{visit}" + return {"tags": [session_tag, proposal_tag, beamline_tag]} + + +def tags_from_own_tag(tag: str, sessions: dict | None = None) -> dict | None: + """The full access blob for a node tagged with just its own composite tag, + or with a bare numeric session id found in `sessions` (authz data's + sessions, keyed by id); None if `tag` is neither.""" + if tag.isdigit(): + if (session := (sessions or {}).get(tag)) is None: + return None + return tags( + session["beamline"], session["proposal_number"], session["visit_number"] + ) + if (match := COMPOSITE_TAG_RE.match(tag)) is None: + return None + number, visit = match["proposal"], match["visit"] + return tags( + match["beamline"], + int(number) if number else None, + int(visit) if visit else None, + ) def plan_moves(runs): @@ -110,7 +153,7 @@ async def get_or_create(conn, parent: int, key: str, access_blob: dict) -> int: return node_id -async def migrate(uri: str, dry_run: bool) -> None: +async def migrate(uri: str, dry_run: bool, sessions: dict | None = None) -> None: engine = create_async_engine(uri) try: async with engine.begin() as conn: @@ -138,10 +181,60 @@ async def migrate(uri: str, dry_run: bool) -> None: for where, n in sorted(counts.items()): print(f" root/{where}: {n} runs") print(f"{len(moves)} runs to move into {len(counts)} sessions.") - if dry_run or not moves: + + # Nodes still carrying only their own composite tag, by that tag. + single = ( + await conn.execute( + text( + "SELECT access_blob -> 'tags' ->> 0, count(*) FROM nodes " + "WHERE jsonb_typeof(access_blob -> 'tags') = 'array' " + "AND jsonb_array_length(access_blob -> 'tags') = 1 " + "GROUP BY 1" + ) + ) + ).all() + retags, to_retag, unknown = [], 0, 0 + for tag, n in single: + if (blob := tags_from_own_tag(tag, sessions)) is not None: + retags.append({"old": json.dumps([tag]), "blob": json.dumps(blob)}) + to_retag += n + elif tag.isdigit(): + unknown += n + print(f"{to_retag} nodes to give their full tags.") + if unknown: + # Includes the runs moved below, which get their tags from the move. + where = "not in --authz-data" if sessions else "(pass --authz-data)" + print(f"{unknown} nodes carry a bare session id {where}.") + if dry_run or not (moves or retags): print("Nothing changed.") return + if retags: + # One join, so nodes is scanned once rather than once per tag + # (the GIN index on access_blob can't serve -> 'tags' = ...). + await conn.execute( + text( + "CREATE TEMP TABLE retags (old jsonb PRIMARY KEY, " + "blob jsonb NOT NULL) ON COMMIT DROP" + ) + ) + await conn.execute( + text( + "INSERT INTO retags VALUES " + "(CAST(:old AS jsonb), CAST(:blob AS jsonb))" + ), + retags, + ) + await conn.execute( + text( + "UPDATE nodes n SET access_blob = r.blob FROM retags r " + "WHERE n.access_blob -> 'tags' = r.old" + ) + ) + print(f"Gave {to_retag} nodes their full tags.") + if not moves: + return + containers: dict[tuple, int] = {} rows = [] for node_id, path in moves: @@ -227,6 +320,12 @@ def main() -> None: default="head", help="tiled catalog schema revision to upgrade to first (default: head)", ) + parser.add_argument( + "--authz-data", + type=argparse.FileType(), + help="the authz bundle's data.json, to retag nodes still carrying a bare " + "numeric session id (e.g. runs left at the root) from that session", + ) parser.add_argument( "--dry-run", action="store_true", @@ -249,7 +348,11 @@ def main() -> None: f"Revision {revision} predates the closure table ({CLOSURE_REVISION}); " "upgrade to it or later before migrating." ) - asyncio.run(migrate(uri, args.dry_run)) + sessions = None + if args.authz_data: + data = json.load(args.authz_data) + sessions = data.get("diamond", {}).get("data", data)["sessions"] + asyncio.run(migrate(uri, args.dry_run, sessions)) if __name__ == "__main__": diff --git a/database_migration/test_integration.py b/database_migration/test_integration.py index 7789404..1bd5d0f 100644 --- a/database_migration/test_integration.py +++ b/database_migration/test_integration.py @@ -369,7 +369,7 @@ def test_user_sees_only_their_session(migrated): assert {n["id"] for n in top.json()["data"]} & beamlines == {beamline} for node_id, path in moves: response = client.get(f"/metadata/{path_of(path)}/{before[node_id]['key']}") - if path[-1][1]["tags"] == [mine]: + if path[-1][1]["tags"][0] == mine: assert response.status_code == 200, response.text else: assert response.status_code == 404, path_of(path) diff --git a/database_migration/test_migrate_node_structure.py b/database_migration/test_migrate_node_structure.py index 082623a..93cb59b 100644 --- a/database_migration/test_migrate_node_structure.py +++ b/database_migration/test_migrate_node_structure.py @@ -1,4 +1,5 @@ -from migrate_node_structure import plan_moves +import pytest +from migrate_node_structure import plan_moves, tags, tags_from_own_tag def test_plan_moves(): @@ -13,13 +14,64 @@ def test_plan_moves(): ( 1, [ - ("i22", {"tags": ["beamline:i22"]}), - ("cm12345", {"tags": ["beamline:i22,proposal:12345"]}), + ("i22", {"tags": ["beamline:i22", "visible:beamline:i22"]}), + ( + "cm12345", + { + "tags": [ + "beamline:i22,proposal:12345", + "beamline:i22", + "visible:beamline:i22,proposal:12345", + ] + }, + ), ( "cm12345-1", - {"tags": ["beamline:i22,proposal:12345,session:12345-1"]}, + { + "tags": [ + "beamline:i22,proposal:12345,session:12345-1", + "beamline:i22,proposal:12345", + "beamline:i22", + ] + }, ), ], ) ] assert [key for key, *_ in skipped] == ["run-b", "i22"] + + +@pytest.mark.parametrize("level", [("i15-1",), ("i15-1", 12345), ("i15-1", 12345, 2)]) +def test_own_tag_gives_full_tags(level): + blob = tags(*level) + assert tags_from_own_tag(blob["tags"][0]) == blob + + +@pytest.mark.parametrize( + "tag", + [ + "11", # bare session id from the old policy; moved with its run instead + "visible:beamline:i22", + "beamline:i22,proposal:12345,session:999-1", # session of another proposal + "beamline:i22,session:12345-1", + ], +) +def test_other_tags_left_alone(tag): + assert tags_from_own_tag(tag) is None + + +SESSIONS = {"11": {"beamline": "b18", "proposal_number": 2024, "visit_number": 1}} + + +def test_session_id_gives_that_sessions_tags(): + assert tags_from_own_tag("11", SESSIONS) == tags("b18", 2024, 1) + + +@pytest.mark.parametrize("sessions", [None, SESSIONS]) +def test_unknown_session_id_left_alone(sessions): + assert tags_from_own_tag("12", sessions) is None + + +def test_own_tag_needs_no_sessions(): + blob = tags("b18", 2024, 1) + assert tags_from_own_tag(blob["tags"][0], SESSIONS) == blob From ef26d1de7f81943c0514dd27dcc0631a83609d6a Mon Sep 17 00:00:00 2001 From: Zoheb Shaikh <26975142+ZohebShaikh@users.noreply.github.com> Date: Sat, 17 Oct 2026 16:52:07 +0100 Subject: [PATCH 7/8] Pin the migration test stack's OPA to 1.21.2 edge-static-debug floated with every OPA build. Pin 1.21.2, the version authz now builds and tests its policy with; the debug variant is kept because entrypoint.sh needs a shell. Co-Authored-By: Claude Opus 5.5 --- database_migration/integration/compose.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/database_migration/integration/compose.yaml b/database_migration/integration/compose.yaml index 5c63b15..bb21373 100644 --- a/database_migration/integration/compose.yaml +++ b/database_migration/integration/compose.yaml @@ -37,7 +37,7 @@ services: # Started by the test once it has written opa_data/data.json from the dump. opa: - image: openpolicyagent/opa:edge-static-debug + image: openpolicyagent/opa:1.21.2-static-debug # debug: entrypoint.sh needs a shell network_mode: host # Port 8181 volumes: - ./opa:/mnt From 5a39180eb6e5a44f237fe39b9af92a63a8426a58 Mon Sep 17 00:00:00 2001 From: Zoheb Shaikh <26975142+ZohebShaikh@users.noreply.github.com> Date: Sat, 17 Oct 2026 16:58:14 +0100 Subject: [PATCH 8/8] Use authz-policy 0.0.28-alpha in the migration test stack 0.0.28-alpha is built from authz 452bd95: the same decisions as 0.0.27-alpha, with the policy written for OPA 1.21.2 (string interpolation, so it needs OPA 1.x to load). Co-Authored-By: Claude Opus 5.5 --- database_migration/integration/opa/config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/database_migration/integration/opa/config.yaml b/database_migration/integration/opa/config.yaml index d751f94..286fe7e 100644 --- a/database_migration/integration/opa/config.yaml +++ b/database_migration/integration/opa/config.yaml @@ -5,4 +5,4 @@ services: bundles: diamond-policies: service: ghcr - resource: ghcr.io/zohebshaikh/authz-policy:0.0.27-alpha + resource: ghcr.io/zohebshaikh/authz-policy:0.0.28-alpha