From b62cc9da451acd3d9e8fcbda82456478f7b975a5 Mon Sep 17 00:00:00 2001 From: dadachann <185672915+dadachann@users.noreply.github.com> Date: Sun, 16 Aug 2026 17:15:57 +0000 Subject: [PATCH] fix(migrations): preserve legacy workspace ownership --- src/langbot/pkg/persistence/mgr.py | 13 +++- .../persistence/test_workspace_migration.py | 73 +++++++++++++++---- 2 files changed, 72 insertions(+), 14 deletions(-) diff --git a/src/langbot/pkg/persistence/mgr.py b/src/langbot/pkg/persistence/mgr.py index e9062c066..b29d6b44e 100644 --- a/src/langbot/pkg/persistence/mgr.py +++ b/src/langbot/pkg/persistence/mgr.py @@ -177,7 +177,6 @@ class PersistenceManager: await self._validate_cloud_runtime() return - self._enable_sqlite_foreign_keys() if self.mode == PersistenceMode.RELEASE_MIGRATION: async with self._release_migration_lock(): await self._initialize_managed_schema() @@ -185,6 +184,7 @@ class PersistenceManager: return await self._initialize_managed_schema() + await self._enable_sqlite_foreign_keys_after_migration() if self.mode == PersistenceMode.OSS_COMPAT: await self.write_space_model_providers() @@ -373,6 +373,17 @@ class PersistenceManager: sqlalchemy.event.listen(self.get_db_engine().sync_engine, 'begin', set_oss_tenant_scope) self._oss_tenant_scope_listener_installed = True + async def _enable_sqlite_foreign_keys_after_migration(self) -> None: + """Enable SQLite FK enforcement only after table-rebuilding migrations.""" + engine = self.get_db_engine() + if engine.dialect.name != 'sqlite': + return + await engine.dispose() + self._enable_sqlite_foreign_keys() + # Dispose again so every runtime connection is opened through the new + # listener instead of reusing a pre-migration pooled connection. + await engine.dispose() + def _enable_sqlite_foreign_keys(self) -> None: """Enable SQLite FK enforcement for every pooled runtime connection.""" engine = self.get_db_engine() diff --git a/tests/integration/persistence/test_workspace_migration.py b/tests/integration/persistence/test_workspace_migration.py index 70c46f89a..2ebb0f1e8 100644 --- a/tests/integration/persistence/test_workspace_migration.py +++ b/tests/integration/persistence/test_workspace_migration.py @@ -179,13 +179,17 @@ async def test_existing_oss_workspace_is_rekeyed_to_instance_identity(tmp_path): ) async with engine.begin() as conn: await conn.run_sync(schema.create_all) - await conn.execute(sa.text("INSERT INTO metadata (key, value) VALUES ('instance_uuid', :value)"), {'value': instance_id}) await conn.execute( - sa.text("INSERT INTO workspaces (uuid, instance_uuid, slug, source) VALUES (:uuid, :instance, 'default', 'local')"), + sa.text("INSERT INTO metadata (key, value) VALUES ('instance_uuid', :value)"), {'value': instance_id} + ) + await conn.execute( + sa.text( + "INSERT INTO workspaces (uuid, instance_uuid, slug, source) VALUES (:uuid, :instance, 'default', 'local')" + ), {'uuid': old_workspace_uuid, 'instance': instance_id}, ) await conn.execute( - sa.text("INSERT INTO tenant_rows (id, workspace_uuid) VALUES (1, :uuid)"), + sa.text('INSERT INTO tenant_rows (id, workspace_uuid) VALUES (1, :uuid)'), {'uuid': old_workspace_uuid}, ) await run_alembic_stamp(engine, '0016_support_admin_sessions') @@ -193,8 +197,8 @@ async def test_existing_oss_workspace_is_rekeyed_to_instance_identity(tmp_path): await run_alembic_upgrade(engine, 'head') async with engine.connect() as conn: - assert (await conn.execute(sa.text("SELECT uuid FROM workspaces"))).scalar_one() == canonical_uuid - assert (await conn.execute(sa.text("SELECT workspace_uuid FROM tenant_rows"))).scalar_one() == canonical_uuid + assert (await conn.execute(sa.text('SELECT uuid FROM workspaces'))).scalar_one() == canonical_uuid + assert (await conn.execute(sa.text('SELECT workspace_uuid FROM tenant_rows'))).scalar_one() == canonical_uuid await engine.dispose() @@ -411,6 +415,45 @@ async def test_persistence_startup_defers_workspace_tables_until_account_upgrade await engine.dispose() +async def test_persistence_startup_preserves_legacy_workspace_membership_with_foreign_keys( + tmp_path, + monkeypatch, +): + database_path = tmp_path / 'startup-foreign-keys.db' + engine = create_async_engine(f'sqlite+aiosqlite:///{database_path}') + try: + await _create_legacy_schema(engine) + await run_alembic_stamp(engine, '0008_mcp_resource_prefs') + finally: + await engine.dispose() + + monkeypatch.setattr(constants, 'instance_id', 'instance_migration_test') + application = type('Application', (), {})() + application.logger = logging.getLogger('workspace-startup-foreign-keys-test') + application.instance_config = type( + 'InstanceConfig', + (), + {'data': {'database': {'use': 'sqlite', 'sqlite': {'path': str(database_path)}}}}, + )() + manager = PersistenceManager(application) + + await manager.initialize() + try: + async with manager.get_db_engine().connect() as conn: + workspace = ( + (await conn.execute(sa.text("SELECT * FROM workspaces WHERE source = 'local'"))).mappings().one() + ) + membership = (await conn.execute(sa.text('SELECT * FROM workspace_memberships'))).mappings().one() + foreign_keys = await conn.scalar(sa.text('PRAGMA foreign_keys')) + + assert workspace['created_by_account_uuid'] == membership['account_uuid'] + assert membership['role'] == 'owner' + assert membership['status'] == 'active' + assert foreign_keys == 1 + finally: + await manager.shutdown() + + async def test_oss_workspace_identity_rekeys_fk_graph_and_metadata(tmp_path): engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-rekey.db"}') try: @@ -425,7 +468,7 @@ async def test_oss_workspace_identity_rekeys_fk_graph_and_metadata(tmp_path): assert instance_uuid await conn.execute( sa.text( - "INSERT INTO workspace_metadata (workspace_uuid, key, value) " + 'INSERT INTO workspace_metadata (workspace_uuid, key, value) ' "VALUES (:workspace_uuid, 'migration_probe', 'present')" ), {'workspace_uuid': old_uuid}, @@ -433,7 +476,7 @@ async def test_oss_workspace_identity_rekeys_fk_graph_and_metadata(tmp_path): await conn.execute( sa.text( "INSERT INTO metadata (key, value) VALUES ('oss_workspace_uuid', :workspace_uuid) " - "ON CONFLICT(key) DO UPDATE SET value = excluded.value" + 'ON CONFLICT(key) DO UPDATE SET value = excluded.value' ), {'workspace_uuid': old_uuid}, ) @@ -442,12 +485,16 @@ async def test_oss_workspace_identity_rekeys_fk_graph_and_metadata(tmp_path): expected_uuid = workspace_uuid_from_instance_id(instance_uuid) async with engine.connect() as conn: assert await conn.scalar(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'")) == expected_uuid - assert await conn.scalar( - sa.text("SELECT workspace_uuid FROM workspace_metadata WHERE key = 'migration_probe'") - ) == expected_uuid - assert await conn.scalar( - sa.text("SELECT value FROM metadata WHERE key = 'oss_workspace_uuid'") - ) == expected_uuid + assert ( + await conn.scalar( + sa.text("SELECT workspace_uuid FROM workspace_metadata WHERE key = 'migration_probe'") + ) + == expected_uuid + ) + assert ( + await conn.scalar(sa.text("SELECT value FROM metadata WHERE key = 'oss_workspace_uuid'")) + == expected_uuid + ) finally: await engine.dispose()