mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-21 20:06:06 +00:00
feat(tenancy): implement workspace isolation
This commit is contained in:
Junyan Qin
@@ -0,0 +1,253 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
|
||||
TENANT_TABLES = (
|
||||
'api_keys',
|
||||
'bots',
|
||||
'bot_admins',
|
||||
'binary_storages',
|
||||
'mcp_servers',
|
||||
'model_providers',
|
||||
'llm_models',
|
||||
'embedding_models',
|
||||
'rerank_models',
|
||||
'legacy_pipelines',
|
||||
'pipeline_run_records',
|
||||
'plugin_settings',
|
||||
'knowledge_bases',
|
||||
'knowledge_base_files',
|
||||
'knowledge_base_chunks',
|
||||
'webhooks',
|
||||
'monitoring_messages',
|
||||
'monitoring_llm_calls',
|
||||
'monitoring_tool_calls',
|
||||
'monitoring_sessions',
|
||||
'monitoring_errors',
|
||||
'monitoring_embedding_calls',
|
||||
'monitoring_feedback',
|
||||
)
|
||||
|
||||
|
||||
def _uuid_table(metadata: sa.MetaData, name: str, *columns: sa.Column) -> sa.Table:
|
||||
return sa.Table(name, metadata, sa.Column('uuid', sa.String(255), primary_key=True), *columns)
|
||||
|
||||
|
||||
async def create_legacy_resource_schema(engine, *, instance_uuid: str) -> None:
|
||||
"""Create the smallest representative pre-0010 schema with one row/table."""
|
||||
metadata = sa.MetaData()
|
||||
system_metadata = sa.Table(
|
||||
'metadata',
|
||||
metadata,
|
||||
sa.Column('key', sa.String(255), primary_key=True),
|
||||
sa.Column('value', sa.String(255)),
|
||||
)
|
||||
users = sa.Table(
|
||||
'users',
|
||||
metadata,
|
||||
sa.Column('id', sa.Integer, primary_key=True),
|
||||
sa.Column('user', sa.String(255), nullable=False),
|
||||
sa.Column('password', sa.String(255), nullable=False),
|
||||
)
|
||||
api_keys = sa.Table(
|
||||
'api_keys',
|
||||
metadata,
|
||||
sa.Column('id', sa.Integer, primary_key=True, autoincrement=True),
|
||||
sa.Column('name', sa.String(255), nullable=False),
|
||||
sa.Column('key', sa.String(255), nullable=False, unique=True),
|
||||
)
|
||||
bots = _uuid_table(
|
||||
metadata,
|
||||
'bots',
|
||||
sa.Column('name', sa.String(255), nullable=False),
|
||||
sa.Column('updated_at', sa.DateTime, nullable=False),
|
||||
)
|
||||
bot_admins = sa.Table(
|
||||
'bot_admins',
|
||||
metadata,
|
||||
sa.Column('id', sa.Integer, primary_key=True, autoincrement=True),
|
||||
sa.Column('bot_uuid', sa.String(255), nullable=False),
|
||||
sa.Column('launcher_type', sa.String(64), nullable=False),
|
||||
sa.Column('launcher_id', sa.String(255), nullable=False),
|
||||
sa.UniqueConstraint('bot_uuid', 'launcher_type', 'launcher_id', name='uq_bot_admin'),
|
||||
)
|
||||
binary_storages = sa.Table(
|
||||
'binary_storages',
|
||||
metadata,
|
||||
sa.Column('unique_key', sa.String(255), primary_key=True),
|
||||
sa.Column('key', sa.String(255), nullable=False),
|
||||
sa.Column('owner_type', sa.String(255), nullable=False),
|
||||
sa.Column('owner', sa.String(255), nullable=False),
|
||||
)
|
||||
mcp_servers = _uuid_table(
|
||||
metadata,
|
||||
'mcp_servers',
|
||||
sa.Column('name', sa.String(255), nullable=False),
|
||||
sa.Column('enable', sa.Boolean, nullable=False),
|
||||
sa.Column('updated_at', sa.DateTime, nullable=False),
|
||||
)
|
||||
model_providers = _uuid_table(
|
||||
metadata,
|
||||
'model_providers',
|
||||
sa.Column('name', sa.String(255), nullable=False),
|
||||
sa.Column('requester', sa.String(255), nullable=False),
|
||||
)
|
||||
llm_models = _uuid_table(
|
||||
metadata,
|
||||
'llm_models',
|
||||
sa.Column('name', sa.String(255), nullable=False),
|
||||
sa.Column('provider_uuid', sa.String(255), nullable=False),
|
||||
)
|
||||
embedding_models = _uuid_table(
|
||||
metadata,
|
||||
'embedding_models',
|
||||
sa.Column('name', sa.String(255), nullable=False),
|
||||
sa.Column('provider_uuid', sa.String(255), nullable=False),
|
||||
)
|
||||
rerank_models = _uuid_table(
|
||||
metadata,
|
||||
'rerank_models',
|
||||
sa.Column('name', sa.String(255), nullable=False),
|
||||
sa.Column('provider_uuid', sa.String(255), nullable=False),
|
||||
)
|
||||
legacy_pipelines = _uuid_table(
|
||||
metadata,
|
||||
'legacy_pipelines',
|
||||
sa.Column('name', sa.String(255), nullable=False),
|
||||
sa.Column('is_default', sa.Boolean, nullable=False),
|
||||
sa.Column('updated_at', sa.DateTime, nullable=False),
|
||||
)
|
||||
pipeline_run_records = _uuid_table(
|
||||
metadata,
|
||||
'pipeline_run_records',
|
||||
sa.Column('pipeline_uuid', sa.String(255), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime, nullable=False),
|
||||
)
|
||||
plugin_settings = sa.Table(
|
||||
'plugin_settings',
|
||||
metadata,
|
||||
sa.Column('plugin_author', sa.String(255), primary_key=True),
|
||||
sa.Column('plugin_name', sa.String(255), primary_key=True),
|
||||
sa.Column('enabled', sa.Boolean, nullable=False),
|
||||
)
|
||||
knowledge_bases = _uuid_table(
|
||||
metadata,
|
||||
'knowledge_bases',
|
||||
sa.Column('name', sa.String(255), nullable=False),
|
||||
sa.Column('collection_id', sa.String(255), nullable=True),
|
||||
)
|
||||
knowledge_base_files = _uuid_table(
|
||||
metadata,
|
||||
'knowledge_base_files',
|
||||
sa.Column('kb_id', sa.String(255), nullable=True),
|
||||
)
|
||||
knowledge_base_chunks = _uuid_table(
|
||||
metadata,
|
||||
'knowledge_base_chunks',
|
||||
sa.Column('file_id', sa.String(255), nullable=True),
|
||||
)
|
||||
webhooks = sa.Table(
|
||||
'webhooks',
|
||||
metadata,
|
||||
sa.Column('id', sa.Integer, primary_key=True, autoincrement=True),
|
||||
sa.Column('name', sa.String(255), nullable=False),
|
||||
sa.Column('enabled', sa.Boolean, nullable=False),
|
||||
sa.Column('created_at', sa.DateTime, nullable=False),
|
||||
)
|
||||
|
||||
monitoring_tables: dict[str, sa.Table] = {}
|
||||
for table_name in (
|
||||
'monitoring_messages',
|
||||
'monitoring_llm_calls',
|
||||
'monitoring_tool_calls',
|
||||
'monitoring_errors',
|
||||
'monitoring_embedding_calls',
|
||||
):
|
||||
monitoring_tables[table_name] = sa.Table(
|
||||
table_name,
|
||||
metadata,
|
||||
sa.Column('id', sa.String(255), primary_key=True),
|
||||
sa.Column('timestamp', sa.DateTime, nullable=False),
|
||||
sa.Column('session_id', sa.String(255), nullable=True),
|
||||
sa.Column('message_id', sa.String(255), nullable=True),
|
||||
)
|
||||
monitoring_tables['monitoring_sessions'] = sa.Table(
|
||||
'monitoring_sessions',
|
||||
metadata,
|
||||
sa.Column('session_id', sa.String(255), primary_key=True),
|
||||
sa.Column('bot_id', sa.String(255), nullable=False),
|
||||
sa.Column('last_activity', sa.DateTime, nullable=False),
|
||||
sa.Column('is_active', sa.Boolean, nullable=False),
|
||||
)
|
||||
monitoring_tables['monitoring_feedback'] = sa.Table(
|
||||
'monitoring_feedback',
|
||||
metadata,
|
||||
sa.Column('id', sa.String(255), primary_key=True),
|
||||
sa.Column('feedback_id', sa.String(255), nullable=False, unique=True),
|
||||
sa.Column('timestamp', sa.DateTime, nullable=False),
|
||||
sa.Column('session_id', sa.String(255), nullable=True),
|
||||
sa.Column('message_id', sa.String(255), nullable=True),
|
||||
)
|
||||
|
||||
now = datetime.datetime(2026, 1, 1)
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(metadata.create_all)
|
||||
await conn.execute(
|
||||
system_metadata.insert(),
|
||||
[
|
||||
{'key': 'database_version', 'value': '25'},
|
||||
{'key': 'instance_uuid', 'value': instance_uuid},
|
||||
{'key': 'wizard_status', 'value': 'completed'},
|
||||
{'key': 'wizard_progress', 'value': '3'},
|
||||
{'key': 'rag_plugin_migration_needed', 'value': 'true'},
|
||||
],
|
||||
)
|
||||
await conn.execute(users.insert().values(user='Owner@Example.COM', password='hash'))
|
||||
await conn.execute(api_keys.insert().values(name='legacy', key='lbk_legacy-secret'))
|
||||
await conn.execute(bots.insert().values(uuid='bot-1', name='bot', updated_at=now))
|
||||
await conn.execute(bot_admins.insert().values(bot_uuid='bot-1', launcher_type='person', launcher_id='owner'))
|
||||
await conn.execute(
|
||||
binary_storages.insert().values(unique_key='plugin:demo:key', key='key', owner_type='plugin', owner='demo')
|
||||
)
|
||||
await conn.execute(mcp_servers.insert().values(uuid='mcp-1', name='shared-name', enable=True, updated_at=now))
|
||||
await conn.execute(model_providers.insert().values(uuid='provider-1', name='provider', requester='openai'))
|
||||
for table in (llm_models, embedding_models, rerank_models):
|
||||
await conn.execute(table.insert().values(uuid=f'{table.name}-1', name='model', provider_uuid='provider-1'))
|
||||
await conn.execute(
|
||||
legacy_pipelines.insert().values(uuid='pipeline-1', name='pipeline', is_default=True, updated_at=now)
|
||||
)
|
||||
await conn.execute(
|
||||
pipeline_run_records.insert().values(uuid='run-1', pipeline_uuid='pipeline-1', created_at=now)
|
||||
)
|
||||
await conn.execute(plugin_settings.insert().values(plugin_author='author', plugin_name='plugin', enabled=True))
|
||||
await conn.execute(knowledge_bases.insert().values(uuid='kb-1', name='knowledge', collection_id='collection-1'))
|
||||
await conn.execute(knowledge_base_files.insert().values(uuid='file-1', kb_id='kb-1'))
|
||||
await conn.execute(knowledge_base_chunks.insert().values(uuid='chunk-1', file_id='file-1'))
|
||||
await conn.execute(webhooks.insert().values(name='hook', enabled=True, created_at=now))
|
||||
for table_name, table in monitoring_tables.items():
|
||||
if table_name == 'monitoring_sessions':
|
||||
values = {
|
||||
'session_id': 'session-1',
|
||||
'bot_id': 'bot-1',
|
||||
'last_activity': now,
|
||||
'is_active': True,
|
||||
}
|
||||
elif table_name == 'monitoring_feedback':
|
||||
values = {
|
||||
'id': 'feedback-row-1',
|
||||
'feedback_id': 'feedback-1',
|
||||
'timestamp': now,
|
||||
'session_id': 'session-1',
|
||||
'message_id': 'message-1',
|
||||
}
|
||||
else:
|
||||
values = {
|
||||
'id': f'{table_name}-1',
|
||||
'timestamp': now,
|
||||
'session_id': 'session-1',
|
||||
'message_id': 'message-1',
|
||||
}
|
||||
await conn.execute(table.insert().values(**values))
|
||||
@@ -13,12 +13,17 @@ CI runs automatically with PostgreSQL service container.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
from sqlalchemy import text
|
||||
|
||||
from langbot.pkg.entity.persistence.base import Base
|
||||
from langbot.pkg.entity.persistence.user import User
|
||||
from langbot.pkg.persistence.mgr import PersistenceManager
|
||||
from langbot.pkg.persistence.alembic_runner import (
|
||||
run_alembic_upgrade,
|
||||
run_alembic_stamp,
|
||||
@@ -27,6 +32,10 @@ from langbot.pkg.persistence.alembic_runner import (
|
||||
)
|
||||
from alembic.config import Config
|
||||
from alembic.script import ScriptDirectory
|
||||
from langbot.pkg.utils import constants
|
||||
from langbot.pkg.workspace.collaboration import normalize_email
|
||||
|
||||
from .resource_migration_support import TENANT_TABLES, create_legacy_resource_schema
|
||||
|
||||
|
||||
def _get_script_head() -> str:
|
||||
@@ -130,6 +139,27 @@ class TestPostgreSQLMigrationBaseline:
|
||||
rev = await get_alembic_current(postgres_engine)
|
||||
assert rev == '0001_baseline'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fresh_postgres_schema_accepts_application_casefold_identity(
|
||||
self,
|
||||
postgres_engine,
|
||||
clean_tables,
|
||||
clean_alembic_version,
|
||||
):
|
||||
canonical_email = normalize_email('Ꭰ@Example.COM')
|
||||
async with postgres_engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
await conn.execute(
|
||||
sa.insert(User).values(
|
||||
uuid='00000000-0000-0000-0000-000000000099',
|
||||
user=canonical_email,
|
||||
normalized_email=canonical_email,
|
||||
password='hash',
|
||||
)
|
||||
)
|
||||
async with postgres_engine.connect() as conn:
|
||||
assert await conn.scalar(sa.select(User.normalized_email)) == 'Ꭰ@example.com'
|
||||
|
||||
|
||||
class TestPostgreSQLMigrationUpgrade:
|
||||
"""Tests for upgrade to head workflow on PostgreSQL."""
|
||||
@@ -223,3 +253,175 @@ class TestPostgreSQLMigrationGetCurrent:
|
||||
|
||||
rev = await get_alembic_current(postgres_engine)
|
||||
assert rev == '0001_baseline'
|
||||
|
||||
|
||||
class TestPostgreSQLWorkspaceMigration:
|
||||
"""Focused coverage for upgrading a pre-tenancy PostgreSQL instance."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_postgres_legacy_instance_gets_default_workspace(
|
||||
self,
|
||||
postgres_engine,
|
||||
clean_tables,
|
||||
clean_alembic_version,
|
||||
monkeypatch,
|
||||
):
|
||||
legacy_metadata = sa.MetaData()
|
||||
metadata_table = sa.Table(
|
||||
'metadata',
|
||||
legacy_metadata,
|
||||
sa.Column('key', sa.String(255), primary_key=True),
|
||||
sa.Column('value', sa.String(255)),
|
||||
)
|
||||
users = sa.Table(
|
||||
'users',
|
||||
legacy_metadata,
|
||||
sa.Column('id', sa.Integer, primary_key=True),
|
||||
sa.Column('user', sa.String(255), nullable=False),
|
||||
sa.Column('password', sa.String(255), nullable=False),
|
||||
sa.Column(
|
||||
'account_type',
|
||||
sa.String(32),
|
||||
nullable=False,
|
||||
server_default='local',
|
||||
),
|
||||
sa.Column('created_at', sa.DateTime, server_default=text('now()')),
|
||||
sa.Column('updated_at', sa.DateTime, server_default=text('now()')),
|
||||
)
|
||||
async with postgres_engine.begin() as conn:
|
||||
await conn.run_sync(legacy_metadata.create_all)
|
||||
await conn.execute(metadata_table.insert().values(key='database_version', value='25'))
|
||||
await conn.execute(metadata_table.insert().values(key='instance_uuid', value='instance_postgres_test'))
|
||||
await conn.execute(users.insert().values(user='owner@example.com', password='owner-hash'))
|
||||
|
||||
await run_alembic_stamp(postgres_engine, '0008_mcp_resource_prefs')
|
||||
monkeypatch.setattr(constants, 'instance_id', 'instance_postgres_test')
|
||||
database = type('Database', (), {'get_engine': lambda self: postgres_engine})()
|
||||
application = type('Application', (), {})()
|
||||
application.logger = logging.getLogger('postgres-workspace-startup-test')
|
||||
manager = PersistenceManager(application)
|
||||
manager.db = database
|
||||
|
||||
await manager.create_tables()
|
||||
async with postgres_engine.connect() as conn:
|
||||
tables_before_migration = set(
|
||||
await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names())
|
||||
)
|
||||
assert 'workspaces' not in tables_before_migration
|
||||
|
||||
await manager._run_alembic_migrations()
|
||||
|
||||
async with postgres_engine.connect() as conn:
|
||||
account = (await conn.execute(text('SELECT uuid, status, source FROM users'))).mappings().one()
|
||||
workspace = (
|
||||
(await conn.execute(text('SELECT * FROM workspaces WHERE source = :source'), {'source': 'local'}))
|
||||
.mappings()
|
||||
.one()
|
||||
)
|
||||
membership = (await conn.execute(text('SELECT * FROM workspace_memberships'))).mappings().one()
|
||||
execution_state = (await conn.execute(text('SELECT * FROM workspace_execution_states'))).mappings().one()
|
||||
|
||||
assert account['status'] == 'active'
|
||||
assert account['source'] == 'local'
|
||||
assert workspace['instance_uuid'] == 'instance_postgres_test'
|
||||
assert workspace['created_by_account_uuid'] == account['uuid']
|
||||
assert membership['account_uuid'] == account['uuid']
|
||||
assert membership['role'] == 'owner'
|
||||
assert execution_state['active_generation'] == 1
|
||||
assert execution_state['write_fenced'] is False
|
||||
|
||||
|
||||
class TestPostgreSQLResourceTenancyMigration:
|
||||
"""Legacy backfill and scoped-key enforcement on real PostgreSQL."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_postgres_resources_are_backfilled_and_scoped(
|
||||
self,
|
||||
postgres_engine,
|
||||
clean_tables,
|
||||
clean_alembic_version,
|
||||
):
|
||||
await create_legacy_resource_schema(
|
||||
postgres_engine,
|
||||
instance_uuid='postgres-resource-migration-test',
|
||||
)
|
||||
async with postgres_engine.begin() as conn:
|
||||
await conn.execute(text('UPDATE users SET "user" = \'Straße@Example.COM\''))
|
||||
await conn.execute(
|
||||
text('INSERT INTO users ("user", password) VALUES (:email, :password)'),
|
||||
{'email': 'Ꭰ@Example.COM', 'password': 'cherokee-hash'},
|
||||
)
|
||||
await run_alembic_stamp(postgres_engine, '0008_mcp_resource_prefs')
|
||||
await run_alembic_upgrade(postgres_engine, 'head')
|
||||
|
||||
async with postgres_engine.connect() as conn:
|
||||
workspace_uuid = await conn.scalar(text("SELECT uuid FROM workspaces WHERE source = 'local'"))
|
||||
for table_name in TENANT_TABLES:
|
||||
count, distinct_workspaces = (
|
||||
await conn.execute(text(f'SELECT COUNT(*), COUNT(DISTINCT workspace_uuid) FROM {table_name}'))
|
||||
).one()
|
||||
assert (count, distinct_workspaces) == (1, 1), table_name
|
||||
columns = await conn.run_sync(
|
||||
lambda sync_conn, name=table_name: {
|
||||
column['name']: column for column in sa.inspect(sync_conn).get_columns(name)
|
||||
}
|
||||
)
|
||||
assert columns['workspace_uuid']['nullable'] is False, table_name
|
||||
|
||||
api_columns = await conn.run_sync(
|
||||
lambda sync_conn: {column['name'] for column in sa.inspect(sync_conn).get_columns('api_keys')}
|
||||
)
|
||||
assert 'key' not in api_columns
|
||||
assert await conn.scalar(text('SELECT scopes FROM api_keys')) == ['*']
|
||||
assert (await conn.execute(text('SELECT normalized_email FROM users ORDER BY id'))).scalars().all() == [
|
||||
'strasse@example.com',
|
||||
'Ꭰ@example.com',
|
||||
]
|
||||
|
||||
second_workspace_uuid = '00000000-0000-0000-0000-000000000002'
|
||||
async with postgres_engine.begin() as conn:
|
||||
await conn.execute(
|
||||
text(
|
||||
'INSERT INTO workspaces '
|
||||
'(uuid, instance_uuid, name, slug, type, status, source, projection_revision) '
|
||||
"VALUES (:uuid, 'postgres-resource-migration-test', 'Second', 'second', "
|
||||
"'team', 'active', 'cloud_projection', 0)"
|
||||
),
|
||||
{'uuid': second_workspace_uuid},
|
||||
)
|
||||
await conn.execute(
|
||||
text(
|
||||
'INSERT INTO mcp_servers (uuid, workspace_uuid, name, enable, updated_at) '
|
||||
"VALUES ('mcp-2', :workspace_uuid, 'shared-name', true, now())"
|
||||
),
|
||||
{'workspace_uuid': second_workspace_uuid},
|
||||
)
|
||||
await conn.execute(
|
||||
text(
|
||||
'INSERT INTO plugin_settings '
|
||||
'(workspace_uuid, plugin_author, plugin_name, enabled) '
|
||||
"VALUES (:workspace_uuid, 'author', 'plugin', true)"
|
||||
),
|
||||
{'workspace_uuid': second_workspace_uuid},
|
||||
)
|
||||
|
||||
with pytest.raises(IntegrityError):
|
||||
async with postgres_engine.begin() as conn:
|
||||
await conn.execute(
|
||||
text(
|
||||
'INSERT INTO mcp_servers '
|
||||
'(uuid, workspace_uuid, name, enable, updated_at) '
|
||||
"VALUES ('mcp-duplicate', :workspace_uuid, 'shared-name', true, now())"
|
||||
),
|
||||
{'workspace_uuid': workspace_uuid},
|
||||
)
|
||||
|
||||
with pytest.raises(IntegrityError):
|
||||
async with postgres_engine.begin() as conn:
|
||||
await conn.execute(
|
||||
text(
|
||||
'INSERT INTO llm_models (uuid, workspace_uuid, name, provider_uuid) '
|
||||
"VALUES ('cross-workspace-model', :workspace_uuid, 'model', 'provider-1')"
|
||||
),
|
||||
{'workspace_uuid': second_workspace_uuid},
|
||||
)
|
||||
|
||||
@@ -0,0 +1,286 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from langbot.pkg.entity import persistence
|
||||
from langbot.pkg.entity.persistence.base import Base
|
||||
from langbot.pkg.persistence.alembic_runner import run_alembic_stamp, run_alembic_upgrade
|
||||
from langbot.pkg.utils import importutil
|
||||
|
||||
from .resource_migration_support import TENANT_TABLES, create_legacy_resource_schema
|
||||
|
||||
|
||||
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
|
||||
|
||||
|
||||
async def _inspect(engine, callback):
|
||||
async with engine.connect() as conn:
|
||||
return await conn.run_sync(callback)
|
||||
|
||||
|
||||
async def test_legacy_sqlite_resources_are_backfilled_and_contracted(tmp_path):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "legacy-resources.db"}')
|
||||
try:
|
||||
await create_legacy_resource_schema(engine, instance_uuid='resource-migration-test')
|
||||
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
|
||||
await run_alembic_upgrade(engine, 'head')
|
||||
|
||||
async with engine.connect() as conn:
|
||||
workspace_uuid = await conn.scalar(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'"))
|
||||
assert workspace_uuid is not None
|
||||
for table_name in TENANT_TABLES:
|
||||
count, distinct_workspaces = (
|
||||
await conn.execute(sa.text(f'SELECT COUNT(*), COUNT(DISTINCT workspace_uuid) FROM {table_name}'))
|
||||
).one()
|
||||
assert count == 1, table_name
|
||||
assert distinct_workspaces == 1, table_name
|
||||
assert (
|
||||
await conn.scalar(
|
||||
sa.text(f'SELECT COUNT(*) FROM {table_name} WHERE workspace_uuid != :workspace_uuid'),
|
||||
{'workspace_uuid': workspace_uuid},
|
||||
)
|
||||
== 0
|
||||
)
|
||||
|
||||
api_key = (
|
||||
(await conn.execute(sa.text('SELECT key_hash, scopes, status, created_by_account_uuid FROM api_keys')))
|
||||
.mappings()
|
||||
.one()
|
||||
)
|
||||
assert api_key['key_hash'] == hashlib.sha256(b'lbk_legacy-secret').hexdigest()
|
||||
stored_scopes = api_key['scopes']
|
||||
if isinstance(stored_scopes, str):
|
||||
stored_scopes = json.loads(stored_scopes)
|
||||
assert stored_scopes == ['*']
|
||||
assert api_key['status'] == 'active'
|
||||
assert api_key['created_by_account_uuid'] is not None
|
||||
assert await conn.scalar(sa.text('SELECT normalized_email FROM users')) == 'owner@example.com'
|
||||
legacy_kb = (
|
||||
(
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'SELECT collection_id, legacy_vector_collection FROM knowledge_bases WHERE uuid = :uuid'
|
||||
),
|
||||
{'uuid': 'kb-1'},
|
||||
)
|
||||
)
|
||||
.mappings()
|
||||
.one()
|
||||
)
|
||||
assert legacy_kb['collection_id'] == 'collection-1'
|
||||
assert legacy_kb['legacy_vector_collection'] == 1
|
||||
assert (
|
||||
await conn.scalar(
|
||||
sa.text(
|
||||
'SELECT COUNT(*) FROM metadata '
|
||||
"WHERE key IN ('wizard_status', 'wizard_progress', 'rag_plugin_migration_needed')"
|
||||
)
|
||||
)
|
||||
== 0
|
||||
)
|
||||
assert (
|
||||
await conn.scalar(
|
||||
sa.text('SELECT COUNT(*) FROM workspace_metadata WHERE workspace_uuid = :workspace_uuid'),
|
||||
{'workspace_uuid': workspace_uuid},
|
||||
)
|
||||
== 3
|
||||
)
|
||||
|
||||
api_columns = await _inspect(
|
||||
engine,
|
||||
lambda conn: {column['name'] for column in sa.inspect(conn).get_columns('api_keys')},
|
||||
)
|
||||
assert 'key' not in api_columns
|
||||
assert {'uuid', 'key_hash', 'scopes', 'status', 'expires_at', 'last_used_at'} <= api_columns
|
||||
for table_name in TENANT_TABLES:
|
||||
columns = await _inspect(
|
||||
engine,
|
||||
lambda conn, name=table_name: {column['name']: column for column in sa.inspect(conn).get_columns(name)},
|
||||
)
|
||||
assert columns['workspace_uuid']['nullable'] is False, table_name
|
||||
if table_name == 'knowledge_bases':
|
||||
assert columns['legacy_vector_collection']['nullable'] is False
|
||||
|
||||
pk_columns = {
|
||||
table_name: tuple(
|
||||
(
|
||||
await _inspect(
|
||||
engine,
|
||||
lambda conn, name=table_name: sa.inspect(conn).get_pk_constraint(name),
|
||||
)
|
||||
)['constrained_columns']
|
||||
)
|
||||
for table_name in ('binary_storages', 'plugin_settings', 'monitoring_sessions')
|
||||
}
|
||||
assert pk_columns == {
|
||||
'binary_storages': ('workspace_uuid', 'unique_key'),
|
||||
'plugin_settings': ('workspace_uuid', 'plugin_author', 'plugin_name'),
|
||||
'monitoring_sessions': ('workspace_uuid', 'session_id'),
|
||||
}
|
||||
|
||||
pipeline_run_foreign_keys = await _inspect(
|
||||
engine,
|
||||
lambda conn: sa.inspect(conn).get_foreign_keys('pipeline_run_records'),
|
||||
)
|
||||
assert any(
|
||||
tuple(foreign_key['constrained_columns']) == ('workspace_uuid', 'pipeline_uuid')
|
||||
and foreign_key['referred_table'] == 'legacy_pipelines'
|
||||
and tuple(foreign_key['referred_columns']) == ('workspace_uuid', 'uuid')
|
||||
for foreign_key in pipeline_run_foreign_keys
|
||||
)
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_legacy_vector_marker_backfill_resumes_from_nullable_expand_step(tmp_path):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "legacy-vector-retry.db"}')
|
||||
try:
|
||||
await create_legacy_resource_schema(engine, instance_uuid='legacy-vector-retry')
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.text('ALTER TABLE knowledge_bases ADD COLUMN legacy_vector_collection BOOLEAN NULL'))
|
||||
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
|
||||
await run_alembic_upgrade(engine, 'head')
|
||||
|
||||
async with engine.connect() as conn:
|
||||
assert (
|
||||
await conn.scalar(
|
||||
sa.text('SELECT legacy_vector_collection FROM knowledge_bases WHERE uuid = :uuid'),
|
||||
{'uuid': 'kb-1'},
|
||||
)
|
||||
== 1
|
||||
)
|
||||
columns = await _inspect(
|
||||
engine,
|
||||
lambda conn: {column['name']: column for column in sa.inspect(conn).get_columns('knowledge_bases')},
|
||||
)
|
||||
assert columns['legacy_vector_collection']['nullable'] is False
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_sqlite_scoped_keys_allow_cross_workspace_but_reject_same_workspace(tmp_path):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "scoped-keys.db"}')
|
||||
try:
|
||||
await create_legacy_resource_schema(engine, instance_uuid='scoped-key-test')
|
||||
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
|
||||
await run_alembic_upgrade(engine, 'head')
|
||||
|
||||
second_workspace_uuid = str(uuid.uuid4())
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.text('PRAGMA foreign_keys=ON'))
|
||||
first_workspace_uuid = await conn.scalar(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'"))
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'INSERT INTO workspaces '
|
||||
'(uuid, instance_uuid, name, slug, type, status, source, projection_revision) '
|
||||
"VALUES (:uuid, 'scoped-key-test', 'Second', 'second', 'team', 'active', "
|
||||
"'cloud_projection', 0)"
|
||||
),
|
||||
{'uuid': second_workspace_uuid},
|
||||
)
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'INSERT INTO mcp_servers (uuid, workspace_uuid, name, enable, updated_at) '
|
||||
"VALUES ('mcp-2', :workspace_uuid, 'shared-name', 1, CURRENT_TIMESTAMP)"
|
||||
),
|
||||
{'workspace_uuid': second_workspace_uuid},
|
||||
)
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'INSERT INTO plugin_settings '
|
||||
'(workspace_uuid, plugin_author, plugin_name, enabled) '
|
||||
"VALUES (:workspace_uuid, 'author', 'plugin', 1)"
|
||||
),
|
||||
{'workspace_uuid': second_workspace_uuid},
|
||||
)
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'INSERT INTO binary_storages '
|
||||
'(workspace_uuid, unique_key, key, owner_type, owner) '
|
||||
"VALUES (:workspace_uuid, 'plugin:demo:key', 'key', 'plugin', 'demo')"
|
||||
),
|
||||
{'workspace_uuid': second_workspace_uuid},
|
||||
)
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'INSERT INTO monitoring_sessions '
|
||||
'(workspace_uuid, session_id, bot_id, last_activity, is_active) '
|
||||
"VALUES (:workspace_uuid, 'session-1', 'bot-2', CURRENT_TIMESTAMP, 1)"
|
||||
),
|
||||
{'workspace_uuid': second_workspace_uuid},
|
||||
)
|
||||
|
||||
with pytest.raises(IntegrityError):
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.text('PRAGMA foreign_keys=ON'))
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'INSERT INTO mcp_servers (uuid, workspace_uuid, name, enable, updated_at) '
|
||||
"VALUES ('mcp-duplicate', :workspace_uuid, 'shared-name', 1, CURRENT_TIMESTAMP)"
|
||||
),
|
||||
{'workspace_uuid': first_workspace_uuid},
|
||||
)
|
||||
|
||||
with pytest.raises(IntegrityError):
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.text('PRAGMA foreign_keys=ON'))
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'INSERT INTO llm_models (uuid, workspace_uuid, name, provider_uuid) '
|
||||
"VALUES ('cross-workspace-model', :workspace_uuid, 'model', 'provider-1')"
|
||||
),
|
||||
{'workspace_uuid': second_workspace_uuid},
|
||||
)
|
||||
|
||||
with pytest.raises(IntegrityError):
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.text('PRAGMA foreign_keys=ON'))
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'INSERT INTO pipeline_run_records '
|
||||
'(uuid, workspace_uuid, pipeline_uuid, created_at) '
|
||||
"VALUES ('cross-workspace-run', :workspace_uuid, 'pipeline-1', CURRENT_TIMESTAMP)"
|
||||
),
|
||||
{'workspace_uuid': second_workspace_uuid},
|
||||
)
|
||||
|
||||
with pytest.raises(IntegrityError):
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'INSERT INTO mcp_servers (uuid, name, enable, updated_at) '
|
||||
"VALUES ('unscoped-mcp', 'unscoped', 1, CURRENT_TIMESTAMP)"
|
||||
)
|
||||
)
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_fresh_sqlite_schema_matches_resource_tenancy_contract(tmp_path):
|
||||
importutil.import_modules_in_pkg(persistence)
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "fresh-resources.db"}')
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
await run_alembic_stamp(engine, '0001_baseline')
|
||||
await run_alembic_upgrade(engine, 'head')
|
||||
|
||||
tables = await _inspect(engine, lambda conn: set(sa.inspect(conn).get_table_names()))
|
||||
assert set(TENANT_TABLES) | {'workspace_metadata'} <= tables
|
||||
for table_name in TENANT_TABLES:
|
||||
columns = await _inspect(
|
||||
engine,
|
||||
lambda conn, name=table_name: {column['name']: column for column in sa.inspect(conn).get_columns(name)},
|
||||
)
|
||||
assert columns['workspace_uuid']['nullable'] is False, table_name
|
||||
if table_name == 'knowledge_bases':
|
||||
assert columns['legacy_vector_collection']['nullable'] is False
|
||||
finally:
|
||||
await engine.dispose()
|
||||
@@ -0,0 +1,107 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import pathlib
|
||||
import sqlite3
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from langbot.pkg.persistence import alembic_runner
|
||||
from langbot.pkg.persistence.mgr import PersistenceManager
|
||||
|
||||
from .resource_migration_support import create_legacy_resource_schema
|
||||
|
||||
|
||||
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
|
||||
|
||||
|
||||
def _manager(engine) -> PersistenceManager:
|
||||
database = type('Database', (), {'get_engine': lambda self: engine})()
|
||||
application = type('Application', (), {})()
|
||||
application.logger = logging.getLogger('sqlite-migration-backup-test')
|
||||
manager = PersistenceManager(application)
|
||||
manager.db = database
|
||||
return manager
|
||||
|
||||
|
||||
def _manifest_payloads(backup_directory) -> list[dict]:
|
||||
return [json.loads(path.read_text(encoding='utf-8')) for path in sorted(backup_directory.glob('*.json'))]
|
||||
|
||||
|
||||
def _assert_verified_backup(payload: dict) -> None:
|
||||
backup_path = pathlib.Path(payload['backup_path'])
|
||||
with sqlite3.connect(f'{backup_path.as_uri()}?mode=ro', uri=True) as connection:
|
||||
assert connection.execute('PRAGMA quick_check').fetchall() == [('ok',)]
|
||||
assert connection.execute('SELECT version_num FROM alembic_version').fetchone()[0] == payload['source_revision']
|
||||
|
||||
|
||||
async def test_tenancy_migrations_retain_verified_boundary_backups(tmp_path):
|
||||
database_path = tmp_path / 'legacy-with-backups.db'
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{database_path}')
|
||||
try:
|
||||
await create_legacy_resource_schema(engine, instance_uuid='backup-success')
|
||||
await alembic_runner.run_alembic_stamp(engine, '0008_mcp_resource_prefs')
|
||||
|
||||
await _manager(engine)._run_alembic_migrations()
|
||||
|
||||
assert await alembic_runner.get_alembic_current(engine) == '0010_scope_resources'
|
||||
payloads = _manifest_payloads(tmp_path / 'migration-backups')
|
||||
assert len(payloads) == 2
|
||||
assert {
|
||||
(payload['source_revision'], payload['target_revision'], payload['status']) for payload in payloads
|
||||
} == {
|
||||
('0008_mcp_resource_prefs', '0009_workspace_tenancy', 'migration_succeeded'),
|
||||
('0009_workspace_tenancy', '0010_scope_resources', 'migration_succeeded'),
|
||||
}
|
||||
for payload in payloads:
|
||||
_assert_verified_backup(payload)
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_failed_tenancy_migration_restores_backup_and_revision(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
):
|
||||
database_path = tmp_path / 'legacy-fault-injection.db'
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{database_path}')
|
||||
real_upgrade = alembic_runner.run_alembic_upgrade
|
||||
|
||||
async def injected_upgrade(async_engine, revision='head'):
|
||||
if revision != '0010_scope_resources':
|
||||
return await real_upgrade(async_engine, revision)
|
||||
async with async_engine.begin() as connection:
|
||||
await connection.execute(sa.text('CREATE TABLE injected_partial_migration (value TEXT NOT NULL)'))
|
||||
await alembic_runner.run_alembic_stamp(async_engine, '0010_scope_resources')
|
||||
raise RuntimeError('injected migration failure after a fake revision stamp')
|
||||
|
||||
try:
|
||||
await create_legacy_resource_schema(engine, instance_uuid='backup-failure')
|
||||
await alembic_runner.run_alembic_stamp(engine, '0008_mcp_resource_prefs')
|
||||
monkeypatch.setattr(alembic_runner, 'run_alembic_upgrade', injected_upgrade)
|
||||
|
||||
with pytest.raises(RuntimeError, match='injected migration failure'):
|
||||
await _manager(engine)._run_alembic_migrations()
|
||||
|
||||
assert await alembic_runner.get_alembic_current(engine) == '0009_workspace_tenancy'
|
||||
async with engine.connect() as connection:
|
||||
tables = set(
|
||||
await connection.run_sync(lambda sync_connection: sa.inspect(sync_connection).get_table_names())
|
||||
)
|
||||
assert 'injected_partial_migration' not in tables
|
||||
|
||||
payloads = _manifest_payloads(tmp_path / 'migration-backups')
|
||||
restored = [payload for payload in payloads if payload['target_revision'] == '0010_scope_resources']
|
||||
assert len(restored) == 1
|
||||
assert restored[0]['status'] == 'restored_after_failure'
|
||||
assert restored[0]['source_revision'] == '0009_workspace_tenancy'
|
||||
_assert_verified_backup(restored[0])
|
||||
|
||||
monkeypatch.setattr(alembic_runner, 'run_alembic_upgrade', real_upgrade)
|
||||
await _manager(engine)._run_alembic_migrations()
|
||||
assert await alembic_runner.get_alembic_current(engine) == '0010_scope_resources'
|
||||
finally:
|
||||
await engine.dispose()
|
||||
@@ -0,0 +1,379 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from langbot.pkg.entity import persistence
|
||||
from langbot.pkg.entity.persistence.base import Base
|
||||
from langbot.pkg.entity.persistence.user import User
|
||||
from langbot.pkg.persistence.mgr import PersistenceManager
|
||||
from langbot.pkg.persistence.alembic_runner import (
|
||||
get_alembic_current,
|
||||
run_alembic_downgrade,
|
||||
run_alembic_stamp,
|
||||
run_alembic_upgrade,
|
||||
)
|
||||
from langbot.pkg.utils import constants
|
||||
from langbot.pkg.utils import importutil
|
||||
from langbot.pkg.workspace.collaboration import normalize_email
|
||||
|
||||
|
||||
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
|
||||
|
||||
|
||||
async def _create_legacy_schema(
|
||||
engine,
|
||||
*,
|
||||
include_instance_uuid: bool = True,
|
||||
include_users: bool = True,
|
||||
) -> None:
|
||||
legacy_metadata = sa.MetaData()
|
||||
metadata_table = sa.Table(
|
||||
'metadata',
|
||||
legacy_metadata,
|
||||
sa.Column('key', sa.String(255), primary_key=True),
|
||||
sa.Column('value', sa.String(255)),
|
||||
)
|
||||
users = sa.Table(
|
||||
'users',
|
||||
legacy_metadata,
|
||||
sa.Column('id', sa.Integer, primary_key=True),
|
||||
sa.Column('user', sa.String(255), nullable=False),
|
||||
sa.Column('password', sa.String(255), nullable=False),
|
||||
sa.Column('account_type', sa.String(32), nullable=False, server_default='local'),
|
||||
sa.Column('space_account_uuid', sa.String(255), nullable=True),
|
||||
sa.Column('space_access_token', sa.Text, nullable=True),
|
||||
sa.Column('space_refresh_token', sa.Text, nullable=True),
|
||||
sa.Column('space_access_token_expires_at', sa.DateTime, nullable=True),
|
||||
sa.Column('space_api_key', sa.String(255), nullable=True),
|
||||
sa.Column('created_at', sa.DateTime, nullable=False, server_default=sa.func.now()),
|
||||
sa.Column('updated_at', sa.DateTime, nullable=False, server_default=sa.func.now()),
|
||||
)
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(legacy_metadata.create_all)
|
||||
await conn.execute(metadata_table.insert().values(key='database_version', value='25'))
|
||||
if include_instance_uuid:
|
||||
await conn.execute(metadata_table.insert().values(key='instance_uuid', value='instance_migration_test'))
|
||||
if include_users:
|
||||
await conn.execute(
|
||||
users.insert(),
|
||||
[
|
||||
{'user': 'owner@example.com', 'password': 'owner-hash'},
|
||||
{'user': 'member@example.com', 'password': 'member-hash'},
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def legacy_engine(tmp_path):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "legacy-workspace.db"}')
|
||||
await _create_legacy_schema(engine)
|
||||
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
|
||||
yield engine
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_legacy_instance_gets_stable_accounts_and_default_workspace(legacy_engine):
|
||||
await run_alembic_upgrade(legacy_engine, 'head')
|
||||
|
||||
async with legacy_engine.connect() as conn:
|
||||
tables = set(await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names()))
|
||||
assert {
|
||||
'workspaces',
|
||||
'workspace_memberships',
|
||||
'workspace_invitations',
|
||||
'workspace_execution_states',
|
||||
}.issubset(tables)
|
||||
|
||||
accounts = (
|
||||
(await conn.execute(sa.text('SELECT id, uuid, status, source, projection_revision FROM users ORDER BY id')))
|
||||
.mappings()
|
||||
.all()
|
||||
)
|
||||
assert len(accounts) == 2
|
||||
assert len({account['uuid'] for account in accounts}) == 2
|
||||
for account in accounts:
|
||||
uuid.UUID(account['uuid'])
|
||||
assert account['status'] == 'active'
|
||||
assert account['source'] == 'local'
|
||||
assert account['projection_revision'] == 0
|
||||
|
||||
workspace = (
|
||||
(await conn.execute(sa.text('SELECT * FROM workspaces WHERE source = :source'), {'source': 'local'}))
|
||||
.mappings()
|
||||
.one()
|
||||
)
|
||||
assert workspace['instance_uuid'] == 'instance_migration_test'
|
||||
assert workspace['slug'] == 'default'
|
||||
assert workspace['status'] == 'active'
|
||||
assert workspace['created_by_account_uuid'] == accounts[0]['uuid']
|
||||
|
||||
membership = (await conn.execute(sa.text('SELECT * FROM workspace_memberships'))).mappings().one()
|
||||
assert membership['workspace_uuid'] == workspace['uuid']
|
||||
assert membership['account_uuid'] == accounts[0]['uuid']
|
||||
assert membership['role'] == 'owner'
|
||||
assert membership['status'] == 'active'
|
||||
|
||||
execution_state = (await conn.execute(sa.text('SELECT * FROM workspace_execution_states'))).mappings().one()
|
||||
assert execution_state['workspace_uuid'] == workspace['uuid']
|
||||
assert execution_state['instance_uuid'] == 'instance_migration_test'
|
||||
assert execution_state['active_generation'] == 1
|
||||
assert execution_state['state'] == 'active'
|
||||
assert execution_state['write_fenced'] in (False, 0)
|
||||
|
||||
assert await get_alembic_current(legacy_engine) == '0010_scope_resources'
|
||||
|
||||
|
||||
async def test_workspace_upgrade_is_idempotent_and_preserves_identifiers(legacy_engine):
|
||||
await run_alembic_upgrade(legacy_engine, 'head')
|
||||
async with legacy_engine.connect() as conn:
|
||||
account_uuids_before = (await conn.execute(sa.text('SELECT uuid FROM users ORDER BY id'))).scalars().all()
|
||||
workspace_uuid_before = (
|
||||
await conn.execute(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'"))
|
||||
).scalar_one()
|
||||
|
||||
await run_alembic_upgrade(legacy_engine, 'head')
|
||||
|
||||
async with legacy_engine.connect() as conn:
|
||||
account_uuids_after = (await conn.execute(sa.text('SELECT uuid FROM users ORDER BY id'))).scalars().all()
|
||||
workspace_uuid_after = (
|
||||
await conn.execute(sa.text("SELECT uuid FROM workspaces WHERE source = 'local'"))
|
||||
).scalar_one()
|
||||
assert account_uuids_after == account_uuids_before
|
||||
assert workspace_uuid_after == workspace_uuid_before
|
||||
|
||||
|
||||
async def test_workspace_kernel_upgrade_downgrade_upgrade_round_trip(tmp_path):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-round-trip.db"}')
|
||||
try:
|
||||
await _create_legacy_schema(engine)
|
||||
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
|
||||
await run_alembic_upgrade(engine, '0009_workspace_tenancy')
|
||||
assert await get_alembic_current(engine) == '0009_workspace_tenancy'
|
||||
|
||||
await run_alembic_downgrade(engine, '0008_mcp_resource_prefs')
|
||||
assert await get_alembic_current(engine) == '0008_mcp_resource_prefs'
|
||||
async with engine.connect() as conn:
|
||||
tables = set(await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names()))
|
||||
user_columns = {
|
||||
column['name']
|
||||
for column in await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_columns('users'))
|
||||
}
|
||||
accounts = (await conn.execute(sa.text('SELECT user, password FROM users ORDER BY id'))).all()
|
||||
assert (
|
||||
not {
|
||||
'workspaces',
|
||||
'workspace_memberships',
|
||||
'workspace_invitations',
|
||||
'workspace_execution_states',
|
||||
}
|
||||
& tables
|
||||
)
|
||||
assert not {'uuid', 'status', 'source', 'projection_revision'} & user_columns
|
||||
assert accounts == [
|
||||
('owner@example.com', 'owner-hash'),
|
||||
('member@example.com', 'member-hash'),
|
||||
]
|
||||
|
||||
await run_alembic_upgrade(engine, '0009_workspace_tenancy')
|
||||
assert await get_alembic_current(engine) == '0009_workspace_tenancy'
|
||||
async with engine.connect() as conn:
|
||||
assert await conn.scalar(sa.text('SELECT COUNT(*) FROM workspaces')) == 1
|
||||
assert await conn.scalar(sa.text('SELECT COUNT(*) FROM workspace_memberships')) == 1
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('raw_email', 'expected_email'),
|
||||
[
|
||||
('Straße@Example.COM', 'strasse@example.com'),
|
||||
('Ꭰ@Example.COM', 'Ꭰ@example.com'),
|
||||
],
|
||||
)
|
||||
async def test_workspace_upgrade_uses_runtime_unicode_email_normalization(
|
||||
tmp_path,
|
||||
raw_email,
|
||||
expected_email,
|
||||
):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "unicode-email.db"}')
|
||||
try:
|
||||
await _create_legacy_schema(engine, include_users=False)
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(
|
||||
sa.text('INSERT INTO users (user, password, account_type) VALUES (:email, :password, :type)'),
|
||||
{'email': raw_email, 'password': 'owner-hash', 'type': 'local'},
|
||||
)
|
||||
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
|
||||
await run_alembic_upgrade(engine, 'head')
|
||||
|
||||
async with engine.connect() as conn:
|
||||
assert await conn.scalar(sa.text('SELECT normalized_email FROM users')) == expected_email
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_fresh_sqlite_schema_accepts_application_casefold_identity(tmp_path):
|
||||
importutil.import_modules_in_pkg(persistence)
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "fresh-unicode-email.db"}')
|
||||
canonical_email = normalize_email('Ꭰ@Example.COM')
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
await conn.execute(
|
||||
sa.insert(User).values(
|
||||
uuid='00000000-0000-0000-0000-000000000099',
|
||||
user=canonical_email,
|
||||
normalized_email=canonical_email,
|
||||
password='hash',
|
||||
)
|
||||
)
|
||||
async with engine.connect() as conn:
|
||||
assert await conn.scalar(sa.select(User.normalized_email)) == 'Ꭰ@example.com'
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_workspace_upgrade_rejects_unicode_casefold_duplicate_accounts(tmp_path):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "unicode-email-duplicate.db"}')
|
||||
try:
|
||||
await _create_legacy_schema(engine, include_users=False)
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'INSERT INTO users (user, password, account_type) VALUES '
|
||||
"('Straße@Example.COM', 'first-hash', 'local'), "
|
||||
"('STRASSE@example.com', 'second-hash', 'local')"
|
||||
)
|
||||
)
|
||||
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
|
||||
with pytest.raises(RuntimeError, match='both normalize'):
|
||||
await run_alembic_upgrade(engine, 'head')
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_uninitialized_instance_gets_ownerless_default_workspace(tmp_path):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "uninitialized-instance.db"}')
|
||||
try:
|
||||
await _create_legacy_schema(engine, include_users=False)
|
||||
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
|
||||
await run_alembic_upgrade(engine, 'head')
|
||||
|
||||
async with engine.connect() as conn:
|
||||
workspace = (await conn.execute(sa.text('SELECT * FROM workspaces'))).mappings().one()
|
||||
membership_count = await conn.scalar(sa.text('SELECT COUNT(*) FROM workspace_memberships'))
|
||||
execution_state = (await conn.execute(sa.text('SELECT * FROM workspace_execution_states'))).mappings().one()
|
||||
|
||||
assert workspace['created_by_account_uuid'] is None
|
||||
assert membership_count == 0
|
||||
assert execution_state['workspace_uuid'] == workspace['uuid']
|
||||
assert execution_state['active_generation'] == 1
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_local_workspace_unique_index_allows_cloud_projections(legacy_engine):
|
||||
await run_alembic_upgrade(legacy_engine, 'head')
|
||||
|
||||
async with legacy_engine.begin() as conn:
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'INSERT INTO workspaces '
|
||||
'(uuid, instance_uuid, name, slug, type, status, source, projection_revision) '
|
||||
'VALUES (:uuid, :instance_uuid, :name, :slug, :type, :status, :source, 0)'
|
||||
),
|
||||
{
|
||||
'uuid': str(uuid.uuid4()),
|
||||
'instance_uuid': 'instance_migration_test',
|
||||
'name': 'Cloud Projection',
|
||||
'slug': 'cloud-projection',
|
||||
'type': 'team',
|
||||
'status': 'active',
|
||||
'source': 'cloud_projection',
|
||||
},
|
||||
)
|
||||
|
||||
with pytest.raises(IntegrityError):
|
||||
async with legacy_engine.begin() as conn:
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
'INSERT INTO workspaces '
|
||||
'(uuid, instance_uuid, name, slug, type, status, source, projection_revision) '
|
||||
'VALUES (:uuid, :instance_uuid, :name, :slug, :type, :status, :source, 0)'
|
||||
),
|
||||
{
|
||||
'uuid': str(uuid.uuid4()),
|
||||
'instance_uuid': 'instance_migration_test',
|
||||
'name': 'Second Local',
|
||||
'slug': 'second-local',
|
||||
'type': 'team',
|
||||
'status': 'active',
|
||||
'source': 'local',
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def test_legacy_instance_without_bound_instance_uuid_fails_closed(tmp_path):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "missing-instance.db"}')
|
||||
try:
|
||||
await _create_legacy_schema(engine, include_instance_uuid=False)
|
||||
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
|
||||
with pytest.raises(RuntimeError, match='instance_uuid'):
|
||||
await run_alembic_upgrade(engine, 'head')
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_persistence_startup_defers_workspace_tables_until_account_upgrade(tmp_path, monkeypatch):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "startup-order.db"}')
|
||||
try:
|
||||
await _create_legacy_schema(engine)
|
||||
await run_alembic_stamp(engine, '0008_mcp_resource_prefs')
|
||||
monkeypatch.setattr(constants, 'instance_id', 'instance_migration_test')
|
||||
|
||||
database = type('Database', (), {'get_engine': lambda self: engine})()
|
||||
application = type('Application', (), {})()
|
||||
application.logger = logging.getLogger('workspace-startup-test')
|
||||
manager = PersistenceManager(application)
|
||||
manager.db = database
|
||||
|
||||
await manager.create_tables()
|
||||
async with engine.connect() as conn:
|
||||
tables_before_migration = set(
|
||||
await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names())
|
||||
)
|
||||
assert 'workspaces' not in tables_before_migration
|
||||
|
||||
await manager._run_alembic_migrations()
|
||||
|
||||
async with engine.connect() as conn:
|
||||
workspace = (
|
||||
(await conn.execute(sa.text("SELECT * FROM workspaces WHERE source = 'local'"))).mappings().one()
|
||||
)
|
||||
assert workspace['instance_uuid'] == 'instance_migration_test'
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_persistence_startup_rejects_instance_uuid_drift(tmp_path, monkeypatch):
|
||||
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "instance-drift.db"}')
|
||||
try:
|
||||
await _create_legacy_schema(engine)
|
||||
monkeypatch.setattr(constants, 'instance_id', 'different_instance')
|
||||
|
||||
database = type('Database', (), {'get_engine': lambda self: engine})()
|
||||
application = type('Application', (), {})()
|
||||
application.logger = logging.getLogger('workspace-instance-drift-test')
|
||||
manager = PersistenceManager(application)
|
||||
manager.db = database
|
||||
|
||||
with pytest.raises(RuntimeError, match='does not match'):
|
||||
await manager.create_tables()
|
||||
finally:
|
||||
await engine.dispose()
|
||||
Reference in New Issue
Block a user