From 56f6572086f73ebbdbcc0f525b7cfac5d192e435 Mon Sep 17 00:00:00 2001 From: 1445043649 <> Date: Wed, 13 May 2026 09:10:49 +0800 Subject: [PATCH] =?UTF-8?q?feat(persistence):=20backfill=20Step=202=20?= =?UTF-8?q?=E2=80=94=20UPDATE=204=20tables=20from=20users.default=5Fworksp?= =?UTF-8?q?ace=5Fid?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-table correlated subquery UPDATE (portable across SQLite + Postgres). Filters workspace_id IS NULL AND user_id IS NOT NULL so already-tagged rows and orphan rows are skipped. Dry-run mode counts via a JOIN, never writes. Returns rows-updated for the orchestrator report. 2 new tests: single-user 4-table fan-out + multi-user isolation. Co-Authored-By: Claude Opus 4.7 (1M context) --- backend/scripts/backfill_workspace_id.py | 51 ++++++++++++-- backend/tests/test_backfill_workspace_id.py | 76 ++++++++++++++++++++- 2 files changed, 120 insertions(+), 7 deletions(-) diff --git a/backend/scripts/backfill_workspace_id.py b/backend/scripts/backfill_workspace_id.py index 04e68b81..7ab84d59 100644 --- a/backend/scripts/backfill_workspace_id.py +++ b/backend/scripts/backfill_workspace_id.py @@ -30,10 +30,15 @@ import asyncio import logging from typing import Any -from sqlalchemy import select, update +from sqlalchemy import func, select, update from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from app.gateway.auth.workspace_slug import auto_slug_from_email, next_available_slug +from deerflow.persistence.base import Base +from deerflow.persistence.feedback.model import FeedbackRow +from deerflow.persistence.models.run_event import RunEventRow +from deerflow.persistence.run.model import RunRow +from deerflow.persistence.thread_meta.model import ThreadMetaRow from deerflow.persistence.user.model import UserRow from deerflow.persistence.workspace import WorkspaceRepository from deerflow.persistence.workspace.sql import SLUG_BLACKLIST @@ -50,6 +55,16 @@ LEGACY_WORKSPACE_NAME = "Legacy Workspace" # The four business tables that gained ``workspace_id`` in alembic 0002. _BUSINESS_TABLES: tuple[str, ...] = ("threads_meta", "runs", "feedback", "run_events") +# Map table-name to ORM class so we can build a portable correlated UPDATE +# using SQLAlchemy expression language (SQLite < 3.33 lacks UPDATE-FROM +# but supports correlated subqueries on every version we ship). +_TABLE_MODELS: dict[str, type[Base]] = { + "threads_meta": ThreadMetaRow, + "runs": RunRow, + "feedback": FeedbackRow, + "run_events": RunEventRow, +} + async def _step1_create_workspaces_for_users( session_factory: async_sessionmaker[AsyncSession], @@ -111,12 +126,36 @@ async def _step2_update_table_from_users( ) -> int: """UPDATE *table* setting workspace_id from owner's users.default_workspace_id. - Filled in by T5.6. + Uses a correlated subquery (portable across SQLite + Postgres). Filters + ``workspace_id IS NULL AND user_id IS NOT NULL`` so already-set rows + and truly orphan rows are skipped (Step 3 handles the latter). + + Returns the number of rows updated (or that *would* be updated under + ``dry_run``). """ - _ = session_factory - _ = table - _ = dry_run - return 0 + model = _TABLE_MODELS[table] + workspace_col = model.workspace_id + user_col = model.user_id + + # Subquery: pull the user's default_workspace_id for each row. + correlated_default = select(UserRow.default_workspace_id).where(UserRow.id == user_col).scalar_subquery() + + if dry_run: + # Count rows whose owner has a default_workspace_id assigned — only + # those would get touched by the actual UPDATE. + count_stmt = select(func.count()).select_from(model).join(UserRow, UserRow.id == user_col).where(workspace_col.is_(None), user_col.is_not(None), UserRow.default_workspace_id.is_not(None)) + async with session_factory() as session: + count = (await session.execute(count_stmt)).scalar_one() or 0 + logger.info("WOULD update %d rows in %s from users.default_workspace_id", count, table) + return int(count) + + stmt = update(model).where(workspace_col.is_(None), user_col.is_not(None)).values(workspace_id=correlated_default) + async with session_factory() as session: + result = await session.execute(stmt) + await session.commit() + rowcount = result.rowcount or 0 + logger.info("Updated %d rows in %s from users.default_workspace_id", rowcount, table) + return int(rowcount) async def _step3_assign_legacy_workspace( diff --git a/backend/tests/test_backfill_workspace_id.py b/backend/tests/test_backfill_workspace_id.py index 93cc8e00..c2a49b5a 100644 --- a/backend/tests/test_backfill_workspace_id.py +++ b/backend/tests/test_backfill_workspace_id.py @@ -16,10 +16,17 @@ import uuid import pytest from sqlalchemy import select +from deerflow.persistence.feedback.model import FeedbackRow +from deerflow.persistence.models.run_event import RunEventRow +from deerflow.persistence.run.model import RunRow +from deerflow.persistence.thread_meta.model import ThreadMetaRow from deerflow.persistence.user.model import UserRow from deerflow.persistence.workspace.model import WorkspaceRow from deerflow.persistence.workspace_membership.model import WorkspaceMembershipRow -from scripts.backfill_workspace_id import _step1_create_workspaces_for_users +from scripts.backfill_workspace_id import ( + _step1_create_workspaces_for_users, + _step2_update_table_from_users, +) pytestmark = pytest.mark.anyio @@ -51,6 +58,12 @@ async def _seed_user(sf, *, email: str, default_workspace_id: str | None = None) return user_id +async def _seed_business_row(sf, model, **fields) -> None: + async with sf() as session: + session.add(model(**fields)) + await session.commit() + + # --------------------------------------------------------------------------- # Step 1: per-user workspace creation # --------------------------------------------------------------------------- @@ -108,6 +121,67 @@ async def test_step1_is_idempotent(tmp_path): await _close() +async def test_backfill_updates_4_tables_from_users(tmp_path): + """Step 2 propagates each user's default_workspace_id into 4 business tables.""" + sf = await _init_engine(tmp_path) + try: + user_id = await _seed_user(sf, email="dave@example.com") + # Pre-seed business rows owned by the user with NULL workspace_id. + await _seed_business_row(sf, ThreadMetaRow, thread_id="t-1", user_id=user_id) + await _seed_business_row(sf, RunRow, run_id="r-1", thread_id="t-1", user_id=user_id) + await _seed_business_row(sf, FeedbackRow, feedback_id="f-1", thread_id="t-1", run_id="r-1", user_id=user_id, rating=1) + await _seed_business_row(sf, RunEventRow, thread_id="t-1", run_id="r-1", user_id=user_id, event_type="lifecycle_started", category="lifecycle", seq=1) + + # Step 1 first so users.default_workspace_id is populated. + await _step1_create_workspaces_for_users(sf, dry_run=False) + + async with sf() as session: + ws_id = (await session.execute(select(WorkspaceRow.id))).scalar_one() + + # Step 2 updates each table. + for table in ("threads_meta", "runs", "feedback", "run_events"): + count = await _step2_update_table_from_users(sf, table, dry_run=False) + assert count == 1, table + + async with sf() as session: + tm = (await session.execute(select(ThreadMetaRow))).scalar_one() + run = (await session.execute(select(RunRow))).scalar_one() + fb = (await session.execute(select(FeedbackRow))).scalar_one() + ev = (await session.execute(select(RunEventRow))).scalar_one() + assert tm.workspace_id == ws_id + assert run.workspace_id == ws_id + assert fb.workspace_id == ws_id + assert ev.workspace_id == ws_id + + # Re-running Step 2 is a no-op (filtered by workspace_id IS NULL). + for table in ("threads_meta", "runs", "feedback", "run_events"): + assert await _step2_update_table_from_users(sf, table, dry_run=False) == 0 + finally: + await _close() + + +async def test_backfill_step2_isolates_per_user(tmp_path): + """Two users with different default workspaces get their own threads tagged independently.""" + sf = await _init_engine(tmp_path) + try: + u_eve = await _seed_user(sf, email="eve@example.com") + u_frank = await _seed_user(sf, email="frank@example.com") + await _seed_business_row(sf, ThreadMetaRow, thread_id="t-eve", user_id=u_eve) + await _seed_business_row(sf, ThreadMetaRow, thread_id="t-frank", user_id=u_frank) + + await _step1_create_workspaces_for_users(sf, dry_run=False) + await _step2_update_table_from_users(sf, "threads_meta", dry_run=False) + + async with sf() as session: + rows = {r.thread_id: r.workspace_id for r in (await session.execute(select(ThreadMetaRow))).scalars().all()} + users = {u.id: u.default_workspace_id for u in (await session.execute(select(UserRow))).scalars().all()} + assert rows["t-eve"] == users[u_eve] + assert rows["t-frank"] == users[u_frank] + assert rows["t-eve"] != rows["t-frank"] + finally: + await _close() + + async def test_step1_skips_blacklisted_base_slug(tmp_path): """A user with email like admin@... gets bumped past the slug blacklist via the walker.""" sf = await _init_engine(tmp_path)