xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-07-21 20:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(tenancy): harden shared cloud runtime boundaries

Browse Source
This commit is contained in:
Junyan Qin
2026-07-20 01:47:42 +08:00
parent 41772920ef
commit a47bfe8167
121 changed files with 18152 additions and 5588 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/integration/api/test_box_security.py
+18
View File
@@ -9,6 +9,7 @@ import pytest
import quart
from langbot.pkg.api.http.controller.groups.box import BoxRouterGroup
from langbot.pkg.cloud.entitlements import EntitlementUnavailableError
pytestmark = pytest.mark.integration
@@ -34,6 +35,8 @@ async def box_security_api():
'owner-token': SimpleNamespace(uuid='owner-account', user='owner@example.com'),
}
application = Mock()
application.deployment = SimpleNamespace(multi_workspace_enabled=False)
application.persistence_mgr = SimpleNamespace(tenant_uow=None)
application.user_service.get_authenticated_account = AsyncMock(side_effect=lambda token: accounts[token])
application.workspace_collaboration_service.resolve_account_workspace = AsyncMock(
side_effect=lambda account_uuid, _workspace_uuid: _access(account_uuid)
@@ -41,6 +44,7 @@ async def box_security_api():
application.box_service.get_status = AsyncMock(return_value={'enabled': True})
application.box_service.get_sessions = AsyncMock(return_value=[{'session_id': 'private-session'}])
application.box_service.get_recent_errors = Mock(return_value=[{'error': 'private error'}])
application.box_service.managed_admission_required = False
quart_app = quart.Quart(__name__)
router = BoxRouterGroup(application, quart_app)
@@ -81,3 +85,17 @@ async def test_owner_can_audit_box_sessions_and_errors(box_security_api):
assert errors.status_code == 200
application.box_service.get_sessions.assert_awaited_once()
application.box_service.get_recent_errors.assert_called_once()
@pytest.mark.asyncio
async def test_box_status_returns_explicit_403_when_workspace_has_no_managed_sandbox(box_security_api):
application, client = box_security_api
application.box_service.get_status.side_effect = EntitlementUnavailableError(
'Workspace entitlement does not grant managed_sandbox'
)
response = await client.get('/api/v1/box/status', headers=_headers('viewer-token'))
assert response.status_code == 403
payload = await response.get_json()
assert payload['code'] == 'managed_sandbox_unavailable'
tests/integration/api/test_fresh_oss_workspace_journey.py
+214
View File
@@ -0,0 +1,214 @@
from __future__ import annotations
import json
import logging
from types import SimpleNamespace
import pytest
import sqlalchemy as sa
from quart import Quart
from langbot.pkg.api.http.controller.groups.system import SystemRouterGroup
from langbot.pkg.api.http.controller.groups.user import UserRouterGroup
from langbot.pkg.api.http.controller.groups.workspaces import WorkspacesRouterGroup
from langbot.pkg.api.http.service.user import UserService
from langbot.pkg.entity.persistence.metadata import WorkspaceMetadata
from langbot.pkg.persistence.mgr import PersistenceManager
from langbot.pkg.utils import constants
from langbot.pkg.workspace.collaboration import WorkspaceCollaborationService
from langbot.pkg.workspace.service import WorkspaceService
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
def _authorization(token: str, workspace_uuid: str | None = None) -> dict[str, str]:
headers = {'Authorization': f'Bearer {token}'}
if workspace_uuid is not None:
headers['X-Workspace-Id'] = workspace_uuid
return headers
async def test_fresh_oss_workspace_http_journey_uses_real_sqlite_persistence(
tmp_path,
monkeypatch,
):
"""Exercise the first-run Workspace journey through real HTTP handlers."""
instance_uuid = 'fresh-oss-workspace-journey'
monkeypatch.setattr(constants, 'instance_id', instance_uuid)
application = SimpleNamespace(
logger=logging.getLogger('fresh-oss-workspace-journey-test'),
instance_config=SimpleNamespace(
data={
'database': {
'use': 'sqlite',
'sqlite': {'path': str(tmp_path / 'langbot.db')},
},
'system': {
'jwt': {'secret': 'fresh-oss-workspace-secret', 'expire': 3600},
'allow_modify_login_info': True,
'limitation': {},
'outbound_ips': [],
},
'api': {'global_api_key': ''},
'plugin': {'enable_marketplace': True},
'space': {
'url': 'https://space.langbot.app',
'models_gateway_api_url': 'https://api.langbot.cloud/v1',
'disable_models_service': False,
},
'mcp': {'stdio': {'enabled': True}},
}
),
)
persistence = PersistenceManager(application)
application.persistence_mgr = persistence
await persistence.initialize()
try:
application.workspace_service = WorkspaceService(
application,
instance_uuid=instance_uuid,
)
application.workspace_collaboration_service = WorkspaceCollaborationService(
application,
application.workspace_service,
)
application.user_service = UserService(application)
quart_app = Quart(__name__)
await UserRouterGroup(application, quart_app).initialize()
await WorkspacesRouterGroup(application, quart_app).initialize()
await SystemRouterGroup(application, quart_app).initialize()
client = quart_app.test_client()
initialization = await client.get('/api/v1/user/init')
assert initialization.status_code == 200
assert (await initialization.get_json())['data'] == {'initialized': False}
initialized = await client.post(
'/api/v1/user/init',
json={'user': 'owner@example.com', 'password': 'owner-password'},
)
assert initialized.status_code == 200
authenticated = await client.post(
'/api/v1/user/auth',
json={'user': 'owner@example.com', 'password': 'owner-password'},
)
assert authenticated.status_code == 200
token = (await authenticated.get_json())['data']['token']
bootstrap = await client.get(
'/api/v1/workspaces/bootstrap',
headers=_authorization(token),
)
assert bootstrap.status_code == 200
bootstrap_workspaces = (await bootstrap.get_json())['data']['workspaces']
assert len(bootstrap_workspaces) == 1
bootstrap_access = bootstrap_workspaces[0]
workspace_uuid = bootstrap_access['workspace']['uuid']
assert bootstrap_access['workspace'] == {
'uuid': workspace_uuid,
'instance_uuid': instance_uuid,
'name': 'Default Workspace',
'slug': 'default',
'type': 'team',
'status': 'active',
'source': 'local',
}
assert bootstrap_access['membership']['email'] == 'owner@example.com'
assert bootstrap_access['membership']['role'] == 'owner'
assert 'workspace.update' in bootstrap_access['permissions']
current = await client.get(
'/api/v1/workspaces/current',
headers=_authorization(token, workspace_uuid),
)
assert current.status_code == 200
current_data = (await current.get_json())['data']
assert current_data['workspace']['uuid'] == workspace_uuid
assert current_data['membership']['account_uuid'] == bootstrap_access['membership']['account_uuid']
assert current_data['membership']['role'] == 'owner'
user_info = await client.get(
'/api/v1/user/info',
headers=_authorization(token, workspace_uuid),
)
assert user_info.status_code == 200
assert (await user_info.get_json())['data'] == {
'account_uuid': bootstrap_access['membership']['account_uuid'],
'user': 'owner@example.com',
'account_type': 'local',
'has_password': True,
}
initial_system_info = await client.get(
'/api/v1/system/info',
headers=_authorization(token, workspace_uuid),
)
assert initial_system_info.status_code == 200
assert (await initial_system_info.get_json())['data']['wizard_status'] == 'none'
assert (await initial_system_info.get_json())['data']['wizard_progress'] is None
progress = {'step': 2, 'selected_adapter': 'telegram', 'bot_saved': False}
updated_progress = await client.put(
'/api/v1/system/wizard/progress',
headers=_authorization(token, workspace_uuid),
json=progress,
)
assert updated_progress.status_code == 200
persisted_progress = await persistence.execute_async(
sa.select(WorkspaceMetadata.value).where(
WorkspaceMetadata.workspace_uuid == workspace_uuid,
WorkspaceMetadata.key == 'wizard_progress',
)
)
assert json.loads(persisted_progress.scalar_one()) == progress
system_info_with_progress = await client.get(
'/api/v1/system/info',
headers=_authorization(token, workspace_uuid),
)
assert system_info_with_progress.status_code == 200
progress_data = (await system_info_with_progress.get_json())['data']
assert progress_data['wizard_status'] == 'none'
assert progress_data['wizard_progress'] == progress
completed = await client.post(
'/api/v1/system/wizard/completed',
headers=_authorization(token, workspace_uuid),
json={'status': 'completed'},
)
assert completed.status_code == 200
completed_system_info = await client.get(
'/api/v1/system/info',
headers=_authorization(token, workspace_uuid),
)
assert completed_system_info.status_code == 200
completed_data = (await completed_system_info.get_json())['data']
assert completed_data['wizard_status'] == 'completed'
assert completed_data['wizard_progress'] is None
rejected_workspace = await client.post(
'/api/v1/workspaces',
headers=_authorization(token, workspace_uuid),
json={'name': 'Second Workspace'},
)
assert rejected_workspace.status_code == 403
rejected_data = await rejected_workspace.get_json()
assert rejected_data['code'] == 'edition_limit'
persisted_wizard_status = await persistence.execute_async(
sa.select(WorkspaceMetadata.value).where(
WorkspaceMetadata.workspace_uuid == workspace_uuid,
WorkspaceMetadata.key == 'wizard_status',
)
)
assert persisted_wizard_status.scalar_one() == 'completed'
finally:
await persistence.get_db_engine().dispose()
tests/integration/api/test_plugins_security.py
+24
View File
@@ -86,6 +86,7 @@ async def plugin_security_api(plugin_module):
}
application = Mock()
application.deployment = SimpleNamespace(multi_workspace_enabled=False)
application.user_service.get_authenticated_account = AsyncMock(side_effect=lambda token: accounts[token])
application.workspace_collaboration_service.resolve_account_workspace = AsyncMock(
side_effect=lambda account_uuid, _workspace_uuid: _access(account_uuid)
@@ -112,6 +113,7 @@ async def plugin_security_api(plugin_module):
persistence_result = Mock()
persistence_result.scalar_one_or_none.return_value = RAW_CONFIG
application.persistence_mgr.execute_async = AsyncMock(return_value=persistence_result)
application.persistence_mgr.tenant_uow = None
quart_app = quart.Quart(__name__)
router = plugin_module.PluginsRouterGroup(application, quart_app)
@@ -246,3 +248,25 @@ async def test_viewer_cannot_read_plugin_runtime_logs(plugin_security_api):
assert response.status_code == 403
assert (await response.get_json())['code'] == 'permission_denied'
application.plugin_connector.get_plugin_logs.assert_not_awaited()
@pytest.mark.asyncio
async def test_github_install_rejects_internal_asset_url_before_task_creation(
plugin_security_api,
):
application, client, _ = plugin_security_api
response = await client.post(
'/api/v1/plugins/install/github',
headers=_headers('manager-token'),
json={
'asset_url': 'http://169.254.169.254/latest/meta-data',
'owner': 'langbot-app',
'repo': 'demo-plugin',
'release_tag': 'v1.0.0',
},
)
assert response.status_code == 400
assert 'HTTPS GitHub release asset URL' in (await response.get_json())['msg']
application.task_mgr.create_user_task.assert_not_called()
tests/integration/api/test_user_space_oauth.py
+2
View File
@@ -25,6 +25,8 @@ async def space_oauth_api():
execution=SimpleNamespace(instance_uuid='instance-a', placement_generation=1),
)
application = Mock()
application.deployment = SimpleNamespace(multi_workspace_enabled=False)
application.persistence_mgr = None
application.user_service.get_authenticated_account = AsyncMock(return_value=account)
application.user_service.issue_space_oauth_state = AsyncMock(
side_effect=lambda purpose, **_: f'opaque-{purpose}-state'
tests/integration/api/test_workspaces.py
+71 -30
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import json
import logging
from types import SimpleNamespace
@@ -13,11 +14,13 @@ from langbot.pkg.api.http.controller.groups.workspaces import (
InvitationsRouterGroup,
WorkspacesRouterGroup,
)
from langbot.pkg.api.http.controller.groups.system import SystemRouterGroup
from langbot.pkg.api.http.controller.groups.apikeys import ApiKeysRouterGroup
from langbot.pkg.api.http.controller.groups.user import UserRouterGroup
from langbot.pkg.api.http.service.apikey import ApiKeyService
from langbot.pkg.api.http.service.user import ControlPlaneDirectoryRequiredError, UserService
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.metadata import WorkspaceMetadata
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.entity.persistence.workspace import (
Workspace,
@@ -25,6 +28,7 @@ from langbot.pkg.entity.persistence.workspace import (
WorkspaceInvitation,
WorkspaceMembership,
)
from langbot.pkg.persistence.mgr import PersistenceManager
from langbot.pkg.workspace.collaboration import WorkspaceCollaborationService
from langbot.pkg.workspace.service import WorkspaceService
from langbot.pkg.workspace.policy import CloudWorkspacePolicy
@@ -33,28 +37,6 @@ from langbot.pkg.workspace.policy import CloudWorkspacePolicy
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
class _PersistenceManager:
def __init__(self, engine):
self.engine = engine
def get_db_engine(self):
return self.engine
async def execute_async(self, *args, **kwargs):
async with self.engine.connect() as connection:
result = await connection.execute(*args, **kwargs)
await connection.commit()
return result
@staticmethod
def serialize_model(model, row, masked_columns=()):
return {
column.name: getattr(row, column.name)
for column in model.__table__.columns
if column.name not in masked_columns
}
@pytest.fixture
async def workspace_api(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-api.db"}')
@@ -62,7 +44,8 @@ async def workspace_api(tmp_path):
await connection.run_sync(Base.metadata.create_all)
application = SimpleNamespace()
application.persistence_mgr = _PersistenceManager(engine)
application.persistence_mgr = PersistenceManager(application)
application.persistence_mgr.db = SimpleNamespace(get_engine=lambda: engine)
application.instance_config = SimpleNamespace(
data={
'system': {
@@ -85,19 +68,27 @@ async def workspace_api(tmp_path):
application.user_service = UserService(application)
application.apikey_service = ApiKeyService(application)
owner = await application.user_service.create_initial_account(
'owner@example.com',
'owner-password',
)
owner_token = await application.user_service.generate_jwt_token(owner)
quart_app = Quart(__name__)
await WorkspacesRouterGroup(application, quart_app).initialize()
await InvitationsRouterGroup(application, quart_app).initialize()
await ApiKeysRouterGroup(application, quart_app).initialize()
await UserRouterGroup(application, quart_app).initialize()
await SystemRouterGroup(application, quart_app).initialize()
yield application, quart_app.test_client(), engine, owner_token
client = quart_app.test_client()
init_response = await client.post(
'/api/v1/user/init',
json={'user': 'owner@example.com', 'password': 'owner-password'},
)
assert init_response.status_code == 200
auth_response = await client.post(
'/api/v1/user/auth',
json={'user': 'owner@example.com', 'password': 'owner-password'},
)
assert auth_response.status_code == 200
owner_token = (await auth_response.get_json())['data']['token']
yield application, client, engine, owner_token
await engine.dispose()
@@ -108,6 +99,56 @@ def _auth(token: str, workspace_uuid: str | None = None) -> dict[str, str]:
return headers
async def test_fresh_sqlite_login_returns_current_workspace_and_user_info(workspace_api):
_, client, _, owner_token = workspace_api
current_response = await client.get('/api/v1/workspaces/current', headers=_auth(owner_token))
assert current_response.status_code == 200
current = (await current_response.get_json())['data']
assert current['workspace']['uuid']
assert current['membership']['email'] == 'owner@example.com'
assert current['membership']['role'] == 'owner'
info_response = await client.get('/api/v1/user/info', headers=_auth(owner_token))
assert info_response.status_code == 200
info = (await info_response.get_json())['data']
assert info['account_uuid'] == current['membership']['account_uuid']
assert info['user'] == 'owner@example.com'
async def test_authenticated_system_info_reads_workspace_wizard_metadata(workspace_api):
application, client, _, owner_token = workspace_api
current_response = await client.get('/api/v1/workspaces/current', headers=_auth(owner_token))
workspace_uuid = (await current_response.get_json())['data']['workspace']['uuid']
progress = {'step': 3, 'selected_adapter': 'telegram'}
await application.persistence_mgr.execute_async(
sqlalchemy.insert(WorkspaceMetadata),
[
{
'workspace_uuid': workspace_uuid,
'key': 'wizard_status',
'value': 'completed',
},
{
'workspace_uuid': workspace_uuid,
'key': 'wizard_progress',
'value': json.dumps(progress),
},
],
)
response = await client.get(
'/api/v1/system/info',
headers=_auth(owner_token, workspace_uuid),
)
assert response.status_code == 200
data = (await response.get_json())['data']
assert data['wizard_status'] == 'completed'
assert data['wizard_progress'] == progress
async def test_owner_invites_second_account_and_secret_is_not_persisted(workspace_api):
application, client, engine, owner_token = workspace_api
tests/integration/persistence/test_migrations.py
+22
View File
@@ -10,6 +10,7 @@ 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
@@ -149,6 +150,27 @@ 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
class TestSQLiteMigrationFreshDatabase:
"""Tests for fresh database workflow."""
tests/integration/persistence/test_migrations_postgres.py
+763 -16
View File
@@ -16,19 +16,30 @@ from __future__ import annotations
import logging
import os
import uuid
import asyncio
import contextlib
import datetime
import hashlib
import typing
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import sqlalchemy as sa
from quart import Quart
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, 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, PersistenceMode
from langbot.pkg.persistence.tenant_uow import TENANT_POLICY_NAME, TENANT_TABLE_COLUMNS, TenantUnitOfWork
from langbot.pkg.persistence.tenant_uow import (
TENANT_POLICY_NAME,
TENANT_TABLE_COLUMNS,
TenantUnitOfWork,
TransactionRollbackOnlyError,
)
from langbot.pkg.persistence.alembic_runner import (
run_alembic_upgrade,
run_alembic_stamp,
@@ -39,6 +50,32 @@ 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 langbot.pkg.workspace.collaboration import WorkspaceCollaborationService
from langbot.pkg.workspace.policy import CloudWorkspacePolicy
from langbot.pkg.workspace.service import WorkspaceService
from langbot.pkg.api.http.authz import Permission
from langbot.pkg.api.http.controller import group as http_group
from langbot.pkg.api.http.controller.groups.system import SystemRouterGroup
from langbot.pkg.api.http.controller.groups.webhooks import WebhookRouterGroup
from langbot.pkg.api.http.context import ExecutionContext, RequestContext
from langbot.pkg.api.http.service.apikey import ApiKeyService
from langbot.pkg.api.http.service.monitoring import MonitoringService
from langbot.pkg.api.http.service.user import UserService
from langbot.pkg.api.mcp.context import get_request_context as get_mcp_request_context
from langbot.pkg.api.mcp.mount import MCPMount
from langbot.pkg.entity.persistence.apikey import ApiKey
from langbot.pkg.entity.persistence.metadata import WorkspaceMetadata
from langbot.pkg.entity.persistence.monitoring import MonitoringFeedback
from langbot.pkg.entity.persistence.workspace import (
Workspace,
WorkspaceExecutionState,
WorkspaceMembership,
)
from langbot.pkg.platform.botmgr import PlatformManager
from langbot.pkg.pipeline.pipelinemgr import PipelineManager
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
from langbot.pkg.provider.tools.loaders.mcp import MCPLoader
from langbot.pkg.rag.knowledge.kbmgr import RAGManager
from .resource_migration_support import TENANT_TABLES, create_legacy_resource_schema
@@ -54,6 +91,42 @@ def _get_script_head() -> str:
return ScriptDirectory.from_config(cfg).get_current_head()
async def _grant_runtime_role_business_objects(
conn: AsyncConnection,
role_name: str,
quote: typing.Callable[[str], str],
) -> None:
"""Mirror the release job's object ACLs without overgranting Alembic."""
business_tables = tuple(sorted({table.name for table in Base.metadata.tables.values()} | {'langbot_vectors'}))
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(role_name)}'))
await conn.execute(text(f'GRANT SELECT ON TABLE public.alembic_version TO {quote(role_name)}'))
sequence_query = text(
"""
SELECT DISTINCT sequence.relname
FROM pg_class sequence
JOIN pg_namespace sequence_namespace ON sequence_namespace.oid = sequence.relnamespace
JOIN pg_depend dependency
ON dependency.classid = 'pg_class'::regclass
AND dependency.objid = sequence.oid
AND dependency.refclassid = 'pg_class'::regclass
AND dependency.deptype IN ('a', 'i')
JOIN pg_class business_table ON business_table.oid = dependency.refobjid
JOIN pg_namespace table_namespace ON table_namespace.oid = business_table.relnamespace
WHERE sequence.relkind = 'S'
AND sequence_namespace.nspname = 'public'
AND table_namespace.nspname = 'public'
AND business_table.relname IN :table_names
ORDER BY sequence.relname
"""
).bindparams(sa.bindparam('table_names', expanding=True))
sequence_names = tuple((await conn.execute(sequence_query, {'table_names': business_tables})).scalars().all())
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(role_name)}'))
def _application_for_postgres_url(postgres_url: str, logger_name: str) -> SimpleNamespace:
url = sa.engine.make_url(postgres_url)
return SimpleNamespace(
@@ -115,15 +188,20 @@ async def postgres_engine(postgres_url):
@pytest.fixture
async def clean_tables(postgres_engine):
"""Drop all tables before and after each test for isolation."""
# Drop all tables before test
async with postgres_engine.begin() as conn:
await conn.run_sync(Base.metadata.drop_all)
async def drop_all_tables() -> None:
# Alembic can create tables (notably langbot_vectors) outside the ORM
# metadata, and legacy migration tests intentionally alter constraints.
# Reflect the dedicated test schema instead of relying on stale ORM DDL.
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 drop_all_tables()
yield
# Drop all tables after test
async with postgres_engine.begin() as conn:
await conn.run_sync(Base.metadata.drop_all)
await drop_all_tables()
@pytest.fixture
@@ -442,10 +520,16 @@ class TestPostgreSQLResourceTenancyMigration:
await conn.execute(
text(
'INSERT INTO plugin_settings '
'(workspace_uuid, plugin_author, plugin_name, enabled) '
"VALUES (:workspace_uuid, 'author', 'plugin', true)"
'(workspace_uuid, plugin_author, plugin_name, enabled, '
'installation_uuid, artifact_digest, runtime_revision) '
"VALUES (:workspace_uuid, 'author', 'plugin', true, "
':installation_uuid, :artifact_digest, 1)'
),
{'workspace_uuid': second_workspace_uuid},
{
'workspace_uuid': second_workspace_uuid,
'installation_uuid': str(uuid.uuid4()),
'artifact_digest': hashlib.sha256(b'test-plugin-artifact').hexdigest(),
},
)
with pytest.raises(IntegrityError):
@@ -556,10 +640,7 @@ class TestPostgreSQLTenantRuntime:
)
)
await conn.execute(text(f'GRANT USAGE ON SCHEMA public TO {quote(role_name)}'))
await conn.execute(
text(f'GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO {quote(role_name)}')
)
await conn.execute(text(f'GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO {quote(role_name)}'))
await _grant_runtime_role_business_objects(conn, role_name, quote)
created_roles.append(role_name)
try:
@@ -703,6 +784,31 @@ class TestPostgreSQLTenantRuntime:
cloud_manager.create_tables.assert_not_awaited()
cloud_manager._run_alembic_migrations.assert_not_awaited()
with pytest.raises(TransactionRollbackOnlyError, match='after-commit work was cancelled'):
async with cloud_manager.tenant_uow(workspace_a):
duplicate_statement = sa.insert(WorkspaceMetadata).values(
workspace_uuid=workspace_a,
key='rollback-only-unique',
value='must-not-commit',
)
await cloud_manager.execute_async(duplicate_statement)
after_commit_gate = cloud_manager.create_after_commit_gate()
assert after_commit_gate is not None
try:
await cloud_manager.execute_async(duplicate_statement)
except IntegrityError:
pass
assert after_commit_gate.cancelled()
async with cloud_manager.tenant_uow(workspace_a):
assert (
await cloud_manager.execute_async(
sa.select(sa.func.count())
.select_from(WorkspaceMetadata)
.where(WorkspaceMetadata.key == 'rollback-only-unique')
)
).scalar_one() == 0
superuser_manager = PersistenceManager(
_application_for_postgres_url(postgres_url, 'postgres-superuser-runtime-test'),
mode=PersistenceMode.CLOUD_RUNTIME,
@@ -743,3 +849,644 @@ class TestPostgreSQLTenantRuntime:
for role_name in reversed(created_roles):
await conn.execute(text(f'DROP OWNED BY {quote(role_name)}'))
await conn.execute(text(f'DROP ROLE {quote(role_name)}'))
@pytest.mark.asyncio
async def test_cloud_discovery_http_and_transaction_contract(
self,
postgres_url,
postgres_engine,
clean_tables,
clean_alembic_version,
monkeypatch,
):
"""Exercise the complete request/discovery path as a non-owner role."""
instance_uuid = 'cloud-request-rls-test'
other_instance_uuid = 'other-cloud-instance'
workspace_a = '31000000-0000-0000-0000-000000000001'
workspace_b = '32000000-0000-0000-0000-000000000002'
workspace_other = '33000000-0000-0000-0000-000000000003'
workspace_fenced = '34000000-0000-0000-0000-000000000004'
account_a_uuid = '41000000-0000-0000-0000-000000000001'
shared_account_uuid = '42000000-0000-0000-0000-000000000002'
active_secret = 'lbk_active-request-key'
revoked_secret = 'lbk_revoked-request-key'
expired_secret = 'lbk_expired-request-key'
role_name = f'lb_request_{uuid.uuid4().hex[:12]}'
role_password = f'Lb{uuid.uuid4().hex}'
quote = postgres_engine.dialect.identifier_preparer.quote
managers: list[PersistenceManager] = []
role_created = False
_restore_postgres_manager_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', instance_uuid)
release_manager = PersistenceManager(
_application_for_postgres_url(postgres_url, 'postgres-request-release-test'),
mode=PersistenceMode.RELEASE_MIGRATION,
)
managers.append(release_manager)
def role_url() -> str:
return (
sa.engine.make_url(postgres_url)
.set(username=role_name, password=role_password)
.render_as_string(hide_password=False)
)
async def seed_workspace(
workspace_uuid: str,
*,
target_instance: str,
state: str = 'active',
write_fenced: bool = False,
) -> None:
async with release_manager.tenant_uow(workspace_uuid) as uow:
uow.session.add(
Workspace(
uuid=workspace_uuid,
instance_uuid=target_instance,
name=workspace_uuid[-4:],
slug=f'workspace-{workspace_uuid[-4:]}',
type='team',
status='active',
source='cloud_projection',
projection_revision=1,
)
)
await uow.session.flush()
uow.session.add(
WorkspaceExecutionState(
workspace_uuid=workspace_uuid,
instance_uuid=target_instance,
active_generation=1,
state=state,
write_fenced=write_fenced,
source='cloud',
desired_state_revision=1,
)
)
uow.session.add(
WorkspaceMetadata(
workspace_uuid=workspace_uuid,
key='tenant-marker',
value=workspace_uuid,
)
)
try:
await release_manager.initialize()
async with release_manager.get_db_engine().begin() as conn:
await conn.execute(
sa.insert(User),
[
{
'uuid': account_a_uuid,
'user': 'account-a@example.com',
'normalized_email': 'account-a@example.com',
'password': 'closed-directory',
'account_type': 'local',
'status': 'active',
'source': 'cloud_projection',
'projection_revision': 1,
},
{
'uuid': shared_account_uuid,
'user': 'shared@example.com',
'normalized_email': 'shared@example.com',
'password': 'closed-directory',
'account_type': 'local',
'status': 'active',
'source': 'cloud_projection',
'projection_revision': 1,
},
],
)
await seed_workspace(workspace_a, target_instance=instance_uuid)
await seed_workspace(workspace_b, target_instance=instance_uuid)
await seed_workspace(workspace_other, target_instance=other_instance_uuid)
await seed_workspace(
workspace_fenced,
target_instance=instance_uuid,
write_fenced=True,
)
async with release_manager.tenant_uow(workspace_a) as uow:
uow.session.add_all(
[
WorkspaceMembership(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_a,
account_uuid=account_a_uuid,
role='owner',
status='active',
projection_revision=1,
),
WorkspaceMembership(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_a,
account_uuid=shared_account_uuid,
role='viewer',
status='active',
projection_revision=1,
),
ApiKey(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_a,
name='active-key',
key_hash=hashlib.sha256(active_secret.encode()).hexdigest(),
scopes=[Permission.WORKSPACE_VIEW.value],
status='active',
),
]
)
async with release_manager.tenant_uow(workspace_b) as uow:
uow.session.add_all(
[
WorkspaceMembership(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_b,
account_uuid=shared_account_uuid,
role='viewer',
status='active',
projection_revision=1,
),
ApiKey(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_b,
name='revoked-key',
key_hash=hashlib.sha256(revoked_secret.encode()).hexdigest(),
scopes=[Permission.WORKSPACE_VIEW.value],
status='revoked',
),
ApiKey(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_b,
name='expired-key',
key_hash=hashlib.sha256(expired_secret.encode()).hexdigest(),
scopes=[Permission.WORKSPACE_VIEW.value],
status='active',
expires_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
- datetime.timedelta(minutes=1),
),
]
)
async with postgres_engine.connect() as conn:
await conn.execute(text(f"CREATE ROLE {quote(role_name)} LOGIN PASSWORD '{role_password}'"))
role_created = True
await conn.execute(
text(
f'GRANT CONNECT ON DATABASE {quote(sa.engine.make_url(postgres_url).database)} '
f'TO {quote(role_name)}'
)
)
await conn.execute(text(f'GRANT USAGE ON SCHEMA public TO {quote(role_name)}'))
await _grant_runtime_role_business_objects(conn, role_name, quote)
runtime_application = _application_for_postgres_url(role_url(), 'postgres-request-runtime-test')
runtime_application.instance_config.data.update(
{
'system': {
'jwt': {'secret': 'postgres-request-jwt', 'expire': 3600},
},
'api': {'global_api_key': ''},
}
)
runtime_application.logger = logging.getLogger('postgres-request-runtime-test')
cloud_manager = PersistenceManager(runtime_application, mode=PersistenceMode.CLOUD_RUNTIME)
managers.append(cloud_manager)
await cloud_manager.initialize()
runtime_application.persistence_mgr = cloud_manager
cloud_manager.ap = runtime_application
runtime_application.workspace_service = WorkspaceService(
runtime_application,
policy=CloudWorkspacePolicy(),
instance_uuid=instance_uuid,
)
runtime_application.workspace_collaboration_service = WorkspaceCollaborationService(
runtime_application,
runtime_application.workspace_service,
policy=CloudWorkspacePolicy(),
)
runtime_application.user_service = UserService(runtime_application)
runtime_application.apikey_service = ApiKeyService(runtime_application)
runtime_application.monitoring_service = MonitoringService(runtime_application)
bindings = await runtime_application.workspace_service.list_active_execution_bindings()
assert {binding.workspace_uuid for binding in bindings} == {workspace_a, workspace_b}
# No scope is an application error before SQL reaches PostgreSQL.
with pytest.raises(RuntimeError, match='explicit Workspace or discovery'):
await cloud_manager.execute_async(sa.select(WorkspaceMetadata))
# Discovery exposes only its index rows and cannot write.
with pytest.raises(TransactionRollbackOnlyError, match='transaction was rolled back'):
async with cloud_manager.account_discovery_uow(shared_account_uuid) as discovery:
assert set(
(
await discovery.session.scalars(
sa.select(WorkspaceMembership.workspace_uuid).order_by(
WorkspaceMembership.workspace_uuid
)
)
).all()
) == {workspace_a, workspace_b}
with pytest.raises(sa.exc.DBAPIError):
await discovery.session.execute(
sa.insert(WorkspaceMembership).values(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_a,
account_uuid=shared_account_uuid,
role='viewer',
status='active',
projection_revision=1,
)
)
async with cloud_manager.api_key_discovery_uow(
hashlib.sha256(active_secret.encode()).hexdigest()
) as discovery:
assert await discovery.session.scalar(sa.select(ApiKey.workspace_uuid)) == workspace_a
update_result = await discovery.session.execute(
sa.update(ApiKey)
.where(ApiKey.key_hash == hashlib.sha256(active_secret.encode()).hexdigest())
.values(name='discovery-must-not-write')
)
assert update_result.rowcount == 0
async with cloud_manager.api_key_discovery_uow(
hashlib.sha256(revoked_secret.encode()).hexdigest()
) as discovery:
assert await discovery.session.scalar(sa.select(ApiKey.workspace_uuid)) is None
async with cloud_manager.api_key_discovery_uow(
hashlib.sha256(expired_secret.encode()).hexdigest()
) as discovery:
assert await discovery.session.scalar(sa.select(ApiKey.workspace_uuid)) is None
async with cloud_manager.instance_discovery_uow(instance_uuid) as discovery:
assert set(
(await discovery.session.scalars(sa.select(WorkspaceExecutionState.workspace_uuid))).all()
) == {workspace_a, workspace_b}
update_result = await discovery.session.execute(
sa.update(WorkspaceExecutionState)
.where(WorkspaceExecutionState.workspace_uuid == workspace_a)
.values(write_fenced=True)
)
assert update_result.rowcount == 0
# Instance discovery deliberately cannot see business rows.
assert (await discovery.session.execute(sa.select(WorkspaceMetadata))).all() == []
platform_manager = PlatformManager(runtime_application)
platform_manager._load_workspace_bots = AsyncMock()
await platform_manager.load_bots_from_db()
assert {call.args[0] for call in platform_manager._load_workspace_bots.await_args_list} == {
workspace_a,
workspace_b,
}
# Every startup cache loader traverses the instance index first,
# then reads business resources in one tenant transaction at a time.
await ModelManager(runtime_application).load_models_from_db()
await PipelineManager(runtime_application).load_pipelines_from_db()
await MCPLoader(runtime_application).load_mcp_servers_from_db()
await RAGManager(runtime_application).load_knowledge_bases_from_db()
# An omitted Workspace predicate remains isolated by RLS.
async with cloud_manager.tenant_uow(workspace_a):
values = (await cloud_manager.execute_async(sa.select(WorkspaceMetadata.value))).scalars().all()
assert values == [workspace_a]
async def read_tenant_repeatedly(workspace_uuid: str) -> list[str]:
observed: list[str] = []
for _ in range(5):
async with cloud_manager.tenant_uow(workspace_uuid):
observed.extend(
(await cloud_manager.execute_async(sa.select(WorkspaceMetadata.value))).scalars().all()
)
await asyncio.sleep(0)
return observed
observed_a, observed_b = await asyncio.gather(
read_tenant_repeatedly(workspace_a),
read_tenant_repeatedly(workspace_b),
)
assert observed_a == [workspace_a] * 5
assert observed_b == [workspace_b] * 5
accesses = await runtime_application.workspace_collaboration_service.list_account_workspaces(
shared_account_uuid
)
assert {access.workspace.uuid for access in accesses} == {workspace_a, workspace_b}
assert await runtime_application.apikey_service.authenticate_api_key(revoked_secret) is None
assert await runtime_application.apikey_service.authenticate_api_key(expired_secret) is None
active_identity = await runtime_application.apikey_service.authenticate_api_key(active_secret)
assert active_identity is not None
assert active_identity.workspace_uuid == workspace_a
async def record_feedback(feedback_type: int) -> str | None:
async with cloud_manager.tenant_scope(workspace_a):
return await runtime_application.monitoring_service.record_feedback(
ExecutionContext(
instance_uuid=instance_uuid,
workspace_uuid=workspace_a,
placement_generation=1,
),
feedback_id='concurrent-feedback',
feedback_type=feedback_type,
)
feedback_ids = await asyncio.gather(record_feedback(1), record_feedback(2))
assert feedback_ids[0] == feedback_ids[1]
async with cloud_manager.tenant_uow(workspace_a) as uow:
feedback_rows = (
(
await uow.execute(
sa.select(MonitoringFeedback).where(MonitoringFeedback.feedback_id == 'concurrent-feedback')
)
)
.scalars()
.all()
)
assert len(feedback_rows) == 1
assert feedback_rows[0].feedback_type in {1, 2}
await uow.execute(
sa.delete(MonitoringFeedback).where(MonitoringFeedback.feedback_id == 'concurrent-feedback')
)
class TenantRuntimeRouter(http_group.RouterGroup):
name = 'postgres-tenant-runtime'
path = '/tenant-runtime'
async def initialize(self) -> None:
@self.route('/account', permission=Permission.WORKSPACE_VIEW)
async def account_route(request_context: RequestContext):
assert self.ap.persistence_mgr.current_session() is None
if self.quart_app.config.get('FORCE_HANDLER_FAILURE'):
async with self.ap.persistence_mgr.tenant_uow(request_context.workspace_uuid):
await self.ap.persistence_mgr.execute_async(
sa.insert(WorkspaceMetadata).values(
workspace_uuid=request_context.workspace_uuid,
key='rolled-back-handler',
value='must-not-commit',
)
)
raise RuntimeError('forced handler failure')
values = (
(await self.ap.persistence_mgr.execute_async(sa.select(WorkspaceMetadata.value)))
.scalars()
.all()
)
assert self.ap.persistence_mgr.current_session() is None
return self.success(data={'workspace_uuid': request_context.workspace_uuid, 'values': values})
@self.route('/key', auth_type=http_group.AuthType.API_KEY)
async def key_route(request_context: RequestContext):
values = (
(await self.ap.persistence_mgr.execute_async(sa.select(WorkspaceMetadata.value)))
.scalars()
.all()
)
return self.success(data={'workspace_uuid': request_context.workspace_uuid, 'values': values})
@self.route('/bootstrap', auth_type=http_group.AuthType.ACCOUNT_TOKEN)
async def bootstrap_route(user_email: str):
account = await self.ap.user_service.get_user_by_email(user_email)
accesses = await self.ap.workspace_collaboration_service.list_account_workspaces(account.uuid)
return self.success(data=sorted(access.workspace.uuid for access in accesses))
quart_app = Quart(__name__)
await TenantRuntimeRouter(runtime_application, quart_app).initialize()
webhook_bot_uuid = str(uuid.uuid4())
class TenantAwareWebhookAdapter:
async def handle_unified_webhook(self, **_kwargs):
assert cloud_manager.current_session() is None
await asyncio.sleep(0)
assert cloud_manager.current_session() is None
values = (
(
await cloud_manager.execute_async(
sa.select(WorkspaceMetadata.value).where(WorkspaceMetadata.key == 'tenant-marker')
)
)
.scalars()
.all()
)
assert cloud_manager.current_session() is None
return {'values': values}
runtime_application.platform_mgr = SimpleNamespace(
resolve_public_bot=AsyncMock(
return_value=SimpleNamespace(
workspace_uuid=workspace_a,
placement_generation=1,
enable=True,
adapter=TenantAwareWebhookAdapter(),
)
)
)
await WebhookRouterGroup(runtime_application, quart_app).initialize()
await SystemRouterGroup(runtime_application, quart_app).initialize()
client = quart_app.test_client()
account_a = await runtime_application.user_service.get_user_by_uuid(account_a_uuid)
shared_account = await runtime_application.user_service.get_user_by_uuid(shared_account_uuid)
assert account_a is not None and shared_account is not None
account_token = await runtime_application.user_service.generate_jwt_token(account_a)
shared_token = await runtime_application.user_service.generate_jwt_token(shared_account)
async with cloud_manager.tenant_uow(workspace_a):
await cloud_manager.execute_async(
sa.insert(WorkspaceMetadata).values(
workspace_uuid=workspace_a,
key='wizard_status',
value='completed',
)
)
response = await client.get(
'/api/v1/system/info',
headers={
'Authorization': f'Bearer {account_token}',
'X-Workspace-Id': workspace_a,
},
)
assert response.status_code == 200
assert (await response.get_json())['data']['wizard_status'] == 'completed'
async with cloud_manager.tenant_uow(workspace_a):
await cloud_manager.execute_async(
sa.delete(WorkspaceMetadata).where(
WorkspaceMetadata.workspace_uuid == workspace_a,
WorkspaceMetadata.key == 'wizard_status',
)
)
response = await client.post(f'/bots/{webhook_bot_uuid}')
assert response.status_code == 200
assert await response.get_json() == {'values': [workspace_a]}
response = await client.get(
'/tenant-runtime/account',
headers={
'Authorization': f'Bearer {account_token}',
'X-Workspace-Id': workspace_a,
},
)
assert response.status_code == 200
assert (await response.get_json())['data'] == {
'workspace_uuid': workspace_a,
'values': [workspace_a],
}
response = await client.get(
'/tenant-runtime/account',
headers={
'Authorization': f'Bearer {account_token}',
'X-Workspace-Id': workspace_b,
},
)
assert response.status_code == 404
response = await client.get(
'/tenant-runtime/account',
headers={'Authorization': f'Bearer {shared_token}'},
)
assert response.status_code == 404
response = await client.get(
'/tenant-runtime/bootstrap',
headers={'Authorization': f'Bearer {shared_token}'},
)
assert response.status_code == 200
assert (await response.get_json())['data'] == [workspace_a, workspace_b]
response = await client.get(
'/tenant-runtime/key',
headers={'X-API-Key': active_secret, 'X-Workspace-Id': workspace_b},
)
assert response.status_code == 200
assert (await response.get_json())['data'] == {
'workspace_uuid': workspace_a,
'values': [workspace_a],
}
# The parallel MCP ASGI entrypoint authenticates the same key and
# retains only a trusted Workspace scope for the full tool request.
# Each DB call gets its own short RLS transaction.
mcp_observation: dict[str, typing.Any] = {}
async def fake_mcp_asgi(scope, receive, send):
del scope, receive
context = get_mcp_request_context()
assert cloud_manager.current_session() is None
values = (await cloud_manager.execute_async(sa.select(WorkspaceMetadata.value))).scalars().all()
assert cloud_manager.current_session() is None
await asyncio.sleep(0)
assert cloud_manager.current_session() is None
repeated_values = (
(await cloud_manager.execute_async(sa.select(WorkspaceMetadata.value))).scalars().all()
)
assert repeated_values == values
assert cloud_manager.current_session() is None
mcp_observation.update(
workspace_uuid=context.workspace_uuid,
values=values,
)
await send({'type': 'http.response.start', 'status': 200, 'headers': []})
await send({'type': 'http.response.body', 'body': b'{}'})
async def unused_quart_asgi(scope, receive, send): # pragma: no cover - routing assertion
del scope, receive, send
raise AssertionError('MCP request was routed to Quart')
mount = MCPMount.__new__(MCPMount)
mount.ap = runtime_application
mount._mcp_asgi = fake_mcp_asgi
sent_messages: list[dict[str, typing.Any]] = []
async def receive():
return {'type': 'http.request', 'body': b'', 'more_body': False}
async def send(message):
sent_messages.append(message)
await mount.wrap(unused_quart_asgi)(
{
'type': 'http',
'path': '/mcp',
'headers': [
(b'x-api-key', active_secret.encode()),
(b'x-workspace-id', workspace_b.encode()),
],
},
receive,
send,
)
assert sent_messages[0]['status'] == 200
assert mcp_observation == {'workspace_uuid': workspace_a, 'values': [workspace_a]}
original_api_key_discovery = cloud_manager.api_key_discovery_uow
@contextlib.asynccontextmanager
async def revoke_after_discovery(key_hash: str):
async with original_api_key_discovery(key_hash) as discovery:
yield discovery
async with cloud_manager.tenant_uow(workspace_a):
await cloud_manager.execute_async(
sa.update(ApiKey).where(ApiKey.key_hash == key_hash).values(status='revoked')
)
monkeypatch.setattr(cloud_manager, 'api_key_discovery_uow', revoke_after_discovery)
assert await runtime_application.apikey_service.authenticate_api_key(active_secret) is None
monkeypatch.setattr(cloud_manager, 'api_key_discovery_uow', original_api_key_discovery)
quart_app.config['FORCE_HANDLER_FAILURE'] = True
response = await client.get(
'/tenant-runtime/account',
headers={
'Authorization': f'Bearer {account_token}',
'X-Workspace-Id': workspace_a,
},
)
assert response.status_code == 500
quart_app.config['FORCE_HANDLER_FAILURE'] = False
async with cloud_manager.tenant_uow(workspace_a):
assert (
await cloud_manager.execute_async(
sa.select(sa.func.count())
.select_from(WorkspaceMetadata)
.where(WorkspaceMetadata.key == 'rolled-back-handler')
)
).scalar_one() == 0
# Transaction-local settings are gone when pooled connections are reused.
async with cloud_manager.get_db_engine().connect() as conn:
assert await conn.scalar(text("SELECT current_setting('langbot.workspace_uuid', true)")) in (
None,
'',
)
assert await conn.scalar(text('SELECT COUNT(*) FROM workspace_metadata')) == 0
# Runtime validation rejects both extra permissive policies and a
# modified expression even if the expected policy name remains.
async with postgres_engine.connect() as conn:
await conn.execute(text('CREATE POLICY injected_policy ON bots FOR SELECT USING (true)'))
with pytest.raises(RuntimeError, match='policy set does not match'):
await cloud_manager._validate_postgres_tenant_schema(validate_runtime_role=True)
async with postgres_engine.connect() as conn:
await conn.execute(text('DROP POLICY injected_policy ON bots'))
await conn.execute(text('DROP POLICY langbot_workspace_isolation ON workspace_metadata'))
await conn.execute(
text(
'CREATE POLICY langbot_workspace_isolation ON workspace_metadata '
'FOR ALL TO PUBLIC USING (true) WITH CHECK (true)'
)
)
with pytest.raises(RuntimeError, match='policy definitions are invalid'):
await cloud_manager._validate_postgres_tenant_schema(validate_runtime_role=True)
finally:
for manager in reversed(managers):
await _dispose_manager(manager)
if role_created:
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP OWNED BY {quote(role_name)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(role_name)}'))
tests/integration/persistence/test_pgvector_postgres.py
+546
View File
@@ -0,0 +1,546 @@
"""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.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, 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(
text(
"""
INSERT INTO workspaces
(uuid, instance_uuid, name, slug, type, status, source, projection_revision)
VALUES
(:uuid, :instance_uuid, :name, :slug, 'team', 'active', 'cloud_projection', 0)
"""
),
{
'uuid': workspace_uuid,
'instance_uuid': instance_uuid,
'name': slug,
'slug': slug,
},
)
await uow.execute(
text(
"""
INSERT INTO knowledge_bases
(uuid, workspace_uuid, name, embedding_dimension)
VALUES (:uuid, :workspace_uuid, :name, 384)
"""
),
{'uuid': kb_uuid, 'workspace_uuid': workspace_uuid, 'name': slug},
)
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(
text("SELECT workspace_uuid, vector_id FROM langbot_vectors WHERE vector_id = 'same-vector'")
)
).all()
assert rows == [(workspace_a, 'same-vector')]
await uow.execute(text('SET LOCAL enable_seqscan = off'))
plan = '\n'.join(
(
await uow.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
runtime_engine = runtime_manager.get_db_engine()
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
+723
View File
@@ -0,0 +1,723 @@
"""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:
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:
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
+9 -3
View File
@@ -195,10 +195,16 @@ async def test_sqlite_scoped_keys_allow_cross_workspace_but_reject_same_workspac
await conn.execute(
sa.text(
'INSERT INTO plugin_settings '
'(workspace_uuid, plugin_author, plugin_name, enabled) '
"VALUES (:workspace_uuid, 'author', 'plugin', 1)"
'(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},
{
'workspace_uuid': second_workspace_uuid,
'installation_uuid': str(uuid.uuid4()),
'artifact_digest': hashlib.sha256(b'test-plugin-artifact').hexdigest(),
},
)
await conn.execute(
sa.text(