From a5a959ad50459e9dea9c8b6c5315ccadd67414db Mon Sep 17 00:00:00 2001 From: Rob Sherman Date: Mon, 29 Jun 2026 19:48:53 -0700 Subject: [PATCH] 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