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 new file mode 100644 index 0000000..b8650dd --- /dev/null +++ b/database_migration/README.md @@ -0,0 +1,236 @@ +# 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 (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 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 --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 and retags +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. + +## 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 < +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 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. 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. + +Example: + uv run --script 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+)$" +) + + +# 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 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): + """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, sessions: dict | None = None) -> 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.") + + # 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: + 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( + "--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", + 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." + ) + 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__": + main() diff --git a/database_migration/postgres_tools.py b/database_migration/postgres_tools.py new file mode 100755 index 0000000..5e6c482 --- /dev/null +++ b/database_migration/postgres_tools.py @@ -0,0 +1,149 @@ +#!/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 +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 + 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. + 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 +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..1bd5d0f --- /dev/null +++ b/database_migration/test_integration.py @@ -0,0 +1,375 @@ +"""Migrate a real tiled dump and read it back through dls-tiled, Keycloak and OPA. + +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 + +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") +# 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 psql(query: str) -> str: + return compose( + "exec", + "-T", + "postgres", + "psql", + "-U", + "postgres", + "-d", + "tiled", + "-AtXc", + query, + capture_output=True, + text=True, + ).stdout + + +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]: + """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") + 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], + 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" + + # The way the README runs it, with the script's own pinned dependencies. + subprocess.run( + ["uv", "run", "--script", 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"][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 new file mode 100644 index 0000000..93cb59b --- /dev/null +++ b/database_migration/test_migrate_node_structure.py @@ -0,0 +1,77 @@ +import pytest +from migrate_node_structure import plan_moves, tags, tags_from_own_tag + + +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", "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", + "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 diff --git a/src/dls_tiled/policy.py b/src/dls_tiled/policy.py index b27fc3e..8682a9d 100644 --- a/src/dls_tiled/policy.py +++ b/src/dls_tiled/policy.py @@ -4,11 +4,13 @@ (https://github.com/DiamondLightSource/diamond-tiled). """ +import base64 import json import logging +import time from fastapi import HTTPException -from pydantic import BaseModel, HttpUrl, TypeAdapter +from pydantic import BaseModel, HttpUrl, TypeAdapter, ValidationError from starlette.status import ( HTTP_401_UNAUTHORIZED, ) @@ -25,11 +27,63 @@ 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): - 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 _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): @@ -62,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, @@ -86,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( @@ -137,9 +195,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) @@ -154,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 d50ff90..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 @@ -6,22 +7,25 @@ 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 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, ) 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 +93,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 +136,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): @@ -204,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 ----------------------------------------------------------------- @@ -223,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",