From c0e88c32e3fc86e30018e5812712f9d7ce9cbd50 Mon Sep 17 00:00:00 2001 From: Alexj9837 <52531949+Alexj9837@users.noreply.github.com> Date: Tue, 13 Oct 2026 10:08:52 +0000 Subject: [PATCH] Add tiled data restructure scripts and docs Ported from DiamondLightSource/diamond-tiled#1. Changes from the original are lint-only: ruff format, import sorting and noqa N806 on the sessionmaker assignment. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01CCtPFejfd3WP589miPqaL3 --- data_restructure/tiled_data_restructure.md | 148 ++++++++ .../tiled_data_restructure_script.py | 282 ++++++++++++++ ...ata_restructure_split_existing_sessions.py | 349 ++++++++++++++++++ 3 files changed, 779 insertions(+) create mode 100644 data_restructure/tiled_data_restructure.md create mode 100644 data_restructure/tiled_data_restructure_script.py create mode 100644 data_restructure/tiled_data_restructure_split_existing_sessions.py diff --git a/data_restructure/tiled_data_restructure.md b/data_restructure/tiled_data_restructure.md new file mode 100644 index 0000000..9dc9914 --- /dev/null +++ b/data_restructure/tiled_data_restructure.md @@ -0,0 +1,148 @@ +# ACQP-568: Tiled catalog node restructuring + +Runbook for `tiled_data_restructure_script.py`, written 2026-08-04 after +running this for real against a p46 replica seeded with a real production +snapshot. If you're reading this because it's a few years later and the +target structure needs to change again, start here. + +## What this was for + +Diamond's production Tiled catalog (`tiled.diamond.ac.uk`) stored every +experimental run as a flat node directly under the catalog root. This +script reorganizes them into: + +``` +root/// +``` + +e.g. `root/p51/cm44254-1/`. + +The `instrument` and `instrument_session` values come from each run's own +metadata (`metadata.start.instrument` / `metadata.start.instrument_session` +— this is standard bluesky run-start-document metadata, not something +Diamond-specific), not from any existing path structure — the catalog was +completely flat before this ran, with ~3,000 top-level nodes each keyed by +an opaque UUID. + +**This was deliberately built as a one-off script, not reusable +infrastructure.** The team lead's call at the time: the target structure +was expected to be stable for 2-3 years, so it wasn't worth building this +as a separate ops repo/package. + +## Prerequisites / access setup + +Two separate Diamond Kubernetes clusters are involved, and switching +between them requires re-authenticating (they're not simultaneously +available in one terminal's kubeconfig): + +- **`pollux`** — hosts `p46-beamline` namespace, the staging/rehearsal + replica (`p46-tiled.diamond.ac.uk`). `module load pollux` (fedid + + password login, `klogout` to force a fresh login if you hit stale-auth + errors). + +- **`argus`** — hosts the `tiled` namespace, i.e. **production** + (`tiled.diamond.ac.uk`). `module load argus`, same login pattern. + +Run `kubectl config current-context` if something that worked a moment ago +suddenly returns "not found" — it's almost always because the context +silently switched clusters. + +Postgres access on both: connect as the **`postgres`** superuser (not the +bootstrapped `tiled`/`USERDB_USER` app account, which lacks table-level +grants), password is the `ADMIN_PASS` key in the `postgres-secret` +(`SealedSecret`) in that namespace, this can be found on the pollux. +Two databases matter: `tiled-db` (thecatalog — this is what this script operates on) + and `data` (writable storage backend, unrelated to node structure). + +To connect a local tool (psql, a GUI client, or this script) to a +cluster's postgres without exec-ing into the pod every time: + +``` +kubectl port-forward -n svc/ 5432:5432 +``` + +Leave that running in its own terminal; connect to `localhost:5432` from +wherever you're running the script/client. + +## The two implementations + +**`tiled_data_restructure_script.py`** (this directory) — the version +actually used. Python + SQLAlchemy async, does everything in one +transaction, set-based (not row-by-row) for performance. + +``` +python scripts/tiled_data_restructure_script.py [--dry-run] +``` + +``` +cd /scratch//path/to/tiled # or wherever you have it +python3.11 -m venv .venv +source .venv/bin/activate +pip install -e ".[all]" +``` + +**Always dry-run first**, on the replica, not production: + +``` +python scripts/tiled_data_restructure_script.py \ + "postgresql+asyncpg://postgres:@localhost:5432/tiled-db" --dry-run +``` + +This prints the instrument/session groupings and counts with zero writes. +Sanity-check the numbers against a manual query before trusting it: + +```sql +SELECT n2.metadata->'start'->>'instrument' AS instrument, + n2.metadata->'start'->>'instrument_session' AS session, + count(*) +FROM nodes n1 JOIN nodes n2 ON n2.parent = n1.id +WHERE n1.parent IS NULL +GROUP BY 1, 2 ORDER BY 1, 2; +``` + +Also worth checking before running for real: how many top-level nodes are +missing either metadata field (the script refuses to proceed if any are +found, but knowing the number up front avoids surprises): + +```sql +SELECT count(*) FILTER (WHERE n2.metadata->'start'->>'instrument' IS NULL) AS missing_instrument, + count(*) FILTER (WHERE n2.metadata->'start'->>'instrument_session' IS NULL) AS missing_session +FROM nodes n1 JOIN nodes n2 ON n2.parent = n1.id +WHERE n1.parent IS NULL; +``` + +## Rehearsal process used for the p46 run (repeat this for any future run) +these can be found also in the old Jira ticket, in the comments https://jira.diamond.ac.uk/browse/ACQP-568 + +1. Get a real snapshot to test against — don't invent synthetic data. + Dump production's `tiled-db` (and `data`, if relevant) via `pg_dump -F c` + from inside the postgres pod, `kubectl cp` it out, sanity-check with + `pg_restore -l` before trusting it. +2. Seed the target replica for real: `dropdb --force` + `createdb` the + replica's existing `tiled-db` (needs `--force` since the live app pod + holds open connections), then `pg_restore` the production dump into it. + Restart the replica's app deployment afterward + (`kubectl rollout restart deployment ...`) so it picks up a clean + connection rather than reusing stale state. +3. Confirm the app itself is healthy against the reseeded data — check + `kubectl logs` on the new pod for connection/startup errors, not just + that the database rows exist. +4. `--dry-run` the migration script, cross-check its printed counts + against a manual SQL query (above). +5. Run it for real (drop `--dry-run`). +6. Verify: spot-check that a specific instrument/session actually has the + expected run count as *children* of the new container (not just trust + the script's own summary), e.g.: + ```sql + SELECT count(*) FROM nodes n JOIN nodes sess ON n.parent = sess.id + WHERE sess.key = ''; + ``` + Compare against the dry-run's printed count for that session. + +script also stored here, as we're unsure if we need a tiled/tiled repo currently and I dont want to lose this. +""" + +Example: + python scripts/acqp568_migrate_node_structure.py \ + postgresql+asyncpg://postgres:@localhost:5432/tiled-db --dry-run +""" diff --git a/data_restructure/tiled_data_restructure_script.py b/data_restructure/tiled_data_restructure_script.py new file mode 100644 index 0000000..487e6f7 --- /dev/null +++ b/data_restructure/tiled_data_restructure_script.py @@ -0,0 +1,282 @@ +""" + +Example: + python scripts/acqp568_migrate_node_structure.py \ + postgresql+asyncpg://postgres:@localhost:5432/tiled-db --dry-run +""" + +import argparse +import asyncio +from collections import defaultdict + +from sqlalchemy import func, select, text +from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine +from sqlalchemy.orm import sessionmaker +from tiled.catalog.orm import Node +from tiled.structures.core import StructureFamily + + +async def get_root_id(session): + result = await session.execute(select(Node.id).where(Node.parent.is_(None))) + return result.scalar_one() + + +async def get_or_create_container(session, parent_id, key): + """Insert a container node if one doesn't already exist under parent_id + with this key. Returns (id, created). Uses ON CONFLICT rather than + check-then-insert, matching the pattern tiled's own create_node uses, + so this is race-free even if run concurrently. + """ + stmt = ( + pg_insert(Node) + .values( + key=key, + parent=parent_id, + metadata_={}, + structure_family=StructureFamily.container, + specs=[], + access_blob={}, + ) + .on_conflict_do_nothing(constraint="key_parent_unique_constraint") + .returning(Node.id) + ) + result = await session.execute(stmt) + row = result.first() + if row is not None: + return row.id, True + result = await session.execute( + select(Node.id).where(Node.parent == parent_id, Node.key == key) + ) + return result.scalar_one(), False + + +async def main(database_uri, dry_run): + engine = create_async_engine(database_uri) + Session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) # noqa: N806 + + async with Session() as session: + root_id = await get_root_id(session) + + result = await session.execute( + select( + Node.id, + Node.key, + Node.metadata_[["start", "instrument"]].as_string().label("instrument"), + Node.metadata_[["start", "instrument_session"]] + .as_string() + .label("session"), + ).where(Node.parent == root_id) + ) + runs = result.all() + + missing = [r for r in runs if r.instrument is None or r.session is None] + if missing: + raise ValueError( + f"{len(missing)} run(s) missing instrument/instrument_session " + f"metadata, refusing to proceed: {[r.key for r in missing][:10]}" + ) + + # instrument_session is "-" (e.g. "cm44254-3"); split + # on the last '-' so a proposal code containing '-' is handled correctly. + splits = {} + malformed = [] + for r in runs: + proposal, sep, visit = r.session.rpartition("-") + if not sep or not proposal or not visit: + malformed.append(r) + else: + splits[r.id] = (proposal, visit) + if malformed: + raise ValueError( + f"{len(malformed)} run(s) have an instrument_session that doesn't " + f"split into '-', refusing to proceed: " + f"{[r.key for r in malformed][:10]}" + ) + + by_target = defaultdict(list) + for run in runs: + proposal, visit = splits[run.id] + by_target[(run.instrument, proposal, visit)].append(run) + + print( + f"Found {len(runs)} runs across {len({r.instrument for r in runs})} " + f"instruments and {len(by_target)} instrument/proposal/visit triples." + ) + for (instrument, proposal, visit), group in sorted(by_target.items()): + print(f" root/{instrument}/{proposal}/{visit}: {len(group)} runs") + + if dry_run: + print("Dry run -- no changes made.") + return + + before_count = ( + await session.execute(select(func.count()).select_from(Node)) + ).scalar_one() + + instrument_ids = {} + proposal_ids = {} + visit_ids = {} + containers_created = 0 + for instrument, proposal, visit in by_target: + if instrument not in instrument_ids: + instrument_ids[instrument], created = await get_or_create_container( + session, root_id, instrument + ) + if created: + containers_created += 1 + print(f"Created container root/{instrument}") + if (instrument, proposal) not in proposal_ids: + ( + proposal_ids[(instrument, proposal)], + created, + ) = await get_or_create_container( + session, instrument_ids[instrument], proposal + ) + if created: + containers_created += 1 + print(f"Created container root/{instrument}/{proposal}") + ( + visit_ids[(instrument, proposal, visit)], + created, + ) = await get_or_create_container( + session, proposal_ids[(instrument, proposal)], visit + ) + if created: + containers_created += 1 + print(f"Created container root/{instrument}/{proposal}/{visit}") + + # Build the move list once, then operate on all of it in a handful + # of set-based statements instead of looping node-by-node. + moves = [(run.id, visit_ids[(run.instrument, *splits[run.id])]) for run in runs] + + await session.execute( + text( + "CREATE TEMP TABLE moves " + "(node_id integer PRIMARY KEY, new_parent_id integer NOT NULL) " + "ON COMMIT DROP" + ) + ) + await session.execute( + text( + "INSERT INTO moves (node_id, new_parent_id) " + "VALUES (:node_id, :new_parent_id)" + ), + [{"node_id": n, "new_parent_id": p} for n, p in moves], + ) + + # Batched safety checks (structurally these should never fire, since + # every target is a container that didn't exist a moment ago, and + # every run has a unique key -- but cheap insurance regardless). + cycle_count = ( + await session.execute( + text( + """ + SELECT count(*) FROM moves m + WHERE EXISTS ( + SELECT 1 FROM nodes_closure + WHERE ancestor = m.node_id AND descendant = m.new_parent_id + ) + """ + ) + ) + ).scalar_one() + if cycle_count: + raise ValueError(f"{cycle_count} move(s) would create a cycle -- aborting.") + + collision_count = ( + await session.execute( + text( + """ + SELECT count(*) FROM moves m + WHERE EXISTS ( + SELECT 1 FROM nodes n2 + WHERE n2.parent = m.new_parent_id + AND n2.key = (SELECT key FROM nodes WHERE id = m.node_id) + AND n2.id != m.node_id + ) + """ + ) + ) + ).scalar_one() + if collision_count: + raise ValueError( + f"{collision_count} move(s) would collide on key -- aborting." + ) + + # Detach: remove closure rows linking each run's OLD ancestor chain + # to its subtree. Rows internal to each run's own subtree are + # untouched, since none of a run's descendants are themselves being + # moved. + await session.execute( + text( + """ + DELETE FROM nodes_closure + WHERE descendant IN ( + SELECT nc.descendant FROM nodes_closure nc + JOIN moves m ON nc.ancestor = m.node_id + ) + AND ancestor IN ( + SELECT nc2.ancestor FROM nodes_closure nc2 + JOIN moves m ON nc2.descendant = m.node_id + WHERE nc2.ancestor != m.node_id + ) + """ + ) + ) + + # Reattach: connect each run's NEW ancestor chain to its subtree. + await session.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_id + """ + ) + ) + + # Update the adjacency pointer itself, for all moved runs at once. + await session.execute( + text( + """ + UPDATE nodes n + SET parent = m.new_parent_id + FROM moves m + WHERE n.id = m.node_id + """ + ) + ) + + after_count = ( + await session.execute(select(func.count()).select_from(Node)) + ).scalar_one() + expected = before_count + containers_created + if after_count != expected: + raise ValueError( + f"Post-move node count ({after_count}) does not match expected " + f"({expected} = {before_count} before + {containers_created} new " + "containers) -- rolling back rather than committing." + ) + + await session.commit() + print( + f"Moved {len(moves)} runs into {containers_created} new container(s). " + f"Node count: {before_count} -> {after_count}." + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "database_uri", help="e.g. postgresql+asyncpg://user:pass@host:5432/dbname" + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Preview the mapping without making changes", + ) + args = parser.parse_args() + asyncio.run(main(args.database_uri, args.dry_run)) diff --git a/data_restructure/tiled_data_restructure_split_existing_sessions.py b/data_restructure/tiled_data_restructure_split_existing_sessions.py new file mode 100644 index 0000000..737626a --- /dev/null +++ b/data_restructure/tiled_data_restructure_split_existing_sessions.py @@ -0,0 +1,349 @@ +"""Variant of tiled_data_restructure_script.py for catalogs that have +already been migrated once to root/// +(the two-level structure) and now need the extra proposal/visit split: +root////. + +Unlike the original script (which reads instrument/instrument_session off +each run's own metadata and assumes runs sit flat under root), this one +reads the proposal/visit split from each existing SESSION CONTAINER's own +key (e.g. "cm44254-1") and moves that container's children -- it doesn't +touch run metadata at all. + +Example: + python tiled_data_restructure_split_existing_sessions.py \ + postgresql+asyncpg://postgres:@localhost:5432/tiled-db --dry-run +""" + +import argparse +import asyncio + +from sqlalchemy import func, select, text +from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine +from sqlalchemy.orm import sessionmaker +from tiled.catalog.orm import Node +from tiled.structures.core import StructureFamily + + +async def get_root_id(session): + result = await session.execute(select(Node.id).where(Node.parent.is_(None))) + return result.scalar_one() + + +async def get_or_create_container(session, parent_id, key): + """Insert a container node if one doesn't already exist under parent_id + with this key. Returns (id, created). Uses ON CONFLICT rather than + check-then-insert, matching the pattern tiled's own create_node uses, + so this is race-free even if run concurrently. + """ + stmt = ( + pg_insert(Node) + .values( + key=key, + parent=parent_id, + metadata_={}, + structure_family=StructureFamily.container, + specs=[], + access_blob={}, + ) + .on_conflict_do_nothing(constraint="key_parent_unique_constraint") + .returning(Node.id) + ) + result = await session.execute(stmt) + row = result.first() + if row is not None: + return row.id, True + result = await session.execute( + select(Node.id).where(Node.parent == parent_id, Node.key == key) + ) + return result.scalar_one(), False + + +async def main(database_uri, dry_run): + engine = create_async_engine(database_uri) + Session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) # noqa: N806 + + async with Session() as session: + root_id = await get_root_id(session) + + instruments = ( + await session.execute( + select(Node.id, Node.key, Node.structure_family).where( + Node.parent == root_id + ) + ) + ).all() + non_containers = [ + i for i in instruments if i.structure_family != StructureFamily.container + ] + if non_containers: + raise ValueError( + f"{len(non_containers)} node(s) directly under root aren't " + f"containers -- expected only instrument containers here " + f"(has the two-level migration already run?): " + f"{[i.key for i in non_containers][:10]}" + ) + instrument_by_id = {i.id: i.key for i in instruments} + + sessions = ( + await session.execute( + select(Node.id, Node.key, Node.parent, Node.structure_family).where( + Node.parent.in_(list(instrument_by_id)) + ) + ) + ).all() + non_containers = [ + s for s in sessions if s.structure_family != StructureFamily.container + ] + if non_containers: + raise ValueError( + f"{len(non_containers)} node(s) directly under an instrument " + f"aren't containers -- expected only session containers here: " + f"{[s.key for s in non_containers][:10]}" + ) + + malformed = [] + splits = {} + for s in sessions: + proposal, sep, visit = s.key.rpartition("-") + if not sep or not proposal or not visit: + malformed.append(s) + else: + splits[s.id] = (proposal, visit) + if malformed: + raise ValueError( + f"{len(malformed)} session container(s) don't split into " + f"'-', refusing to proceed: " + f"{[s.key for s in malformed][:10]}" + ) + + child_counts = dict( + ( + await session.execute( + select(Node.parent, func.count()) + .where(Node.parent.in_([s.id for s in sessions])) + .group_by(Node.parent) + ) + ).all() + ) + + to_split = [s for s in sessions if child_counts.get(s.id, 0) > 0] + empty = [s for s in sessions if child_counts.get(s.id, 0) == 0] + + print( + f"Found {len(sessions)} session container(s) across " + f"{len(instrument_by_id)} instrument(s); {len(to_split)} have " + f"runs to move, {len(empty)} are already empty and will be left " + "alone." + ) + for s in sorted(to_split, key=lambda s: (instrument_by_id[s.parent], s.key)): + proposal, visit = splits[s.id] + print( + f" root/{instrument_by_id[s.parent]}/{s.key} " + f"-> root/{instrument_by_id[s.parent]}/{proposal}/{visit}: " + f"{child_counts[s.id]} runs" + ) + if empty: + print(f" (skipping, no children): {[s.key for s in empty][:10]}") + + if dry_run: + print("Dry run -- no changes made.") + return + + if not to_split: + print("Nothing to do.") + return + + before_count = ( + await session.execute(select(func.count()).select_from(Node)) + ).scalar_one() + + proposal_ids = {} + visit_ids = {} + containers_created = 0 + for s in to_split: + instrument_id = s.parent + proposal, visit = splits[s.id] + if (instrument_id, proposal) not in proposal_ids: + ( + proposal_ids[(instrument_id, proposal)], + created, + ) = await get_or_create_container(session, instrument_id, proposal) + if created: + containers_created += 1 + print( + f"Created container root/{instrument_by_id[instrument_id]}" + f"/{proposal}" + ) + visit_ids[s.id], created = await get_or_create_container( + session, proposal_ids[(instrument_id, proposal)], visit + ) + if created: + containers_created += 1 + print( + f"Created container root/{instrument_by_id[instrument_id]}" + f"/{proposal}/{visit}" + ) + + runs = ( + await session.execute( + select(Node.id, Node.parent).where( + Node.parent.in_([s.id for s in to_split]) + ) + ) + ).all() + moves = [(r.id, visit_ids[r.parent]) for r in runs] + + await session.execute( + text( + "CREATE TEMP TABLE moves " + "(node_id integer PRIMARY KEY, new_parent_id integer NOT NULL) " + "ON COMMIT DROP" + ) + ) + await session.execute( + text( + "INSERT INTO moves (node_id, new_parent_id) " + "VALUES (:node_id, :new_parent_id)" + ), + [{"node_id": n, "new_parent_id": p} for n, p in moves], + ) + + # Batched safety checks -- structurally these should never fire + # (every target is a container that either pre-dates this run or + # was just created, and every moved node has a unique key) but + # cheap insurance regardless. + cycle_count = ( + await session.execute( + text( + """ + SELECT count(*) FROM moves m + WHERE EXISTS ( + SELECT 1 FROM nodes_closure + WHERE ancestor = m.node_id AND descendant = m.new_parent_id + ) + """ + ) + ) + ).scalar_one() + if cycle_count: + raise ValueError(f"{cycle_count} move(s) would create a cycle -- aborting.") + + collision_count = ( + await session.execute( + text( + """ + SELECT count(*) FROM moves m + WHERE EXISTS ( + SELECT 1 FROM nodes n2 + WHERE n2.parent = m.new_parent_id + AND n2.key = (SELECT key FROM nodes WHERE id = m.node_id) + AND n2.id != m.node_id + ) + """ + ) + ) + ).scalar_one() + if collision_count: + raise ValueError( + f"{collision_count} move(s) would collide on key -- aborting." + ) + + # Detach: remove closure rows linking each run's OLD ancestor chain + # (through the old session container) to its subtree. + await session.execute( + text( + """ + DELETE FROM nodes_closure + WHERE descendant IN ( + SELECT nc.descendant FROM nodes_closure nc + JOIN moves m ON nc.ancestor = m.node_id + ) + AND ancestor IN ( + SELECT nc2.ancestor FROM nodes_closure nc2 + JOIN moves m ON nc2.descendant = m.node_id + WHERE nc2.ancestor != m.node_id + ) + """ + ) + ) + + # Reattach: connect each run's NEW ancestor chain (through the new + # proposal/visit containers) to its subtree. + await session.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_id + """ + ) + ) + + # Update the adjacency pointer itself, for all moved runs at once. + await session.execute( + text( + """ + UPDATE nodes n + SET parent = m.new_parent_id + FROM moves m + WHERE n.id = m.node_id + """ + ) + ) + + # Old session containers should now be empty -- confirm that before + # deleting them, since a delete cascades to any remaining children. + session_ids_to_delete = [s.id for s in to_split] + remaining = ( + await session.execute( + select(func.count()) + .select_from(Node) + .where(Node.parent.in_(session_ids_to_delete)) + ) + ).scalar_one() + if remaining: + raise ValueError( + f"{remaining} node(s) still parented under an old session " + "container after the move -- refusing to delete it, aborting." + ) + + await session.execute( + text("DELETE FROM nodes WHERE id = ANY(:ids)"), + {"ids": session_ids_to_delete}, + ) + + after_count = ( + await session.execute(select(func.count()).select_from(Node)) + ).scalar_one() + expected = before_count + containers_created - len(session_ids_to_delete) + if after_count != expected: + raise ValueError( + f"Post-move node count ({after_count}) does not match expected " + f"({expected} = {before_count} before + {containers_created} new " + f"containers - {len(session_ids_to_delete)} emptied session " + "containers removed) -- rolling back rather than committing." + ) + + await session.commit() + print( + f"Moved {len(moves)} runs into {containers_created} new container(s) " + f"and removed {len(session_ids_to_delete)} emptied session " + f"container(s). Node count: {before_count} -> {after_count}." + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "database_uri", help="e.g. postgresql+asyncpg://user:pass@host:5432/dbname" + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Preview the mapping without making changes", + ) + args = parser.parse_args() + asyncio.run(main(args.database_uri, args.dry_run))