mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 12:40:59 +00:00
feat(tenancy): add Workspace multi-tenant foundation (#2353)
* Document multi-tenant workspace architecture * Add OSS and commercial workspace boundaries * docs: redesign multi-tenant workspace architecture * feat(tenancy): implement workspace isolation * docs(tenancy): record verification evidence * docs(tenancy): revise single-instance SaaS topology * docs(tenancy): refine architecture options * docs: finalize cloud v2 multi-tenant decisions * feat(tenancy): establish cloud isolation foundations * feat(tenancy): harden shared cloud runtime boundaries * docs(tenancy): record final isolation verification * fix(tenancy): close isolation and permission gaps * docs(tenancy): record final isolation verification * feat(tenancy): connect cloud workspace control plane * fix(build): install git for pinned SDK * docs(cloud): update control plane verification * chore: update multi-tenant SDK pin * fix(cloud): skip legacy model sync during startup * test(cloud): preserve minimal model manager fixtures * fix(cloud): preserve authenticated account context * fix(cloud): reuse authenticated account for user info * feat(cloud): complete Workspace settings navigation * test(web): cover Workspace dropdown menu * feat(web): place workspace controls in sidebar * refactor(web): streamline workspace controls * style(web): format workspace layout test * fix(cloud): surface runtime and workspace plan status * fix(plugin): keep runtime identity stable across restarts * fix(ui): widen and center workspace switcher * fix(ui): hide roles from workspace switcher * fix(ui): align workspace switcher with sidebar entries * feat(workspace): add in-product collaboration and direct Cloud launch * style: format collaboration changes * fix(workspace): bind collaboration APIs to tenant UoW * fix(cloud): preserve Core-owned collaboration state * test(cloud): require Space identity for invite registration * feat(cloud): complete secure invitation experience * style(web): format invitation flows * fix(cloud): recover box runtime without unscoped skill reload * feat(oss): enforce invitation account and owner billing flows * style: format OSS account service * test(oss): cover invitation logout handoff * fix(oss): resolve workspace owner in scoped session * feat(cloud): harden multi-tenant runtime resources * fix(cloud): bound runtime restart storms * fix(cloud): eliminate periodic runtime CPU spikes * fix(cloud): enforce instance capacity ceilings * fix(cloud): scope public login capability discovery * fix(cloud): bound tenant maintenance and monitoring work * fix(runtime): bound tenant resource amplification * fix(deps): pin green multi-tenant plugin SDK * fix(cloud): handle unavailable skill capability * fix(security): require authentication for image file endpoint (H-2) - Changed /api/v1/files/image from AuthType.NONE to USER_TOKEN_OR_API_KEY - Added Permission.RESOURCE_VIEW requirement - Prevents unauthenticated cross-tenant file access via leaked keys - Fixes HIGH severity finding from multi-tenant security review docs: add comprehensive database migration guide - Complete migration steps for OSS → multi-tenant - Backup, execution, verification procedures - Rollback scenarios and recovery plans - Performance tuning recommendations * test: add comprehensive cross-tenant isolation tests Added 7 critical test scenarios for multi-tenant boundaries: - Cross-tenant bot access prevention - Viewer role read-only enforcement - Removed member immediate access revocation - Model provider credential isolation - WebSocket message isolation - Invitation token workspace scoping - Multi-workspace context validation These tests address P0-2 coverage gaps for: - workspaces.py (membership & invitation flows) - user.py (authentication & authorization) - websocket_chat.py (real-time isolation) - plugins.py (resource access control) docs: finalize database migration guide * fix(security): resolve M-1, M-2, M-3 security findings M-1: WebSocket authorization TOCTOU race (FIXED) - Changed _revalidate_websocket_authorization to return RequestContext - Ensures validated context is used immediately without race window - Prevents removed members from sending messages during revalidation gap M-2: Model Manager cache workspace isolation (VERIFIED) - Confirmed _CacheKey already uses 4-tuple: (instance, workspace, generation, resource) - Cache is properly scoped per workspace, no cross-tenant leakage possible - No code change needed, documented as working correctly M-3: Invitation lock workspace scoping (FIXED) - Changed lock key from token_digest to workspace_uuid:token_digest - Prevents DoS where attacker locks token in Workspace A to block Workspace B - Locks now isolated per workspace All MEDIUM severity findings from security review now resolved. * fix(cloud): unblock tenant CI and enforce knowledge quotas * fix(tenancy): scope rerank model sync --------- Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
This commit is contained in:
RockChinQ
@@ -0,0 +1,542 @@
|
||||
"""add the workspace tenancy persistence kernel
|
||||
|
||||
Revision ID: 0009_workspace_tenancy
|
||||
Revises: 0008_mcp_resource_prefs
|
||||
Create Date: 2026-07-18
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
import uuid
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision = '0009_workspace_tenancy'
|
||||
down_revision = '0008_mcp_resource_prefs'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def _table_names(conn: sa.Connection) -> set[str]:
|
||||
return set(sa.inspect(conn).get_table_names())
|
||||
|
||||
|
||||
def _column_map(conn: sa.Connection, table_name: str) -> dict[str, dict]:
|
||||
return {column['name']: column for column in sa.inspect(conn).get_columns(table_name)}
|
||||
|
||||
|
||||
def _constraint_names(conn: sa.Connection, table_name: str) -> set[str]:
|
||||
inspector = sa.inspect(conn)
|
||||
names = {
|
||||
constraint['name']
|
||||
for constraint in inspector.get_check_constraints(table_name)
|
||||
if constraint.get('name') is not None
|
||||
}
|
||||
names.update(
|
||||
constraint['name']
|
||||
for constraint in inspector.get_unique_constraints(table_name)
|
||||
if constraint.get('name') is not None
|
||||
)
|
||||
return names
|
||||
|
||||
|
||||
def _index_names(conn: sa.Connection, table_name: str) -> set[str]:
|
||||
return {index['name'] for index in sa.inspect(conn).get_indexes(table_name)}
|
||||
|
||||
|
||||
def _upgrade_users(conn: sa.Connection) -> None:
|
||||
if 'users' not in _table_names(conn):
|
||||
return
|
||||
|
||||
columns = _column_map(conn, 'users')
|
||||
if 'uuid' not in columns:
|
||||
op.add_column('users', sa.Column('uuid', sa.String(36), nullable=True))
|
||||
if 'status' not in columns:
|
||||
op.add_column('users', sa.Column('status', sa.String(32), nullable=True, server_default='active'))
|
||||
if 'source' not in columns:
|
||||
op.add_column('users', sa.Column('source', sa.String(32), nullable=True, server_default='local'))
|
||||
if 'projection_revision' not in columns:
|
||||
op.add_column(
|
||||
'users',
|
||||
sa.Column('projection_revision', sa.BigInteger(), nullable=True, server_default='0'),
|
||||
)
|
||||
|
||||
users = sa.table(
|
||||
'users',
|
||||
sa.column('id', sa.Integer()),
|
||||
sa.column('uuid', sa.String(36)),
|
||||
sa.column('status', sa.String(32)),
|
||||
sa.column('source', sa.String(32)),
|
||||
sa.column('projection_revision', sa.BigInteger()),
|
||||
)
|
||||
|
||||
seen_uuids: set[str] = set()
|
||||
for user_id, account_uuid in conn.execute(sa.select(users.c.id, users.c.uuid).order_by(users.c.id)).all():
|
||||
normalized_uuid = account_uuid.strip() if isinstance(account_uuid, str) else ''
|
||||
try:
|
||||
normalized_uuid = str(uuid.UUID(normalized_uuid))
|
||||
except (ValueError, AttributeError):
|
||||
normalized_uuid = ''
|
||||
if not normalized_uuid or normalized_uuid in seen_uuids:
|
||||
normalized_uuid = str(uuid.uuid4())
|
||||
if normalized_uuid != account_uuid:
|
||||
conn.execute(users.update().where(users.c.id == user_id).values(uuid=normalized_uuid))
|
||||
seen_uuids.add(normalized_uuid)
|
||||
|
||||
conn.execute(users.update().where(users.c.status.is_(None)).values(status='active'))
|
||||
conn.execute(users.update().where(users.c.source.is_(None)).values(source='local'))
|
||||
conn.execute(users.update().where(users.c.projection_revision.is_(None)).values(projection_revision=0))
|
||||
|
||||
columns = _column_map(conn, 'users')
|
||||
constraint_names = _constraint_names(conn, 'users')
|
||||
needs_batch_alter = any(
|
||||
columns[column_name]['nullable'] for column_name in ('uuid', 'status', 'source', 'projection_revision')
|
||||
) or not {'ck_users_status', 'ck_users_source'}.issubset(constraint_names)
|
||||
|
||||
if needs_batch_alter:
|
||||
with op.batch_alter_table('users') as batch_op:
|
||||
if columns['uuid']['nullable']:
|
||||
batch_op.alter_column('uuid', existing_type=sa.String(36), nullable=False)
|
||||
if columns['status']['nullable']:
|
||||
batch_op.alter_column(
|
||||
'status',
|
||||
existing_type=sa.String(32),
|
||||
nullable=False,
|
||||
server_default='active',
|
||||
)
|
||||
if columns['source']['nullable']:
|
||||
batch_op.alter_column(
|
||||
'source',
|
||||
existing_type=sa.String(32),
|
||||
nullable=False,
|
||||
server_default='local',
|
||||
)
|
||||
if columns['projection_revision']['nullable']:
|
||||
batch_op.alter_column(
|
||||
'projection_revision',
|
||||
existing_type=sa.BigInteger(),
|
||||
nullable=False,
|
||||
server_default='0',
|
||||
)
|
||||
if 'ck_users_status' not in constraint_names:
|
||||
batch_op.create_check_constraint(
|
||||
'ck_users_status',
|
||||
"status IN ('active', 'disabled', 'deleted')",
|
||||
)
|
||||
if 'ck_users_source' not in constraint_names:
|
||||
batch_op.create_check_constraint(
|
||||
'ck_users_source',
|
||||
"source IN ('local', 'cloud_projection')",
|
||||
)
|
||||
|
||||
if 'uq_users_uuid' not in _index_names(conn, 'users'):
|
||||
op.create_index('uq_users_uuid', 'users', ['uuid'], unique=True)
|
||||
|
||||
|
||||
def _create_workspace_tables(conn: sa.Connection) -> None:
|
||||
tables = _table_names(conn)
|
||||
if 'users' not in tables:
|
||||
# LangBot's supported startup path creates the baseline schema before
|
||||
# Alembic runs. Keep direct Alembic probes on an empty database safe.
|
||||
return
|
||||
|
||||
if 'workspaces' not in tables:
|
||||
op.create_table(
|
||||
'workspaces',
|
||||
sa.Column('uuid', sa.String(36), primary_key=True),
|
||||
sa.Column('instance_uuid', sa.String(255), nullable=False),
|
||||
sa.Column('name', sa.String(255), nullable=False),
|
||||
sa.Column('slug', sa.String(255), nullable=False),
|
||||
sa.Column('type', sa.String(32), nullable=False, server_default='team'),
|
||||
sa.Column('status', sa.String(32), nullable=False, server_default='active'),
|
||||
sa.Column('created_by_account_uuid', sa.String(36), nullable=True),
|
||||
sa.Column('source', sa.String(32), nullable=False, server_default='local'),
|
||||
sa.Column('projection_revision', sa.BigInteger(), nullable=False, server_default='0'),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=False, server_default=sa.func.now()),
|
||||
sa.Column('updated_at', sa.DateTime(), nullable=False, server_default=sa.func.now()),
|
||||
sa.ForeignKeyConstraint(
|
||||
['created_by_account_uuid'],
|
||||
['users.uuid'],
|
||||
name='fk_workspaces_created_by_account',
|
||||
ondelete='SET NULL',
|
||||
),
|
||||
sa.UniqueConstraint('instance_uuid', 'slug', name='uq_workspaces_instance_slug'),
|
||||
sa.CheckConstraint("type IN ('personal', 'team')", name='ck_workspaces_type'),
|
||||
sa.CheckConstraint(
|
||||
"status IN ('provisioning', 'active', 'suspended', 'archived', 'deleted')",
|
||||
name='ck_workspaces_status',
|
||||
),
|
||||
sa.CheckConstraint(
|
||||
"source IN ('local', 'cloud_projection')",
|
||||
name='ck_workspaces_source',
|
||||
),
|
||||
)
|
||||
workspace_indexes = _index_names(conn, 'workspaces')
|
||||
if 'ix_workspaces_instance_status' not in workspace_indexes:
|
||||
op.create_index(
|
||||
'ix_workspaces_instance_status',
|
||||
'workspaces',
|
||||
['instance_uuid', 'status'],
|
||||
)
|
||||
if 'uq_workspaces_local_instance' not in workspace_indexes:
|
||||
op.create_index(
|
||||
'uq_workspaces_local_instance',
|
||||
'workspaces',
|
||||
['instance_uuid'],
|
||||
unique=True,
|
||||
sqlite_where=sa.text("source = 'local'"),
|
||||
postgresql_where=sa.text("source = 'local'"),
|
||||
)
|
||||
|
||||
tables = _table_names(conn)
|
||||
if 'workspace_memberships' not in tables:
|
||||
op.create_table(
|
||||
'workspace_memberships',
|
||||
sa.Column('uuid', sa.String(36), primary_key=True),
|
||||
sa.Column('workspace_uuid', sa.String(36), nullable=False),
|
||||
sa.Column('account_uuid', sa.String(36), nullable=False),
|
||||
sa.Column('role', sa.String(32), nullable=False),
|
||||
sa.Column('status', sa.String(32), nullable=False, server_default='active'),
|
||||
sa.Column('invited_by_account_uuid', sa.String(36), nullable=True),
|
||||
sa.Column('joined_at', sa.DateTime(), nullable=True),
|
||||
sa.Column('projection_revision', sa.BigInteger(), nullable=False, server_default='0'),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=False, server_default=sa.func.now()),
|
||||
sa.Column('updated_at', sa.DateTime(), nullable=False, server_default=sa.func.now()),
|
||||
sa.ForeignKeyConstraint(
|
||||
['workspace_uuid'],
|
||||
['workspaces.uuid'],
|
||||
name='fk_workspace_memberships_workspace',
|
||||
ondelete='CASCADE',
|
||||
),
|
||||
sa.ForeignKeyConstraint(
|
||||
['account_uuid'],
|
||||
['users.uuid'],
|
||||
name='fk_workspace_memberships_account',
|
||||
ondelete='CASCADE',
|
||||
),
|
||||
sa.ForeignKeyConstraint(
|
||||
['invited_by_account_uuid'],
|
||||
['users.uuid'],
|
||||
name='fk_workspace_memberships_invited_by_account',
|
||||
ondelete='SET NULL',
|
||||
),
|
||||
sa.UniqueConstraint(
|
||||
'workspace_uuid',
|
||||
'account_uuid',
|
||||
name='uq_workspace_membership_account',
|
||||
),
|
||||
sa.CheckConstraint(
|
||||
"role IN ('owner', 'admin', 'developer', 'operator', 'viewer')",
|
||||
name='ck_workspace_memberships_role',
|
||||
),
|
||||
sa.CheckConstraint(
|
||||
"status IN ('active', 'disabled', 'removed')",
|
||||
name='ck_workspace_memberships_status',
|
||||
),
|
||||
)
|
||||
membership_indexes = _index_names(conn, 'workspace_memberships')
|
||||
if 'ix_workspace_memberships_account_status' not in membership_indexes:
|
||||
op.create_index(
|
||||
'ix_workspace_memberships_account_status',
|
||||
'workspace_memberships',
|
||||
['account_uuid', 'status'],
|
||||
)
|
||||
|
||||
tables = _table_names(conn)
|
||||
if 'workspace_invitations' not in tables:
|
||||
op.create_table(
|
||||
'workspace_invitations',
|
||||
sa.Column('uuid', sa.String(36), primary_key=True),
|
||||
sa.Column('workspace_uuid', sa.String(36), nullable=False),
|
||||
sa.Column('normalized_email', sa.String(320), nullable=False),
|
||||
sa.Column('role', sa.String(32), nullable=False),
|
||||
sa.Column('token_hash', sa.String(255), nullable=False),
|
||||
sa.Column('status', sa.String(32), nullable=False, server_default='pending'),
|
||||
sa.Column('expires_at', sa.DateTime(), nullable=False),
|
||||
sa.Column('accepted_at', sa.DateTime(), nullable=True),
|
||||
sa.Column('revoked_at', sa.DateTime(), nullable=True),
|
||||
sa.Column('created_by_account_uuid', sa.String(36), nullable=False),
|
||||
sa.Column('created_at', sa.DateTime(), nullable=False, server_default=sa.func.now()),
|
||||
sa.Column('updated_at', sa.DateTime(), nullable=False, server_default=sa.func.now()),
|
||||
sa.ForeignKeyConstraint(
|
||||
['workspace_uuid'],
|
||||
['workspaces.uuid'],
|
||||
name='fk_workspace_invitations_workspace',
|
||||
ondelete='CASCADE',
|
||||
),
|
||||
sa.ForeignKeyConstraint(
|
||||
['created_by_account_uuid'],
|
||||
['users.uuid'],
|
||||
name='fk_workspace_invitations_created_by_account',
|
||||
ondelete='CASCADE',
|
||||
),
|
||||
sa.CheckConstraint(
|
||||
"role IN ('admin', 'developer', 'operator', 'viewer')",
|
||||
name='ck_workspace_invitations_role',
|
||||
),
|
||||
sa.CheckConstraint(
|
||||
"status IN ('pending', 'accepted', 'revoked', 'expired')",
|
||||
name='ck_workspace_invitations_status',
|
||||
),
|
||||
)
|
||||
invitation_indexes = _index_names(conn, 'workspace_invitations')
|
||||
if 'uq_workspace_invitations_token_hash' not in invitation_indexes:
|
||||
op.create_index(
|
||||
'uq_workspace_invitations_token_hash',
|
||||
'workspace_invitations',
|
||||
['token_hash'],
|
||||
unique=True,
|
||||
)
|
||||
if 'uq_workspace_invitations_pending_email' not in invitation_indexes:
|
||||
op.create_index(
|
||||
'uq_workspace_invitations_pending_email',
|
||||
'workspace_invitations',
|
||||
['workspace_uuid', 'normalized_email'],
|
||||
unique=True,
|
||||
sqlite_where=sa.text("status = 'pending'"),
|
||||
postgresql_where=sa.text("status = 'pending'"),
|
||||
)
|
||||
|
||||
tables = _table_names(conn)
|
||||
if 'workspace_execution_states' not in tables:
|
||||
op.create_table(
|
||||
'workspace_execution_states',
|
||||
sa.Column('workspace_uuid', sa.String(36), primary_key=True),
|
||||
sa.Column('instance_uuid', sa.String(255), nullable=False),
|
||||
sa.Column('active_generation', sa.BigInteger(), nullable=False, server_default='1'),
|
||||
sa.Column('state', sa.String(32), nullable=False, server_default='active'),
|
||||
sa.Column('write_fenced', sa.Boolean(), nullable=False, server_default=sa.false()),
|
||||
sa.Column('source', sa.String(32), nullable=False, server_default='local'),
|
||||
sa.Column('desired_state_revision', sa.BigInteger(), nullable=False, server_default='0'),
|
||||
sa.Column('updated_at', sa.DateTime(), nullable=False, server_default=sa.func.now()),
|
||||
sa.ForeignKeyConstraint(
|
||||
['workspace_uuid'],
|
||||
['workspaces.uuid'],
|
||||
name='fk_workspace_execution_states_workspace',
|
||||
ondelete='CASCADE',
|
||||
),
|
||||
sa.CheckConstraint('active_generation > 0', name='ck_workspace_execution_generation'),
|
||||
sa.CheckConstraint(
|
||||
"state IN ('provisioning', 'active', 'migrating', 'draining', 'inactive')",
|
||||
name='ck_workspace_execution_state',
|
||||
),
|
||||
sa.CheckConstraint(
|
||||
"source IN ('local', 'cloud')",
|
||||
name='ck_workspace_execution_source',
|
||||
),
|
||||
)
|
||||
execution_indexes = _index_names(conn, 'workspace_execution_states')
|
||||
if 'ix_workspace_execution_states_instance_state' not in execution_indexes:
|
||||
op.create_index(
|
||||
'ix_workspace_execution_states_instance_state',
|
||||
'workspace_execution_states',
|
||||
['instance_uuid', 'state'],
|
||||
)
|
||||
|
||||
|
||||
def _load_instance_uuid(conn: sa.Connection) -> str | None:
|
||||
if 'metadata' not in _table_names(conn):
|
||||
return None
|
||||
|
||||
metadata = sa.table(
|
||||
'metadata',
|
||||
sa.column('key', sa.String(255)),
|
||||
sa.column('value', sa.String(255)),
|
||||
)
|
||||
value = conn.execute(sa.select(metadata.c.value).where(metadata.c.key == 'instance_uuid')).scalar_one_or_none()
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
return None
|
||||
return value.strip()
|
||||
|
||||
|
||||
def _bootstrap_default_workspace(conn: sa.Connection) -> None:
|
||||
required_tables = {'users', 'workspaces', 'workspace_memberships', 'workspace_execution_states'}
|
||||
if not required_tables.issubset(_table_names(conn)):
|
||||
return
|
||||
|
||||
instance_uuid = _load_instance_uuid(conn)
|
||||
users_exist = 'users' in _table_names(conn) and bool(conn.execute(sa.text('SELECT 1 FROM users LIMIT 1')).first())
|
||||
if instance_uuid is None:
|
||||
if users_exist:
|
||||
raise RuntimeError("Cannot bootstrap the default workspace without metadata['instance_uuid']")
|
||||
return
|
||||
|
||||
workspaces = sa.table(
|
||||
'workspaces',
|
||||
sa.column('uuid', sa.String(36)),
|
||||
sa.column('instance_uuid', sa.String(255)),
|
||||
sa.column('name', sa.String(255)),
|
||||
sa.column('slug', sa.String(255)),
|
||||
sa.column('type', sa.String(32)),
|
||||
sa.column('status', sa.String(32)),
|
||||
sa.column('created_by_account_uuid', sa.String(36)),
|
||||
sa.column('source', sa.String(32)),
|
||||
sa.column('projection_revision', sa.BigInteger()),
|
||||
)
|
||||
local_rows = conn.execute(
|
||||
sa.select(workspaces.c.uuid).where(
|
||||
workspaces.c.instance_uuid == instance_uuid,
|
||||
workspaces.c.source == 'local',
|
||||
)
|
||||
).all()
|
||||
if len(local_rows) > 1:
|
||||
raise RuntimeError(f'Multiple local workspaces already exist for instance {instance_uuid!r}')
|
||||
|
||||
users = sa.table(
|
||||
'users',
|
||||
sa.column('id', sa.Integer()),
|
||||
sa.column('uuid', sa.String(36)),
|
||||
)
|
||||
owner_account_uuid = None
|
||||
if 'users' in _table_names(conn):
|
||||
owner_account_uuid = conn.execute(sa.select(users.c.uuid).order_by(users.c.id).limit(1)).scalar_one_or_none()
|
||||
|
||||
if local_rows:
|
||||
workspace_uuid = local_rows[0][0]
|
||||
if owner_account_uuid is not None:
|
||||
conn.execute(
|
||||
workspaces.update()
|
||||
.where(workspaces.c.uuid == workspace_uuid)
|
||||
.where(workspaces.c.created_by_account_uuid.is_(None))
|
||||
.values(created_by_account_uuid=owner_account_uuid)
|
||||
)
|
||||
else:
|
||||
workspace_uuid = str(uuid.uuid4())
|
||||
conn.execute(
|
||||
workspaces.insert().values(
|
||||
uuid=workspace_uuid,
|
||||
instance_uuid=instance_uuid,
|
||||
name='Default Workspace',
|
||||
slug='default',
|
||||
type='team',
|
||||
status='active',
|
||||
created_by_account_uuid=owner_account_uuid,
|
||||
source='local',
|
||||
projection_revision=0,
|
||||
)
|
||||
)
|
||||
|
||||
execution_states = sa.table(
|
||||
'workspace_execution_states',
|
||||
sa.column('workspace_uuid', sa.String(36)),
|
||||
sa.column('instance_uuid', sa.String(255)),
|
||||
sa.column('active_generation', sa.BigInteger()),
|
||||
sa.column('state', sa.String(32)),
|
||||
sa.column('write_fenced', sa.Boolean()),
|
||||
sa.column('source', sa.String(32)),
|
||||
sa.column('desired_state_revision', sa.BigInteger()),
|
||||
)
|
||||
execution_state = conn.execute(
|
||||
sa.select(
|
||||
execution_states.c.instance_uuid,
|
||||
execution_states.c.active_generation,
|
||||
execution_states.c.state,
|
||||
execution_states.c.write_fenced,
|
||||
execution_states.c.source,
|
||||
).where(execution_states.c.workspace_uuid == workspace_uuid)
|
||||
).first()
|
||||
if execution_state is None:
|
||||
conn.execute(
|
||||
execution_states.insert().values(
|
||||
workspace_uuid=workspace_uuid,
|
||||
instance_uuid=instance_uuid,
|
||||
active_generation=1,
|
||||
state='active',
|
||||
write_fenced=False,
|
||||
source='local',
|
||||
desired_state_revision=0,
|
||||
)
|
||||
)
|
||||
elif (
|
||||
execution_state.instance_uuid != instance_uuid
|
||||
or execution_state.active_generation != 1
|
||||
or execution_state.state != 'active'
|
||||
or execution_state.write_fenced
|
||||
or execution_state.source != 'local'
|
||||
):
|
||||
raise RuntimeError(f'Default workspace {workspace_uuid!r} has an invalid local execution state')
|
||||
|
||||
if owner_account_uuid is None:
|
||||
return
|
||||
|
||||
memberships = sa.table(
|
||||
'workspace_memberships',
|
||||
sa.column('uuid', sa.String(36)),
|
||||
sa.column('workspace_uuid', sa.String(36)),
|
||||
sa.column('account_uuid', sa.String(36)),
|
||||
sa.column('role', sa.String(32)),
|
||||
sa.column('status', sa.String(32)),
|
||||
sa.column('joined_at', sa.DateTime()),
|
||||
sa.column('projection_revision', sa.BigInteger()),
|
||||
)
|
||||
membership = conn.execute(
|
||||
sa.select(memberships.c.uuid, memberships.c.joined_at).where(
|
||||
memberships.c.workspace_uuid == workspace_uuid,
|
||||
memberships.c.account_uuid == owner_account_uuid,
|
||||
)
|
||||
).first()
|
||||
now = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
|
||||
if membership is None:
|
||||
conn.execute(
|
||||
memberships.insert().values(
|
||||
uuid=str(uuid.uuid4()),
|
||||
workspace_uuid=workspace_uuid,
|
||||
account_uuid=owner_account_uuid,
|
||||
role='owner',
|
||||
status='active',
|
||||
joined_at=now,
|
||||
projection_revision=0,
|
||||
)
|
||||
)
|
||||
else:
|
||||
conn.execute(
|
||||
memberships.update()
|
||||
.where(memberships.c.uuid == membership.uuid)
|
||||
.values(
|
||||
role='owner',
|
||||
status='active',
|
||||
joined_at=membership.joined_at or now,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
_upgrade_users(conn)
|
||||
_create_workspace_tables(conn)
|
||||
_bootstrap_default_workspace(conn)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
tables = _table_names(conn)
|
||||
for table_name in (
|
||||
'workspace_execution_states',
|
||||
'workspace_invitations',
|
||||
'workspace_memberships',
|
||||
'workspaces',
|
||||
):
|
||||
if table_name in tables:
|
||||
op.drop_table(table_name)
|
||||
|
||||
if 'users' not in _table_names(conn):
|
||||
return
|
||||
|
||||
indexes = _index_names(conn, 'users')
|
||||
if 'uq_users_uuid' in indexes:
|
||||
op.drop_index('uq_users_uuid', table_name='users')
|
||||
|
||||
columns = _column_map(conn, 'users')
|
||||
constraint_names = _constraint_names(conn, 'users')
|
||||
with op.batch_alter_table('users') as batch_op:
|
||||
# SQLite batch recreation otherwise preserves the named checks while
|
||||
# dropping their referenced columns, producing ``no such column`` only
|
||||
# after the Workspace directory tables have already been removed.
|
||||
for constraint_name in ('ck_users_source', 'ck_users_status'):
|
||||
if constraint_name in constraint_names:
|
||||
batch_op.drop_constraint(constraint_name, type_='check')
|
||||
for column_name in ('projection_revision', 'source', 'status', 'uuid'):
|
||||
if column_name in columns:
|
||||
batch_op.drop_column(column_name)
|
||||
@@ -0,0 +1,884 @@
|
||||
"""scope every tenant-owned resource to a workspace
|
||||
|
||||
Revision ID: 0010_scope_resources
|
||||
Revises: 0009_workspace_tenancy
|
||||
Create Date: 2026-07-19
|
||||
|
||||
This migration is intentionally expand/backfill/contract. Existing rows are
|
||||
bound to the single local Workspace created by revision 0009 before any
|
||||
non-null or scoped-key constraint is installed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import uuid
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision = '0010_scope_resources'
|
||||
down_revision = '0009_workspace_tenancy'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
_TENANT_TABLES = (
|
||||
'api_keys',
|
||||
'bots',
|
||||
'bot_admins',
|
||||
'binary_storages',
|
||||
'mcp_servers',
|
||||
'model_providers',
|
||||
'llm_models',
|
||||
'embedding_models',
|
||||
'rerank_models',
|
||||
'legacy_pipelines',
|
||||
'pipeline_run_records',
|
||||
'plugin_settings',
|
||||
'knowledge_bases',
|
||||
'knowledge_base_files',
|
||||
'knowledge_base_chunks',
|
||||
'webhooks',
|
||||
'monitoring_messages',
|
||||
'monitoring_llm_calls',
|
||||
'monitoring_tool_calls',
|
||||
'monitoring_sessions',
|
||||
'monitoring_errors',
|
||||
'monitoring_embedding_calls',
|
||||
'monitoring_feedback',
|
||||
)
|
||||
|
||||
_COMPOSITE_PRIMARY_KEYS = {
|
||||
'binary_storages': ('workspace_uuid', 'unique_key'),
|
||||
'plugin_settings': ('workspace_uuid', 'plugin_author', 'plugin_name'),
|
||||
'monitoring_sessions': ('workspace_uuid', 'session_id'),
|
||||
}
|
||||
|
||||
_COMPOSITE_FOREIGN_KEYS = {
|
||||
'bot_admins': (
|
||||
(
|
||||
'fk_bot_admins_workspace_bot',
|
||||
('workspace_uuid', 'bot_uuid'),
|
||||
'bots',
|
||||
('workspace_uuid', 'uuid'),
|
||||
'CASCADE',
|
||||
),
|
||||
),
|
||||
'llm_models': (
|
||||
(
|
||||
'fk_llm_models_workspace_provider',
|
||||
('workspace_uuid', 'provider_uuid'),
|
||||
'model_providers',
|
||||
('workspace_uuid', 'uuid'),
|
||||
None,
|
||||
),
|
||||
),
|
||||
'embedding_models': (
|
||||
(
|
||||
'fk_embedding_models_workspace_provider',
|
||||
('workspace_uuid', 'provider_uuid'),
|
||||
'model_providers',
|
||||
('workspace_uuid', 'uuid'),
|
||||
None,
|
||||
),
|
||||
),
|
||||
'rerank_models': (
|
||||
(
|
||||
'fk_rerank_models_workspace_provider',
|
||||
('workspace_uuid', 'provider_uuid'),
|
||||
'model_providers',
|
||||
('workspace_uuid', 'uuid'),
|
||||
None,
|
||||
),
|
||||
),
|
||||
'pipeline_run_records': (
|
||||
(
|
||||
'fk_pipeline_run_records_workspace_pipeline',
|
||||
('workspace_uuid', 'pipeline_uuid'),
|
||||
'legacy_pipelines',
|
||||
('workspace_uuid', 'uuid'),
|
||||
'CASCADE',
|
||||
),
|
||||
),
|
||||
'knowledge_base_files': (
|
||||
(
|
||||
'fk_knowledge_base_files_workspace_kb',
|
||||
('workspace_uuid', 'kb_id'),
|
||||
'knowledge_bases',
|
||||
('workspace_uuid', 'uuid'),
|
||||
'CASCADE',
|
||||
),
|
||||
),
|
||||
'knowledge_base_chunks': (
|
||||
(
|
||||
'fk_knowledge_base_chunks_workspace_file',
|
||||
('workspace_uuid', 'file_id'),
|
||||
'knowledge_base_files',
|
||||
('workspace_uuid', 'uuid'),
|
||||
'CASCADE',
|
||||
),
|
||||
),
|
||||
}
|
||||
|
||||
_SCOPED_INDEXES: dict[str, tuple[tuple[str, tuple[str, ...], bool, sa.TextClause | None], ...]] = {
|
||||
'api_keys': (
|
||||
('uq_api_keys_uuid', ('uuid',), True, None),
|
||||
('uq_api_keys_key_hash', ('key_hash',), True, None),
|
||||
('ix_api_keys_workspace_name', ('workspace_uuid', 'name'), False, None),
|
||||
('ix_api_keys_workspace_status', ('workspace_uuid', 'status'), False, None),
|
||||
),
|
||||
'bots': (
|
||||
('uq_bots_workspace_uuid', ('workspace_uuid', 'uuid'), True, None),
|
||||
('ix_bots_workspace_name', ('workspace_uuid', 'name'), False, None),
|
||||
('ix_bots_workspace_updated', ('workspace_uuid', 'updated_at'), False, None),
|
||||
),
|
||||
'bot_admins': (
|
||||
(
|
||||
'uq_bot_admin',
|
||||
('workspace_uuid', 'bot_uuid', 'launcher_type', 'launcher_id'),
|
||||
True,
|
||||
None,
|
||||
),
|
||||
('ix_bot_admins_workspace_bot', ('workspace_uuid', 'bot_uuid'), False, None),
|
||||
),
|
||||
'binary_storages': (
|
||||
(
|
||||
'ix_binary_storages_workspace_owner',
|
||||
('workspace_uuid', 'owner_type', 'owner'),
|
||||
False,
|
||||
None,
|
||||
),
|
||||
),
|
||||
'mcp_servers': (
|
||||
('uq_mcp_servers_workspace_name', ('workspace_uuid', 'name'), True, None),
|
||||
('ix_mcp_servers_workspace_enable', ('workspace_uuid', 'enable'), False, None),
|
||||
('ix_mcp_servers_workspace_updated', ('workspace_uuid', 'updated_at'), False, None),
|
||||
),
|
||||
'model_providers': (
|
||||
('uq_model_providers_workspace_uuid', ('workspace_uuid', 'uuid'), True, None),
|
||||
('ix_model_providers_workspace_name', ('workspace_uuid', 'name'), False, None),
|
||||
('ix_model_providers_workspace_requester', ('workspace_uuid', 'requester'), False, None),
|
||||
),
|
||||
'llm_models': (
|
||||
('ix_llm_models_workspace_provider', ('workspace_uuid', 'provider_uuid'), False, None),
|
||||
('ix_llm_models_workspace_name', ('workspace_uuid', 'name'), False, None),
|
||||
),
|
||||
'embedding_models': (
|
||||
('ix_embedding_models_workspace_provider', ('workspace_uuid', 'provider_uuid'), False, None),
|
||||
('ix_embedding_models_workspace_name', ('workspace_uuid', 'name'), False, None),
|
||||
),
|
||||
'rerank_models': (
|
||||
('ix_rerank_models_workspace_provider', ('workspace_uuid', 'provider_uuid'), False, None),
|
||||
('ix_rerank_models_workspace_name', ('workspace_uuid', 'name'), False, None),
|
||||
),
|
||||
'legacy_pipelines': (
|
||||
('uq_legacy_pipelines_workspace_uuid', ('workspace_uuid', 'uuid'), True, None),
|
||||
('ix_legacy_pipelines_workspace_name', ('workspace_uuid', 'name'), False, None),
|
||||
('ix_legacy_pipelines_workspace_default', ('workspace_uuid', 'is_default'), False, None),
|
||||
('ix_legacy_pipelines_workspace_updated', ('workspace_uuid', 'updated_at'), False, None),
|
||||
),
|
||||
'pipeline_run_records': (
|
||||
(
|
||||
'ix_pipeline_run_records_workspace_pipeline',
|
||||
('workspace_uuid', 'pipeline_uuid'),
|
||||
False,
|
||||
None,
|
||||
),
|
||||
(
|
||||
'ix_pipeline_run_records_workspace_created',
|
||||
('workspace_uuid', 'created_at'),
|
||||
False,
|
||||
None,
|
||||
),
|
||||
),
|
||||
'plugin_settings': (('ix_plugin_settings_workspace_enabled', ('workspace_uuid', 'enabled'), False, None),),
|
||||
'knowledge_bases': (
|
||||
('uq_knowledge_bases_workspace_uuid', ('workspace_uuid', 'uuid'), True, None),
|
||||
('ix_knowledge_bases_workspace_name', ('workspace_uuid', 'name'), False, None),
|
||||
(
|
||||
'uq_knowledge_bases_workspace_collection',
|
||||
('workspace_uuid', 'collection_id'),
|
||||
True,
|
||||
sa.text('collection_id IS NOT NULL'),
|
||||
),
|
||||
),
|
||||
'knowledge_base_files': (
|
||||
('uq_knowledge_base_files_workspace_uuid', ('workspace_uuid', 'uuid'), True, None),
|
||||
('ix_knowledge_base_files_workspace_kb', ('workspace_uuid', 'kb_id'), False, None),
|
||||
),
|
||||
'knowledge_base_chunks': (('ix_knowledge_base_chunks_workspace_file', ('workspace_uuid', 'file_id'), False, None),),
|
||||
'webhooks': (
|
||||
('ix_webhooks_workspace_name', ('workspace_uuid', 'name'), False, None),
|
||||
('ix_webhooks_workspace_enabled', ('workspace_uuid', 'enabled'), False, None),
|
||||
('ix_webhooks_workspace_created', ('workspace_uuid', 'created_at'), False, None),
|
||||
),
|
||||
'monitoring_messages': (
|
||||
('ix_monitoring_messages_workspace_timestamp', ('workspace_uuid', 'timestamp'), False, None),
|
||||
('ix_monitoring_messages_workspace_bot', ('workspace_uuid', 'bot_id', 'timestamp'), False, None),
|
||||
(
|
||||
'ix_monitoring_messages_workspace_pipeline',
|
||||
('workspace_uuid', 'pipeline_id', 'timestamp'),
|
||||
False,
|
||||
None,
|
||||
),
|
||||
('ix_monitoring_messages_workspace_session', ('workspace_uuid', 'session_id'), False, None),
|
||||
),
|
||||
'monitoring_llm_calls': (
|
||||
('ix_monitoring_llm_calls_workspace_timestamp', ('workspace_uuid', 'timestamp'), False, None),
|
||||
('ix_monitoring_llm_calls_workspace_session', ('workspace_uuid', 'session_id'), False, None),
|
||||
('ix_monitoring_llm_calls_workspace_message', ('workspace_uuid', 'message_id'), False, None),
|
||||
),
|
||||
'monitoring_tool_calls': (
|
||||
('ix_monitoring_tool_calls_workspace_timestamp', ('workspace_uuid', 'timestamp'), False, None),
|
||||
('ix_monitoring_tool_calls_workspace_session', ('workspace_uuid', 'session_id'), False, None),
|
||||
('ix_monitoring_tool_calls_workspace_message', ('workspace_uuid', 'message_id'), False, None),
|
||||
),
|
||||
'monitoring_sessions': (
|
||||
('ix_monitoring_sessions_workspace_activity', ('workspace_uuid', 'last_activity'), False, None),
|
||||
('ix_monitoring_sessions_workspace_active', ('workspace_uuid', 'is_active'), False, None),
|
||||
('ix_monitoring_sessions_workspace_bot', ('workspace_uuid', 'bot_id', 'last_activity'), False, None),
|
||||
),
|
||||
'monitoring_errors': (
|
||||
('ix_monitoring_errors_workspace_timestamp', ('workspace_uuid', 'timestamp'), False, None),
|
||||
('ix_monitoring_errors_workspace_session', ('workspace_uuid', 'session_id'), False, None),
|
||||
('ix_monitoring_errors_workspace_message', ('workspace_uuid', 'message_id'), False, None),
|
||||
),
|
||||
'monitoring_embedding_calls': (
|
||||
(
|
||||
'ix_monitoring_embedding_calls_workspace_timestamp',
|
||||
('workspace_uuid', 'timestamp'),
|
||||
False,
|
||||
None,
|
||||
),
|
||||
(
|
||||
'ix_monitoring_embedding_calls_workspace_kb',
|
||||
('workspace_uuid', 'knowledge_base_id'),
|
||||
False,
|
||||
None,
|
||||
),
|
||||
(
|
||||
'ix_monitoring_embedding_calls_workspace_session',
|
||||
('workspace_uuid', 'session_id'),
|
||||
False,
|
||||
None,
|
||||
),
|
||||
),
|
||||
'monitoring_feedback': (
|
||||
(
|
||||
'uq_monitoring_feedback_workspace_feedback_id',
|
||||
('workspace_uuid', 'feedback_id'),
|
||||
True,
|
||||
None,
|
||||
),
|
||||
('ix_monitoring_feedback_workspace_timestamp', ('workspace_uuid', 'timestamp'), False, None),
|
||||
('ix_monitoring_feedback_workspace_session', ('workspace_uuid', 'session_id'), False, None),
|
||||
('ix_monitoring_feedback_workspace_message', ('workspace_uuid', 'message_id'), False, None),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _inspector(conn: sa.Connection) -> sa.Inspector:
|
||||
return sa.inspect(conn)
|
||||
|
||||
|
||||
def _table_names(conn: sa.Connection) -> set[str]:
|
||||
return set(_inspector(conn).get_table_names())
|
||||
|
||||
|
||||
def _columns(conn: sa.Connection, table_name: str) -> dict[str, dict]:
|
||||
return {column['name']: column for column in _inspector(conn).get_columns(table_name)}
|
||||
|
||||
|
||||
def _index_names(conn: sa.Connection, table_name: str) -> set[str]:
|
||||
return {index['name'] for index in _inspector(conn).get_indexes(table_name)}
|
||||
|
||||
|
||||
def _unique_column_sets(conn: sa.Connection, table_name: str) -> set[tuple[str, ...]]:
|
||||
inspector = _inspector(conn)
|
||||
result = {
|
||||
tuple(constraint.get('column_names') or ()) for constraint in inspector.get_unique_constraints(table_name)
|
||||
}
|
||||
result.update(
|
||||
tuple(index.get('column_names') or ()) for index in inspector.get_indexes(table_name) if index.get('unique')
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def _foreign_key_exists(
|
||||
conn: sa.Connection,
|
||||
table_name: str,
|
||||
local_columns: tuple[str, ...],
|
||||
referred_table: str,
|
||||
referred_columns: tuple[str, ...],
|
||||
) -> bool:
|
||||
return any(
|
||||
tuple(foreign_key.get('constrained_columns') or ()) == local_columns
|
||||
and foreign_key.get('referred_table') == referred_table
|
||||
and tuple(foreign_key.get('referred_columns') or ()) == referred_columns
|
||||
for foreign_key in _inspector(conn).get_foreign_keys(table_name)
|
||||
)
|
||||
|
||||
|
||||
def _metadata_value(conn: sa.Connection, key: str) -> str | None:
|
||||
if 'metadata' not in _table_names(conn):
|
||||
return None
|
||||
metadata = sa.table(
|
||||
'metadata',
|
||||
sa.column('key', sa.String(255)),
|
||||
sa.column('value', sa.String(255)),
|
||||
)
|
||||
value = conn.execute(sa.select(metadata.c.value).where(metadata.c.key == key)).scalar_one_or_none()
|
||||
return value.strip() if isinstance(value, str) and value.strip() else None
|
||||
|
||||
|
||||
def _default_workspace_uuid(conn: sa.Connection) -> str | None:
|
||||
if 'workspaces' not in _table_names(conn):
|
||||
return None
|
||||
workspaces = sa.table(
|
||||
'workspaces',
|
||||
sa.column('uuid', sa.String(36)),
|
||||
sa.column('instance_uuid', sa.String(255)),
|
||||
sa.column('source', sa.String(32)),
|
||||
)
|
||||
instance_uuid = _metadata_value(conn, 'instance_uuid')
|
||||
query = sa.select(workspaces.c.uuid).where(workspaces.c.source == 'local')
|
||||
if instance_uuid is not None:
|
||||
query = query.where(workspaces.c.instance_uuid == instance_uuid)
|
||||
rows = conn.execute(query).all()
|
||||
if len(rows) > 1:
|
||||
raise RuntimeError('Cannot backfill tenant resources: multiple local Workspaces exist')
|
||||
return rows[0][0] if rows else None
|
||||
|
||||
|
||||
def _upgrade_normalized_email(conn: sa.Connection) -> None:
|
||||
if 'users' not in _table_names(conn):
|
||||
return
|
||||
columns = _columns(conn, 'users')
|
||||
if 'normalized_email' not in columns:
|
||||
op.add_column('users', sa.Column('normalized_email', sa.String(320), nullable=True))
|
||||
|
||||
users = sa.table(
|
||||
'users',
|
||||
sa.column('id', sa.Integer()),
|
||||
sa.column('user', sa.String(255)),
|
||||
sa.column('normalized_email', sa.String(320)),
|
||||
)
|
||||
# Use the exact same normalization algorithm as the runtime. Database
|
||||
# ``lower()`` is ASCII-only on SQLite and is not equivalent to Python
|
||||
# ``casefold()`` (for example, Straße -> strasse). Recompute every row so
|
||||
# an interrupted expand/backfill attempt using an older migration body can
|
||||
# be resumed safely.
|
||||
seen_emails: dict[str, int] = {}
|
||||
for user_id, email in conn.execute(sa.select(users.c.id, users.c.user).order_by(users.c.id)).all():
|
||||
normalized_email = str(email or '').strip().casefold()
|
||||
if not normalized_email:
|
||||
raise RuntimeError(f'Cannot normalize empty account identity for user row {user_id}')
|
||||
if len(normalized_email) > 320:
|
||||
raise RuntimeError(
|
||||
f'Cannot normalize account identity for user row {user_id}: canonical value exceeds 320 characters'
|
||||
)
|
||||
duplicate_user_id = seen_emails.get(normalized_email)
|
||||
if duplicate_user_id is not None:
|
||||
raise RuntimeError(
|
||||
f'Cannot create normalized account identity: user rows '
|
||||
f'{duplicate_user_id} and {user_id} both normalize to {normalized_email!r}'
|
||||
)
|
||||
seen_emails[normalized_email] = user_id
|
||||
conn.execute(users.update().where(users.c.id == user_id).values(normalized_email=normalized_email))
|
||||
|
||||
columns = _columns(conn, 'users')
|
||||
checks = {
|
||||
constraint.get('name'): constraint
|
||||
for constraint in _inspector(conn).get_check_constraints('users')
|
||||
if constraint.get('name') is not None
|
||||
}
|
||||
identity_check = checks.get('ck_users_normalized_email')
|
||||
identity_check_sql = str((identity_check or {}).get('sqltext') or '').casefold()
|
||||
# Python casefold is the canonical identity algorithm. SQL ``lower`` is
|
||||
# dialect/locale dependent (notably Cherokee folds to uppercase in Python
|
||||
# but PostgreSQL lowercases it), so the database validates only portable
|
||||
# structural invariants and uniqueness.
|
||||
replace_legacy_identity_check = identity_check is not None and 'lower' in identity_check_sql
|
||||
needs_contract = columns['normalized_email']['nullable'] or identity_check is None or replace_legacy_identity_check
|
||||
if needs_contract:
|
||||
with op.batch_alter_table('users') as batch_op:
|
||||
if replace_legacy_identity_check:
|
||||
batch_op.drop_constraint('ck_users_normalized_email', type_='check')
|
||||
if columns['normalized_email']['nullable']:
|
||||
batch_op.alter_column(
|
||||
'normalized_email',
|
||||
existing_type=columns['normalized_email']['type'],
|
||||
nullable=False,
|
||||
)
|
||||
if identity_check is None or replace_legacy_identity_check:
|
||||
batch_op.create_check_constraint(
|
||||
'ck_users_normalized_email',
|
||||
'normalized_email = trim(normalized_email) '
|
||||
'AND length(normalized_email) > 0 '
|
||||
'AND length(normalized_email) <= 320',
|
||||
)
|
||||
if 'uq_users_normalized_email' not in _index_names(conn, 'users'):
|
||||
op.create_index('uq_users_normalized_email', 'users', ['normalized_email'], unique=True)
|
||||
|
||||
|
||||
def _api_key_owner(conn: sa.Connection, workspace_uuid: str | None) -> str | None:
|
||||
if workspace_uuid is None or 'workspace_memberships' not in _table_names(conn):
|
||||
return None
|
||||
memberships = sa.table(
|
||||
'workspace_memberships',
|
||||
sa.column('workspace_uuid', sa.String(36)),
|
||||
sa.column('account_uuid', sa.String(36)),
|
||||
sa.column('role', sa.String(32)),
|
||||
)
|
||||
return conn.execute(
|
||||
sa.select(memberships.c.account_uuid)
|
||||
.where(
|
||||
memberships.c.workspace_uuid == workspace_uuid,
|
||||
memberships.c.role == 'owner',
|
||||
)
|
||||
.limit(1)
|
||||
).scalar_one_or_none()
|
||||
|
||||
|
||||
def _expand_and_hash_api_keys(conn: sa.Connection, workspace_uuid: str | None) -> None:
|
||||
if 'api_keys' not in _table_names(conn):
|
||||
return
|
||||
columns = _columns(conn, 'api_keys')
|
||||
additions = (
|
||||
('uuid', sa.Column('uuid', sa.String(36), nullable=True)),
|
||||
('created_by_account_uuid', sa.Column('created_by_account_uuid', sa.String(36), nullable=True)),
|
||||
('key_hash', sa.Column('key_hash', sa.String(64), nullable=True)),
|
||||
('scopes', sa.Column('scopes', sa.JSON(), nullable=True)),
|
||||
('status', sa.Column('status', sa.String(32), nullable=True, server_default='active')),
|
||||
('expires_at', sa.Column('expires_at', sa.DateTime(), nullable=True)),
|
||||
('last_used_at', sa.Column('last_used_at', sa.DateTime(), nullable=True)),
|
||||
)
|
||||
for name, column in additions:
|
||||
if name not in columns:
|
||||
op.add_column('api_keys', column)
|
||||
|
||||
columns = _columns(conn, 'api_keys')
|
||||
api_key_columns = [sa.column('id', sa.Integer())]
|
||||
for name in ('uuid', 'created_by_account_uuid', 'key_hash', 'scopes', 'status', 'key'):
|
||||
if name in columns:
|
||||
api_key_columns.append(sa.column(name, columns[name]['type']))
|
||||
api_keys = sa.table('api_keys', *api_key_columns)
|
||||
owner_uuid = _api_key_owner(conn, workspace_uuid)
|
||||
selected_columns = [api_keys.c.id, api_keys.c.uuid, api_keys.c.key_hash]
|
||||
if 'key' in api_keys.c:
|
||||
selected_columns.append(api_keys.c.key)
|
||||
rows = conn.execute(sa.select(*selected_columns).order_by(api_keys.c.id)).mappings().all()
|
||||
for row in rows:
|
||||
values: dict[str, object] = {}
|
||||
if not row['uuid']:
|
||||
values['uuid'] = str(uuid.uuid4())
|
||||
if not row['key_hash']:
|
||||
plaintext = row.get('key')
|
||||
if not isinstance(plaintext, str) or not plaintext:
|
||||
raise RuntimeError(f'API key row {row["id"]} has no secret to hash')
|
||||
values['key_hash'] = hashlib.sha256(plaintext.encode()).hexdigest()
|
||||
if values:
|
||||
conn.execute(api_keys.update().where(api_keys.c.id == row['id']).values(**values))
|
||||
# Pre-tenancy API keys historically had unrestricted instance access. A
|
||||
# wildcard preserves that behavior while binding it to the backfilled
|
||||
# Workspace; new keys must store their requested explicit scopes.
|
||||
conn.execute(api_keys.update().where(api_keys.c.scopes.is_(None)).values(scopes=['*']))
|
||||
conn.execute(api_keys.update().where(api_keys.c.status.is_(None)).values(status='active'))
|
||||
if owner_uuid is not None:
|
||||
conn.execute(
|
||||
api_keys.update()
|
||||
.where(api_keys.c.created_by_account_uuid.is_(None))
|
||||
.values(created_by_account_uuid=owner_uuid)
|
||||
)
|
||||
|
||||
columns = _columns(conn, 'api_keys')
|
||||
check_names = {constraint.get('name') for constraint in _inspector(conn).get_check_constraints('api_keys')}
|
||||
existing_fks = _inspector(conn).get_foreign_keys('api_keys')
|
||||
creator_fk_exists = any(
|
||||
tuple(foreign_key.get('constrained_columns') or ()) == ('created_by_account_uuid',)
|
||||
and foreign_key.get('referred_table') == 'users'
|
||||
and tuple(foreign_key.get('referred_columns') or ()) == ('uuid',)
|
||||
for foreign_key in existing_fks
|
||||
)
|
||||
has_legacy_key = 'key' in columns
|
||||
needs_contract = (
|
||||
any(columns[name]['nullable'] for name in ('uuid', 'key_hash', 'scopes', 'status'))
|
||||
or 'ck_api_keys_status' not in check_names
|
||||
or not creator_fk_exists
|
||||
or has_legacy_key
|
||||
)
|
||||
if needs_contract:
|
||||
naming = {'uq': 'uq_%(table_name)s_%(column_0_name)s'}
|
||||
with op.batch_alter_table('api_keys', naming_convention=naming) as batch_op:
|
||||
for name in ('uuid', 'key_hash', 'scopes', 'status'):
|
||||
if columns[name]['nullable']:
|
||||
batch_op.alter_column(name, existing_type=columns[name]['type'], nullable=False)
|
||||
if 'ck_api_keys_status' not in check_names:
|
||||
batch_op.create_check_constraint('ck_api_keys_status', "status IN ('active', 'revoked')")
|
||||
if not creator_fk_exists:
|
||||
batch_op.create_foreign_key(
|
||||
'fk_api_keys_created_by_account',
|
||||
'users',
|
||||
['created_by_account_uuid'],
|
||||
['uuid'],
|
||||
ondelete='SET NULL',
|
||||
)
|
||||
if has_legacy_key:
|
||||
# Dropping the column also removes its old global plaintext
|
||||
# unique constraint/index during SQLite's batch rebuild.
|
||||
batch_op.drop_column('key')
|
||||
|
||||
|
||||
def _expand_workspace_columns(conn: sa.Connection, workspace_uuid: str | None) -> None:
|
||||
tables = _table_names(conn)
|
||||
for table_name in _TENANT_TABLES:
|
||||
if table_name not in tables:
|
||||
continue
|
||||
columns = _columns(conn, table_name)
|
||||
if 'workspace_uuid' not in columns:
|
||||
op.add_column(table_name, sa.Column('workspace_uuid', sa.String(36), nullable=True))
|
||||
tenant_table = sa.table(table_name, sa.column('workspace_uuid', sa.String(36)))
|
||||
null_count = conn.scalar(
|
||||
sa.select(sa.func.count()).select_from(tenant_table).where(tenant_table.c.workspace_uuid.is_(None))
|
||||
)
|
||||
if null_count:
|
||||
if workspace_uuid is None:
|
||||
raise RuntimeError(f'Cannot backfill {table_name}: the instance has no unique local Workspace')
|
||||
conn.execute(
|
||||
tenant_table.update()
|
||||
.where(tenant_table.c.workspace_uuid.is_(None))
|
||||
.values(workspace_uuid=workspace_uuid)
|
||||
)
|
||||
|
||||
|
||||
def _mark_legacy_vector_collections(conn: sa.Connection, workspace_uuid: str | None) -> None:
|
||||
"""Persist which pre-tenancy KBs must keep using ``collection_id``.
|
||||
|
||||
The marker is backfilled only when this migration introduces the column,
|
||||
or resumes while that newly added column is still nullable. A fresh
|
||||
schema already contains the non-null column with ``false`` as its default,
|
||||
so knowledge bases created under the scoped-vector contract can never be
|
||||
mistaken for legacy data during a later migration retry.
|
||||
"""
|
||||
|
||||
if 'knowledge_bases' not in _table_names(conn):
|
||||
return
|
||||
columns = _columns(conn, 'knowledge_bases')
|
||||
introduced = 'legacy_vector_collection' not in columns
|
||||
if introduced:
|
||||
op.add_column(
|
||||
'knowledge_bases',
|
||||
sa.Column('legacy_vector_collection', sa.Boolean(), nullable=True),
|
||||
)
|
||||
columns = _columns(conn, 'knowledge_bases')
|
||||
needs_legacy_backfill = introduced or columns['legacy_vector_collection']['nullable']
|
||||
|
||||
knowledge_bases = sa.table(
|
||||
'knowledge_bases',
|
||||
sa.column('collection_id', columns['collection_id']['type']),
|
||||
sa.column('legacy_vector_collection', sa.Boolean()),
|
||||
*((sa.column('workspace_uuid', columns['workspace_uuid']['type']),) if 'workspace_uuid' in columns else ()),
|
||||
)
|
||||
if needs_legacy_backfill and workspace_uuid is not None:
|
||||
legacy_filter = sa.and_(
|
||||
knowledge_bases.c.collection_id.is_not(None),
|
||||
sa.func.length(sa.func.trim(knowledge_bases.c.collection_id)) > 0,
|
||||
)
|
||||
if 'workspace_uuid' in knowledge_bases.c:
|
||||
# A partially migrated database may already have Workspace
|
||||
# columns. Never mark a projected cloud row as legacy.
|
||||
legacy_filter = sa.and_(
|
||||
legacy_filter,
|
||||
sa.or_(
|
||||
knowledge_bases.c.workspace_uuid.is_(None),
|
||||
knowledge_bases.c.workspace_uuid == workspace_uuid,
|
||||
),
|
||||
)
|
||||
conn.execute(knowledge_bases.update().where(legacy_filter).values(legacy_vector_collection=True))
|
||||
|
||||
conn.execute(
|
||||
knowledge_bases.update()
|
||||
.where(knowledge_bases.c.legacy_vector_collection.is_(None))
|
||||
.values(legacy_vector_collection=False)
|
||||
)
|
||||
columns = _columns(conn, 'knowledge_bases')
|
||||
if columns['legacy_vector_collection']['nullable']:
|
||||
with op.batch_alter_table('knowledge_bases') as batch_op:
|
||||
batch_op.alter_column(
|
||||
'legacy_vector_collection',
|
||||
existing_type=columns['legacy_vector_collection']['type'],
|
||||
nullable=False,
|
||||
server_default=sa.false(),
|
||||
)
|
||||
|
||||
|
||||
def _drop_legacy_uniqueness(conn: sa.Connection) -> None:
|
||||
if 'bot_admins' in _table_names(conn):
|
||||
for constraint in _inspector(conn).get_unique_constraints('bot_admins'):
|
||||
if tuple(constraint.get('column_names') or ()) == ('bot_uuid', 'launcher_type', 'launcher_id'):
|
||||
with op.batch_alter_table('bot_admins') as batch_op:
|
||||
batch_op.drop_constraint(constraint['name'], type_='unique')
|
||||
break
|
||||
|
||||
if 'monitoring_feedback' in _table_names(conn):
|
||||
dropped_constraint = False
|
||||
for constraint in _inspector(conn).get_unique_constraints('monitoring_feedback'):
|
||||
if tuple(constraint.get('column_names') or ()) == ('feedback_id',):
|
||||
convention = {'uq': 'uq_%(table_name)s_%(column_0_name)s'}
|
||||
constraint_name = constraint.get('name') or 'uq_monitoring_feedback_feedback_id'
|
||||
with op.batch_alter_table(
|
||||
'monitoring_feedback',
|
||||
naming_convention=convention,
|
||||
) as batch_op:
|
||||
batch_op.drop_constraint(constraint_name, type_='unique')
|
||||
dropped_constraint = True
|
||||
break
|
||||
if not dropped_constraint:
|
||||
for index in _inspector(conn).get_indexes('monitoring_feedback'):
|
||||
if index.get('unique') and tuple(index.get('column_names') or ()) == ('feedback_id',):
|
||||
op.drop_index(index['name'], table_name='monitoring_feedback')
|
||||
|
||||
|
||||
def _create_index_if_missing(
|
||||
conn: sa.Connection,
|
||||
table_name: str,
|
||||
name: str,
|
||||
columns: tuple[str, ...],
|
||||
unique: bool,
|
||||
predicate: sa.TextClause | None,
|
||||
) -> None:
|
||||
if name in _index_names(conn, table_name):
|
||||
return
|
||||
if unique and predicate is None and columns in _unique_column_sets(conn, table_name):
|
||||
return
|
||||
kwargs = {}
|
||||
if predicate is not None:
|
||||
kwargs = {'sqlite_where': predicate, 'postgresql_where': predicate}
|
||||
op.create_index(name, table_name, list(columns), unique=unique, **kwargs)
|
||||
|
||||
|
||||
def _validate_scoped_unique_data(conn: sa.Connection) -> None:
|
||||
checks = (
|
||||
('mcp_servers', ('workspace_uuid', 'name')),
|
||||
('knowledge_bases', ('workspace_uuid', 'collection_id')),
|
||||
)
|
||||
for table_name, column_names in checks:
|
||||
if table_name not in _table_names(conn):
|
||||
continue
|
||||
columns = _columns(conn, table_name)
|
||||
if not all(column_name in columns for column_name in column_names):
|
||||
continue
|
||||
table = sa.table(
|
||||
table_name,
|
||||
*(sa.column(column_name, columns[column_name]['type']) for column_name in column_names),
|
||||
)
|
||||
group_columns = [table.c[column_name] for column_name in column_names]
|
||||
query = sa.select(*group_columns, sa.func.count()).group_by(*group_columns).having(sa.func.count() > 1)
|
||||
if column_names[-1] == 'collection_id':
|
||||
query = query.where(group_columns[-1].is_not(None))
|
||||
duplicate = conn.execute(query.limit(1)).first()
|
||||
if duplicate is not None:
|
||||
raise RuntimeError(
|
||||
f'Cannot create scoped unique key on {table_name}{column_names}: duplicate {duplicate!r}'
|
||||
)
|
||||
|
||||
|
||||
def _create_parent_and_scoped_indexes(conn: sa.Connection) -> None:
|
||||
_validate_scoped_unique_data(conn)
|
||||
tables = _table_names(conn)
|
||||
for table_name, indexes in _SCOPED_INDEXES.items():
|
||||
if table_name not in tables:
|
||||
continue
|
||||
available_columns = _columns(conn, table_name)
|
||||
for name, columns, unique, predicate in indexes:
|
||||
if all(column in available_columns for column in columns):
|
||||
_create_index_if_missing(conn, table_name, name, columns, unique, predicate)
|
||||
|
||||
|
||||
def _contract_table(conn: sa.Connection, table_name: str) -> None:
|
||||
columns = _columns(conn, table_name)
|
||||
if 'workspace_uuid' not in columns:
|
||||
return
|
||||
direct_workspace_fk = _foreign_key_exists(
|
||||
conn,
|
||||
table_name,
|
||||
('workspace_uuid',),
|
||||
'workspaces',
|
||||
('uuid',),
|
||||
)
|
||||
current_pk = tuple(_inspector(conn).get_pk_constraint(table_name).get('constrained_columns') or ())
|
||||
desired_pk = _COMPOSITE_PRIMARY_KEYS.get(table_name)
|
||||
missing_composite_fks = [
|
||||
foreign_key
|
||||
for foreign_key in _COMPOSITE_FOREIGN_KEYS.get(table_name, ())
|
||||
if not _foreign_key_exists(
|
||||
conn,
|
||||
table_name,
|
||||
foreign_key[1],
|
||||
foreign_key[2],
|
||||
foreign_key[3],
|
||||
)
|
||||
]
|
||||
needs_contract = (
|
||||
columns['workspace_uuid']['nullable']
|
||||
or not direct_workspace_fk
|
||||
or (desired_pk is not None and current_pk != desired_pk)
|
||||
or bool(missing_composite_fks)
|
||||
)
|
||||
if not needs_contract:
|
||||
return
|
||||
|
||||
naming = {
|
||||
'pk': 'pk_%(table_name)s',
|
||||
'fk': 'fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s',
|
||||
}
|
||||
pk_name = _inspector(conn).get_pk_constraint(table_name).get('name') or f'pk_{table_name}'
|
||||
with op.batch_alter_table(table_name, naming_convention=naming) as batch_op:
|
||||
if columns['workspace_uuid']['nullable']:
|
||||
batch_op.alter_column(
|
||||
'workspace_uuid',
|
||||
existing_type=columns['workspace_uuid']['type'],
|
||||
nullable=False,
|
||||
)
|
||||
if desired_pk is not None and current_pk != desired_pk:
|
||||
batch_op.drop_constraint(pk_name, type_='primary')
|
||||
batch_op.create_primary_key(f'pk_{table_name}', list(desired_pk))
|
||||
if not direct_workspace_fk:
|
||||
batch_op.create_foreign_key(
|
||||
f'fk_{table_name}_workspace',
|
||||
'workspaces',
|
||||
['workspace_uuid'],
|
||||
['uuid'],
|
||||
ondelete='CASCADE',
|
||||
)
|
||||
for name, local_columns, referred_table, referred_columns, ondelete in missing_composite_fks:
|
||||
batch_op.create_foreign_key(
|
||||
name,
|
||||
referred_table,
|
||||
list(local_columns),
|
||||
list(referred_columns),
|
||||
ondelete=ondelete,
|
||||
)
|
||||
|
||||
|
||||
def _contract_workspace_columns(conn: sa.Connection) -> None:
|
||||
tables = _table_names(conn)
|
||||
# Parents must be contracted before their children so SQLite can validate
|
||||
# the exact composite target key during a batch-table rebuild.
|
||||
order = (
|
||||
'api_keys',
|
||||
'bots',
|
||||
'bot_admins',
|
||||
'binary_storages',
|
||||
'mcp_servers',
|
||||
'model_providers',
|
||||
'llm_models',
|
||||
'embedding_models',
|
||||
'rerank_models',
|
||||
'legacy_pipelines',
|
||||
'pipeline_run_records',
|
||||
'plugin_settings',
|
||||
'knowledge_bases',
|
||||
'knowledge_base_files',
|
||||
'knowledge_base_chunks',
|
||||
'webhooks',
|
||||
'monitoring_messages',
|
||||
'monitoring_llm_calls',
|
||||
'monitoring_tool_calls',
|
||||
'monitoring_sessions',
|
||||
'monitoring_errors',
|
||||
'monitoring_embedding_calls',
|
||||
'monitoring_feedback',
|
||||
)
|
||||
for table_name in order:
|
||||
if table_name in tables:
|
||||
_contract_table(conn, table_name)
|
||||
|
||||
|
||||
def _migrate_workspace_metadata(conn: sa.Connection, workspace_uuid: str | None) -> None:
|
||||
tables = _table_names(conn)
|
||||
if 'workspaces' not in tables:
|
||||
return
|
||||
if 'workspace_metadata' not in tables:
|
||||
op.create_table(
|
||||
'workspace_metadata',
|
||||
sa.Column('workspace_uuid', sa.String(36), nullable=False),
|
||||
sa.Column('key', sa.String(255), nullable=False),
|
||||
sa.Column('value', sa.String(255), nullable=True),
|
||||
sa.ForeignKeyConstraint(
|
||||
['workspace_uuid'],
|
||||
['workspaces.uuid'],
|
||||
name='fk_workspace_metadata_workspace',
|
||||
ondelete='CASCADE',
|
||||
),
|
||||
sa.PrimaryKeyConstraint('workspace_uuid', 'key', name='pk_workspace_metadata'),
|
||||
)
|
||||
if workspace_uuid is None or 'metadata' not in tables:
|
||||
return
|
||||
metadata = sa.table(
|
||||
'metadata',
|
||||
sa.column('key', sa.String(255)),
|
||||
sa.column('value', sa.String(255)),
|
||||
)
|
||||
workspace_metadata = sa.table(
|
||||
'workspace_metadata',
|
||||
sa.column('workspace_uuid', sa.String(36)),
|
||||
sa.column('key', sa.String(255)),
|
||||
sa.column('value', sa.String(255)),
|
||||
)
|
||||
tenant_keys = ('wizard_status', 'wizard_progress', 'rag_plugin_migration_needed')
|
||||
rows = conn.execute(sa.select(metadata.c.key, metadata.c.value).where(metadata.c.key.in_(tenant_keys))).all()
|
||||
for key, value in rows:
|
||||
exists = conn.execute(
|
||||
sa.select(workspace_metadata.c.key).where(
|
||||
workspace_metadata.c.workspace_uuid == workspace_uuid,
|
||||
workspace_metadata.c.key == key,
|
||||
)
|
||||
).first()
|
||||
if exists is None:
|
||||
conn.execute(
|
||||
workspace_metadata.insert().values(
|
||||
workspace_uuid=workspace_uuid,
|
||||
key=key,
|
||||
value=value,
|
||||
)
|
||||
)
|
||||
if rows:
|
||||
conn.execute(metadata.delete().where(metadata.c.key.in_(tenant_keys)))
|
||||
|
||||
|
||||
def _validate_contract(conn: sa.Connection) -> None:
|
||||
for table_name in _TENANT_TABLES:
|
||||
if table_name not in _table_names(conn):
|
||||
continue
|
||||
columns = _columns(conn, table_name)
|
||||
if 'workspace_uuid' not in columns or columns['workspace_uuid']['nullable']:
|
||||
raise RuntimeError(f'{table_name}.workspace_uuid was not contracted to NOT NULL')
|
||||
table = sa.table(table_name, sa.column('workspace_uuid', sa.String(36)))
|
||||
if conn.scalar(sa.select(sa.func.count()).select_from(table).where(table.c.workspace_uuid.is_(None))):
|
||||
raise RuntimeError(f'{table_name} still contains unscoped rows')
|
||||
if conn.dialect.name == 'sqlite':
|
||||
violations = conn.execute(sa.text('PRAGMA foreign_key_check')).all()
|
||||
if violations:
|
||||
raise RuntimeError(f'SQLite foreign key validation failed: {violations[:5]!r}')
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
_upgrade_normalized_email(conn)
|
||||
workspace_uuid = _default_workspace_uuid(conn)
|
||||
_mark_legacy_vector_collections(conn, workspace_uuid)
|
||||
_expand_workspace_columns(conn, workspace_uuid)
|
||||
_expand_and_hash_api_keys(conn, workspace_uuid)
|
||||
_drop_legacy_uniqueness(conn)
|
||||
_create_parent_and_scoped_indexes(conn)
|
||||
_contract_workspace_columns(conn)
|
||||
_migrate_workspace_metadata(conn, workspace_uuid)
|
||||
_validate_contract(conn)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
raise RuntimeError(
|
||||
'0010_scope_resources is intentionally irreversible because plaintext API key secrets were securely removed'
|
||||
)
|
||||
@@ -0,0 +1,201 @@
|
||||
"""enforce PostgreSQL tenant isolation with exact discovery contracts
|
||||
|
||||
Revision ID: 0011_postgres_tenant_rls
|
||||
Revises: 0010_scope_resources
|
||||
Create Date: 2026-07-19
|
||||
|
||||
The table and policy lists are deliberately duplicated from the runtime
|
||||
contract. Alembic revisions must remain self-contained after application code
|
||||
evolves. Discovery policies are SELECT-only and reveal the minimum rows needed
|
||||
to turn an authenticated credential into one Workspace transaction.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision = '0011_postgres_tenant_rls'
|
||||
down_revision = '0010_scope_resources'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
_POLICY_NAME = 'langbot_workspace_isolation'
|
||||
_ACCOUNT_POLICY_NAME = 'langbot_account_discovery'
|
||||
_API_KEY_POLICY_NAME = 'langbot_api_key_discovery'
|
||||
_INVITATION_POLICY_NAME = 'langbot_invitation_discovery'
|
||||
_INSTANCE_POLICY_NAME = 'langbot_instance_discovery'
|
||||
|
||||
_TENANT_SETTING = 'langbot.workspace_uuid'
|
||||
_ACCOUNT_SETTING = 'langbot.account_uuid'
|
||||
_API_KEY_HASH_SETTING = 'langbot.api_key_hash'
|
||||
_INVITATION_HASH_SETTING = 'langbot.invitation_hash'
|
||||
_INSTANCE_SETTING = 'langbot.instance_uuid'
|
||||
_OSS_WORKSPACE_METADATA_KEY = 'oss_workspace_uuid'
|
||||
|
||||
_TENANT_TABLE_COLUMNS: dict[str, str] = {
|
||||
'workspaces': 'uuid',
|
||||
'workspace_memberships': 'workspace_uuid',
|
||||
'workspace_invitations': 'workspace_uuid',
|
||||
'workspace_execution_states': 'workspace_uuid',
|
||||
'workspace_metadata': 'workspace_uuid',
|
||||
'api_keys': 'workspace_uuid',
|
||||
'bots': 'workspace_uuid',
|
||||
'bot_admins': 'workspace_uuid',
|
||||
'binary_storages': 'workspace_uuid',
|
||||
'mcp_servers': 'workspace_uuid',
|
||||
'model_providers': 'workspace_uuid',
|
||||
'llm_models': 'workspace_uuid',
|
||||
'embedding_models': 'workspace_uuid',
|
||||
'rerank_models': 'workspace_uuid',
|
||||
'legacy_pipelines': 'workspace_uuid',
|
||||
'pipeline_run_records': 'workspace_uuid',
|
||||
'plugin_settings': 'workspace_uuid',
|
||||
'knowledge_bases': 'workspace_uuid',
|
||||
'knowledge_base_files': 'workspace_uuid',
|
||||
'knowledge_base_chunks': 'workspace_uuid',
|
||||
'webhooks': 'workspace_uuid',
|
||||
'monitoring_messages': 'workspace_uuid',
|
||||
'monitoring_llm_calls': 'workspace_uuid',
|
||||
'monitoring_tool_calls': 'workspace_uuid',
|
||||
'monitoring_sessions': 'workspace_uuid',
|
||||
'monitoring_errors': 'workspace_uuid',
|
||||
'monitoring_embedding_calls': 'workspace_uuid',
|
||||
'monitoring_feedback': 'workspace_uuid',
|
||||
}
|
||||
|
||||
|
||||
def _setting(name: str) -> str:
|
||||
return f"NULLIF(current_setting('{name}', true), '')"
|
||||
|
||||
|
||||
def _tenant_expression(column: str) -> str:
|
||||
return f'{column}::text = {_setting(_TENANT_SETTING)}'
|
||||
|
||||
|
||||
_DISCOVERY_POLICIES: dict[str, dict[str, str]] = {
|
||||
'workspace_memberships': {
|
||||
_ACCOUNT_POLICY_NAME: (f"account_uuid::text = {_setting(_ACCOUNT_SETTING)} AND status = 'active'"),
|
||||
},
|
||||
'workspace_execution_states': {
|
||||
_INSTANCE_POLICY_NAME: (
|
||||
f"instance_uuid = {_setting(_INSTANCE_SETTING)} AND state = 'active' AND write_fenced = false"
|
||||
),
|
||||
},
|
||||
'api_keys': {
|
||||
_API_KEY_POLICY_NAME: (
|
||||
f"key_hash = {_setting(_API_KEY_HASH_SETTING)} AND status = 'active' "
|
||||
'AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)'
|
||||
),
|
||||
},
|
||||
'workspace_invitations': {
|
||||
_INVITATION_POLICY_NAME: f'token_hash = {_setting(_INVITATION_HASH_SETTING)}',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _quote_identifier(conn: sa.Connection, identifier: str) -> str:
|
||||
return conn.dialect.identifier_preparer.quote(identifier)
|
||||
|
||||
|
||||
def _record_oss_workspace_scope(conn: sa.Connection) -> None:
|
||||
"""Keep PostgreSQL OSS usable after FORCE RLS is enabled."""
|
||||
|
||||
local_workspaces = (
|
||||
conn.execute(sa.text("SELECT uuid FROM workspaces WHERE source = 'local' ORDER BY uuid")).scalars().all()
|
||||
)
|
||||
if len(local_workspaces) != 1:
|
||||
return
|
||||
|
||||
existing = conn.execute(
|
||||
sa.text('SELECT value FROM metadata WHERE key = :key'),
|
||||
{'key': _OSS_WORKSPACE_METADATA_KEY},
|
||||
).scalar_one_or_none()
|
||||
if existing is None:
|
||||
conn.execute(
|
||||
sa.text('INSERT INTO metadata (key, value) VALUES (:key, :value)'),
|
||||
{'key': _OSS_WORKSPACE_METADATA_KEY, 'value': local_workspaces[0]},
|
||||
)
|
||||
elif existing != local_workspaces[0]:
|
||||
raise RuntimeError('Stored OSS Workspace scope does not match the local Workspace')
|
||||
|
||||
|
||||
def _drop_all_policies(conn: sa.Connection, table_name: str) -> None:
|
||||
policy_names = conn.execute(
|
||||
sa.text(
|
||||
"""
|
||||
SELECT p.polname
|
||||
FROM pg_policy p
|
||||
JOIN pg_class c ON c.oid = p.polrelid
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE n.nspname = current_schema() AND c.relname = :table_name
|
||||
"""
|
||||
),
|
||||
{'table_name': table_name},
|
||||
).scalars()
|
||||
table = _quote_identifier(conn, table_name)
|
||||
for policy in policy_names:
|
||||
op.execute(sa.text(f'DROP POLICY {_quote_identifier(conn, policy)} ON {table}'))
|
||||
|
||||
|
||||
def _create_policy(
|
||||
conn: sa.Connection,
|
||||
table_name: str,
|
||||
policy_name: str,
|
||||
expression: str,
|
||||
*,
|
||||
command: str,
|
||||
) -> None:
|
||||
table = _quote_identifier(conn, table_name)
|
||||
policy = _quote_identifier(conn, policy_name)
|
||||
if command == 'ALL':
|
||||
sql = f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR ALL TO PUBLIC USING ({expression}) WITH CHECK ({expression})'
|
||||
elif command == 'SELECT':
|
||||
sql = f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR SELECT TO PUBLIC USING ({expression})'
|
||||
else: # pragma: no cover - migration-local invariant
|
||||
raise AssertionError(f'Unsupported RLS command: {command}')
|
||||
op.execute(sa.text(sql))
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
if conn.dialect.name != 'postgresql':
|
||||
return
|
||||
|
||||
existing_tables = set(sa.inspect(conn).get_table_names())
|
||||
missing_tables = set(_TENANT_TABLE_COLUMNS) - existing_tables
|
||||
if missing_tables:
|
||||
raise RuntimeError(f'Cannot enable tenant RLS before all tenant-owned tables exist: {sorted(missing_tables)!r}')
|
||||
|
||||
_record_oss_workspace_scope(conn)
|
||||
|
||||
for table_name, tenant_column in _TENANT_TABLE_COLUMNS.items():
|
||||
table = _quote_identifier(conn, table_name)
|
||||
_drop_all_policies(conn, table_name)
|
||||
op.execute(sa.text(f'ALTER TABLE {table} ENABLE ROW LEVEL SECURITY'))
|
||||
op.execute(sa.text(f'ALTER TABLE {table} FORCE ROW LEVEL SECURITY'))
|
||||
_create_policy(
|
||||
conn,
|
||||
table_name,
|
||||
_POLICY_NAME,
|
||||
_tenant_expression(_quote_identifier(conn, tenant_column)),
|
||||
command='ALL',
|
||||
)
|
||||
for policy_name, expression in _DISCOVERY_POLICIES.get(table_name, {}).items():
|
||||
_create_policy(conn, table_name, policy_name, expression, command='SELECT')
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
if conn.dialect.name != 'postgresql':
|
||||
return
|
||||
|
||||
existing_tables = set(sa.inspect(conn).get_table_names())
|
||||
for table_name in _TENANT_TABLE_COLUMNS:
|
||||
if table_name not in existing_tables:
|
||||
continue
|
||||
table = _quote_identifier(conn, table_name)
|
||||
_drop_all_policies(conn, table_name)
|
||||
op.execute(sa.text(f'ALTER TABLE {table} NO FORCE ROW LEVEL SECURITY'))
|
||||
op.execute(sa.text(f'ALTER TABLE {table} DISABLE ROW LEVEL SECURITY'))
|
||||
@@ -0,0 +1,179 @@
|
||||
"""add immutable plugin installation identity
|
||||
|
||||
Revision ID: 0012_plugin_identity
|
||||
Revises: 0011_postgres_tenant_rls
|
||||
Create Date: 2026-07-19
|
||||
|
||||
The migration gives every legacy row a random, stable installation UUID. A
|
||||
legacy artifact digest is only a valid SHA-256-shaped recovery marker; Core
|
||||
replaces it with the package digest and increments ``runtime_revision`` before
|
||||
the next package apply.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import uuid
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
|
||||
revision = '0012_plugin_identity'
|
||||
down_revision = '0011_postgres_tenant_rls'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
_TABLE = 'plugin_settings'
|
||||
_INSTALLATION_INDEX = 'ix_plugin_settings_workspace_installation'
|
||||
_INSTALLATION_UNIQUE = 'uq_plugin_settings_installation_uuid'
|
||||
_REVISION_CHECK = 'ck_plugin_settings_runtime_revision_positive'
|
||||
_DIGEST_CHECK = 'ck_plugin_settings_artifact_digest_length'
|
||||
|
||||
|
||||
def _column_names(conn: sa.Connection) -> set[str]:
|
||||
inspector = sa.inspect(conn)
|
||||
if _TABLE not in inspector.get_table_names():
|
||||
return set()
|
||||
return {column['name'] for column in inspector.get_columns(_TABLE)}
|
||||
|
||||
|
||||
def _legacy_digest(installation_uuid: str) -> str:
|
||||
return hashlib.sha256(f'legacy-installation:{installation_uuid}'.encode()).hexdigest()
|
||||
|
||||
|
||||
def _suspend_postgres_rls(conn: sa.Connection) -> tuple[bool, bool]:
|
||||
"""Let the release migration backfill every tenant row after revision 0011."""
|
||||
|
||||
if conn.dialect.name != 'postgresql':
|
||||
return False, False
|
||||
row = conn.execute(
|
||||
sa.text('SELECT relrowsecurity, relforcerowsecurity FROM pg_class WHERE oid = to_regclass(:table_name)'),
|
||||
{'table_name': _TABLE},
|
||||
).one()
|
||||
rls_enabled, rls_forced = bool(row.relrowsecurity), bool(row.relforcerowsecurity)
|
||||
if rls_forced:
|
||||
op.execute(sa.text(f'ALTER TABLE {_TABLE} NO FORCE ROW LEVEL SECURITY'))
|
||||
if rls_enabled:
|
||||
op.execute(sa.text(f'ALTER TABLE {_TABLE} DISABLE ROW LEVEL SECURITY'))
|
||||
return rls_enabled, rls_forced
|
||||
|
||||
|
||||
def _restore_postgres_rls(conn: sa.Connection, state: tuple[bool, bool]) -> None:
|
||||
if conn.dialect.name != 'postgresql':
|
||||
return
|
||||
rls_enabled, rls_forced = state
|
||||
if rls_enabled:
|
||||
op.execute(sa.text(f'ALTER TABLE {_TABLE} ENABLE ROW LEVEL SECURITY'))
|
||||
if rls_forced:
|
||||
op.execute(sa.text(f'ALTER TABLE {_TABLE} FORCE ROW LEVEL SECURITY'))
|
||||
|
||||
|
||||
def _backfill(conn: sa.Connection) -> None:
|
||||
table = sa.table(
|
||||
_TABLE,
|
||||
sa.column('workspace_uuid', sa.String(36)),
|
||||
sa.column('plugin_author', sa.String(255)),
|
||||
sa.column('plugin_name', sa.String(255)),
|
||||
sa.column('installation_uuid', sa.String(36)),
|
||||
sa.column('artifact_digest', sa.String(64)),
|
||||
sa.column('runtime_revision', sa.Integer()),
|
||||
)
|
||||
rows = conn.execute(
|
||||
sa.select(
|
||||
table.c.workspace_uuid,
|
||||
table.c.plugin_author,
|
||||
table.c.plugin_name,
|
||||
table.c.installation_uuid,
|
||||
table.c.artifact_digest,
|
||||
table.c.runtime_revision,
|
||||
)
|
||||
).all()
|
||||
for row in rows:
|
||||
installation_uuid = str(row.installation_uuid or uuid.uuid4())
|
||||
values: dict[str, object] = {}
|
||||
if not row.installation_uuid:
|
||||
values['installation_uuid'] = installation_uuid
|
||||
if not row.artifact_digest or len(str(row.artifact_digest)) != 64:
|
||||
values['artifact_digest'] = _legacy_digest(installation_uuid)
|
||||
if row.runtime_revision is None or row.runtime_revision < 1:
|
||||
values['runtime_revision'] = 1
|
||||
if values:
|
||||
conn.execute(
|
||||
table.update()
|
||||
.where(table.c.workspace_uuid == row.workspace_uuid)
|
||||
.where(table.c.plugin_author == row.plugin_author)
|
||||
.where(table.c.plugin_name == row.plugin_name)
|
||||
.values(**values)
|
||||
)
|
||||
|
||||
|
||||
def _constraint_names(conn: sa.Connection, kind: str) -> set[str]:
|
||||
inspector = sa.inspect(conn)
|
||||
getter = inspector.get_unique_constraints if kind == 'unique' else inspector.get_check_constraints
|
||||
return {str(item.get('name')) for item in getter(_TABLE) if item.get('name')}
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
columns = _column_names(conn)
|
||||
if not columns:
|
||||
return
|
||||
|
||||
if 'installation_uuid' not in columns:
|
||||
op.add_column(_TABLE, sa.Column('installation_uuid', sa.String(36), nullable=True))
|
||||
if 'artifact_digest' not in columns:
|
||||
op.add_column(_TABLE, sa.Column('artifact_digest', sa.String(64), nullable=True))
|
||||
if 'runtime_revision' not in columns:
|
||||
op.add_column(_TABLE, sa.Column('runtime_revision', sa.Integer(), nullable=True))
|
||||
|
||||
rls_state = _suspend_postgres_rls(conn)
|
||||
try:
|
||||
_backfill(conn)
|
||||
finally:
|
||||
_restore_postgres_rls(conn, rls_state)
|
||||
|
||||
# Fresh databases are created from current metadata before Alembic runs;
|
||||
# guards make the revision safe for that path and for interrupted upgrades.
|
||||
indexes = {index['name'] for index in sa.inspect(conn).get_indexes(_TABLE)}
|
||||
unique_constraints = _constraint_names(conn, 'unique')
|
||||
check_constraints = _constraint_names(conn, 'check')
|
||||
with op.batch_alter_table(_TABLE) as batch:
|
||||
batch.alter_column('installation_uuid', existing_type=sa.String(36), nullable=False)
|
||||
batch.alter_column('artifact_digest', existing_type=sa.String(64), nullable=False)
|
||||
batch.alter_column('runtime_revision', existing_type=sa.Integer(), nullable=False)
|
||||
if _REVISION_CHECK not in check_constraints:
|
||||
batch.create_check_constraint(_REVISION_CHECK, 'runtime_revision >= 1')
|
||||
if _DIGEST_CHECK not in check_constraints:
|
||||
batch.create_check_constraint(_DIGEST_CHECK, 'length(artifact_digest) = 64')
|
||||
if _INSTALLATION_UNIQUE not in unique_constraints:
|
||||
batch.create_unique_constraint(_INSTALLATION_UNIQUE, ['installation_uuid'])
|
||||
if _INSTALLATION_INDEX not in indexes and _INSTALLATION_INDEX not in unique_constraints:
|
||||
batch.create_index(
|
||||
_INSTALLATION_INDEX,
|
||||
['workspace_uuid', 'installation_uuid'],
|
||||
unique=True,
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
columns = _column_names(conn)
|
||||
if not columns:
|
||||
return
|
||||
indexes = {index['name'] for index in sa.inspect(conn).get_indexes(_TABLE)}
|
||||
checks = _constraint_names(conn, 'check')
|
||||
uniques = _constraint_names(conn, 'unique')
|
||||
with op.batch_alter_table(_TABLE) as batch:
|
||||
if _INSTALLATION_INDEX in indexes:
|
||||
batch.drop_index(_INSTALLATION_INDEX)
|
||||
if _DIGEST_CHECK in checks:
|
||||
batch.drop_constraint(_DIGEST_CHECK, type_='check')
|
||||
if _REVISION_CHECK in checks:
|
||||
batch.drop_constraint(_REVISION_CHECK, type_='check')
|
||||
if _INSTALLATION_UNIQUE in uniques:
|
||||
batch.drop_constraint(_INSTALLATION_UNIQUE, type_='unique')
|
||||
for column_name in ('runtime_revision', 'artifact_digest', 'installation_uuid'):
|
||||
if column_name in columns:
|
||||
batch.drop_column(column_name)
|
||||
@@ -0,0 +1,382 @@
|
||||
"""create tenant-scoped pgvector storage in the business database
|
||||
|
||||
Revision ID: 0013_tenant_pgvector
|
||||
Revises: 0012_plugin_identity
|
||||
Create Date: 2026-07-19
|
||||
|
||||
The Cloud application role never executes this DDL. A release migration role
|
||||
installs pgvector once, creates an untyped vector column, and builds a bounded
|
||||
set of expression/partial ANN indexes. Existing legacy rows are migrated only
|
||||
when each row maps unambiguously to one knowledge base.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import typing
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from pgvector.sqlalchemy import Vector
|
||||
|
||||
|
||||
revision = '0013_tenant_pgvector'
|
||||
down_revision = '0012_plugin_identity'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
_VECTOR_TABLE = 'langbot_vectors'
|
||||
_LEGACY_TABLE = 'langbot_vectors_legacy_0013'
|
||||
_TENANT_POLICY = 'langbot_workspace_isolation'
|
||||
_TENANT_SETTING = 'langbot.workspace_uuid'
|
||||
_KB_DIMENSION_CHECK = 'ck_knowledge_bases_embedding_dimension_positive'
|
||||
_VECTOR_DIMENSION_CHECK = 'ck_langbot_vectors_embedding_dimension'
|
||||
_VECTOR_ALLOWED_CHECK = 'ck_langbot_vectors_embedding_dimension_enabled'
|
||||
_ALLOWED_DIMENSIONS = (384, 512, 768, 1024, 1536)
|
||||
_LEGACY_SOURCE_TABLES = (
|
||||
'knowledge_bases',
|
||||
'knowledge_base_files',
|
||||
'knowledge_base_chunks',
|
||||
)
|
||||
|
||||
|
||||
def _quote(conn: sa.Connection, identifier: str) -> str:
|
||||
return conn.dialect.identifier_preparer.quote(identifier)
|
||||
|
||||
|
||||
def _columns(conn: sa.Connection, table_name: str) -> set[str]:
|
||||
inspector = sa.inspect(conn)
|
||||
if table_name not in inspector.get_table_names():
|
||||
return set()
|
||||
return {column['name'] for column in inspector.get_columns(table_name)}
|
||||
|
||||
|
||||
def _checks(conn: sa.Connection, table_name: str) -> set[str]:
|
||||
return {str(item['name']) for item in sa.inspect(conn).get_check_constraints(table_name) if item.get('name')}
|
||||
|
||||
|
||||
def _ensure_knowledge_base_dimension(conn: sa.Connection) -> None:
|
||||
columns = _columns(conn, 'knowledge_bases')
|
||||
if 'embedding_dimension' not in columns:
|
||||
op.add_column('knowledge_bases', sa.Column('embedding_dimension', sa.Integer(), nullable=True))
|
||||
if _KB_DIMENSION_CHECK not in _checks(conn, 'knowledge_bases'):
|
||||
with op.batch_alter_table('knowledge_bases') as batch:
|
||||
batch.create_check_constraint(
|
||||
_KB_DIMENSION_CHECK,
|
||||
'embedding_dimension IS NULL OR embedding_dimension > 0',
|
||||
)
|
||||
|
||||
|
||||
def _create_vector_table() -> None:
|
||||
enabled = ', '.join(str(item) for item in _ALLOWED_DIMENSIONS)
|
||||
op.create_table(
|
||||
_VECTOR_TABLE,
|
||||
sa.Column('workspace_uuid', sa.String(36), nullable=False),
|
||||
sa.Column('knowledge_base_uuid', sa.String(255), nullable=False),
|
||||
sa.Column('vector_id', sa.String(255), nullable=False),
|
||||
sa.Column('embedding_dimension', sa.Integer(), nullable=False),
|
||||
sa.Column('embedding', Vector(), nullable=False),
|
||||
sa.Column('text', sa.Text(), nullable=True),
|
||||
sa.Column('file_id', sa.String(255), nullable=True),
|
||||
sa.Column('chunk_uuid', sa.String(255), nullable=True),
|
||||
sa.PrimaryKeyConstraint(
|
||||
'workspace_uuid',
|
||||
'knowledge_base_uuid',
|
||||
'vector_id',
|
||||
name='pk_langbot_vectors',
|
||||
),
|
||||
sa.ForeignKeyConstraint(
|
||||
['workspace_uuid', 'knowledge_base_uuid'],
|
||||
['knowledge_bases.workspace_uuid', 'knowledge_bases.uuid'],
|
||||
name='fk_langbot_vectors_workspace_kb',
|
||||
ondelete='CASCADE',
|
||||
),
|
||||
sa.CheckConstraint(
|
||||
'vector_dims(embedding) = embedding_dimension',
|
||||
name=_VECTOR_DIMENSION_CHECK,
|
||||
),
|
||||
sa.CheckConstraint(
|
||||
f'embedding_dimension IN ({enabled})',
|
||||
name=_VECTOR_ALLOWED_CHECK,
|
||||
),
|
||||
)
|
||||
op.create_index(
|
||||
'ix_langbot_vectors_workspace_kb_file',
|
||||
_VECTOR_TABLE,
|
||||
['workspace_uuid', 'knowledge_base_uuid', 'file_id'],
|
||||
)
|
||||
op.create_index(
|
||||
'ix_langbot_vectors_workspace_kb_chunk',
|
||||
_VECTOR_TABLE,
|
||||
['workspace_uuid', 'knowledge_base_uuid', 'chunk_uuid'],
|
||||
)
|
||||
|
||||
|
||||
def _legacy_mapping_predicate() -> str:
|
||||
return """
|
||||
kb.collection_id = legacy.collection
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM knowledge_base_files AS files
|
||||
LEFT JOIN knowledge_base_chunks AS chunks
|
||||
ON chunks.workspace_uuid = files.workspace_uuid
|
||||
AND chunks.file_id = files.uuid
|
||||
WHERE files.workspace_uuid = kb.workspace_uuid
|
||||
AND files.kb_id = kb.uuid
|
||||
AND (files.uuid = legacy.file_id OR chunks.uuid = legacy.chunk_uuid)
|
||||
)
|
||||
"""
|
||||
|
||||
|
||||
def _legacy_source_rls_states(conn: sa.Connection) -> dict[str, tuple[bool, bool]]:
|
||||
rows = (
|
||||
conn.execute(
|
||||
sa.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
|
||||
AND c.relkind IN ('r', 'p')
|
||||
"""
|
||||
).bindparams(sa.bindparam('table_names', expanding=True)),
|
||||
{'table_names': _LEGACY_SOURCE_TABLES},
|
||||
)
|
||||
.mappings()
|
||||
.all()
|
||||
)
|
||||
states = {str(row['relname']): (bool(row['relrowsecurity']), bool(row['relforcerowsecurity'])) for row in rows}
|
||||
missing = set(_LEGACY_SOURCE_TABLES) - set(states)
|
||||
if missing:
|
||||
raise RuntimeError(f'Legacy pgvector source tables are missing: {sorted(missing)!r}')
|
||||
return states
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _suspend_legacy_source_rls(conn: sa.Connection) -> typing.Iterator[None]:
|
||||
"""Temporarily let the table-owning migrator map all legacy tenant rows.
|
||||
|
||||
Revision 0011 enables and forces RLS on each source table. The release
|
||||
migrator intentionally has neither superuser nor BYPASSRLS, so even a table
|
||||
owner cannot read those rows until FORCE RLS is paused. Preserve both flags
|
||||
independently and restore them in ``finally`` so mixed pre-existing states
|
||||
survive successful, rejected, and interrupted legacy migrations.
|
||||
"""
|
||||
|
||||
states = _legacy_source_rls_states(conn)
|
||||
try:
|
||||
for table_name in _LEGACY_SOURCE_TABLES:
|
||||
table = _quote(conn, table_name)
|
||||
conn.execute(sa.text(f'ALTER TABLE {table} NO FORCE ROW LEVEL SECURITY'))
|
||||
conn.execute(sa.text(f'ALTER TABLE {table} DISABLE ROW LEVEL SECURITY'))
|
||||
yield
|
||||
finally:
|
||||
for table_name in _LEGACY_SOURCE_TABLES:
|
||||
table = _quote(conn, table_name)
|
||||
rls_enabled, rls_forced = states[table_name]
|
||||
enabled_clause = 'ENABLE' if rls_enabled else 'DISABLE'
|
||||
forced_clause = 'FORCE' if rls_forced else 'NO FORCE'
|
||||
conn.execute(sa.text(f'ALTER TABLE {table} {enabled_clause} ROW LEVEL SECURITY'))
|
||||
conn.execute(sa.text(f'ALTER TABLE {table} {forced_clause} ROW LEVEL SECURITY'))
|
||||
|
||||
|
||||
def _migrate_legacy_rows(conn: sa.Connection) -> None:
|
||||
predicate = _legacy_mapping_predicate()
|
||||
ambiguous = conn.execute(
|
||||
sa.text(
|
||||
f"""
|
||||
WITH candidates AS (
|
||||
SELECT legacy.id, kb.workspace_uuid, kb.uuid AS knowledge_base_uuid
|
||||
FROM {_LEGACY_TABLE} AS legacy
|
||||
JOIN knowledge_bases AS kb ON ({predicate})
|
||||
), candidate_counts AS (
|
||||
SELECT id, COUNT(*) AS count
|
||||
FROM candidates
|
||||
GROUP BY id
|
||||
)
|
||||
SELECT legacy.id, COALESCE(candidate_counts.count, 0) AS candidate_count
|
||||
FROM {_LEGACY_TABLE} AS legacy
|
||||
LEFT JOIN candidate_counts ON candidate_counts.id = legacy.id
|
||||
WHERE COALESCE(candidate_counts.count, 0) <> 1
|
||||
LIMIT 1
|
||||
"""
|
||||
)
|
||||
).first()
|
||||
if ambiguous is not None:
|
||||
raise RuntimeError(
|
||||
'Legacy pgvector row cannot be mapped to exactly one Workspace/knowledge base: '
|
||||
f'{ambiguous.id!r} has {ambiguous.candidate_count} candidates'
|
||||
)
|
||||
|
||||
conn.execute(
|
||||
sa.text(
|
||||
f"""
|
||||
INSERT INTO {_VECTOR_TABLE} (
|
||||
workspace_uuid,
|
||||
knowledge_base_uuid,
|
||||
vector_id,
|
||||
embedding_dimension,
|
||||
embedding,
|
||||
text,
|
||||
file_id,
|
||||
chunk_uuid
|
||||
)
|
||||
SELECT
|
||||
kb.workspace_uuid,
|
||||
kb.uuid,
|
||||
legacy.id,
|
||||
vector_dims(legacy.embedding),
|
||||
legacy.embedding,
|
||||
legacy.text,
|
||||
legacy.file_id,
|
||||
legacy.chunk_uuid
|
||||
FROM {_LEGACY_TABLE} AS legacy
|
||||
JOIN knowledge_bases AS kb ON ({predicate})
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
mixed_dimension = conn.execute(
|
||||
sa.text(
|
||||
f"""
|
||||
SELECT workspace_uuid, knowledge_base_uuid
|
||||
FROM {_VECTOR_TABLE}
|
||||
GROUP BY workspace_uuid, knowledge_base_uuid
|
||||
HAVING MIN(embedding_dimension) <> MAX(embedding_dimension)
|
||||
LIMIT 1
|
||||
"""
|
||||
)
|
||||
).first()
|
||||
if mixed_dimension is not None:
|
||||
raise RuntimeError('Legacy knowledge base contains mixed embedding dimensions')
|
||||
|
||||
conn.execute(
|
||||
sa.text(
|
||||
f"""
|
||||
UPDATE knowledge_bases AS kb
|
||||
SET embedding_dimension = dimensions.embedding_dimension
|
||||
FROM (
|
||||
SELECT workspace_uuid, knowledge_base_uuid, MIN(embedding_dimension) AS embedding_dimension
|
||||
FROM {_VECTOR_TABLE}
|
||||
GROUP BY workspace_uuid, knowledge_base_uuid
|
||||
) AS dimensions
|
||||
WHERE kb.workspace_uuid = dimensions.workspace_uuid
|
||||
AND kb.uuid = dimensions.knowledge_base_uuid
|
||||
AND kb.embedding_dimension IS NULL
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _drop_all_policies(conn: sa.Connection) -> None:
|
||||
table = _quote(conn, _VECTOR_TABLE)
|
||||
policies = conn.execute(
|
||||
sa.text(
|
||||
"""
|
||||
SELECT p.polname
|
||||
FROM pg_policy p
|
||||
JOIN pg_class c ON c.oid = p.polrelid
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE n.nspname = current_schema() AND c.relname = :table_name
|
||||
"""
|
||||
),
|
||||
{'table_name': _VECTOR_TABLE},
|
||||
).scalars()
|
||||
for policy_name in policies:
|
||||
op.execute(sa.text(f'DROP POLICY {_quote(conn, policy_name)} ON {table}'))
|
||||
|
||||
|
||||
def _enable_rls(conn: sa.Connection) -> None:
|
||||
table = _quote(conn, _VECTOR_TABLE)
|
||||
policy = _quote(conn, _TENANT_POLICY)
|
||||
expression = f"workspace_uuid::text = NULLIF(current_setting('{_TENANT_SETTING}', true), '')"
|
||||
_drop_all_policies(conn)
|
||||
op.execute(sa.text(f'ALTER TABLE {table} ENABLE ROW LEVEL SECURITY'))
|
||||
op.execute(sa.text(f'ALTER TABLE {table} FORCE ROW LEVEL SECURITY'))
|
||||
op.execute(
|
||||
sa.text(
|
||||
f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR ALL TO PUBLIC '
|
||||
f'USING ({expression}) WITH CHECK ({expression})'
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _create_ann_indexes(conn: sa.Connection) -> None:
|
||||
table = _quote(conn, _VECTOR_TABLE)
|
||||
for dimension in _ALLOWED_DIMENSIONS:
|
||||
index = _quote(conn, f'ix_langbot_vectors_hnsw_cosine_{dimension}')
|
||||
op.execute(
|
||||
sa.text(
|
||||
f'CREATE INDEX {index} ON {table} USING hnsw '
|
||||
f'((embedding::vector({dimension})) vector_cosine_ops) '
|
||||
f'WHERE embedding_dimension = {dimension}'
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
if 'knowledge_bases' not in sa.inspect(conn).get_table_names():
|
||||
if conn.dialect.name == 'postgresql':
|
||||
# The supported PostgreSQL release path creates the business
|
||||
# tables before stamping 0010 and reaching this migration. Missing
|
||||
# knowledge_bases therefore means the operator bypassed the
|
||||
# release bootstrap or the schema is incomplete; stamping head in
|
||||
# that state would make Cloud runtime validation unrecoverable.
|
||||
raise RuntimeError('PostgreSQL release migration requires the knowledge_bases table')
|
||||
# A direct empty SQLite Alembic walk is still used by migration tooling;
|
||||
# Core creates the complete ORM schema on its following compatibility
|
||||
# pass, including the portable embedding_dimension field.
|
||||
return
|
||||
|
||||
_ensure_knowledge_base_dimension(conn)
|
||||
|
||||
# pgvector storage is PostgreSQL-only, but ``embedding_dimension`` is an
|
||||
# ORM field used by both deployment modes. Existing OSS SQLite databases
|
||||
# must receive the column before this revision becomes a no-op.
|
||||
if conn.dialect.name != 'postgresql':
|
||||
return
|
||||
|
||||
op.execute(sa.text('CREATE EXTENSION IF NOT EXISTS vector'))
|
||||
columns = _columns(conn, _VECTOR_TABLE)
|
||||
if columns:
|
||||
scoped_columns = {
|
||||
'workspace_uuid',
|
||||
'knowledge_base_uuid',
|
||||
'vector_id',
|
||||
'embedding_dimension',
|
||||
'embedding',
|
||||
}
|
||||
if scoped_columns.issubset(columns):
|
||||
raise RuntimeError('Tenant pgvector table exists before its owning release migration')
|
||||
legacy_columns = {'id', 'collection', 'embedding'}
|
||||
if not legacy_columns.issubset(columns):
|
||||
raise RuntimeError('Existing pgvector table has an unsupported schema')
|
||||
if _LEGACY_TABLE in sa.inspect(conn).get_table_names():
|
||||
raise RuntimeError(f'Interrupted pgvector migration left {_LEGACY_TABLE!r} behind')
|
||||
op.rename_table(_VECTOR_TABLE, _LEGACY_TABLE)
|
||||
|
||||
_create_vector_table()
|
||||
if _LEGACY_TABLE in sa.inspect(conn).get_table_names():
|
||||
with _suspend_legacy_source_rls(conn):
|
||||
_migrate_legacy_rows(conn)
|
||||
op.drop_table(_LEGACY_TABLE)
|
||||
_create_ann_indexes(conn)
|
||||
_enable_rls(conn)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
if conn.dialect.name == 'postgresql' and _VECTOR_TABLE in sa.inspect(conn).get_table_names():
|
||||
_drop_all_policies(conn)
|
||||
op.drop_table(_VECTOR_TABLE)
|
||||
|
||||
columns = _columns(conn, 'knowledge_bases')
|
||||
if 'embedding_dimension' in columns:
|
||||
checks = _checks(conn, 'knowledge_bases')
|
||||
with op.batch_alter_table('knowledge_bases') as batch:
|
||||
if _KB_DIMENSION_CHECK in checks:
|
||||
batch.drop_constraint(_KB_DIMENSION_CHECK, type_='check')
|
||||
batch.drop_column('embedding_dimension')
|
||||
@@ -0,0 +1,268 @@
|
||||
"""add the Cloud directory projection persistence boundary
|
||||
|
||||
Revision ID: 0014_cloud_directory
|
||||
Revises: 0013_tenant_pgvector
|
||||
Create Date: 2026-07-24
|
||||
|
||||
The open Core projector receives already-verified control-plane data and is the
|
||||
only runtime path allowed to mutate projected Workspace directory rows. Its
|
||||
transaction-local instance setting is intentionally distinct from both normal
|
||||
Workspace scope and the read-only instance discovery scope.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
|
||||
revision = '0014_cloud_directory'
|
||||
down_revision = '0013_tenant_pgvector'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
_STATE_TABLE = 'directory_projection_states'
|
||||
_INBOX_TABLE = 'directory_projection_inbox'
|
||||
_DIRECTORY_POLICY_NAME = 'langbot_directory_projection'
|
||||
_TENANT_POLICY_NAME = 'langbot_workspace_isolation'
|
||||
_LOCAL_WRITE_POLICY_NAME = 'langbot_workspace_local_directory_write'
|
||||
_DIRECTORY_SETTING = 'langbot.directory_instance_uuid'
|
||||
_TENANT_SETTING = 'langbot.workspace_uuid'
|
||||
_PROJECTED_TENANT_TABLES = (
|
||||
'workspaces',
|
||||
'workspace_memberships',
|
||||
'workspace_execution_states',
|
||||
)
|
||||
|
||||
|
||||
def _setting(name: str) -> str:
|
||||
return f"NULLIF(current_setting('{name}', true), '')"
|
||||
|
||||
|
||||
def _quote(conn: sa.Connection, identifier: str) -> str:
|
||||
return conn.dialect.identifier_preparer.quote(identifier)
|
||||
|
||||
|
||||
def _create_tables(conn: sa.Connection) -> None:
|
||||
existing_tables = set(sa.inspect(conn).get_table_names())
|
||||
if _STATE_TABLE not in existing_tables:
|
||||
op.create_table(
|
||||
_STATE_TABLE,
|
||||
sa.Column('instance_uuid', sa.String(255), nullable=False),
|
||||
sa.Column('cursor', sa.BigInteger(), server_default='0', nullable=False),
|
||||
sa.Column('snapshot_coverage_cursor', sa.BigInteger(), server_default='0', nullable=False),
|
||||
sa.Column('snapshot_fingerprint', sa.Text(), nullable=False),
|
||||
sa.Column('last_applied_at', sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column('lease_expires_at', sa.DateTime(timezone=True), nullable=True),
|
||||
sa.CheckConstraint(
|
||||
'cursor >= 0',
|
||||
name='ck_directory_projection_state_cursor',
|
||||
),
|
||||
sa.CheckConstraint(
|
||||
'snapshot_coverage_cursor >= 0 AND snapshot_coverage_cursor <= cursor',
|
||||
name='ck_directory_projection_state_snapshot_coverage',
|
||||
),
|
||||
sa.CheckConstraint(
|
||||
'length(snapshot_fingerprint) = 64',
|
||||
name='ck_directory_projection_state_fingerprint',
|
||||
),
|
||||
sa.PrimaryKeyConstraint('instance_uuid'),
|
||||
)
|
||||
if _INBOX_TABLE not in existing_tables:
|
||||
op.create_table(
|
||||
_INBOX_TABLE,
|
||||
sa.Column('instance_uuid', sa.String(255), nullable=False),
|
||||
sa.Column('event_uuid', sa.String(36), nullable=False),
|
||||
sa.Column('cursor', sa.BigInteger(), nullable=False),
|
||||
sa.Column('event_type', sa.String(128), nullable=False),
|
||||
sa.Column('revision', sa.BigInteger(), nullable=False),
|
||||
sa.Column('fingerprint', sa.Text(), nullable=False),
|
||||
sa.Column(
|
||||
'received_at',
|
||||
sa.DateTime(timezone=True),
|
||||
server_default=sa.func.now(),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column('applied_at', sa.DateTime(timezone=True), nullable=True),
|
||||
sa.CheckConstraint(
|
||||
'cursor > 0',
|
||||
name='ck_directory_projection_inbox_cursor',
|
||||
),
|
||||
sa.CheckConstraint(
|
||||
'revision > 0',
|
||||
name='ck_directory_projection_inbox_revision',
|
||||
),
|
||||
sa.CheckConstraint(
|
||||
'length(fingerprint) = 64',
|
||||
name='ck_directory_projection_inbox_fingerprint',
|
||||
),
|
||||
sa.PrimaryKeyConstraint('instance_uuid', 'event_uuid'),
|
||||
sa.UniqueConstraint(
|
||||
'instance_uuid',
|
||||
'cursor',
|
||||
name='uq_directory_projection_inbox_cursor',
|
||||
),
|
||||
)
|
||||
op.create_index(
|
||||
'ix_directory_projection_inbox_pending',
|
||||
_INBOX_TABLE,
|
||||
['instance_uuid', 'applied_at', 'cursor'],
|
||||
unique=False,
|
||||
)
|
||||
|
||||
|
||||
def _drop_policy(conn: sa.Connection, table_name: str, policy_name: str) -> None:
|
||||
table = _quote(conn, table_name)
|
||||
policy = _quote(conn, policy_name)
|
||||
op.execute(sa.text(f'DROP POLICY IF EXISTS {policy} ON {table}'))
|
||||
|
||||
|
||||
def _create_policy(
|
||||
conn: sa.Connection,
|
||||
table_name: str,
|
||||
policy_name: str,
|
||||
expression: str,
|
||||
*,
|
||||
command: str = 'ALL',
|
||||
) -> None:
|
||||
table = _quote(conn, table_name)
|
||||
policy = _quote(conn, policy_name)
|
||||
_drop_policy(conn, table_name, policy_name)
|
||||
op.execute(sa.text(f'ALTER TABLE {table} ENABLE ROW LEVEL SECURITY'))
|
||||
op.execute(sa.text(f'ALTER TABLE {table} FORCE ROW LEVEL SECURITY'))
|
||||
if command == 'SELECT':
|
||||
sql = f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR SELECT TO PUBLIC USING ({expression})'
|
||||
elif command == 'ALL':
|
||||
sql = (
|
||||
f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR ALL TO PUBLIC '
|
||||
f'USING ({expression}) WITH CHECK ({expression})'
|
||||
)
|
||||
else: # pragma: no cover - migration-local invariant.
|
||||
raise AssertionError(f'Unsupported RLS policy command: {command}')
|
||||
op.execute(sa.text(sql))
|
||||
|
||||
|
||||
def _install_postgres_policies(conn: sa.Connection) -> None:
|
||||
existing_tables = set(sa.inspect(conn).get_table_names())
|
||||
required_tables = set(_PROJECTED_TENANT_TABLES) | {_STATE_TABLE, _INBOX_TABLE}
|
||||
missing_tables = required_tables - existing_tables
|
||||
if missing_tables:
|
||||
raise RuntimeError(
|
||||
f'Cannot enable Cloud directory projection RLS before all required tables exist: {sorted(missing_tables)!r}'
|
||||
)
|
||||
|
||||
directory_setting = _setting(_DIRECTORY_SETTING)
|
||||
tenant_setting = _setting(_TENANT_SETTING)
|
||||
directory_expressions = {
|
||||
'workspaces': (f"instance_uuid::text = {directory_setting} AND source = 'cloud_projection'"),
|
||||
'workspace_memberships': (
|
||||
'EXISTS ('
|
||||
'SELECT 1 FROM workspaces AS directory_workspace '
|
||||
'WHERE directory_workspace.uuid = workspace_memberships.workspace_uuid '
|
||||
f'AND directory_workspace.instance_uuid::text = {directory_setting} '
|
||||
"AND directory_workspace.source = 'cloud_projection'"
|
||||
')'
|
||||
),
|
||||
'workspace_execution_states': (
|
||||
f"instance_uuid::text = {directory_setting} AND source = 'cloud' AND EXISTS ("
|
||||
'SELECT 1 FROM workspaces AS directory_workspace '
|
||||
'WHERE directory_workspace.uuid = workspace_execution_states.workspace_uuid '
|
||||
f'AND directory_workspace.instance_uuid::text = {directory_setting} '
|
||||
"AND directory_workspace.source = 'cloud_projection'"
|
||||
')'
|
||||
),
|
||||
_STATE_TABLE: f'instance_uuid::text = {directory_setting}',
|
||||
_INBOX_TABLE: f'instance_uuid::text = {directory_setting}',
|
||||
}
|
||||
tenant_expressions = {
|
||||
'workspaces': f'uuid::text = {tenant_setting}',
|
||||
'workspace_memberships': f'workspace_uuid::text = {tenant_setting}',
|
||||
'workspace_execution_states': f'workspace_uuid::text = {tenant_setting}',
|
||||
}
|
||||
local_write_expressions = {
|
||||
'workspaces': f"uuid::text = {tenant_setting} AND source = 'local'",
|
||||
'workspace_memberships': (
|
||||
f'workspace_uuid::text = {tenant_setting} AND EXISTS ('
|
||||
'SELECT 1 FROM workspaces AS local_workspace '
|
||||
'WHERE local_workspace.uuid = workspace_memberships.workspace_uuid '
|
||||
"AND local_workspace.source = 'local'"
|
||||
')'
|
||||
),
|
||||
'workspace_execution_states': (
|
||||
f'workspace_uuid::text = {tenant_setting} AND EXISTS ('
|
||||
'SELECT 1 FROM workspaces AS local_workspace '
|
||||
'WHERE local_workspace.uuid = workspace_execution_states.workspace_uuid '
|
||||
"AND local_workspace.source = 'local'"
|
||||
')'
|
||||
),
|
||||
}
|
||||
for table_name in _PROJECTED_TENANT_TABLES:
|
||||
_create_policy(
|
||||
conn,
|
||||
table_name,
|
||||
_TENANT_POLICY_NAME,
|
||||
tenant_expressions[table_name],
|
||||
command='SELECT',
|
||||
)
|
||||
_create_policy(
|
||||
conn,
|
||||
table_name,
|
||||
_LOCAL_WRITE_POLICY_NAME,
|
||||
local_write_expressions[table_name],
|
||||
)
|
||||
_create_policy(
|
||||
conn,
|
||||
table_name,
|
||||
_DIRECTORY_POLICY_NAME,
|
||||
directory_expressions[table_name],
|
||||
)
|
||||
for table_name in (_STATE_TABLE, _INBOX_TABLE):
|
||||
_create_policy(
|
||||
conn,
|
||||
table_name,
|
||||
_DIRECTORY_POLICY_NAME,
|
||||
directory_expressions[table_name],
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
_create_tables(conn)
|
||||
if conn.dialect.name == 'postgresql':
|
||||
_install_postgres_policies(conn)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
existing_tables = set(sa.inspect(conn).get_table_names())
|
||||
if conn.dialect.name == 'postgresql':
|
||||
tenant_setting = _setting(_TENANT_SETTING)
|
||||
tenant_columns = {
|
||||
'workspaces': 'uuid',
|
||||
'workspace_memberships': 'workspace_uuid',
|
||||
'workspace_execution_states': 'workspace_uuid',
|
||||
}
|
||||
for table_name in _PROJECTED_TENANT_TABLES:
|
||||
if table_name not in existing_tables:
|
||||
continue
|
||||
_drop_policy(conn, table_name, _DIRECTORY_POLICY_NAME)
|
||||
_drop_policy(conn, table_name, _LOCAL_WRITE_POLICY_NAME)
|
||||
_create_policy(
|
||||
conn,
|
||||
table_name,
|
||||
_TENANT_POLICY_NAME,
|
||||
f'{tenant_columns[table_name]}::text = {tenant_setting}',
|
||||
)
|
||||
for table_name in (_STATE_TABLE, _INBOX_TABLE):
|
||||
if table_name not in existing_tables:
|
||||
continue
|
||||
_drop_policy(conn, table_name, _DIRECTORY_POLICY_NAME)
|
||||
table = _quote(conn, table_name)
|
||||
op.execute(sa.text(f'ALTER TABLE {table} NO FORCE ROW LEVEL SECURITY'))
|
||||
op.execute(sa.text(f'ALTER TABLE {table} DISABLE ROW LEVEL SECURITY'))
|
||||
|
||||
if _INBOX_TABLE in existing_tables:
|
||||
op.drop_table(_INBOX_TABLE)
|
||||
if _STATE_TABLE in existing_tables:
|
||||
op.drop_table(_STATE_TABLE)
|
||||
@@ -0,0 +1,75 @@
|
||||
"""allow Core-owned collaboration writes on Cloud Workspaces
|
||||
|
||||
Revision ID: 0015_cloud_core_collab
|
||||
Revises: 0014_cloud_directory
|
||||
Create Date: 2026-07-26
|
||||
|
||||
Cloud-projected Workspace identity remains projected by the directory
|
||||
boundary, but membership role/remove and invitation acceptance are now owned
|
||||
by Core tenant scope.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
|
||||
revision = '0015_cloud_core_collab'
|
||||
down_revision = '0014_cloud_directory'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
_TABLE_NAME = 'workspace_memberships'
|
||||
_POLICY_NAME = 'langbot_workspace_local_directory_write'
|
||||
_TENANT_SETTING = 'langbot.workspace_uuid'
|
||||
|
||||
|
||||
def _setting(name: str) -> str:
|
||||
return f"NULLIF(current_setting('{name}', true), '')"
|
||||
|
||||
|
||||
def _quote(conn: sa.Connection, identifier: str) -> str:
|
||||
return conn.dialect.identifier_preparer.quote(identifier)
|
||||
|
||||
|
||||
def _drop_policy(conn: sa.Connection) -> None:
|
||||
table = _quote(conn, _TABLE_NAME)
|
||||
policy = _quote(conn, _POLICY_NAME)
|
||||
op.execute(sa.text(f'DROP POLICY IF EXISTS {policy} ON {table}'))
|
||||
|
||||
|
||||
def _create_policy(conn: sa.Connection, expression: str) -> None:
|
||||
table = _quote(conn, _TABLE_NAME)
|
||||
policy = _quote(conn, _POLICY_NAME)
|
||||
op.execute(
|
||||
sa.text(
|
||||
f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR ALL TO PUBLIC '
|
||||
f'USING ({expression}) WITH CHECK ({expression})'
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
if conn.dialect.name != 'postgresql':
|
||||
return
|
||||
expression = f'workspace_uuid::text = {_setting(_TENANT_SETTING)}'
|
||||
_drop_policy(conn)
|
||||
_create_policy(conn, expression)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
if conn.dialect.name != 'postgresql':
|
||||
return
|
||||
expression = (
|
||||
f'workspace_uuid::text = {_setting(_TENANT_SETTING)} AND EXISTS ('
|
||||
'SELECT 1 FROM workspaces AS local_workspace '
|
||||
'WHERE local_workspace.uuid = workspace_memberships.workspace_uuid '
|
||||
"AND local_workspace.source = 'local'"
|
||||
')'
|
||||
)
|
||||
_drop_policy(conn)
|
||||
_create_policy(conn, expression)
|
||||
@@ -18,6 +18,7 @@ from typing import TYPE_CHECKING
|
||||
from alembic.config import Config
|
||||
from alembic import command
|
||||
from alembic.runtime.migration import MigrationContext
|
||||
from alembic.script import ScriptDirectory
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
@@ -47,12 +48,28 @@ def _do_stamp(connection: Connection, revision: str = 'head') -> None:
|
||||
command.stamp(cfg, revision)
|
||||
|
||||
|
||||
def _do_downgrade(connection: Connection, revision: str) -> None:
|
||||
"""Synchronous downgrade — runs inside run_sync."""
|
||||
cfg = _build_config(connection)
|
||||
command.downgrade(cfg, revision)
|
||||
|
||||
|
||||
def _do_get_current(connection: Connection) -> str | None:
|
||||
"""Get current alembic revision synchronously."""
|
||||
ctx = MigrationContext.configure(connection)
|
||||
return ctx.get_current_revision()
|
||||
|
||||
|
||||
def get_alembic_head() -> str:
|
||||
"""Resolve the single release head without opening a database connection."""
|
||||
cfg = Config()
|
||||
cfg.set_main_option('script_location', _ALEMBIC_DIR)
|
||||
head = ScriptDirectory.from_config(cfg).get_current_head()
|
||||
if head is None:
|
||||
raise RuntimeError('Alembic has no migration head')
|
||||
return head
|
||||
|
||||
|
||||
def _do_autogenerate(connection: Connection, message: str = 'auto migration') -> None:
|
||||
"""Synchronous autogenerate — runs inside run_sync."""
|
||||
cfg = _build_config(connection)
|
||||
@@ -73,6 +90,13 @@ async def run_alembic_stamp(async_engine: AsyncEngine, revision: str = 'head') -
|
||||
await conn.commit()
|
||||
|
||||
|
||||
async def run_alembic_downgrade(async_engine: AsyncEngine, revision: str) -> None:
|
||||
"""Run Alembic downgrade to the given revision."""
|
||||
async with async_engine.connect() as conn:
|
||||
await conn.run_sync(_do_downgrade, revision)
|
||||
await conn.commit()
|
||||
|
||||
|
||||
async def get_alembic_current(async_engine: AsyncEngine) -> str | None:
|
||||
"""Get current alembic revision, or None if not stamped."""
|
||||
async with async_engine.connect() as conn:
|
||||
@@ -121,6 +145,7 @@ if __name__ == '__main__':
|
||||
print('Commands:')
|
||||
print(' autogenerate "message" — Generate migration from ORM model diff')
|
||||
print(' upgrade [revision] — Upgrade database (default: head)')
|
||||
print(' downgrade <revision> — Downgrade database to a revision')
|
||||
print(' stamp [revision] — Stamp revision without running (default: head)')
|
||||
print(' current — Show current revision')
|
||||
sys.exit(1)
|
||||
@@ -140,6 +165,13 @@ if __name__ == '__main__':
|
||||
rev = sys.argv[2] if len(sys.argv) > 2 else 'head'
|
||||
asyncio.run(run_alembic_stamp(engine, rev))
|
||||
print(f'Stamped: {rev}')
|
||||
elif cmd == 'downgrade':
|
||||
if len(sys.argv) < 3:
|
||||
print('Usage: python -m langbot.pkg.persistence.alembic_runner downgrade <revision>')
|
||||
sys.exit(1)
|
||||
rev = sys.argv[2]
|
||||
asyncio.run(run_alembic_downgrade(engine, rev))
|
||||
print(f'Downgraded to: {rev}')
|
||||
elif cmd == 'current':
|
||||
rev = asyncio.run(get_alembic_current(engine))
|
||||
print(f'Current revision: {rev}')
|
||||
|
||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import abc
|
||||
|
||||
import sqlalchemy
|
||||
import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
|
||||
|
||||
from ..core import app
|
||||
@@ -30,8 +31,15 @@ class BaseDatabaseManager(abc.ABC):
|
||||
|
||||
engine: sqlalchemy_asyncio.AsyncEngine
|
||||
|
||||
def __init__(self, ap: app.Application) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
ap: app.Application,
|
||||
*,
|
||||
url_override: sqlalchemy.engine.URL | None = None,
|
||||
) -> None:
|
||||
self.ap = ap
|
||||
self.url_override = url_override
|
||||
self.persistence_mode: str | None = None
|
||||
|
||||
@abc.abstractmethod
|
||||
async def initialize(self) -> None:
|
||||
|
||||
@@ -1,21 +1,167 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy
|
||||
import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
|
||||
|
||||
from .. import database
|
||||
from ..postgresql_url import normalize_asyncpg_url
|
||||
|
||||
|
||||
MAX_POOL_CONNECTIONS = 100
|
||||
MAX_POOL_TIMEOUT_SECONDS = 300
|
||||
MAX_POOL_RECYCLE_SECONDS = 86_400
|
||||
MAX_STATEMENT_TIMEOUT_MS = 300_000
|
||||
MAX_LOCK_TIMEOUT_MS = 60_000
|
||||
MAX_IDLE_TRANSACTION_TIMEOUT_MS = 300_000
|
||||
|
||||
|
||||
@database.manager_class('postgresql')
|
||||
class PostgreSQLDatabaseManager(database.BaseDatabaseManager):
|
||||
"""PostgreSQL database manager"""
|
||||
|
||||
async def initialize(self) -> None:
|
||||
postgresql_config = self.ap.instance_config.data.get('database', {}).get('postgresql', {})
|
||||
@staticmethod
|
||||
def _pool_integer(
|
||||
config: dict,
|
||||
name: str,
|
||||
default: int,
|
||||
*,
|
||||
minimum: int,
|
||||
maximum: int,
|
||||
) -> int:
|
||||
value = config.get(name, default)
|
||||
if isinstance(value, bool) or not isinstance(value, int) or not minimum <= value <= maximum:
|
||||
comparator = 'non-negative' if minimum == 0 else 'positive'
|
||||
raise ValueError(f'database.postgresql.{name} must be a {comparator} integer no greater than {maximum}')
|
||||
return value
|
||||
|
||||
host = postgresql_config.get('host', '127.0.0.1')
|
||||
port = postgresql_config.get('port', 5432)
|
||||
user = postgresql_config.get('user', 'postgres')
|
||||
password = postgresql_config.get('password', 'postgres')
|
||||
database = postgresql_config.get('database', 'postgres')
|
||||
engine_url = f'postgresql+asyncpg://{user}:{password}@{host}:{port}/{database}'
|
||||
self.engine = sqlalchemy_asyncio.create_async_engine(engine_url)
|
||||
async def initialize(self) -> None:
|
||||
self._pool_timeouts_total = 0
|
||||
postgresql_config = self.ap.instance_config.data.get('database', {}).get('postgresql', {})
|
||||
if not isinstance(postgresql_config, dict):
|
||||
raise ValueError('database.postgresql must be an object')
|
||||
if self.url_override is not None:
|
||||
engine_url = self.url_override
|
||||
else:
|
||||
explicit_url = postgresql_config.get('url')
|
||||
if explicit_url:
|
||||
if not isinstance(explicit_url, str):
|
||||
raise ValueError('database.postgresql.url must be a string')
|
||||
try:
|
||||
engine_url = sqlalchemy.engine.make_url(explicit_url)
|
||||
except Exception:
|
||||
raise ValueError('database.postgresql.url is invalid') from None
|
||||
try:
|
||||
engine_url = normalize_asyncpg_url(engine_url)
|
||||
except ValueError:
|
||||
raise ValueError('database.postgresql.url must use valid PostgreSQL asyncpg options') from None
|
||||
else:
|
||||
engine_url = sqlalchemy.URL.create(
|
||||
'postgresql+asyncpg',
|
||||
username=postgresql_config.get('user', 'postgres'),
|
||||
password=postgresql_config.get('password', 'postgres'),
|
||||
host=postgresql_config.get('host', '127.0.0.1'),
|
||||
port=postgresql_config.get('port', 5432),
|
||||
database=postgresql_config.get('database', 'postgres'),
|
||||
)
|
||||
self.pool_size = self._pool_integer(
|
||||
postgresql_config,
|
||||
'pool_size',
|
||||
10,
|
||||
minimum=1,
|
||||
maximum=MAX_POOL_CONNECTIONS,
|
||||
)
|
||||
self.max_overflow = self._pool_integer(
|
||||
postgresql_config,
|
||||
'max_overflow',
|
||||
10,
|
||||
minimum=0,
|
||||
maximum=MAX_POOL_CONNECTIONS,
|
||||
)
|
||||
if self.pool_size + self.max_overflow > MAX_POOL_CONNECTIONS:
|
||||
raise ValueError(f'database.postgresql pool_size + max_overflow must not exceed {MAX_POOL_CONNECTIONS}')
|
||||
self.pool_timeout_seconds = self._pool_integer(
|
||||
postgresql_config,
|
||||
'pool_timeout_seconds',
|
||||
30,
|
||||
minimum=1,
|
||||
maximum=MAX_POOL_TIMEOUT_SECONDS,
|
||||
)
|
||||
self.pool_recycle_seconds = self._pool_integer(
|
||||
postgresql_config,
|
||||
'pool_recycle_seconds',
|
||||
1800,
|
||||
minimum=1,
|
||||
maximum=MAX_POOL_RECYCLE_SECONDS,
|
||||
)
|
||||
connect_args = {}
|
||||
self.statement_timeout_ms = 0
|
||||
self.lock_timeout_ms = 0
|
||||
self.idle_transaction_timeout_ms = 0
|
||||
if self.persistence_mode == 'cloud_runtime':
|
||||
self.statement_timeout_ms = self._pool_integer(
|
||||
postgresql_config,
|
||||
'statement_timeout_ms',
|
||||
60_000,
|
||||
minimum=1,
|
||||
maximum=MAX_STATEMENT_TIMEOUT_MS,
|
||||
)
|
||||
self.lock_timeout_ms = self._pool_integer(
|
||||
postgresql_config,
|
||||
'lock_timeout_ms',
|
||||
5_000,
|
||||
minimum=1,
|
||||
maximum=MAX_LOCK_TIMEOUT_MS,
|
||||
)
|
||||
self.idle_transaction_timeout_ms = self._pool_integer(
|
||||
postgresql_config,
|
||||
'idle_in_transaction_session_timeout_ms',
|
||||
60_000,
|
||||
minimum=1,
|
||||
maximum=MAX_IDLE_TRANSACTION_TIMEOUT_MS,
|
||||
)
|
||||
connect_args = {
|
||||
'server_settings': {
|
||||
'statement_timeout': str(self.statement_timeout_ms),
|
||||
'lock_timeout': str(self.lock_timeout_ms),
|
||||
'idle_in_transaction_session_timeout': str(self.idle_transaction_timeout_ms),
|
||||
}
|
||||
}
|
||||
self.engine = sqlalchemy_asyncio.create_async_engine(
|
||||
engine_url,
|
||||
pool_size=self.pool_size,
|
||||
max_overflow=self.max_overflow,
|
||||
pool_timeout=self.pool_timeout_seconds,
|
||||
pool_recycle=self.pool_recycle_seconds,
|
||||
pool_pre_ping=True,
|
||||
**({'connect_args': connect_args} if connect_args else {}),
|
||||
)
|
||||
|
||||
def resource_stats(self) -> dict[str, int]:
|
||||
"""Return aggregate pool gauges without exposing connection details."""
|
||||
|
||||
pool = self.engine.pool
|
||||
|
||||
def read(name: str) -> int:
|
||||
method = getattr(pool, name, None)
|
||||
if not callable(method):
|
||||
return 0
|
||||
try:
|
||||
return int(method())
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
return {
|
||||
'configured_size': self.pool_size,
|
||||
'configured_max_overflow': self.max_overflow,
|
||||
'configured_capacity': self.pool_size + self.max_overflow,
|
||||
'statement_timeout_ms': self.statement_timeout_ms,
|
||||
'lock_timeout_ms': self.lock_timeout_ms,
|
||||
'idle_in_transaction_session_timeout_ms': self.idle_transaction_timeout_ms,
|
||||
'timeouts_total': self._pool_timeouts_total,
|
||||
'checked_in': read('checkedin'),
|
||||
'checked_out': read('checkedout'),
|
||||
'overflow': max(read('overflow'), 0),
|
||||
}
|
||||
|
||||
def record_pool_timeout(self) -> None:
|
||||
self._pool_timeouts_total += 1
|
||||
|
||||
+1751
-15
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,25 @@
|
||||
"""Safe PostgreSQL URL normalization shared by runtime and migration jobs."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy
|
||||
|
||||
|
||||
def normalize_asyncpg_url(url: sqlalchemy.engine.URL) -> sqlalchemy.engine.URL:
|
||||
"""Select asyncpg and translate the common libpq TLS query spelling."""
|
||||
|
||||
if url.drivername == 'postgresql':
|
||||
url = url.set(drivername='postgresql+asyncpg')
|
||||
elif url.drivername != 'postgresql+asyncpg':
|
||||
raise ValueError('PostgreSQL URL must use PostgreSQL with the asyncpg driver')
|
||||
|
||||
query = dict(url.query)
|
||||
sslmode = query.pop('sslmode', None)
|
||||
if sslmode is not None:
|
||||
if 'ssl' in query and query['ssl'] != sslmode:
|
||||
raise ValueError('PostgreSQL URL cannot specify conflicting ssl and sslmode options')
|
||||
# SQLAlchemy expands URL query keys into asyncpg keyword arguments.
|
||||
# asyncpg calls this keyword ``ssl`` even though PostgreSQL DSNs
|
||||
# conventionally spell the same mode ``sslmode``.
|
||||
query['ssl'] = sslmode
|
||||
return url.set(query=query)
|
||||
@@ -0,0 +1,190 @@
|
||||
"""One-shot, operator-only Cloud PostgreSQL release migration entrypoint."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
|
||||
import sqlalchemy
|
||||
|
||||
from ..cloud.bootstrap import SUPPORTED_PGVECTOR_DIMENSIONS
|
||||
from ..core import app as core_app
|
||||
from ..core.stages.load_config import LoadConfigStage
|
||||
from ..core.stages.setup_logger import SetupLoggerStage
|
||||
from .mgr import PersistenceManager, PersistenceMode
|
||||
from .postgresql_url import normalize_asyncpg_url
|
||||
|
||||
|
||||
DEFAULT_OPERATOR_DSN_ENV = 'LANGBOT_CLOUD_MIGRATION_DSN'
|
||||
_ENV_NAME = re.compile(r'^[A-Z_][A-Z0-9_]*$')
|
||||
|
||||
|
||||
class CloudReleaseMigrationConfigurationError(RuntimeError):
|
||||
"""Raised before any database operation when migration input is unsafe."""
|
||||
|
||||
|
||||
def _url_endpoint(url: sqlalchemy.engine.URL, *, label: str) -> tuple[str, int]:
|
||||
"""Return a comparison-safe PostgreSQL endpoint without leaking its DSN."""
|
||||
|
||||
try:
|
||||
host = (url.host or '').strip().casefold()
|
||||
port = url.port or 5432
|
||||
except (TypeError, ValueError):
|
||||
raise CloudReleaseMigrationConfigurationError(f'{label} PostgreSQL host or port is invalid') from None
|
||||
if not host or isinstance(port, bool) or not isinstance(port, int) or not 1 <= port <= 65535:
|
||||
raise CloudReleaseMigrationConfigurationError(f'{label} PostgreSQL host or port is invalid')
|
||||
return host, port
|
||||
|
||||
|
||||
def _operator_database_url(
|
||||
instance_config: dict,
|
||||
*,
|
||||
environ: Mapping[str, str],
|
||||
) -> sqlalchemy.engine.URL:
|
||||
database_config = instance_config.get('database')
|
||||
if not isinstance(database_config, dict) or database_config.get('use') != 'postgresql':
|
||||
raise CloudReleaseMigrationConfigurationError(
|
||||
'Cloud release migration requires explicit database.use=postgresql; SQLite fallback is forbidden'
|
||||
)
|
||||
|
||||
runtime_config = database_config.get('postgresql')
|
||||
if not isinstance(runtime_config, dict):
|
||||
raise CloudReleaseMigrationConfigurationError('Cloud runtime PostgreSQL configuration is missing')
|
||||
|
||||
migration_config = database_config.get('cloud_migration', {})
|
||||
if not isinstance(migration_config, dict):
|
||||
raise CloudReleaseMigrationConfigurationError('database.cloud_migration must be a mapping')
|
||||
dsn_env = migration_config.get('operator_dsn_env', DEFAULT_OPERATOR_DSN_ENV)
|
||||
if not isinstance(dsn_env, str) or not _ENV_NAME.fullmatch(dsn_env):
|
||||
raise CloudReleaseMigrationConfigurationError(
|
||||
'database.cloud_migration.operator_dsn_env must name an uppercase environment variable'
|
||||
)
|
||||
|
||||
raw_dsn = environ.get(dsn_env, '').strip()
|
||||
if not raw_dsn:
|
||||
raise CloudReleaseMigrationConfigurationError(
|
||||
f'Cloud release migration requires the operator DSN in environment variable {dsn_env}'
|
||||
)
|
||||
try:
|
||||
operator_url = sqlalchemy.engine.make_url(raw_dsn)
|
||||
except Exception:
|
||||
# Never echo a malformed DSN because it may contain an unescaped secret.
|
||||
raise CloudReleaseMigrationConfigurationError('Cloud release migration operator DSN is invalid') from None
|
||||
|
||||
try:
|
||||
operator_url = normalize_asyncpg_url(operator_url)
|
||||
except ValueError:
|
||||
raise CloudReleaseMigrationConfigurationError(
|
||||
'Cloud release migration operator DSN must use valid PostgreSQL asyncpg options'
|
||||
) from None
|
||||
|
||||
operator_user = (operator_url.username or '').strip()
|
||||
operator_database = (operator_url.database or '').strip()
|
||||
operator_host, operator_port = _url_endpoint(operator_url, label='Cloud release migration operator')
|
||||
runtime_url_value = runtime_config.get('url')
|
||||
if runtime_url_value:
|
||||
if not isinstance(runtime_url_value, str):
|
||||
raise CloudReleaseMigrationConfigurationError('Cloud runtime PostgreSQL URL must be a string')
|
||||
try:
|
||||
runtime_url = sqlalchemy.engine.make_url(runtime_url_value)
|
||||
except Exception:
|
||||
raise CloudReleaseMigrationConfigurationError('Cloud runtime PostgreSQL URL is invalid') from None
|
||||
if runtime_url.drivername not in {'postgresql', 'postgresql+asyncpg'}:
|
||||
raise CloudReleaseMigrationConfigurationError('Cloud runtime database URL must use PostgreSQL')
|
||||
runtime_user = (runtime_url.username or '').strip()
|
||||
runtime_database = (runtime_url.database or '').strip()
|
||||
runtime_host, runtime_port = _url_endpoint(runtime_url, label='Cloud runtime')
|
||||
else:
|
||||
runtime_user = str(runtime_config.get('user', 'postgres') or '').strip()
|
||||
runtime_database = str(runtime_config.get('database', 'postgres') or '').strip()
|
||||
runtime_host = str(runtime_config.get('host', '') or '').strip().casefold()
|
||||
runtime_port = runtime_config.get('port', 5432)
|
||||
if not operator_user or not operator_database:
|
||||
raise CloudReleaseMigrationConfigurationError(
|
||||
'Cloud release migration operator DSN must include a user, host, and database'
|
||||
)
|
||||
if (
|
||||
not runtime_user
|
||||
or not runtime_database
|
||||
or not runtime_host
|
||||
or isinstance(runtime_port, bool)
|
||||
or not isinstance(runtime_port, int)
|
||||
or not 1 <= runtime_port <= 65535
|
||||
):
|
||||
raise CloudReleaseMigrationConfigurationError(
|
||||
'Cloud runtime PostgreSQL user, host, port, and database are required'
|
||||
)
|
||||
if operator_user == runtime_user:
|
||||
raise CloudReleaseMigrationConfigurationError(
|
||||
'Cloud release migration requires a distinct operator role from the runtime PostgreSQL role'
|
||||
)
|
||||
if operator_database != runtime_database:
|
||||
raise CloudReleaseMigrationConfigurationError(
|
||||
'Cloud release migration operator DSN must target the configured runtime database'
|
||||
)
|
||||
if operator_host != runtime_host or operator_port != runtime_port:
|
||||
# The first Cloud release intentionally requires the migrator and
|
||||
# runtime to use the same PostgreSQL endpoint. Supporting a direct
|
||||
# operator endpoint plus a runtime pooler requires a database-backed
|
||||
# immutable cluster identity check; accepting aliases here would turn a
|
||||
# same-named database on another cluster into a silent migration target.
|
||||
raise CloudReleaseMigrationConfigurationError(
|
||||
'Cloud release migration operator DSN must target the configured runtime PostgreSQL endpoint'
|
||||
)
|
||||
|
||||
vdb_config = instance_config.get('vdb')
|
||||
if not isinstance(vdb_config, dict) or vdb_config.get('use') != 'pgvector':
|
||||
raise CloudReleaseMigrationConfigurationError('Cloud release migration requires vdb.use=pgvector')
|
||||
pgvector_config = vdb_config.get('pgvector')
|
||||
if not isinstance(pgvector_config, dict) or pgvector_config.get('use_business_database') is not True:
|
||||
raise CloudReleaseMigrationConfigurationError(
|
||||
'Cloud release migration requires vdb.pgvector.use_business_database=true'
|
||||
)
|
||||
dimensions = pgvector_config.get('allowed_dimensions')
|
||||
if (
|
||||
not isinstance(dimensions, list)
|
||||
or not dimensions
|
||||
or any(isinstance(item, bool) or not isinstance(item, int) for item in dimensions)
|
||||
or not set(dimensions).issubset(SUPPORTED_PGVECTOR_DIMENSIONS)
|
||||
):
|
||||
raise CloudReleaseMigrationConfigurationError(
|
||||
'Cloud release migration pgvector dimensions are outside the release-created index set'
|
||||
)
|
||||
|
||||
return operator_url
|
||||
|
||||
|
||||
async def run_cloud_release_migration(
|
||||
ap: core_app.Application,
|
||||
*,
|
||||
environ: Mapping[str, str] | None = None,
|
||||
) -> None:
|
||||
"""Run and validate one release migration with an isolated operator DSN."""
|
||||
|
||||
operator_url = _operator_database_url(
|
||||
ap.instance_config.data,
|
||||
environ=os.environ if environ is None else environ,
|
||||
)
|
||||
manager = PersistenceManager(
|
||||
ap,
|
||||
mode=PersistenceMode.RELEASE_MIGRATION,
|
||||
database_url=operator_url,
|
||||
)
|
||||
ap.persistence_mgr = manager
|
||||
try:
|
||||
await manager.initialize()
|
||||
ap.logger.info('Cloud PostgreSQL release migration reached and validated the exact release head.')
|
||||
finally:
|
||||
await manager.shutdown()
|
||||
|
||||
|
||||
async def run_cloud_release_migration_from_config(loop: asyncio.AbstractEventLoop) -> None:
|
||||
"""Load only process configuration/logging, then run the one-shot job."""
|
||||
|
||||
ap = core_app.Application()
|
||||
ap.event_loop = loop
|
||||
await LoadConfigStage().run(ap)
|
||||
await SetupLoggerStage().run(ap)
|
||||
await run_cloud_release_migration(ap)
|
||||
@@ -0,0 +1,272 @@
|
||||
"""Durable SQLite backups for destructive Alembic migration boundaries."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import dataclasses
|
||||
import datetime
|
||||
import json
|
||||
import os
|
||||
import pathlib
|
||||
import re
|
||||
import secrets
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import typing
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
|
||||
|
||||
class SQLiteMigrationBackupError(RuntimeError):
|
||||
"""A verified migration backup could not be created or restored."""
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True, slots=True)
|
||||
class SQLiteMigrationBackup:
|
||||
database_path: pathlib.Path
|
||||
backup_path: pathlib.Path
|
||||
manifest_path: pathlib.Path
|
||||
source_revision: str
|
||||
target_revision: str
|
||||
created_at: str
|
||||
|
||||
|
||||
def _safe_label(value: str) -> str:
|
||||
label = re.sub(r'[^A-Za-z0-9_.-]+', '-', value).strip('-')
|
||||
return label or 'unknown'
|
||||
|
||||
|
||||
def _database_path(engine: AsyncEngine) -> pathlib.Path:
|
||||
if engine.dialect.name != 'sqlite':
|
||||
raise SQLiteMigrationBackupError('SQLite migration backups require a SQLite engine')
|
||||
database = engine.url.database
|
||||
if not database or database == ':memory:' or engine.url.query.get('mode') == 'memory':
|
||||
raise SQLiteMigrationBackupError('Tenant schema migrations require a file-backed SQLite database for recovery')
|
||||
database_path = pathlib.Path(database).expanduser()
|
||||
if not database_path.is_absolute():
|
||||
database_path = pathlib.Path.cwd() / database_path
|
||||
database_path = database_path.resolve()
|
||||
if not database_path.is_file():
|
||||
raise SQLiteMigrationBackupError(f'SQLite database does not exist: {database_path}')
|
||||
return database_path
|
||||
|
||||
|
||||
def _open_read_only(path: pathlib.Path) -> sqlite3.Connection:
|
||||
return sqlite3.connect(f'{path.as_uri()}?mode=ro', uri=True, timeout=30)
|
||||
|
||||
|
||||
def _read_revision(connection: sqlite3.Connection) -> str | None:
|
||||
has_version_table = connection.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'alembic_version'"
|
||||
).fetchone()
|
||||
if has_version_table is None:
|
||||
return None
|
||||
rows = connection.execute('SELECT version_num FROM alembic_version').fetchall()
|
||||
if not rows:
|
||||
return None
|
||||
if len(rows) != 1 or not isinstance(rows[0][0], str):
|
||||
raise SQLiteMigrationBackupError('SQLite backup has an invalid Alembic revision table')
|
||||
return rows[0][0]
|
||||
|
||||
|
||||
def _verify_connection(connection: sqlite3.Connection, expected_revision: str) -> None:
|
||||
quick_check = connection.execute('PRAGMA quick_check').fetchall()
|
||||
if quick_check != [('ok',)]:
|
||||
raise SQLiteMigrationBackupError(f'SQLite quick_check failed: {quick_check[:5]!r}')
|
||||
actual_revision = _read_revision(connection)
|
||||
if actual_revision != expected_revision:
|
||||
raise SQLiteMigrationBackupError(
|
||||
f'SQLite backup revision mismatch: {actual_revision!r} != {expected_revision!r}'
|
||||
)
|
||||
|
||||
|
||||
def _verify_file(path: pathlib.Path, expected_revision: str) -> None:
|
||||
with _open_read_only(path) as connection:
|
||||
_verify_connection(connection, expected_revision)
|
||||
|
||||
|
||||
def _write_manifest(backup: SQLiteMigrationBackup, status: str, **extra: typing.Any) -> None:
|
||||
payload: dict[str, typing.Any] = {
|
||||
'version': 1,
|
||||
'status': status,
|
||||
'created_at': backup.created_at,
|
||||
'database_path': str(backup.database_path),
|
||||
'backup_path': str(backup.backup_path),
|
||||
'source_revision': backup.source_revision,
|
||||
'target_revision': backup.target_revision,
|
||||
'quick_check': 'ok',
|
||||
**extra,
|
||||
}
|
||||
backup.manifest_path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||||
descriptor, temporary_name = tempfile.mkstemp(
|
||||
prefix=f'.{backup.manifest_path.name}.',
|
||||
suffix='.tmp',
|
||||
dir=backup.manifest_path.parent,
|
||||
)
|
||||
temporary_path = pathlib.Path(temporary_name)
|
||||
try:
|
||||
with os.fdopen(descriptor, 'w', encoding='utf-8') as file:
|
||||
json.dump(payload, file, ensure_ascii=False, indent=2, sort_keys=True)
|
||||
file.write('\n')
|
||||
file.flush()
|
||||
os.fsync(file.fileno())
|
||||
os.chmod(temporary_path, 0o600)
|
||||
os.replace(temporary_path, backup.manifest_path)
|
||||
_fsync_directory(backup.manifest_path.parent)
|
||||
finally:
|
||||
temporary_path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
def _fsync_file(path: pathlib.Path) -> None:
|
||||
descriptor = os.open(path, os.O_RDONLY)
|
||||
try:
|
||||
os.fsync(descriptor)
|
||||
finally:
|
||||
os.close(descriptor)
|
||||
|
||||
|
||||
def _fsync_directory(path: pathlib.Path) -> None:
|
||||
descriptor = os.open(path, os.O_RDONLY)
|
||||
try:
|
||||
os.fsync(descriptor)
|
||||
finally:
|
||||
os.close(descriptor)
|
||||
|
||||
|
||||
def _create_backup(
|
||||
database_path: pathlib.Path,
|
||||
source_revision: str,
|
||||
target_revision: str,
|
||||
) -> SQLiteMigrationBackup:
|
||||
backup_directory = database_path.parent / 'migration-backups'
|
||||
backup_directory.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||||
os.chmod(backup_directory, 0o700)
|
||||
created_at = datetime.datetime.now(datetime.UTC).strftime('%Y-%m-%dT%H-%M-%S.%fZ')
|
||||
stem = (
|
||||
f'{database_path.stem}-pre-{_safe_label(target_revision)}-'
|
||||
f'from-{_safe_label(source_revision)}-{created_at}-{secrets.token_hex(4)}'
|
||||
)
|
||||
backup_path = backup_directory / f'{stem}.sqlite3'
|
||||
manifest_path = backup_directory / f'{stem}.json'
|
||||
descriptor, temporary_name = tempfile.mkstemp(
|
||||
prefix=f'.{stem}.',
|
||||
suffix='.creating',
|
||||
dir=backup_directory,
|
||||
)
|
||||
os.close(descriptor)
|
||||
temporary_path = pathlib.Path(temporary_name)
|
||||
try:
|
||||
with (
|
||||
_open_read_only(database_path) as source,
|
||||
sqlite3.connect(
|
||||
temporary_path,
|
||||
timeout=30,
|
||||
) as destination,
|
||||
):
|
||||
source.execute('PRAGMA busy_timeout = 30000')
|
||||
source.backup(destination)
|
||||
destination.commit()
|
||||
_verify_connection(destination, source_revision)
|
||||
os.chmod(temporary_path, 0o600)
|
||||
_fsync_file(temporary_path)
|
||||
os.replace(temporary_path, backup_path)
|
||||
_fsync_file(backup_path)
|
||||
_fsync_directory(backup_directory)
|
||||
backup = SQLiteMigrationBackup(
|
||||
database_path=database_path,
|
||||
backup_path=backup_path,
|
||||
manifest_path=manifest_path,
|
||||
source_revision=source_revision,
|
||||
target_revision=target_revision,
|
||||
created_at=created_at,
|
||||
)
|
||||
_write_manifest(backup, 'verified')
|
||||
return backup
|
||||
except Exception:
|
||||
backup_path.unlink(missing_ok=True)
|
||||
manifest_path.unlink(missing_ok=True)
|
||||
raise
|
||||
finally:
|
||||
temporary_path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
async def create_verified_backup(
|
||||
engine: AsyncEngine,
|
||||
*,
|
||||
source_revision: str,
|
||||
target_revision: str,
|
||||
) -> SQLiteMigrationBackup:
|
||||
"""Create and verify an online-consistent backup next to instance data."""
|
||||
|
||||
database_path = _database_path(engine)
|
||||
return await asyncio.to_thread(
|
||||
_create_backup,
|
||||
database_path,
|
||||
source_revision,
|
||||
target_revision,
|
||||
)
|
||||
|
||||
|
||||
def _restore_backup(backup: SQLiteMigrationBackup) -> None:
|
||||
_verify_file(backup.backup_path, backup.source_revision)
|
||||
descriptor, temporary_name = tempfile.mkstemp(
|
||||
prefix=f'.{backup.database_path.name}.',
|
||||
suffix='.restoring',
|
||||
dir=backup.database_path.parent,
|
||||
)
|
||||
os.close(descriptor)
|
||||
temporary_path = pathlib.Path(temporary_name)
|
||||
try:
|
||||
with (
|
||||
_open_read_only(backup.backup_path) as source,
|
||||
sqlite3.connect(
|
||||
temporary_path,
|
||||
timeout=30,
|
||||
) as destination,
|
||||
):
|
||||
source.backup(destination)
|
||||
destination.commit()
|
||||
_verify_connection(destination, backup.source_revision)
|
||||
os.chmod(temporary_path, 0o600)
|
||||
_fsync_file(temporary_path)
|
||||
|
||||
# A stale WAL could replay pages from the failed migration after the
|
||||
# main database file is replaced. The engine is disposed before this
|
||||
# function runs, so these exact sidecars are safe to remove.
|
||||
for suffix in ('-wal', '-shm', '-journal'):
|
||||
pathlib.Path(f'{backup.database_path}{suffix}').unlink(missing_ok=True)
|
||||
os.replace(temporary_path, backup.database_path)
|
||||
_fsync_file(backup.database_path)
|
||||
_fsync_directory(backup.database_path.parent)
|
||||
_verify_file(backup.database_path, backup.source_revision)
|
||||
finally:
|
||||
temporary_path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
async def restore_verified_backup(engine: AsyncEngine, backup: SQLiteMigrationBackup) -> None:
|
||||
"""Atomically restore a verified backup after a migration failure."""
|
||||
|
||||
await engine.dispose()
|
||||
await asyncio.to_thread(_restore_backup, backup)
|
||||
await asyncio.to_thread(
|
||||
_write_manifest,
|
||||
backup,
|
||||
'restored_after_failure',
|
||||
restored_at=datetime.datetime.now(datetime.UTC).isoformat(),
|
||||
)
|
||||
|
||||
|
||||
async def mark_migration_succeeded(
|
||||
backup: SQLiteMigrationBackup,
|
||||
*,
|
||||
completed_revision: str,
|
||||
) -> None:
|
||||
"""Mark a retained verified backup after its migration boundary succeeds."""
|
||||
|
||||
await asyncio.to_thread(
|
||||
_write_manifest,
|
||||
backup,
|
||||
'migration_succeeded',
|
||||
completed_at=datetime.datetime.now(datetime.UTC).isoformat(),
|
||||
completed_revision=completed_revision,
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user