xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-08-09 04:40:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(tenancy): add Workspace multi-tenant foundation (#2353)

Browse Source 
* Document multi-tenant workspace architecture

* Add OSS and commercial workspace boundaries

* docs: redesign multi-tenant workspace architecture

* feat(tenancy): implement workspace isolation

* docs(tenancy): record verification evidence

* docs(tenancy): revise single-instance SaaS topology

* docs(tenancy): refine architecture options

* docs: finalize cloud v2 multi-tenant decisions

* feat(tenancy): establish cloud isolation foundations

* feat(tenancy): harden shared cloud runtime boundaries

* docs(tenancy): record final isolation verification

* fix(tenancy): close isolation and permission gaps

* docs(tenancy): record final isolation verification

* feat(tenancy): connect cloud workspace control plane

* fix(build): install git for pinned SDK

* docs(cloud): update control plane verification

* chore: update multi-tenant SDK pin

* fix(cloud): skip legacy model sync during startup

* test(cloud): preserve minimal model manager fixtures

* fix(cloud): preserve authenticated account context

* fix(cloud): reuse authenticated account for user info

* feat(cloud): complete Workspace settings navigation

* test(web): cover Workspace dropdown menu

* feat(web): place workspace controls in sidebar

* refactor(web): streamline workspace controls

* style(web): format workspace layout test

* fix(cloud): surface runtime and workspace plan status

* fix(plugin): keep runtime identity stable across restarts

* fix(ui): widen and center workspace switcher

* fix(ui): hide roles from workspace switcher

* fix(ui): align workspace switcher with sidebar entries

* feat(workspace): add in-product collaboration and direct Cloud launch

* style: format collaboration changes

* fix(workspace): bind collaboration APIs to tenant UoW

* fix(cloud): preserve Core-owned collaboration state

* test(cloud): require Space identity for invite registration

* feat(cloud): complete secure invitation experience

* style(web): format invitation flows

* fix(cloud): recover box runtime without unscoped skill reload

* feat(oss): enforce invitation account and owner billing flows

* style: format OSS account service

* test(oss): cover invitation logout handoff

* fix(oss): resolve workspace owner in scoped session

* feat(cloud): harden multi-tenant runtime resources

* fix(cloud): bound runtime restart storms

* fix(cloud): eliminate periodic runtime CPU spikes

* fix(cloud): enforce instance capacity ceilings

* fix(cloud): scope public login capability discovery

* fix(cloud): bound tenant maintenance and monitoring work

* fix(runtime): bound tenant resource amplification

* fix(deps): pin green multi-tenant plugin SDK

* fix(cloud): handle unavailable skill capability

* fix(security): require authentication for image file endpoint (H-2)

- Changed /api/v1/files/image from AuthType.NONE to USER_TOKEN_OR_API_KEY
- Added Permission.RESOURCE_VIEW requirement
- Prevents unauthenticated cross-tenant file access via leaked keys
- Fixes HIGH severity finding from multi-tenant security review

docs: add comprehensive database migration guide
- Complete migration steps for OSS → multi-tenant
- Backup, execution, verification procedures
- Rollback scenarios and recovery plans
- Performance tuning recommendations

* test: add comprehensive cross-tenant isolation tests

Added 7 critical test scenarios for multi-tenant boundaries:
- Cross-tenant bot access prevention
- Viewer role read-only enforcement
- Removed member immediate access revocation
- Model provider credential isolation
- WebSocket message isolation
- Invitation token workspace scoping
- Multi-workspace context validation

These tests address P0-2 coverage gaps for:
- workspaces.py (membership & invitation flows)
- user.py (authentication & authorization)
- websocket_chat.py (real-time isolation)
- plugins.py (resource access control)

docs: finalize database migration guide

* fix(security): resolve M-1, M-2, M-3 security findings

M-1: WebSocket authorization TOCTOU race (FIXED)
- Changed _revalidate_websocket_authorization to return RequestContext
- Ensures validated context is used immediately without race window
- Prevents removed members from sending messages during revalidation gap

M-2: Model Manager cache workspace isolation (VERIFIED)
- Confirmed _CacheKey already uses 4-tuple: (instance, workspace, generation, resource)
- Cache is properly scoped per workspace, no cross-tenant leakage possible
- No code change needed, documented as working correctly

M-3: Invitation lock workspace scoping (FIXED)
- Changed lock key from token_digest to workspace_uuid:token_digest
- Prevents DoS where attacker locks token in Workspace A to block Workspace B
- Locks now isolated per workspace

All MEDIUM severity findings from security review now resolved.

* fix(cloud): unblock tenant CI and enforce knowledge quotas

* fix(tenancy): scope rerank model sync

---------

Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
This commit is contained in:
RockChinQ
2026-07-30 21:43:35 +08:00
committed by GitHub
parent 463b120923
commit e1ac5e0fc8
468 changed files with 78320 additions and 13137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/integration/persistence/resource_migration_support.py
+253
View File
@@ -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))
tests/integration/persistence/test_migrations.py
+41
View File
@@ -10,10 +10,12 @@ Run: uv run pytest tests/integration/persistence/test_migrations.py -q
from __future__ import annotations
import pytest
import sqlalchemy
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.persistence.alembic_runner import (
run_alembic_downgrade,
run_alembic_upgrade,
run_alembic_stamp,
get_alembic_current,
@@ -149,6 +151,45 @@ class TestSQLiteMigrationUpgrade:
rev2 = await get_alembic_current(sqlite_engine)
assert rev2 == rev1, f'Expected {rev1}, got {rev2}'
@pytest.mark.asyncio
async def test_upgrade_from_0012_adds_knowledge_base_embedding_dimension(self, sqlite_engine):
"""The PostgreSQL pgvector revision also evolves the OSS ORM schema."""
async with sqlite_engine.begin() as conn:
await conn.exec_driver_sql(
'CREATE TABLE knowledge_bases ('
'uuid VARCHAR(255) PRIMARY KEY, workspace_uuid VARCHAR(36) NOT NULL, name VARCHAR(255) NOT NULL)'
)
await run_alembic_stamp(sqlite_engine, '0012_plugin_identity')
await run_alembic_upgrade(sqlite_engine, 'head')
async with sqlite_engine.connect() as conn:
columns = await conn.run_sync(
lambda sync_conn: {
item['name'] for item in sqlalchemy.inspect(sync_conn).get_columns('knowledge_bases')
}
)
assert 'embedding_dimension' in columns
@pytest.mark.asyncio
async def test_directory_projection_upgrade_downgrade_round_trip(self, sqlite_engine):
await run_alembic_stamp(sqlite_engine, '0013_tenant_pgvector')
await run_alembic_upgrade(sqlite_engine, 'head')
async with sqlite_engine.connect() as conn:
tables = await conn.run_sync(lambda sync_conn: set(sqlalchemy.inspect(sync_conn).get_table_names()))
assert {'directory_projection_states', 'directory_projection_inbox'} <= tables
await run_alembic_downgrade(sqlite_engine, '0013_tenant_pgvector')
async with sqlite_engine.connect() as conn:
tables = await conn.run_sync(lambda sync_conn: set(sqlalchemy.inspect(sync_conn).get_table_names()))
assert 'directory_projection_states' not in tables
assert 'directory_projection_inbox' not in tables
await run_alembic_upgrade(sqlite_engine, 'head')
assert await get_alembic_current(sqlite_engine) == _get_script_head()
class TestSQLiteMigrationFreshDatabase:
"""Tests for fresh database workflow."""
tests/integration/persistence/test_migrations_postgres.py
+1962 -8
View File
File diff suppressed because it is too large Load Diff
tests/integration/persistence/test_pgvector_postgres.py
+553
View File
@@ -0,0 +1,553 @@
"""Real PostgreSQL/pgvector tenant isolation and CRUD verification."""
from __future__ import annotations
import logging
import os
import uuid
from types import SimpleNamespace
import pytest
import sqlalchemy as sa
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.rag import KnowledgeBase
from langbot.pkg.entity.persistence.workspace import Workspace
from langbot.pkg.persistence.alembic_runner import get_alembic_current, run_alembic_stamp, run_alembic_upgrade
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
from langbot.pkg.utils import constants
from langbot.pkg.vector.vdbs.pgvector_db import PgVectorDatabase, PgVectorEntry, PgVectorScope
pytestmark = [pytest.mark.integration, pytest.mark.slow, pytest.mark.asyncio]
def _application(postgres_url: str, logger_name: str) -> SimpleNamespace:
url = sa.engine.make_url(postgres_url)
return SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'use': 'postgresql',
'postgresql': {
'host': url.host,
'port': url.port,
'user': url.username,
'password': url.password,
'database': url.database,
},
}
}
),
logger=logging.getLogger(logger_name),
)
def _restore_postgres_registry(monkeypatch) -> None:
from langbot.pkg.persistence import mgr as persistence_mgr_module
from langbot.pkg.persistence.databases.postgresql import PostgreSQLDatabaseManager
monkeypatch.setattr(
persistence_mgr_module.database,
'preregistered_managers',
[PostgreSQLDatabaseManager],
)
@pytest.fixture
def postgres_url() -> str:
url = os.environ.get('TEST_POSTGRES_URL')
if not url:
pytest.skip('TEST_POSTGRES_URL not set')
return url
@pytest.fixture
async def postgres_engine(postgres_url: str):
engine = create_async_engine(postgres_url, isolation_level='AUTOCOMMIT')
yield engine
await engine.dispose()
@pytest.fixture
async def clean_database(postgres_engine: AsyncEngine):
async def clean() -> None:
async with postgres_engine.begin() as conn:
await conn.execute(text('DROP TABLE IF EXISTS langbot_vectors_legacy_0013 CASCADE'))
await conn.execute(text('DROP TABLE IF EXISTS langbot_vectors CASCADE'))
await conn.run_sync(Base.metadata.drop_all)
await conn.execute(text('DROP TABLE IF EXISTS alembic_version'))
await clean()
yield
await clean()
async def test_legacy_upgrade_temporarily_suspends_and_restores_source_rls_for_unprivileged_owner(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
) -> None:
workspace_uuid = '30000000-0000-0000-0000-000000000303'
knowledge_base_uuid = 'legacy-knowledge-base'
migrator_role = f'lb_vector_migrator_{uuid.uuid4().hex[:12]}'
migrator_password = f'Lb{uuid.uuid4().hex}'
quote = postgres_engine.dialect.identifier_preparer.quote
database_name = sa.engine.make_url(postgres_url).database
migrator_url = (
sa.engine.make_url(postgres_url)
.set(username=migrator_role, password=migrator_password)
.render_as_string(hide_password=False)
)
migrator_engine: AsyncEngine | None = None
role_created = False
admin_role: str | None = None
source_tables = ('knowledge_bases', 'knowledge_base_files', 'knowledge_base_chunks')
async def read_rls_states(conn, table_names: tuple[str, ...]) -> dict[str, tuple[bool, bool]]:
rows = (
(
await conn.execute(
text(
"""
SELECT c.relname, c.relrowsecurity, c.relforcerowsecurity
FROM pg_class AS c
JOIN pg_namespace AS n ON n.oid = c.relnamespace
WHERE n.nspname = current_schema()
AND c.relname IN :table_names
"""
).bindparams(sa.bindparam('table_names', expanding=True)),
{'table_names': table_names},
)
)
.mappings()
.all()
)
return {str(row['relname']): (bool(row['relrowsecurity']), bool(row['relforcerowsecurity'])) for row in rows}
try:
async with postgres_engine.begin() as conn:
admin_role = await conn.scalar(text('SELECT current_user'))
await conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
await conn.run_sync(Base.metadata.create_all)
await conn.execute(
text(
'ALTER TABLE knowledge_bases DROP CONSTRAINT IF EXISTS '
'ck_knowledge_bases_embedding_dimension_positive'
)
)
await conn.execute(text('ALTER TABLE knowledge_bases DROP COLUMN IF EXISTS embedding_dimension'))
await run_alembic_stamp(postgres_engine, '0010_scope_resources')
await run_alembic_upgrade(postgres_engine, '0012_plugin_identity')
assert await get_alembic_current(postgres_engine) == '0012_plugin_identity'
embedding_a = '[' + ','.join(['0.125'] * 384) + ']'
embedding_b = '[' + ','.join(['0.25'] * 384) + ']'
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, 'legacy-vector-instance', 'legacy', 'legacy-vector',
'team', 'active', 'cloud_projection', 0)
"""
),
{'uuid': workspace_uuid},
)
await conn.execute(
text(
"""
INSERT INTO knowledge_bases
(uuid, workspace_uuid, name, collection_id, legacy_vector_collection)
VALUES
(:uuid, :workspace_uuid, 'legacy', 'legacy-collection', true)
"""
),
{'uuid': knowledge_base_uuid, 'workspace_uuid': workspace_uuid},
)
await conn.execute(
text(
"""
INSERT INTO knowledge_base_files
(uuid, workspace_uuid, kb_id, file_name, extension, status)
VALUES
('legacy-file', :workspace_uuid, :kb_uuid, 'legacy.txt', 'txt', 'completed')
"""
),
{'workspace_uuid': workspace_uuid, 'kb_uuid': knowledge_base_uuid},
)
await conn.execute(
text(
"""
INSERT INTO knowledge_base_chunks (uuid, workspace_uuid, file_id, text)
VALUES ('legacy-chunk', :workspace_uuid, 'legacy-file', 'chunk text')
"""
),
{'workspace_uuid': workspace_uuid},
)
await conn.execute(
text(
"""
CREATE TABLE langbot_vectors (
id VARCHAR(255) PRIMARY KEY,
collection VARCHAR(255),
embedding vector NOT NULL,
text TEXT,
file_id VARCHAR(255),
chunk_uuid VARCHAR(255)
)
"""
)
)
await conn.execute(
text(
"""
INSERT INTO langbot_vectors
(id, collection, embedding, text, file_id, chunk_uuid)
VALUES
('legacy-by-collection', 'legacy-collection', CAST(:embedding_a AS vector),
'collection row', NULL, NULL),
('legacy-by-chunk', 'unmatched-collection', CAST(:embedding_b AS vector),
'chunk row', NULL, 'legacy-chunk')
"""
),
{'embedding_a': embedding_a, 'embedding_b': embedding_b},
)
# Exercise exact restoration rather than assuming all source tables
# arrived with identical flags.
await conn.execute(text('ALTER TABLE knowledge_base_files NO FORCE ROW LEVEL SECURITY'))
await conn.execute(text('ALTER TABLE knowledge_base_chunks NO FORCE ROW LEVEL SECURITY'))
await conn.execute(text('ALTER TABLE knowledge_base_chunks DISABLE ROW LEVEL SECURITY'))
expected_source_rls = await read_rls_states(conn, source_tables)
assert expected_source_rls == {
'knowledge_bases': (True, True),
'knowledge_base_files': (True, False),
'knowledge_base_chunks': (False, False),
}
await conn.execute(
text(f"CREATE ROLE {quote(migrator_role)} LOGIN PASSWORD '{migrator_password}' NOSUPERUSER NOBYPASSRLS")
)
role_created = True
await conn.execute(text(f'GRANT CONNECT ON DATABASE {quote(database_name)} TO {quote(migrator_role)}'))
await conn.execute(text(f'GRANT USAGE, CREATE ON SCHEMA public TO {quote(migrator_role)}'))
await conn.execute(text(f'GRANT SELECT, UPDATE ON alembic_version TO {quote(migrator_role)}'))
for table_name in (*source_tables, 'langbot_vectors'):
await conn.execute(text(f'ALTER TABLE {quote(table_name)} OWNER TO {quote(migrator_role)}'))
role = (
await conn.execute(
text('SELECT rolsuper, rolbypassrls FROM pg_roles WHERE rolname = :role'),
{'role': migrator_role},
)
).one()
assert role == (False, False)
owned_tables = set(
(
await conn.execute(
text(
"""
SELECT c.relname
FROM pg_class AS c
JOIN pg_namespace AS n ON n.oid = c.relnamespace
WHERE n.nspname = current_schema()
AND c.relname IN :table_names
AND pg_get_userbyid(c.relowner) = :role
"""
).bindparams(sa.bindparam('table_names', expanding=True)),
{'table_names': (*source_tables, 'langbot_vectors'), 'role': migrator_role},
)
).scalars()
)
assert owned_tables == {*source_tables, 'langbot_vectors'}
migrator_engine = create_async_engine(migrator_url)
await run_alembic_upgrade(migrator_engine, '0013_tenant_pgvector')
assert await get_alembic_current(migrator_engine) == '0013_tenant_pgvector'
async with postgres_engine.connect() as conn:
migrated_rows = (
(
await conn.execute(
text(
"""
SELECT workspace_uuid, knowledge_base_uuid, vector_id,
embedding_dimension, text, file_id, chunk_uuid
FROM langbot_vectors
ORDER BY vector_id
"""
)
)
)
.mappings()
.all()
)
assert [dict(row) for row in migrated_rows] == [
{
'workspace_uuid': workspace_uuid,
'knowledge_base_uuid': knowledge_base_uuid,
'vector_id': 'legacy-by-chunk',
'embedding_dimension': 384,
'text': 'chunk row',
'file_id': None,
'chunk_uuid': 'legacy-chunk',
},
{
'workspace_uuid': workspace_uuid,
'knowledge_base_uuid': knowledge_base_uuid,
'vector_id': 'legacy-by-collection',
'embedding_dimension': 384,
'text': 'collection row',
'file_id': None,
'chunk_uuid': None,
},
]
assert (
await conn.scalar(
text('SELECT embedding_dimension FROM knowledge_bases WHERE uuid = :uuid'),
{'uuid': knowledge_base_uuid},
)
== 384
)
assert await conn.scalar(text("SELECT to_regclass('langbot_vectors_legacy_0013') IS NULL")) is True
assert await read_rls_states(conn, source_tables) == expected_source_rls
assert await read_rls_states(conn, ('langbot_vectors',)) == {'langbot_vectors': (True, True)}
assert (
await conn.scalar(
text(
"""
SELECT COUNT(*)
FROM pg_policy AS p
JOIN pg_class AS c ON c.oid = p.polrelid
JOIN pg_namespace AS n ON n.oid = c.relnamespace
WHERE n.nspname = current_schema()
AND c.relname = 'langbot_vectors'
AND p.polname = 'langbot_workspace_isolation'
"""
)
)
== 1
)
async with migrator_engine.connect() as conn:
assert await conn.scalar(text('SELECT COUNT(*) FROM langbot_vectors')) == 0
async with migrator_engine.begin() as conn:
await conn.execute(
text("SELECT set_config('langbot.workspace_uuid', :workspace_uuid, true)"),
{'workspace_uuid': workspace_uuid},
)
assert await conn.scalar(text('SELECT COUNT(*) FROM langbot_vectors')) == 2
finally:
if migrator_engine is not None:
await migrator_engine.dispose()
if role_created:
async with postgres_engine.connect() as conn:
if admin_role is None: # pragma: no cover - setup cannot create the role without an admin
admin_role = await conn.scalar(text('SELECT current_user'))
await conn.execute(text(f'REASSIGN OWNED BY {quote(migrator_role)} TO {quote(admin_role)}'))
await conn.execute(text(f'DROP OWNED BY {quote(migrator_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(migrator_role)}'))
async def test_pgvector_shared_database_is_scoped_indexed_and_ddl_free_at_runtime(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
monkeypatch,
) -> None:
instance_uuid = 'pgvector-tenant-integration'
workspace_a = '10000000-0000-0000-0000-000000000101'
workspace_b = '20000000-0000-0000-0000-000000000202'
kb_a = 'knowledge-base-a'
kb_b = 'knowledge-base-b'
role_suffix = uuid.uuid4().hex[:12]
runtime_role = f'lb_vector_runtime_{role_suffix}'
role_password = f'Lb{uuid.uuid4().hex}'
release_manager: PersistenceManager | None = None
runtime_manager: PersistenceManager | None = None
role_created = False
_restore_postgres_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', instance_uuid)
quote = postgres_engine.dialect.identifier_preparer.quote
database_name = sa.engine.make_url(postgres_url).database
runtime_url = (
sa.engine.make_url(postgres_url)
.set(username=runtime_role, password=role_password)
.render_as_string(hide_password=False)
)
try:
release_app = _application(postgres_url, 'pgvector-release-migration-test')
release_manager = PersistenceManager(release_app, mode=PersistenceMode.RELEASE_MIGRATION)
release_app.persistence_mgr = release_manager
await release_manager.initialize()
for workspace_uuid, kb_uuid, slug in (
(workspace_a, kb_a, 'vector-a'),
(workspace_b, kb_b, 'vector-b'),
):
async with release_manager.tenant_uow(workspace_uuid) as uow:
await uow.execute(
sa.insert(Workspace).values(
uuid=workspace_uuid,
instance_uuid=instance_uuid,
name=slug,
slug=slug,
type='team',
status='active',
source='cloud_projection',
projection_revision=0,
)
)
await uow.execute(
sa.insert(KnowledgeBase).values(
uuid=kb_uuid,
workspace_uuid=workspace_uuid,
name=slug,
embedding_dimension=384,
)
)
async with postgres_engine.connect() as conn:
await conn.execute(text(f"CREATE ROLE {quote(runtime_role)} LOGIN PASSWORD '{role_password}'"))
await conn.execute(text(f'GRANT CONNECT ON DATABASE {quote(database_name)} TO {quote(runtime_role)}'))
await conn.execute(text(f'GRANT USAGE ON SCHEMA public TO {quote(runtime_role)}'))
business_tables = release_manager._runtime_business_table_names()
quoted_tables = ', '.join(f'public.{quote(table_name)}' for table_name in business_tables)
await conn.execute(
text(f'GRANT SELECT, INSERT, UPDATE, DELETE ON TABLE {quoted_tables} TO {quote(runtime_role)}')
)
await conn.execute(text(f'GRANT SELECT ON TABLE public.alembic_version TO {quote(runtime_role)}'))
sequence_names = await release_manager._runtime_business_sequence_names(conn, business_tables)
if sequence_names:
quoted_sequences = ', '.join(f'public.{quote(sequence_name)}' for sequence_name in sequence_names)
await conn.execute(text(f'GRANT USAGE, SELECT ON SEQUENCE {quoted_sequences} TO {quote(runtime_role)}'))
role_created = True
runtime_app = _application(runtime_url, 'pgvector-runtime-test')
runtime_manager = PersistenceManager(runtime_app, mode=PersistenceMode.CLOUD_RUNTIME)
runtime_app.persistence_mgr = runtime_manager
await runtime_manager.initialize()
adapter = PgVectorDatabase(
runtime_app,
use_business_database=True,
allowed_dimensions=[384],
)
scope_a = PgVectorScope(workspace_a, kb_a, 384)
scope_b = PgVectorScope(workspace_b, kb_b, 384)
# The same vector ID is valid in two Workspaces because the relational
# primary key includes Workspace and knowledge base.
await adapter.add_embeddings(
'opaque-a',
['same-vector'],
[[0.1] * 384],
[{'text': 'workspace-a', 'file_id': 'file-a', 'uuid': 'chunk-a'}],
scope=scope_a,
)
await adapter.add_embeddings(
'opaque-b',
['same-vector'],
[[0.2] * 384],
[{'text': 'workspace-b', 'file_id': 'file-b', 'uuid': 'chunk-b'}],
scope=scope_b,
)
result_a = await adapter.search('opaque-a', [0.1] * 384, scope=scope_a)
result_b = await adapter.search('opaque-b', [0.2] * 384, scope=scope_b)
assert result_a['metadatas'][0][0]['text'] == 'workspace-a'
assert result_b['metadatas'][0][0]['text'] == 'workspace-b'
# Guessing another knowledge-base UUID while retaining A's Workspace
# cannot escape either the explicit conditions or PostgreSQL RLS.
guessed = await adapter.search(
'attacker-controlled-name',
[0.2] * 384,
scope=PgVectorScope(workspace_a, kb_b, 384),
)
assert guessed['ids'] == [[]]
with pytest.raises(ValueError, match='trusted PgVectorScope'):
await adapter.search('opaque-a', [0.1] * 384)
with pytest.raises(ValueError, match='selected dimension'):
await adapter.add_embeddings(
'opaque-a',
['bad-dimension'],
[[0.1] * 383],
[{}],
scope=scope_a,
)
# Deliberately omit the application Workspace predicate. FORCE RLS is
# still the second isolation boundary and returns only A.
async with runtime_manager.tenant_uow(workspace_a) as uow:
rows = (
await uow.execute(
sa.select(PgVectorEntry.workspace_uuid, PgVectorEntry.vector_id).where(
PgVectorEntry.vector_id == 'same-vector'
)
)
).all()
assert rows == [(workspace_a, 'same-vector')]
# Index-plan inspection is deployment diagnostics, not a public tenant
# Session capability. Establish the same transaction-local RLS scope on
# a test-only connection and keep raw EXPLAIN outside TenantUnitOfWork.
runtime_engine = runtime_manager.get_db_engine()
async with runtime_engine.begin() as conn:
await conn.execute(
text('SELECT set_config(:setting_name, :setting_value, true)'),
{'setting_name': 'langbot.workspace_uuid', 'setting_value': workspace_a},
)
await conn.execute(text('SET LOCAL enable_seqscan = off'))
plan = '\n'.join(
(
await conn.execute(
text(
"""
EXPLAIN SELECT vector_id
FROM langbot_vectors
WHERE workspace_uuid = :workspace_uuid
AND knowledge_base_uuid = :knowledge_base_uuid
AND embedding_dimension = 384
ORDER BY (embedding::vector(384)) <=> CAST(:query AS vector(384))
LIMIT 5
"""
),
{
'workspace_uuid': workspace_a,
'knowledge_base_uuid': kb_a,
'query': '[' + ','.join(['0.1'] * 384) + ']',
},
)
).scalars()
)
assert 'ix_langbot_vectors_hnsw_cosine_384' in plan
async with runtime_engine.connect() as conn:
assert await conn.scalar(text('SELECT COUNT(*) FROM langbot_vectors')) == 0
assert await conn.scalar(text("SELECT current_setting('langbot.workspace_uuid', true)")) in (None, '')
items_a, total_a = await adapter.list_by_filter('opaque-a', scope=scope_a)
assert total_a == 1
assert items_a[0]['metadata']['file_id'] == 'file-a'
await adapter.delete_by_file_id('opaque-a', 'file-a', scope=scope_a)
assert (await adapter.list_by_filter('opaque-a', scope=scope_a))[1] == 0
assert (await adapter.list_by_filter('opaque-b', scope=scope_b))[1] == 1
finally:
if runtime_manager is not None and getattr(runtime_manager, 'db', None) is not None:
await runtime_manager.get_db_engine().dispose()
if release_manager is not None and getattr(release_manager, 'db', None) is not None:
await release_manager.get_db_engine().dispose()
if role_created:
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}'))
tests/integration/persistence/test_plugin_identity_migration.py
+94
View File
@@ -0,0 +1,94 @@
from __future__ import annotations
import hashlib
import uuid
import pytest
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.persistence.alembic_runner import (
get_alembic_current,
run_alembic_stamp,
run_alembic_upgrade,
)
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
async def test_legacy_plugin_settings_receive_stable_random_installation_identities(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "plugin-identity.db"}')
metadata = sa.MetaData()
plugin_settings = sa.Table(
'plugin_settings',
metadata,
sa.Column('workspace_uuid', sa.String(36), primary_key=True),
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, server_default=sa.true()),
sa.Column('priority', sa.Integer, nullable=False, server_default='0'),
sa.Column('config', sa.JSON, nullable=False, server_default='{}'),
sa.Column('install_source', sa.String(255), nullable=False, server_default='local'),
sa.Column('install_info', sa.JSON, nullable=False, server_default='{}'),
)
try:
async with engine.begin() as connection:
await connection.run_sync(metadata.create_all)
await connection.execute(
plugin_settings.insert(),
[
{
'workspace_uuid': '11111111-1111-4111-8111-111111111111',
'plugin_author': 'author',
'plugin_name': 'one',
},
{
'workspace_uuid': '22222222-2222-4222-8222-222222222222',
'plugin_author': 'author',
'plugin_name': 'two',
},
],
)
await run_alembic_stamp(engine, '0011_postgres_tenant_rls')
await run_alembic_upgrade(engine, '0012_plugin_identity')
async with engine.connect() as connection:
rows = (
(
await connection.execute(
sa.text(
'SELECT installation_uuid, artifact_digest, runtime_revision '
'FROM plugin_settings ORDER BY workspace_uuid'
)
)
)
.mappings()
.all()
)
columns = await connection.run_sync(
lambda sync_connection: {
column['name']: column for column in sa.inspect(sync_connection).get_columns('plugin_settings')
}
)
indexes = await connection.run_sync(
lambda sync_connection: {
index['name']: index for index in sa.inspect(sync_connection).get_indexes('plugin_settings')
}
)
assert await get_alembic_current(engine) == '0012_plugin_identity'
assert columns['installation_uuid']['nullable'] is False
assert columns['artifact_digest']['nullable'] is False
assert columns['runtime_revision']['nullable'] is False
assert indexes['ix_plugin_settings_workspace_installation']['unique'] == 1
assert len({row['installation_uuid'] for row in rows}) == 2
for row in rows:
uuid.UUID(row['installation_uuid'])
assert row['runtime_revision'] == 1
assert (
row['artifact_digest']
== hashlib.sha256(f'legacy-installation:{row["installation_uuid"]}'.encode()).hexdigest()
)
finally:
await engine.dispose()
tests/integration/persistence/test_release_migration_postgres.py
+734
View File
@@ -0,0 +1,734 @@
"""Real PostgreSQL coverage for the one-shot Cloud release migration job."""
from __future__ import annotations
import logging
import os
import uuid
from types import SimpleNamespace
import pytest
import sqlalchemy as sa
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from langbot.pkg.persistence import release_migration
from langbot.pkg.persistence.alembic_runner import (
get_alembic_current,
get_alembic_head,
run_alembic_stamp,
run_alembic_upgrade,
)
from langbot.pkg.persistence.mgr import (
PersistenceManager,
PersistenceMode,
_RELEASE_MIGRATION_ADVISORY_LOCK_ID,
)
from langbot.pkg.utils import constants
pytestmark = [pytest.mark.integration, pytest.mark.slow, pytest.mark.asyncio]
_RUNTIME_PASSWORD = 'runtime-secret-not-used-by-migration'
@pytest.fixture
def postgres_url() -> str:
url = os.environ.get('TEST_POSTGRES_URL')
if not url:
pytest.skip('TEST_POSTGRES_URL not set')
return url
@pytest.fixture
async def postgres_engine(postgres_url: str):
engine = create_async_engine(postgres_url, isolation_level='AUTOCOMMIT')
yield engine
await engine.dispose()
@pytest.fixture
async def clean_database(postgres_engine: AsyncEngine):
async def clean() -> None:
async with postgres_engine.begin() as conn:
table_names = await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names())
quote = postgres_engine.dialect.identifier_preparer.quote
for table_name in table_names:
await conn.execute(text(f'DROP TABLE {quote(table_name)} CASCADE'))
await clean()
yield
await clean()
def _restore_postgres_registry(monkeypatch) -> None:
from langbot.pkg.persistence import mgr as persistence_mgr_module
from langbot.pkg.persistence.databases.postgresql import PostgreSQLDatabaseManager
monkeypatch.setattr(
persistence_mgr_module.database,
'preregistered_managers',
[PostgreSQLDatabaseManager],
)
def _application(postgres_url: str, *, runtime_role: str = 'langbot_runtime_not_used_by_migration') -> SimpleNamespace:
url = sa.engine.make_url(postgres_url)
return SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'use': 'postgresql',
'postgresql': {
'host': url.host,
'port': url.port,
# This is deliberately not the operator role in the DSN.
'user': runtime_role,
'password': _RUNTIME_PASSWORD,
'database': url.database,
},
'cloud_migration': {'operator_dsn_env': 'TEST_RELEASE_OPERATOR_DSN'},
},
'vdb': {
'use': 'pgvector',
'pgvector': {
'use_business_database': True,
'allowed_dimensions': [384, 512, 768, 1024, 1536],
},
},
}
),
logger=logging.getLogger('cloud-release-migration-entrypoint-test'),
persistence_mgr=None,
)
async def test_release_entrypoint_holds_lock_migrates_validates_and_disposes(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
monkeypatch,
) -> None:
_restore_postgres_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', 'release-migration-entrypoint-test')
original_validate = PersistenceManager._validate_release_schema
validation_observed_lock = False
runtime_role = f'lb_release_runtime_{uuid.uuid4().hex[:12]}'
quote = postgres_engine.dialect.identifier_preparer.quote
async def validate_while_asserting_lock(self: PersistenceManager) -> None:
nonlocal validation_observed_lock
async with postgres_engine.connect() as conn:
acquired = await conn.scalar(
text('SELECT pg_try_advisory_lock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
if acquired:
await conn.scalar(
text('SELECT pg_advisory_unlock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
assert acquired is False
validation_observed_lock = True
await original_validate(self)
monkeypatch.setattr(PersistenceManager, '_validate_release_schema', validate_while_asserting_lock)
async with postgres_engine.connect() as conn:
await conn.execute(text(f"CREATE ROLE {quote(runtime_role)} LOGIN PASSWORD '{_RUNTIME_PASSWORD}'"))
# The release job must make this otherwise bare LOGIN usable without
# relying on pre-provisioned object ACLs.
assert (
await conn.scalar(
text(
"""
SELECT NOT EXISTS (
SELECT 1
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
CROSS JOIN LATERAL aclexplode(c.relacl) acl
JOIN pg_roles grantee ON grantee.oid = acl.grantee
WHERE n.nspname = current_schema()
AND grantee.rolname = :runtime_role
)
"""
),
{'runtime_role': runtime_role},
)
is True
)
ap = _application(postgres_url, runtime_role=runtime_role)
try:
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
assert validation_observed_lock is True
assert await get_alembic_current(postgres_engine) == get_alembic_head()
manager = ap.persistence_mgr
assert isinstance(manager, PersistenceManager)
business_tables = set(manager._runtime_business_table_names())
async with postgres_engine.connect() as conn:
assert await conn.scalar(text("SELECT to_regclass('langbot_vectors') IS NOT NULL")) is True
assert (
await conn.scalar(text("SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'vector')")) is True
)
runtime_role_state = (
(
await conn.execute(
text(
"""
SELECT
rolcanlogin,
rolsuper,
rolbypassrls,
rolcreatedb,
rolcreaterole,
rolreplication
FROM pg_roles
WHERE rolname = :runtime_role
"""
),
{'runtime_role': runtime_role},
)
)
.mappings()
.one()
)
assert dict(runtime_role_state) == {
'rolcanlogin': True,
'rolsuper': False,
'rolbypassrls': False,
'rolcreatedb': False,
'rolcreaterole': False,
'rolreplication': False,
}
database_privileges = set(
(
await conn.execute(
text(
"""
SELECT acl.privilege_type
FROM pg_database database
CROSS JOIN LATERAL aclexplode(database.datacl) acl
JOIN pg_roles grantee ON grantee.oid = acl.grantee
WHERE database.datname = current_database()
AND grantee.rolname = :runtime_role
AND acl.is_grantable IS FALSE
"""
),
{'runtime_role': runtime_role},
)
)
.scalars()
.all()
)
schema_privileges = set(
(
await conn.execute(
text(
"""
SELECT acl.privilege_type
FROM pg_namespace namespace
CROSS JOIN LATERAL aclexplode(namespace.nspacl) acl
JOIN pg_roles grantee ON grantee.oid = acl.grantee
WHERE namespace.nspname = current_schema()
AND grantee.rolname = :runtime_role
AND acl.is_grantable IS FALSE
"""
),
{'runtime_role': runtime_role},
)
)
.scalars()
.all()
)
assert database_privileges == {'CONNECT'}
assert schema_privileges == {'USAGE'}
assert (
await conn.scalar(
text("SELECT has_database_privilege(:runtime_role, current_database(), 'CREATE')"),
{'runtime_role': runtime_role},
)
is False
)
assert (
await conn.scalar(
text("SELECT has_schema_privilege(:runtime_role, current_schema(), 'CREATE')"),
{'runtime_role': runtime_role},
)
is False
)
# PostgreSQL grants TEMP to PUBLIC by default. The first Cloud
# release deliberately tolerates that inherited compatibility
# privilege while granting no direct TEMP ACL to the runtime role.
assert (
await conn.scalar(
text("SELECT has_database_privilege(:runtime_role, current_database(), 'TEMP')"),
{'runtime_role': runtime_role},
)
is True
)
object_grants = (
(
await conn.execute(
text(
"""
SELECT
c.relname,
c.relkind::text AS relkind,
acl.privilege_type,
acl.is_grantable
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
CROSS JOIN LATERAL aclexplode(c.relacl) acl
JOIN pg_roles grantee ON grantee.oid = acl.grantee
WHERE n.nspname = current_schema()
AND c.relkind IN ('r', 'p', 'S')
AND grantee.rolname = :runtime_role
ORDER BY c.relname, acl.privilege_type
"""
),
{'runtime_role': runtime_role},
)
)
.mappings()
.all()
)
direct_table_grants: dict[str, set[str]] = {}
direct_sequence_grants: dict[str, set[str]] = {}
for grant in object_grants:
assert grant['is_grantable'] is False
target = direct_sequence_grants if grant['relkind'] == 'S' else direct_table_grants
target.setdefault(grant['relname'], set()).add(grant['privilege_type'])
assert set(direct_table_grants) == business_tables | {'alembic_version'}
assert all(
privileges == {'SELECT', 'INSERT', 'UPDATE', 'DELETE'}
for table_name, privileges in direct_table_grants.items()
if table_name != 'alembic_version'
)
assert direct_table_grants['alembic_version'] == {'SELECT'}
assert direct_sequence_grants
assert all(privileges == {'USAGE', 'SELECT'} for privileges in direct_sequence_grants.values())
# The session-level lock must be released before the one-shot job exits.
assert (
await conn.scalar(
text('SELECT pg_try_advisory_lock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
is True
)
assert (
await conn.scalar(
text('SELECT pg_advisory_unlock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
is True
)
runtime_url = (
sa.engine.make_url(postgres_url)
.set(username=runtime_role, password=_RUNTIME_PASSWORD)
.render_as_string(hide_password=False)
)
runtime_engine = create_async_engine(runtime_url)
try:
async with runtime_engine.begin() as conn:
assert await conn.scalar(text('SELECT current_user')) == runtime_role
await conn.execute(text("INSERT INTO metadata (key, value) VALUES ('runtime-grant-smoke', 'created')"))
await conn.execute(text("UPDATE metadata SET value = 'updated' WHERE key = 'runtime-grant-smoke'"))
assert (
await conn.scalar(text("SELECT value FROM metadata WHERE key = 'runtime-grant-smoke'")) == 'updated'
)
await conn.execute(text("DELETE FROM metadata WHERE key = 'runtime-grant-smoke'"))
await conn.scalar(
text('SELECT nextval(CAST(:sequence_name AS regclass))'),
{'sequence_name': f'public.{sorted(direct_sequence_grants)[0]}'},
)
finally:
await runtime_engine.dispose()
runtime_application = _application(postgres_url, runtime_role=runtime_role)
runtime_manager = PersistenceManager(runtime_application, mode=PersistenceMode.CLOUD_RUNTIME)
runtime_application.persistence_mgr = runtime_manager
try:
# This reads alembic_version as the actual runtime role and then
# reruns the complete grant/catalog validator before startup.
await runtime_manager.initialize()
finally:
await runtime_manager.get_db_engine().dispose()
# The catalog validator must fail closed if the role is no longer
# deployable, even though the remaining grants still look plausible.
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE DELETE ON TABLE public.metadata FROM {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match="table 'metadata' grants are incomplete"):
await manager._validate_configured_runtime_postgres_role(require_grants=True)
finally:
manager = getattr(ap, 'persistence_mgr', None)
if isinstance(manager, PersistenceManager):
# The one-shot entrypoint disposes its pool before returning. This
# test deliberately reuses the manager for catalog mutation checks,
# which can open a fresh pool and therefore owns a second shutdown.
await manager.shutdown()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}'))
async def test_release_entrypoint_rejects_privileged_or_table_owning_runtime_role(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
monkeypatch,
) -> None:
_restore_postgres_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', 'release-runtime-role-validation-test')
runtime_role = f'lb_release_runtime_{uuid.uuid4().hex[:12]}'
quote = postgres_engine.dialect.identifier_preparer.quote
async with postgres_engine.connect() as conn:
operator_role = await conn.scalar(text('SELECT current_user'))
await conn.execute(text(f'CREATE ROLE {quote(runtime_role)} LOGIN'))
ap = _application(postgres_url, runtime_role=runtime_role)
try:
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} SUPERUSER'))
with pytest.raises(RuntimeError, match='must not be superuser or BYPASSRLS'):
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} NOSUPERUSER BYPASSRLS'))
with pytest.raises(RuntimeError, match='must not be superuser or BYPASSRLS'):
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} NOBYPASSRLS'))
await conn.execute(text(f'ALTER TABLE bots OWNER TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='owns tenant tables'):
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
finally:
async with postgres_engine.connect() as conn:
table_exists = await conn.scalar(text("SELECT to_regclass('bots') IS NOT NULL"))
if table_exists:
await conn.execute(text(f'ALTER TABLE bots OWNER TO {quote(operator_role)}'))
await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}'))
async def test_runtime_role_catalog_validator_rejects_delegation_and_escape_hatches(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
monkeypatch,
) -> None:
_restore_postgres_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', 'release-runtime-role-catalog-test')
suffix = uuid.uuid4().hex[:12]
runtime_role = f'lb_release_runtime_{suffix}'
delegated_role = f'lb_release_delegate_{suffix}'
extra_schema = f'lb_release_schema_{suffix}'
extra_view = f'lb_release_view_{suffix}'
owned_routine = f'lb_release_owned_routine_{suffix}'
security_definer = f'lb_release_definer_{suffix}'
foreign_wrapper = f'lb_release_fdw_{suffix}'
foreign_server = f'lb_release_server_{suffix}'
persistent_setting_secret = f'lb_release_secret_{suffix}'
database_name = sa.engine.make_url(postgres_url).database
assert database_name
quote = postgres_engine.dialect.identifier_preparer.quote
async with postgres_engine.connect() as conn:
await conn.execute(text(f"CREATE ROLE {quote(runtime_role)} LOGIN PASSWORD '{_RUNTIME_PASSWORD}'"))
await conn.execute(text(f'CREATE ROLE {quote(delegated_role)}'))
ap = _application(postgres_url, runtime_role=runtime_role)
try:
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
manager = ap.persistence_mgr
assert isinstance(manager, PersistenceManager)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'GRANT pg_read_all_data TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='must not participate in role memberships'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE pg_read_all_data FROM {quote(runtime_role)}'))
# Reject delegation in the other direction too: no role may be
# allowed to SET ROLE to the runtime identity or administer it.
await conn.execute(text(f'GRANT {quote(runtime_role)} TO {quote(delegated_role)} WITH ADMIN OPTION'))
with pytest.raises(RuntimeError, match='must not participate in role memberships'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE {quote(runtime_role)} FROM {quote(delegated_role)}'))
await conn.execute(
text(f'GRANT SELECT ON TABLE public.metadata TO {quote(runtime_role)} WITH GRANT OPTION')
)
with pytest.raises(RuntimeError, match='GRANT OPTION'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(
text(f'REVOKE GRANT OPTION FOR SELECT ON TABLE public.metadata FROM {quote(runtime_role)}')
)
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} SET search_path TO public'))
with pytest.raises(RuntimeError, match='persistent session overrides'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET search_path'))
await conn.execute(text(f'ALTER DATABASE {quote(database_name)} SET search_path TO public'))
with pytest.raises(RuntimeError, match='persistent session overrides'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER DATABASE {quote(database_name)} RESET search_path'))
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} SET session_replication_role TO replica'))
with pytest.raises(RuntimeError, match='persistent session overrides'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET session_replication_role'))
await conn.execute(
text(f"ALTER ROLE {quote(runtime_role)} SET application_name TO '{persistent_setting_secret}'")
)
with pytest.raises(RuntimeError, match='persistent session overrides') as error:
await manager._validate_configured_runtime_postgres_role()
assert persistent_setting_secret not in str(error.value)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET application_name'))
await conn.execute(text('CREATE EXTENSION dblink'))
with pytest.raises(RuntimeError, match='extensions must include vector and be limited'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text('DROP EXTENSION dblink'))
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'GRANT CREATE ON DATABASE {quote(database_name)} TO {quote(runtime_role)}'))
await conn.execute(text(f'GRANT CREATE ON SCHEMA public TO {quote(runtime_role)}'))
runtime_url = (
sa.engine.make_url(postgres_url)
.set(username=runtime_role, password=_RUNTIME_PASSWORD)
.render_as_string(hide_password=False)
)
runtime_engine = create_async_engine(runtime_url)
try:
async with runtime_engine.begin() as conn:
# hstore is a trusted extension in the production PG16 image,
# so this creates a real runtime-owned extension catalog row.
await conn.execute(text('CREATE EXTENSION hstore'))
finally:
await runtime_engine.dispose()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE CREATE ON SCHEMA public FROM {quote(runtime_role)}'))
await conn.execute(text(f'REVOKE CREATE ON DATABASE {quote(database_name)} FROM {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='must not own extensions'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text('DROP EXTENSION hstore CASCADE'))
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'CREATE FOREIGN DATA WRAPPER {quote(foreign_wrapper)}'))
await conn.execute(
text(f'CREATE SERVER {quote(foreign_server)} FOREIGN DATA WRAPPER {quote(foreign_wrapper)}')
)
await conn.execute(text(f'CREATE USER MAPPING FOR {quote(runtime_role)} SERVER {quote(foreign_server)}'))
with pytest.raises(RuntimeError, match='foreign data wrappers, servers, or user mappings'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP FOREIGN DATA WRAPPER {quote(foreign_wrapper)} CASCADE'))
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'CREATE SCHEMA {quote(extra_schema)} AUTHORIZATION {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='non-business schemas'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP SCHEMA {quote(extra_schema)} CASCADE'))
await conn.execute(text(f'CREATE VIEW public.{quote(extra_view)} AS SELECT key FROM public.metadata'))
await conn.execute(text(f'GRANT SELECT ON public.{quote(extra_view)} TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='non-business objects|table privileges are unsafe'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP VIEW public.{quote(extra_view)}'))
await conn.execute(text(f'GRANT SELECT (key) ON public.metadata TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='column-level ACLs'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE SELECT (key) ON public.metadata FROM {quote(runtime_role)}'))
# System file access functions are not SECURITY DEFINER, so the
# validator must reject their explicit EXECUTE ACL independently.
await conn.execute(
text(f'GRANT EXECUTE ON FUNCTION pg_catalog.pg_read_file(text) TO {quote(runtime_role)}')
)
runtime_application = _application(postgres_url, runtime_role=runtime_role)
runtime_manager = PersistenceManager(runtime_application, mode=PersistenceMode.CLOUD_RUNTIME)
runtime_application.persistence_mgr = runtime_manager
try:
with pytest.raises(RuntimeError, match='explicit EXECUTE privileges on routines'):
await runtime_manager.initialize()
finally:
await runtime_manager.get_db_engine().dispose()
async with postgres_engine.connect() as conn:
await conn.execute(
text(f'REVOKE EXECUTE ON FUNCTION pg_catalog.pg_read_file(text) FROM {quote(runtime_role)}')
)
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
# replica disables ordinary triggers/rules and foreign-key
# enforcement; it must never reach the runtime identity.
await conn.execute(text(f'GRANT SET ON PARAMETER session_replication_role TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='explicit SET or ALTER SYSTEM parameter privileges'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE SET ON PARAMETER session_replication_role FROM {quote(runtime_role)}'))
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(
text(f"CREATE FUNCTION public.{quote(owned_routine)}() RETURNS integer LANGUAGE sql AS 'SELECT 1'")
)
await conn.execute(text(f'ALTER FUNCTION public.{quote(owned_routine)}() OWNER TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='must not own routines'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP FUNCTION public.{quote(owned_routine)}()'))
async with postgres_engine.connect() as conn:
await conn.execute(
text(
f'CREATE FUNCTION public.{quote(security_definer)}() RETURNS integer '
"LANGUAGE sql SECURITY DEFINER AS 'SELECT 1'"
)
)
# Extension membership must not exempt an executable definer from
# the runtime audit, even for an allowlisted extension.
await conn.execute(text(f'ALTER EXTENSION vector ADD FUNCTION public.{quote(security_definer)}()'))
with pytest.raises(RuntimeError, match='SECURITY DEFINER'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER EXTENSION vector DROP FUNCTION public.{quote(security_definer)}()'))
finally:
manager = getattr(ap, 'persistence_mgr', None)
if isinstance(manager, PersistenceManager):
# The entrypoint disposed the release pool before returning; all
# validator calls above happened afterward and can reopen it.
await manager.shutdown()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER DATABASE {quote(database_name)} RESET search_path'))
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET search_path'))
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET session_replication_role'))
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET application_name'))
await conn.execute(text(f'REVOKE pg_read_all_data FROM {quote(runtime_role)}'))
await conn.execute(text(f'REVOKE {quote(runtime_role)} FROM {quote(delegated_role)}'))
await conn.execute(text(f'REVOKE CREATE ON SCHEMA public FROM {quote(runtime_role)}'))
await conn.execute(text(f'REVOKE CREATE ON DATABASE {quote(database_name)} FROM {quote(runtime_role)}'))
await conn.execute(
text(f'REVOKE EXECUTE ON FUNCTION pg_catalog.pg_read_file(text) FROM {quote(runtime_role)}')
)
await conn.execute(text(f'REVOKE SET ON PARAMETER session_replication_role FROM {quote(runtime_role)}'))
await conn.execute(text('DROP EXTENSION IF EXISTS dblink CASCADE'))
await conn.execute(text('DROP EXTENSION IF EXISTS hstore CASCADE'))
await conn.execute(text(f'DROP FOREIGN DATA WRAPPER IF EXISTS {quote(foreign_wrapper)} CASCADE'))
await conn.execute(text(f'DROP FUNCTION IF EXISTS public.{quote(owned_routine)}()'))
security_definer_is_extension_member = await conn.scalar(
text(
"""
SELECT EXISTS (
SELECT 1
FROM pg_depend dependency
JOIN pg_extension extension ON extension.oid = dependency.refobjid
WHERE dependency.classid = 'pg_proc'::regclass
AND dependency.objid = to_regprocedure(:routine)
AND dependency.refclassid = 'pg_extension'::regclass
AND dependency.deptype = 'e'
AND extension.extname = 'vector'
)
"""
),
{'routine': f'public.{security_definer}()'},
)
if security_definer_is_extension_member:
await conn.execute(text(f'ALTER EXTENSION vector DROP FUNCTION public.{quote(security_definer)}()'))
await conn.execute(text(f'DROP FUNCTION IF EXISTS public.{quote(security_definer)}()'))
await conn.execute(text(f'DROP VIEW IF EXISTS public.{quote(extra_view)}'))
await conn.execute(text(f'DROP SCHEMA IF EXISTS {quote(extra_schema)} CASCADE'))
await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(delegated_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}'))
async def test_direct_postgres_head_stamp_fails_without_business_schema(
postgres_engine: AsyncEngine,
clean_database,
) -> None:
await run_alembic_stamp(postgres_engine, '0012_plugin_identity')
with pytest.raises(RuntimeError, match='requires the knowledge_bases table'):
await run_alembic_upgrade(postgres_engine)
assert await get_alembic_current(postgres_engine) == '0012_plugin_identity'
async def test_release_entrypoint_fails_immediately_when_another_job_holds_lock(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
monkeypatch,
) -> None:
_restore_postgres_registry(monkeypatch)
ap = _application(postgres_url)
async with postgres_engine.connect() as lock_connection:
assert (
await lock_connection.scalar(
text('SELECT pg_try_advisory_lock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
is True
)
try:
with pytest.raises(RuntimeError, match='already holds the advisory lock'):
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
finally:
assert (
await lock_connection.scalar(
text('SELECT pg_advisory_unlock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
is True
)
async with postgres_engine.connect() as conn:
assert await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names()) == []
tests/integration/persistence/test_resource_tenancy_migration.py
+292
View File
@@ -0,0 +1,292 @@
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, '
'installation_uuid, artifact_digest, runtime_revision) '
"VALUES (:workspace_uuid, 'author', 'plugin', 1, "
':installation_uuid, :artifact_digest, 1)'
),
{
'workspace_uuid': second_workspace_uuid,
'installation_uuid': str(uuid.uuid4()),
'artifact_digest': hashlib.sha256(b'test-plugin-artifact').hexdigest(),
},
)
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()
tests/integration/persistence/test_sqlite_migration_backup.py
+107
View File
@@ -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) == alembic_runner.get_alembic_head()
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) == alembic_runner.get_alembic_head()
finally:
await engine.dispose()
tests/integration/persistence/test_workspace_migration.py
+380
View File
@@ -0,0 +1,380 @@
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_head,
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) == get_alembic_head()
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()