668 lines
24 KiB
Python
668 lines
24 KiB
Python
#!/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 _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,
|
|
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 = "auto",
|
|
) -> None:
|
|
"""Apply the given plans: for each plan ensure dependencies, clear queue,
|
|
relocate, then flush and assert integrity.
|
|
|
|
``strategy`` selects the relocate implementation:
|
|
- ``"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).
|
|
"""
|
|
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(
|
|
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")
|
|
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
|
|
|
|
|
|
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, resolved_strategy: str) -> None:
|
|
mode = "APPLY" if apply else "DRY-RUN"
|
|
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(
|
|
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
|
|
)
|
|
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:
|
|
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,
|
|
strategy=args.strategy,
|
|
)
|
|
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())
|