From c0d6bc6f3b00f7d0e7f94e042c07da02f88755af Mon Sep 17 00:00:00 2001 From: Rob Sherman Date: Mon, 29 Jun 2026 18:03:59 -0700 Subject: [PATCH 1/2] feat(scripts): add tool to losslessly move sessions between workspaces Adds scripts/move_session_workspace.py, an admin tool that relocates named sessions (with their messages, embeddings, and conclusions) from one workspace to another without data loss. Preserves session IDs in place via transaction-local deferrable constraints, copies any missing target peers/collections full-column, renames on collision, clears transient queue items, and asserts referential integrity inside a single pg_dump-backed transaction. Dry-run by default; a no-privilege create-new-row fallback exists for deployments that cannot defer constraints. Includes 20 tests (tests/scripts/test_move_session_workspace.py). --- scripts/move_session_workspace.py | 618 ++++++++++++++++++ tests/scripts/test_move_session_workspace.py | 638 +++++++++++++++++++ 2 files changed, 1256 insertions(+) create mode 100644 scripts/move_session_workspace.py create mode 100644 tests/scripts/test_move_session_workspace.py diff --git a/scripts/move_session_workspace.py b/scripts/move_session_workspace.py new file mode 100644 index 00000000..342c9443 --- /dev/null +++ b/scripts/move_session_workspace.py @@ -0,0 +1,618 @@ +#!/usr/bin/env uv run python +"""Losslessly move named sessions between Honcho workspaces. See +docs/superpowers/specs/2026-06-29-cross-workspace-session-move-design.md.""" + +from __future__ import annotations + +import argparse +import asyncio +import os +import sys +from dataclasses import dataclass, field +from typing import Any, cast + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from sqlalchemy import CursorResult, delete, func, select, text, update +from sqlalchemy.ext.asyncio import AsyncSession + +from src import models + +# A session-scoped child model: every member carries ``workspace_name`` and +# ``session_name`` columns that the relocation moves in place. +_ChildModelType = ( + type[models.Message] + | type[models.MessageEmbedding] + | type[models.Document] + | type[models.SessionPeer] +) +# A model whose full-column copy ``_copy_row`` produces (parent/dependency rows). +_CopyableModel = type[models.Peer] | type[models.Collection] | type[models.Session] + +_CHILD_MODELS: tuple[_ChildModelType, ...] = ( + models.Message, + models.MessageEmbedding, + models.Document, + models.SessionPeer, +) + + +class MoveError(Exception): + pass + + +@dataclass +class SessionPlan: + source_name: str + target_name: str + renamed: bool + messages: int + embeddings: int = 0 + documents: int = 0 + peers_to_create: list[str] = field(default_factory=list) + collections_to_create: list[tuple[str, str]] = field(default_factory=list) + queue_rows: int = 0 + cross_boundary_premises: list[str] = field(default_factory=list) + + +async def _workspace_exists(session: AsyncSession, name: str) -> bool: + r = await session.scalar( + select(models.Workspace.name).where(models.Workspace.name == name) + ) + return r is not None + + +async def _session_row( + session: AsyncSession, ws: str, name: str +) -> models.Session | None: + return await session.scalar( + select(models.Session).where( + models.Session.workspace_name == ws, models.Session.name == name + ) + ) + + +async def _count( + session: AsyncSession, model: _ChildModelType, ws: str, name: str +) -> int: + return ( + await session.scalar( + select(func.count()) + .select_from(model) + .where(model.workspace_name == ws, model.session_name == name) + ) + or 0 + ) + + +async def _resolve_target_name( + session: AsyncSession, + target_ws: str, + name: str, + on_collision: str, + rename_suffix: str, + source_ws: str, +) -> tuple[str, bool, bool]: + if await _session_row(session, target_ws, name) is None: + return name, False, False + if on_collision == "skip": + return name, False, True + base = name + rename_suffix.format(source=source_ws) + candidate, n = base, 1 + while await _session_row(session, target_ws, candidate) is not None: + n += 1 + candidate = f"{base}-{n}" + return candidate, True, False + + +def _copy_row(src_obj: Any, model: _CopyableModel, **overrides: Any) -> Any: + """Full-column copy of an ORM row into a new instance, with overrides.""" + data = {c.name: getattr(src_obj, c.name) for c in model.__table__.columns} + data.pop("id", None) # let the nanoid default generate a fresh PK + data.update(overrides) + return model(**data) + + +async def _required_peers(session: AsyncSession, ws: str, name: str) -> set[str]: + peers: set[str] = set() + peers.update( + await session.scalars( + select(models.Message.peer_name.distinct()).where( + models.Message.workspace_name == ws, models.Message.session_name == name + ) + ) + ) + peers.update( + await session.scalars( + select(models.MessageEmbedding.peer_name.distinct()).where( + models.MessageEmbedding.workspace_name == ws, + models.MessageEmbedding.session_name == name, + ) + ) + ) + peers.update( + await session.scalars( + select(models.SessionPeer.peer_name.distinct()).where( + models.SessionPeer.workspace_name == ws, + models.SessionPeer.session_name == name, + ) + ) + ) + for col in (models.Document.observer, models.Document.observed): + peers.update( + await session.scalars( + select(col.distinct()).where( + models.Document.workspace_name == ws, + models.Document.session_name == name, + ) + ) + ) + peers.discard(None) + return peers + + +async def _required_collections( + session: AsyncSession, ws: str, name: str +) -> set[tuple[str, str]]: + rows = await session.execute( + select(models.Document.observer, models.Document.observed) + .where( + models.Document.workspace_name == ws, + models.Document.session_name == name, + ) + .distinct() + ) + return {(o, d) for o, d in rows.all()} + + +async def ensure_dependencies( + session: AsyncSession, + source_ws: str, + target_ws: str, + name: str, +) -> tuple[list[str], list[tuple[str, str]]]: + """Create missing target peers/collections as full-column copies from source. + + Existing rows in the target workspace are left untouched. Returns lists of + peer names and (observer, observed) pairs that were created. + """ + created_peers: list[str] = [] + for pname in sorted(await _required_peers(session, source_ws, name)): + exists = await session.scalar( + select(models.Peer).where( + models.Peer.workspace_name == target_ws, + models.Peer.name == pname, + ) + ) + if exists is None: + src = await session.scalar( + select(models.Peer).where( + models.Peer.workspace_name == source_ws, + models.Peer.name == pname, + ) + ) + if src is not None: + session.add(_copy_row(src, models.Peer, workspace_name=target_ws)) + created_peers.append(pname) + + # Flush peers before collections: Collection has FK to peers in same workspace. + if created_peers: + await session.flush() + + created_cols: list[tuple[str, str]] = [] + for obs, observed in sorted(await _required_collections(session, source_ws, name)): + exists = await session.scalar( + select(models.Collection).where( + models.Collection.workspace_name == target_ws, + models.Collection.observer == obs, + models.Collection.observed == observed, + ) + ) + if exists is None: + src = await session.scalar( + select(models.Collection).where( + models.Collection.workspace_name == source_ws, + models.Collection.observer == obs, + models.Collection.observed == observed, + ) + ) + if src is not None: + session.add(_copy_row(src, models.Collection, workspace_name=target_ws)) + created_cols.append((obs, observed)) + + return created_peers, created_cols + + +async def _session_fk_constraints(session: AsyncSession) -> list[tuple[str, str]]: + """Return ``(child_table, conname)`` for every FK whose ``confrelid`` is + ``sessions`` (the composite child FKs plus ``queue.session_id``).""" + rows = await session.execute( + text( + "SELECT conrelid::regclass::text AS child, conname" + + " FROM pg_constraint" + + " WHERE contype='f' AND confrelid='sessions'::regclass" + ) + ) + return [(r.child, r.conname) for r in rows] + + +async def relocate_in_place( + session: AsyncSession, + source_ws: str, + target_ws: str, + source_name: str, + target_name: str, +) -> None: + """Move a session row and its children to a new ``(name, workspace_name)`` + in place, preserving every ``id``/``public_id``. + + Uses transaction-local deferrable constraints so the parent and child + composite FKs are checked together at drain time rather than per-statement. + """ + fks = await _session_fk_constraints(session) + # 1. make the session FKs deferrable for this transaction + for child, conname in fks: + await session.execute( + text(f'ALTER TABLE {child} ALTER CONSTRAINT "{conname}" DEFERRABLE') + ) + await session.execute(text("SET CONSTRAINTS ALL DEFERRED")) + # 2. move the parent row in place (id/created_at/metadata preserved) + await session.execute( + update(models.Session) + .where( + models.Session.workspace_name == source_ws, + models.Session.name == source_name, + ) + .values(workspace_name=target_ws, name=target_name) + ) + # 3. move children in place (public_id/id preserved -> no CASCADE). + # queue is intentionally excluded from _CHILD_MODELS (handled later). + for model in _CHILD_MODELS: + await session.execute( + update(model) + .where( + model.workspace_name == source_ws, + model.session_name == source_name, + ) + .values(workspace_name=target_ws, session_name=target_name) + ) + # 4. drain deferred checks (now consistent) BEFORE restoring NOT DEFERRABLE + await session.execute(text("SET CONSTRAINTS ALL IMMEDIATE")) + for child, conname in fks: + await session.execute( + text(f'ALTER TABLE {child} ALTER CONSTRAINT "{conname}" NOT DEFERRABLE') + ) + + +async def relocate_create_new( + session: AsyncSession, + source_ws: str, + target_ws: str, + source_name: str, + target_name: str, +) -> None: + """Move a session to a new ``(name, workspace_name)`` by creating a fresh + session row (full-column copy, new id), repointing children to the new row, + then deleting the old session row. + + This is the no-privilege fallback for deployments where deferrable + constraints cannot be used. Queue rows MUST be cleared before calling this + (``apply_moves`` does so automatically). + """ + old = await _session_row(session, source_ws, source_name) + if old is None: + return + # 1. new target session row (full-column copy, fresh id) + session.add( + _copy_row(old, models.Session, workspace_name=target_ws, name=target_name) + ) + await session.flush() + # 2. repoint children to the new (name, workspace) — both rows exist, FK resolves + for model in _CHILD_MODELS: + await session.execute( + update(model) + .where(model.workspace_name == source_ws, model.session_name == source_name) + .values(workspace_name=target_ws, session_name=target_name) + ) + # 3. delete the now-unreferenced old session row + # (queue rows were already cleared by apply_moves before this call) + await session.execute(delete(models.Session).where(models.Session.id == old.id)) + + +async def clear_session_queue( + session: AsyncSession, ws: str, name: str, force: bool +) -> int: + """Delete all queue rows for ``session``/``ws``/``name``. + + If any rows are unprocessed and ``force`` is False, raises ``MoveError`` + rather than deleting. Returns the count of rows deleted. Queue rows are + transient work-state whose ``work_unit_key``/``payload`` embed the old + workspace identity, so they are never repointed — only cleared. + """ + sess = await _session_row(session, ws, name) + if sess is None: + return 0 + pending = await session.scalar( + select(func.count()) + .select_from(models.QueueItem) + .where( + models.QueueItem.session_id == sess.id, + models.QueueItem.processed.is_(False), + ) + ) + if pending and not force: + raise MoveError( + f"session '{name}' has {pending} pending queue items;" + + " re-run with --force-clear-queue to delete them" + ) + result = cast( + "CursorResult[Any]", + await session.execute( + delete(models.QueueItem).where(models.QueueItem.session_id == sess.id) + ), + ) + return result.rowcount or 0 + + +async def cross_boundary_premises( + session: AsyncSession, ws: str, moved_names: set[str] +) -> list[str]: + """Return premise doc ids cited by docs in moved sessions that originate + outside the move set (peer-global or in a non-moved session). + + Co-moved premises (``session_name in moved_names``) are NOT flagged. + This is a read-only report. + """ + # Collect all premise ids cited by docs in the moved sessions + rows = await session.scalars( + select(models.Document.source_ids).where( + models.Document.workspace_name == ws, + models.Document.session_name.in_(moved_names), + ) + ) + premise_ids: set[str] = set() + for sid_list in rows: + if sid_list: + premise_ids.update(sid_list) + if not premise_ids: + return [] + # Find premises that are peer-global or in a non-moved session + flagged: list[str] = [] + prem_rows = await session.execute( + select(models.Document.id, models.Document.session_name).where( + models.Document.workspace_name == ws, + models.Document.id.in_(premise_ids), + ) + ) + for doc_id, sess_name in prem_rows.all(): + if sess_name is None or sess_name not in moved_names: + flagged.append(doc_id) + return flagged + + +async def _assert_integrity(session: AsyncSession, _target_ws: str) -> None: + """Raise MoveError if any child row has an unparented (session_name, workspace_name) + or any queue row references a missing session. + + ``_target_ws`` is part of the documented signature for a future scoped check; + the current integrity sweep is workspace-agnostic, so it is unused here. + """ + for model in _CHILD_MODELS: + tname = getattr(model, "__tablename__", None) or model.__table__.name + # A NULL session_name is legitimately session-less (peer-global documents), + # not an orphan — only flag rows whose non-null session_name fails to resolve. + orphan = await session.scalar( + text( + f"SELECT 1 FROM {tname} c" + + " LEFT JOIN sessions s ON s.name=c.session_name AND s.workspace_name=c.workspace_name" + + " WHERE c.session_name IS NOT NULL AND s.name IS NULL LIMIT 1" + ) + ) + if orphan: + raise MoveError(f"integrity: orphaned rows in {tname}") + dangling = await session.scalar( + text( + "SELECT 1 FROM queue q LEFT JOIN sessions s ON s.id=q.session_id" + + " WHERE q.session_id IS NOT NULL AND s.id IS NULL LIMIT 1" + ) + ) + if dangling: + raise MoveError("integrity: queue rows reference a missing session") + + +async def apply_moves( + session: AsyncSession, + source_ws: str, + target_ws: str, + plans: list[SessionPlan], + force_clear_queue: bool, + strategy: str = "in_place", +) -> None: + """Apply the given plans: for each plan ensure dependencies, clear queue, + relocate, then flush and assert integrity. + + ``strategy`` selects the relocate implementation: + - ``"in_place"`` (default): id-preserving, requires deferrable constraints. + - ``"create_new"``: creates a fresh session row; works without deferrable + constraints but the session id changes. + + The caller controls the outer transaction (commit/rollback). + """ + relocate = relocate_in_place if strategy == "in_place" else relocate_create_new + for plan in plans: + await ensure_dependencies(session, source_ws, target_ws, plan.source_name) + await clear_session_queue( + session, source_ws, plan.source_name, force=force_clear_queue + ) + await relocate( + session, source_ws, target_ws, plan.source_name, plan.target_name + ) + await session.flush() + await _assert_integrity(session, target_ws) + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser( + description="Losslessly move sessions between Honcho workspaces." + ) + p.add_argument("--from", dest="source", required=True) + p.add_argument("--to", dest="target", required=True) + p.add_argument("--session", action="append", required=True, help="repeatable") + p.add_argument("--on-collision", choices=["rename", "skip"], default="rename") + p.add_argument("--rename-suffix", default="-from-{source}") + p.add_argument("--apply", action="store_true", help="default: dry-run") + p.add_argument("--force-clear-queue", action="store_true") + p.add_argument("--no-backup", action="store_true") + return p + + +def _pg_dump(out_path: str) -> None: + import subprocess + + from src.config import settings + + uri = settings.DB.CONNECTION_URI + # strip the +psycopg driver suffix for libpq / pg_dump + libpq = uri.replace("postgresql+psycopg", "postgresql") + subprocess.run(["pg_dump", "--dbname", libpq, "--file", out_path], check=True) + + +def _print_plans(plans: list[SessionPlan], apply: bool) -> None: + mode = "APPLY" if apply else "DRY-RUN" + print(f"[{mode}] {len(plans)} session(s):") + for p in plans: + rn = f" -> {p.target_name} (renamed)" if p.renamed else "" + print( + f" {p.source_name}{rn}: {p.messages} msgs, {p.documents} docs," + + f" {p.embeddings} embeddings; create peers={p.peers_to_create}" + + f" collections={p.collections_to_create}; queue={p.queue_rows}" + ) + if p.cross_boundary_premises: + print( + " WARNING cross-boundary premises (will dangle): " + + f"{p.cross_boundary_premises}" + ) + + +async def main_async(args: argparse.Namespace) -> int: + # Planning is read-only; apply uses a separate write session so that + # session.begin() is the FIRST DB operation on it, making it a true outer + # transaction (not a savepoint nested inside an auto-begun T1). + from src.db import ReadSessionLocal, SessionLocal + + async with ReadSessionLocal() as r: + plans = await plan_moves( + r, + args.source, + args.target, + args.session, + args.on_collision, + args.rename_suffix, + ) + moved = {p.source_name for p in plans} + for p in plans: + p.cross_boundary_premises = await cross_boundary_premises( + r, args.source, moved + ) + _print_plans(plans, apply=args.apply) + if not args.apply: + return 0 + if not args.no_backup: + import datetime + + path = f"/tmp/honcho-backup-{datetime.datetime.now(datetime.timezone.utc):%Y%m%dT%H%M%SZ}.sql" + _pg_dump(path) + print(f"backup written: {path}") + # session.begin() is the first DB op on a fresh session → true outer txn + async with SessionLocal() as session, session.begin(): + await apply_moves( + session, args.source, args.target, plans, args.force_clear_queue + ) + print("move applied.") + return 0 + + +def main() -> int: + return asyncio.run(main_async(build_parser().parse_args())) + + +async def plan_moves( + session: AsyncSession, + source_ws: str, + target_ws: str, + names: list[str], + on_collision: str = "rename", + rename_suffix: str = "-from-{source}", +) -> list[SessionPlan]: + if source_ws == target_ws: + raise MoveError("source and target are the same workspace") + if not await _workspace_exists(session, source_ws): + raise MoveError(f"source workspace '{source_ws}' not found") + if not await _workspace_exists(session, target_ws): + raise MoveError(f"target workspace '{target_ws}' not found") + + plans: list[SessionPlan] = [] + for name in names: + src = await _session_row(session, source_ws, name) + if src is None: + raise MoveError(f"session '{name}' not found in workspace '{source_ws}'") + target_name, renamed, skip = await _resolve_target_name( + session, target_ws, name, on_collision, rename_suffix, source_ws + ) + if skip: + continue + # Compute dry-run display fields (all SELECT-only) + required_peers = await _required_peers(session, source_ws, name) + peers_to_create: list[str] = [] + for pname in sorted(required_peers): + exists = await session.scalar( + select(models.Peer).where( + models.Peer.workspace_name == target_ws, + models.Peer.name == pname, + ) + ) + if exists is None: + peers_to_create.append(pname) + + required_cols = await _required_collections(session, source_ws, name) + collections_to_create: list[tuple[str, str]] = [] + for obs, observed in sorted(required_cols): + exists = await session.scalar( + select(models.Collection).where( + models.Collection.workspace_name == target_ws, + models.Collection.observer == obs, + models.Collection.observed == observed, + ) + ) + if exists is None: + collections_to_create.append((obs, observed)) + + queue_rows = ( + await session.scalar( + select(func.count()) + .select_from(models.QueueItem) + .where(models.QueueItem.session_id == src.id) + ) + or 0 + ) + + plans.append( + SessionPlan( + source_name=name, + target_name=target_name, + renamed=renamed, + messages=await _count(session, models.Message, source_ws, name), + embeddings=await _count( + session, models.MessageEmbedding, source_ws, name + ), + documents=await _count(session, models.Document, source_ws, name), + peers_to_create=peers_to_create, + collections_to_create=collections_to_create, + queue_rows=queue_rows, + ) + ) + return plans + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/scripts/test_move_session_workspace.py b/tests/scripts/test_move_session_workspace.py new file mode 100644 index 00000000..d034fe1a --- /dev/null +++ b/tests/scripts/test_move_session_workspace.py @@ -0,0 +1,638 @@ +import pytest +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from scripts.move_session_workspace import ( + MoveError, + _count, # pyright: ignore[reportPrivateUsage] + plan_moves, +) +from scripts.move_session_workspace import ( + _session_row as _session_row_helper, # pyright: ignore[reportPrivateUsage] +) +from src import models + + +async def _mk_workspace(db: AsyncSession, name: str) -> None: + db.add(models.Workspace(name=name)) + await db.flush() + + +async def _mk_session_with_messages( + db: AsyncSession, ws: str, name: str, peer: str, n: int +) -> None: + db.add(models.Peer(name=peer, workspace_name=ws)) + db.add(models.Session(name=name, workspace_name=ws)) + await db.flush() + for i in range(n): + db.add( + models.Message( + session_name=name, + workspace_name=ws, + peer_name=peer, + content=f"m{i}", + seq_in_session=i, + ) + ) + await db.flush() + + +@pytest.mark.asyncio +async def test_plan_clean_session_counts(db_session: AsyncSession): + await _mk_workspace(db_session, "personal") + await _mk_workspace(db_session, "highway") + await _mk_session_with_messages(db_session, "personal", "s1", "robsherman", 3) + + plans = await plan_moves(db_session, "personal", "highway", ["s1"]) + + assert len(plans) == 1 + assert plans[0].source_name == "s1" + assert plans[0].target_name == "s1" + assert plans[0].renamed is False + assert plans[0].messages == 3 + + +@pytest.mark.asyncio +async def test_plan_rejects_same_workspace(db_session: AsyncSession): + await _mk_workspace(db_session, "personal") + with pytest.raises(MoveError, match="same workspace"): + await plan_moves(db_session, "personal", "personal", ["s1"]) + + +@pytest.mark.asyncio +async def test_plan_rejects_missing_session(db_session: AsyncSession): + await _mk_workspace(db_session, "personal") + await _mk_workspace(db_session, "highway") + with pytest.raises(MoveError, match="not found"): + await plan_moves(db_session, "personal", "highway", ["nope"]) + + +@pytest.mark.asyncio +async def test_plan_renames_on_collision(db_session: AsyncSession): + await _mk_workspace(db_session, "personal") + await _mk_workspace(db_session, "highway") + await _mk_session_with_messages(db_session, "personal", "maca", "robsherman", 2) + await _mk_session_with_messages( + db_session, "highway", "maca", "robsherman", 5 + ) # collision + + plans = await plan_moves(db_session, "personal", "highway", ["maca"]) + assert plans[0].target_name == "maca-from-personal" + assert plans[0].renamed is True + + +@pytest.mark.asyncio +async def test_plan_skip_mode_leaves_collision(db_session: AsyncSession): + await _mk_workspace(db_session, "personal") + await _mk_workspace(db_session, "highway") + await _mk_session_with_messages(db_session, "personal", "maca", "robsherman", 2) + await _mk_session_with_messages(db_session, "highway", "maca", "robsherman", 5) + + plans = await plan_moves( + db_session, "personal", "highway", ["maca"], on_collision="skip" + ) + assert plans == [] # skipped, nothing to do + + +@pytest.mark.asyncio +async def test_ensure_dependencies_copies_missing_peer_fullcolumn( + db_session: AsyncSession, +): + await _mk_workspace(db_session, "personal") + await _mk_workspace(db_session, "highway") + # peer with metadata in personal; session uses it + db_session.add( + models.Peer( + name="robsherman", + workspace_name="personal", + internal_metadata={"card": "x"}, + ) + ) + db_session.add(models.Session(name="s1", workspace_name="personal")) + await db_session.flush() + db_session.add( + models.Message( + session_name="s1", + workspace_name="personal", + peer_name="robsherman", + content="hi", + seq_in_session=0, + ) + ) + await db_session.flush() + + from scripts.move_session_workspace import ensure_dependencies + + created_peers, _ = await ensure_dependencies( + db_session, "personal", "highway", "s1" + ) + await db_session.flush() + + assert created_peers == ["robsherman"] + moved = await db_session.scalar( + select(models.Peer).where( + models.Peer.workspace_name == "highway", + models.Peer.name == "robsherman", + ) + ) + assert moved is not None + assert moved.internal_metadata == { + "card": "x" + } # full-column copy preserved peer card + + +@pytest.mark.asyncio +async def test_ensure_dependencies_leaves_existing_peer_untouched( + db_session: AsyncSession, +): + await _mk_workspace(db_session, "personal") + await _mk_workspace(db_session, "highway") + db_session.add( + models.Peer( + name="robsherman", + workspace_name="personal", + internal_metadata={"card": "SOURCE"}, + ) + ) + db_session.add( + models.Peer( + name="robsherman", + workspace_name="highway", + internal_metadata={"card": "TARGET"}, + ) + ) + db_session.add(models.Session(name="s1", workspace_name="personal")) + await db_session.flush() + db_session.add( + models.Message( + session_name="s1", + workspace_name="personal", + peer_name="robsherman", + content="hi", + seq_in_session=0, + ) + ) + await db_session.flush() + + from scripts.move_session_workspace import ensure_dependencies + + created_peers, _ = await ensure_dependencies( + db_session, "personal", "highway", "s1" + ) + await db_session.flush() + + assert created_peers == [] # already present, not created + existing = await db_session.scalar( + select(models.Peer).where( + models.Peer.workspace_name == "highway", + models.Peer.name == "robsherman", + ) + ) + assert existing is not None + assert existing.internal_metadata == {"card": "TARGET"} # NOT clobbered + + +@pytest.mark.asyncio +async def test_relocate_preserves_id_and_moves_children(db_session: AsyncSession): + await _mk_workspace(db_session, "personal") + await _mk_workspace(db_session, "highway") + db_session.add( + models.Peer(name="robsherman", workspace_name="highway") + ) # target peer exists + db_session.add(models.Session(name="s1", workspace_name="personal")) + await db_session.flush() + src_sess = await db_session.scalar( + select(models.Session).where( + models.Session.workspace_name == "personal", + models.Session.name == "s1", + ) + ) + assert src_sess is not None + src_id, src_created = src_sess.id, src_sess.created_at + db_session.add(models.Peer(name="robsherman", workspace_name="personal")) + await db_session.flush() + db_session.add( + models.Message( + session_name="s1", + workspace_name="personal", + peer_name="robsherman", + content="hi", + seq_in_session=0, + ) + ) + await db_session.flush() + + from scripts.move_session_workspace import ensure_dependencies, relocate_in_place + + await ensure_dependencies(db_session, "personal", "highway", "s1") + await relocate_in_place(db_session, "personal", "highway", "s1", "s1") + await db_session.flush() + + moved = await db_session.scalar( + select(models.Session).where( + models.Session.workspace_name == "highway", + models.Session.name == "s1", + ) + ) + assert moved is not None + assert moved.id == src_id # id preserved + assert moved.created_at == src_created + assert await _count(db_session, models.Message, "highway", "s1") == 1 + assert await _count(db_session, models.Message, "personal", "s1") == 0 + # no orphaned source session row + assert await _session_row_helper(db_session, "personal", "s1") is None + + +@pytest.mark.asyncio +async def test_clear_queue_deletes_rows(db_session: AsyncSession): + await _mk_workspace(db_session, "personal") + db_session.add(models.Session(name="s1", workspace_name="personal")) + await db_session.flush() + # seed a processed queue row for the session + sess = await db_session.scalar( + select(models.Session).where( + models.Session.workspace_name == "personal", models.Session.name == "s1" + ) + ) + assert sess is not None + db_session.add( + models.QueueItem( + session_id=sess.id, + workspace_name="personal", + work_unit_key="test-key-1", + task_type="representation", + payload={}, + processed=True, + ) + ) + await db_session.flush() + + from scripts.move_session_workspace import clear_session_queue + + deleted = await clear_session_queue(db_session, "personal", "s1", force=False) + await db_session.flush() + assert deleted == 1 + + +@pytest.mark.asyncio +async def test_clear_queue_raises_on_pending_without_force(db_session: AsyncSession): + await _mk_workspace(db_session, "personal") + db_session.add(models.Session(name="s1", workspace_name="personal")) + await db_session.flush() + sess = await db_session.scalar( + select(models.Session).where( + models.Session.workspace_name == "personal", models.Session.name == "s1" + ) + ) + assert sess is not None + db_session.add( + models.QueueItem( + session_id=sess.id, + workspace_name="personal", + work_unit_key="test-key-2", + task_type="representation", + payload={}, + processed=False, + ) + ) + await db_session.flush() + + from scripts.move_session_workspace import clear_session_queue + + with pytest.raises(MoveError, match="pending queue items"): + await clear_session_queue(db_session, "personal", "s1", force=False) + + +@pytest.mark.asyncio +async def test_clear_queue_force_deletes_pending(db_session: AsyncSession): + await _mk_workspace(db_session, "personal") + db_session.add(models.Session(name="s1", workspace_name="personal")) + await db_session.flush() + sess = await db_session.scalar( + select(models.Session).where( + models.Session.workspace_name == "personal", models.Session.name == "s1" + ) + ) + assert sess is not None + db_session.add( + models.QueueItem( + session_id=sess.id, + workspace_name="personal", + work_unit_key="test-key-3", + task_type="representation", + payload={}, + processed=False, + ) + ) + await db_session.flush() + + from scripts.move_session_workspace import clear_session_queue + + deleted = await clear_session_queue(db_session, "personal", "s1", force=True) + await db_session.flush() + assert deleted == 1 + + +@pytest.mark.asyncio +async def test_apply_then_integrity_clean(db_session: AsyncSession): + await _mk_workspace(db_session, "personal") + await _mk_workspace(db_session, "highway") + await _mk_session_with_messages(db_session, "personal", "s1", "robsherman", 2) + + from scripts.move_session_workspace import apply_moves, plan_moves + + plans = await plan_moves(db_session, "personal", "highway", ["s1"]) + await apply_moves(db_session, "personal", "highway", plans, force_clear_queue=True) + await db_session.flush() + + assert await _count(db_session, models.Message, "highway", "s1") == 2 + assert await _count(db_session, models.Message, "personal", "s1") == 0 + + +@pytest.mark.asyncio +async def test_dry_run_writes_nothing(db_session: AsyncSession): + await _mk_workspace(db_session, "personal") + await _mk_workspace(db_session, "highway") + await _mk_session_with_messages(db_session, "personal", "s1", "robsherman", 2) + + from scripts.move_session_workspace import plan_moves + + before = await _count(db_session, models.Message, "personal", "s1") + await plan_moves(db_session, "personal", "highway", ["s1"]) # plan only, no apply + after = await _count(db_session, models.Message, "personal", "s1") + assert before == after == 2 # plan_moves is read-only + + +def test_build_parser_defaults(): + from scripts.move_session_workspace import build_parser + + args = build_parser().parse_args( + ["--from", "personal", "--to", "highway", "--session", "s1"] + ) + assert args.source == "personal" and args.target == "highway" + assert args.session == ["s1"] + assert args.apply is False # dry-run default + assert args.on_collision == "rename" + + +@pytest.mark.asyncio +async def test_cross_boundary_premises_flags_only_outside_move_set( + db_session: AsyncSession, +): + await _mk_workspace(db_session, "personal") + db_session.add(models.Peer(name="p", workspace_name="personal")) + for s in ("a", "b", "other"): + db_session.add(models.Session(name=s, workspace_name="personal")) + await db_session.flush() + # Collection required by Document FK (observer, observed, workspace_name) + db_session.add( + models.Collection(observer="p", observed="p", workspace_name="personal") + ) + await db_session.flush() + # premise docs + prem_co = models.Document( + workspace_name="personal", + session_name="b", + observer="p", + observed="p", + content="co", + source_ids=[], + ) + prem_out = models.Document( + workspace_name="personal", + session_name="other", + observer="p", + observed="p", + content="out", + source_ids=[], + ) + db_session.add_all([prem_co, prem_out]) + await db_session.flush() + # a conclusion in session "a" citing both premises + db_session.add( + models.Document( + workspace_name="personal", + session_name="a", + observer="p", + observed="p", + content="concl", + source_ids=[prem_co.id, prem_out.id], + ) + ) + await db_session.flush() + + from scripts.move_session_workspace import cross_boundary_premises + + flagged = await cross_boundary_premises(db_session, "personal", {"a", "b"}) + assert prem_out.id in flagged # "other" is outside the move set → flagged + assert prem_co.id not in flagged # "b" is co-moved → not flagged + + +@pytest.mark.asyncio +async def test_relocate_create_new_repoints_and_deletes_old(db_session: AsyncSession): + await _mk_workspace(db_session, "personal") + await _mk_workspace(db_session, "highway") + db_session.add(models.Peer(name="robsherman", workspace_name="highway")) + db_session.add(models.Peer(name="robsherman", workspace_name="personal")) + db_session.add(models.Session(name="s1", workspace_name="personal")) + await db_session.flush() + old = await db_session.scalar( + select(models.Session).where( + models.Session.workspace_name == "personal", models.Session.name == "s1" + ) + ) + assert old is not None + old_id = old.id + db_session.add( + models.Message( + session_name="s1", + workspace_name="personal", + peer_name="robsherman", + content="hi", + seq_in_session=0, + ) + ) + await db_session.flush() + + from scripts.move_session_workspace import relocate_create_new + + await relocate_create_new(db_session, "personal", "highway", "s1", "s1") + await db_session.flush() + + moved = await db_session.scalar( + select(models.Session).where( + models.Session.workspace_name == "highway", models.Session.name == "s1" + ) + ) + assert moved is not None and moved.id != old_id # id churns on fallback + assert await _count(db_session, models.Message, "highway", "s1") == 1 + assert await _session_row_helper(db_session, "personal", "s1") is None + + +@pytest.mark.asyncio +async def test_plan_populates_create_lists_and_queue_count(db_session: AsyncSession): + """plan_moves populates peers_to_create and queue_rows on the SessionPlan.""" + await _mk_workspace(db_session, "personal") + await _mk_workspace(db_session, "highway") + + # Create the source peer+session in personal; highway has NO peer yet (absent) + db_session.add(models.Peer(name="robsherman", workspace_name="personal")) + db_session.add(models.Session(name="s1", workspace_name="personal")) + await db_session.flush() + + db_session.add( + models.Message( + session_name="s1", + workspace_name="personal", + peer_name="robsherman", + content="hi", + seq_in_session=0, + ) + ) + await db_session.flush() + + # Seed a queue row for the source session + sess = await db_session.scalar( + select(models.Session).where( + models.Session.workspace_name == "personal", models.Session.name == "s1" + ) + ) + assert sess is not None + db_session.add( + models.QueueItem( + session_id=sess.id, + workspace_name="personal", + work_unit_key="plan-test-key", + task_type="representation", + payload={}, + processed=True, + ) + ) + await db_session.flush() + + plans = await plan_moves(db_session, "personal", "highway", ["s1"]) + + assert len(plans) == 1 + plan = plans[0] + assert "robsherman" in plan.peers_to_create # absent in target → needs creation + assert plan.queue_rows == 1 # the queue row we seeded + + +@pytest.mark.asyncio +async def test_apply_rename_path_end_to_end(db_session: AsyncSession): + """apply_moves renames the moved session when a same-named session exists in target.""" + await _mk_workspace(db_session, "personal") + await _mk_workspace(db_session, "highway") + + # Both workspaces have the peer + db_session.add(models.Peer(name="robsherman", workspace_name="personal")) + db_session.add(models.Peer(name="robsherman", workspace_name="highway")) + + # "maca" session in personal (source) with 3 messages + db_session.add(models.Session(name="maca", workspace_name="personal")) + # Pre-existing "maca" in highway (the collision) with 5 messages + db_session.add(models.Session(name="maca", workspace_name="highway")) + await db_session.flush() + + for i in range(3): + db_session.add( + models.Message( + session_name="maca", + workspace_name="personal", + peer_name="robsherman", + content=f"personal-m{i}", + seq_in_session=i, + ) + ) + for i in range(5): + db_session.add( + models.Message( + session_name="maca", + workspace_name="highway", + peer_name="robsherman", + content=f"highway-m{i}", + seq_in_session=i, + ) + ) + await db_session.flush() + + from scripts.move_session_workspace import apply_moves + + plans = await plan_moves(db_session, "personal", "highway", ["maca"]) + + # The plan should rename due to collision + assert len(plans) == 1 + assert plans[0].renamed is True + assert plans[0].target_name == "maca-from-personal" + + await apply_moves(db_session, "personal", "highway", plans, force_clear_queue=True) + await db_session.flush() + + # Moved session exists under renamed name in highway + renamed_sess = await _session_row_helper( + db_session, "highway", "maca-from-personal" + ) + assert renamed_sess is not None + + # Moved messages are now under the renamed session in highway + assert ( + await _count(db_session, models.Message, "highway", "maca-from-personal") == 3 + ) + + # Original personal session is gone + assert await _session_row_helper(db_session, "personal", "maca") is None + assert await _count(db_session, models.Message, "personal", "maca") == 0 + + # Pre-existing highway "maca" session is untouched with its original 5 messages + original_highway = await _session_row_helper(db_session, "highway", "maca") + assert original_highway is not None + assert await _count(db_session, models.Message, "highway", "maca") == 5 + + +def test_cli_runs_as_script_without_name_error(): + """Regression: running the module as a script must not NameError because a + function (plan_moves) is defined after the __main__ guard. --from==--to hits + the same-workspace guard before any DB query, so this needs no database.""" + import os + import subprocess + import sys + + repo = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + script = os.path.join(repo, "scripts", "move_session_workspace.py") + result = subprocess.run( + [sys.executable, script, "--from", "w", "--to", "w", "--session", "s"], + capture_output=True, + text=True, + cwd=repo, + ) + combined = result.stdout + result.stderr + assert "NameError" not in combined, combined + assert "same workspace" in combined, combined + + +@pytest.mark.asyncio +async def test_assert_integrity_ignores_peer_global_documents(db_session: AsyncSession): + """Regression: peer-global documents (session_name IS NULL) are legitimately + session-less and must NOT be flagged as orphans by _assert_integrity.""" + from scripts.move_session_workspace import ( + _assert_integrity, # pyright: ignore[reportPrivateUsage] + ) + + await _mk_workspace(db_session, "wsgi") + db_session.add(models.Peer(name="p", workspace_name="wsgi")) + await db_session.flush() # peer must exist before the collection FK + db_session.add(models.Collection(observer="p", observed="p", workspace_name="wsgi")) + await db_session.flush() + db_session.add( + models.Document( + workspace_name="wsgi", + session_name=None, # peer-global + observer="p", + observed="p", + content="global fact", + source_ids=[], + ) + ) + await db_session.flush() + # Must NOT raise (would raise "orphaned rows in documents" before the fix). + await _assert_integrity(db_session, "wsgi") From a5a959ad50459e9dea9c8b6c5315ccadd67414db Mon Sep 17 00:00:00 2001 From: Rob Sherman Date: Mon, 29 Jun 2026 19:48:53 -0700 Subject: [PATCH 2/2] feat(scripts): add --strategy flag with auto privilege detection for relocate fallback Adds _can_defer_constraints() to probe whether the current role can ALTER the sessions table's constraints, and _resolve_strategy() to resolve 'auto' (the new default) to 'in_place' or 'create_new' based on that probe. apply_moves() now defaults to strategy='auto' and raises MoveError on unknown values instead of silently routing to create_new. The CLI gains --strategy {auto,in_place,create_new} and the dry-run output shows which concrete strategy will run. Claude-Session: https://claude.ai/code/session_017E61SyPWyCgohR5wUwNXG7 --- scripts/move_session_workspace.py | 67 +++++++++++-- tests/scripts/test_move_session_workspace.py | 98 ++++++++++++++++++++ 2 files changed, 156 insertions(+), 9 deletions(-) diff --git a/scripts/move_session_workspace.py b/scripts/move_session_workspace.py index 342c9443..9654f42f 100644 --- a/scripts/move_session_workspace.py +++ b/scripts/move_session_workspace.py @@ -236,6 +236,36 @@ async def _session_fk_constraints(session: AsyncSession) -> list[tuple[str, str] return [(r.child, r.conname) for r in rows] +async def _can_defer_constraints(session: AsyncSession) -> bool: + """Return True if the current role can ALTER the ``sessions`` table's constraints. + + Heuristic: if the probe is wrong, the apply transaction rolls back (pg_dump-backed), + so a false positive is safe — the ALTER will fail and the outer transaction aborts. + """ + result = await session.scalar( + text( + "SELECT COALESCE(" + + " (SELECT rolsuper FROM pg_roles WHERE rolname = current_user)" + + " OR pg_catalog.pg_get_userbyid(relowner) = current_user, false)" + + " FROM pg_class WHERE relname = 'sessions' AND relkind = 'r'" + ) + ) + return bool(result) + + +async def _resolve_strategy(session: AsyncSession, strategy: str) -> str: + """Resolve ``strategy`` to a concrete ``"in_place"`` or ``"create_new"`` value. + + Raises ``MoveError`` for unrecognised values. ``"auto"`` probes the current + role's ALTER privilege and picks the best available strategy. + """ + if strategy not in {"auto", "in_place", "create_new"}: + raise MoveError(f"unknown strategy: {strategy!r}") + if strategy == "auto": + return "in_place" if await _can_defer_constraints(session) else "create_new" + return strategy + + async def relocate_in_place( session: AsyncSession, source_ws: str, @@ -426,19 +456,22 @@ async def apply_moves( target_ws: str, plans: list[SessionPlan], force_clear_queue: bool, - strategy: str = "in_place", + strategy: str = "auto", ) -> None: """Apply the given plans: for each plan ensure dependencies, clear queue, relocate, then flush and assert integrity. ``strategy`` selects the relocate implementation: - - ``"in_place"`` (default): id-preserving, requires deferrable constraints. - - ``"create_new"``: creates a fresh session row; works without deferrable - constraints but the session id changes. + - ``"auto"`` (default): probe ALTER privilege and pick the best available strategy. + - ``"in_place"``: id-preserving, requires table-owner or superuser role. + - ``"create_new"``: creates a fresh session row; works without special privilege + but the session id changes. + Raises ``MoveError`` for unrecognised strategy values. The caller controls the outer transaction (commit/rollback). """ - relocate = relocate_in_place if strategy == "in_place" else relocate_create_new + concrete = await _resolve_strategy(session, strategy) + relocate = relocate_in_place if concrete == "in_place" else relocate_create_new for plan in plans: await ensure_dependencies(session, source_ws, target_ws, plan.source_name) await clear_session_queue( @@ -463,6 +496,16 @@ def build_parser() -> argparse.ArgumentParser: p.add_argument("--apply", action="store_true", help="default: dry-run") p.add_argument("--force-clear-queue", action="store_true") p.add_argument("--no-backup", action="store_true") + p.add_argument( + "--strategy", + choices=["auto", "in_place", "create_new"], + default="auto", + help=( + "relocate strategy: auto=detect privilege (default)," + + " in_place=id-preserving (requires table-owner/superuser)," + + " create_new=no special privilege required (session id changes)" + ), + ) return p @@ -477,9 +520,9 @@ def _pg_dump(out_path: str) -> None: subprocess.run(["pg_dump", "--dbname", libpq, "--file", out_path], check=True) -def _print_plans(plans: list[SessionPlan], apply: bool) -> None: +def _print_plans(plans: list[SessionPlan], apply: bool, resolved_strategy: str) -> None: mode = "APPLY" if apply else "DRY-RUN" - print(f"[{mode}] {len(plans)} session(s):") + print(f"[{mode}] {len(plans)} session(s) strategy: {resolved_strategy}:") for p in plans: rn = f" -> {p.target_name} (renamed)" if p.renamed else "" print( @@ -514,7 +557,8 @@ async def main_async(args: argparse.Namespace) -> int: p.cross_boundary_premises = await cross_boundary_premises( r, args.source, moved ) - _print_plans(plans, apply=args.apply) + resolved = await _resolve_strategy(r, args.strategy) + _print_plans(plans, apply=args.apply, resolved_strategy=resolved) if not args.apply: return 0 if not args.no_backup: @@ -526,7 +570,12 @@ async def main_async(args: argparse.Namespace) -> int: # session.begin() is the first DB op on a fresh session → true outer txn async with SessionLocal() as session, session.begin(): await apply_moves( - session, args.source, args.target, plans, args.force_clear_queue + session, + args.source, + args.target, + plans, + args.force_clear_queue, + strategy=args.strategy, ) print("move applied.") return 0 diff --git a/tests/scripts/test_move_session_workspace.py b/tests/scripts/test_move_session_workspace.py index d034fe1a..ee03eb02 100644 --- a/tests/scripts/test_move_session_workspace.py +++ b/tests/scripts/test_move_session_workspace.py @@ -4,7 +4,9 @@ from sqlalchemy.ext.asyncio import AsyncSession from scripts.move_session_workspace import ( MoveError, + _can_defer_constraints, # pyright: ignore[reportPrivateUsage] _count, # pyright: ignore[reportPrivateUsage] + _resolve_strategy, # pyright: ignore[reportPrivateUsage] plan_moves, ) from scripts.move_session_workspace import ( @@ -636,3 +638,99 @@ async def test_assert_integrity_ignores_peer_global_documents(db_session: AsyncS await db_session.flush() # Must NOT raise (would raise "orphaned rows in documents" before the fix). await _assert_integrity(db_session, "wsgi") + + +def test_build_parser_strategy_default(): + """--strategy defaults to 'auto'; explicit --strategy create_new is accepted.""" + from scripts.move_session_workspace import build_parser + + args_default = build_parser().parse_args( + ["--from", "personal", "--to", "highway", "--session", "s1"] + ) + assert args_default.strategy == "auto" + + args_explicit = build_parser().parse_args( + [ + "--from", + "personal", + "--to", + "highway", + "--session", + "s1", + "--strategy", + "create_new", + ] + ) + assert args_explicit.strategy == "create_new" + + +@pytest.mark.asyncio +async def test_can_defer_constraints_true_in_test_env(db_session: AsyncSession): + """The test role can alter constraints (owner or superuser) so this returns True. + + The in-place relocation tests already prove the role can ALTER TABLE … ALTER CONSTRAINT, + so the probe must agree. + """ + result = await _can_defer_constraints(db_session) + assert result is True + + +@pytest.mark.asyncio +async def test_resolve_strategy_auto_picks_in_place(db_session: AsyncSession): + """With a privileged test role, auto-detection resolves to 'in_place'.""" + resolved = await _resolve_strategy(db_session, "auto") + assert resolved == "in_place" + + +@pytest.mark.asyncio +async def test_resolve_strategy_rejects_unknown(db_session: AsyncSession): + """An unrecognised strategy string must raise MoveError immediately.""" + with pytest.raises(MoveError, match="unknown strategy"): + await _resolve_strategy(db_session, "bogus") + + +@pytest.mark.asyncio +async def test_apply_moves_create_new_strategy(db_session: AsyncSession): + """apply_moves with strategy='create_new' moves session to target with a new id.""" + await _mk_workspace(db_session, "personal") + await _mk_workspace(db_session, "highway") + await _mk_session_with_messages(db_session, "personal", "s1", "robsherman", 2) + + # Capture the original session id before the move + old_sess = await db_session.scalar( + select(models.Session).where( + models.Session.workspace_name == "personal", + models.Session.name == "s1", + ) + ) + assert old_sess is not None + old_id = old_sess.id + + from scripts.move_session_workspace import apply_moves + + plans = await plan_moves(db_session, "personal", "highway", ["s1"]) + await apply_moves( + db_session, + "personal", + "highway", + plans, + force_clear_queue=False, + strategy="create_new", + ) + await db_session.flush() + + # Session exists in target + moved = await db_session.scalar( + select(models.Session).where( + models.Session.workspace_name == "highway", + models.Session.name == "s1", + ) + ) + assert moved is not None + assert moved.id != old_id # id must churn with create_new + + # Messages followed + assert await _count(db_session, models.Message, "highway", "s1") == 2 + assert await _count(db_session, models.Message, "personal", "s1") == 0 + # Source session row is gone + assert await _session_row_helper(db_session, "personal", "s1") is None