mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-16 23:40:59 +00:00
feat(tenancy): implement workspace isolation
This commit is contained in:
Junyan Qin
@@ -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'
|
||||
)
|
||||
@@ -47,6 +47,12 @@ 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)
|
||||
@@ -73,6 +79,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 +134,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 +154,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}')
|
||||
|
||||
@@ -1,14 +1,16 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
import sqlite3
|
||||
import typing
|
||||
|
||||
|
||||
import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
|
||||
import sqlalchemy
|
||||
|
||||
from . import database, migration
|
||||
from . import database, migration, sqlite_migration_backup
|
||||
from ..entity.persistence import base, metadata, model as persistence_model
|
||||
from ..entity.persistence import workspace as persistence_workspace
|
||||
from ..entity import persistence
|
||||
from ..core import app
|
||||
from ..utils import constants, importutil
|
||||
@@ -19,6 +21,51 @@ importutil.import_modules_in_pkg(migrations)
|
||||
importutil.import_modules_in_pkg(persistence)
|
||||
|
||||
|
||||
_ALEMBIC_TENANT_TABLES = {
|
||||
'workspaces',
|
||||
'workspace_memberships',
|
||||
'workspace_invitations',
|
||||
'workspace_execution_states',
|
||||
'workspace_metadata',
|
||||
'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',
|
||||
}
|
||||
|
||||
_PRE_WORKSPACE_ALEMBIC_REVISIONS = {
|
||||
'0001_baseline',
|
||||
'0002_sample',
|
||||
'0003_add_rerank_models',
|
||||
'0004_add_mcp_readme',
|
||||
'0005_add_llm_context_length',
|
||||
'0006_normalize_mcp_remote_mode',
|
||||
'0007_add_bot_admins',
|
||||
'0008_mcp_resource_prefs',
|
||||
}
|
||||
_WORKSPACE_ALEMBIC_REVISION = '0009_workspace_tenancy'
|
||||
_RESOURCE_SCOPE_ALEMBIC_REVISION = '0010_scope_resources'
|
||||
|
||||
|
||||
class PersistenceManager:
|
||||
"""Persistence module manager"""
|
||||
|
||||
@@ -42,6 +89,8 @@ class PersistenceManager:
|
||||
await self.db.initialize()
|
||||
break
|
||||
|
||||
self._enable_sqlite_foreign_keys()
|
||||
|
||||
await self.create_tables()
|
||||
|
||||
# run migrations
|
||||
@@ -79,12 +128,33 @@ class PersistenceManager:
|
||||
# Run Alembic migrations (new migration system)
|
||||
await self._run_alembic_migrations()
|
||||
|
||||
# A legacy database may not contain tenant tables introduced by a
|
||||
# newer release. They were deliberately deferred before 0009 because
|
||||
# their Workspace/account FK targets did not exist yet; create them
|
||||
# now that the tenancy contract is in place.
|
||||
await self.create_tables()
|
||||
|
||||
await self.write_space_model_providers()
|
||||
|
||||
async def create_tables(self):
|
||||
# create tables
|
||||
async with self.get_db_engine().connect() as conn:
|
||||
await conn.run_sync(self.meta.create_all)
|
||||
|
||||
def create_compatible_tables(sync_conn: sqlalchemy.Connection) -> None:
|
||||
inspector = sqlalchemy.inspect(sync_conn)
|
||||
existing_tables = set(inspector.get_table_names())
|
||||
legacy_users = 'users' in existing_tables and (
|
||||
'uuid' not in {column['name'] for column in inspector.get_columns('users')}
|
||||
or 'workspaces' not in existing_tables
|
||||
)
|
||||
# On a legacy installation, resource tables already exist
|
||||
# without workspace_uuid and Workspace itself references the
|
||||
# account UUID introduced by 0009. Alembic must expand those
|
||||
# tables before SQLAlchemy may create any new tenant table.
|
||||
excluded_tables = _ALEMBIC_TENANT_TABLES if legacy_users else set()
|
||||
tables_to_create = [table for table in self.meta.sorted_tables if table.name not in excluded_tables]
|
||||
self.meta.create_all(sync_conn, tables=tables_to_create)
|
||||
|
||||
await conn.run_sync(create_compatible_tables)
|
||||
|
||||
await conn.commit()
|
||||
|
||||
@@ -101,15 +171,80 @@ class PersistenceManager:
|
||||
if row is None:
|
||||
await self.execute_async(sqlalchemy.insert(metadata.Metadata).values(item))
|
||||
|
||||
await self._ensure_instance_uuid_metadata()
|
||||
|
||||
def _enable_sqlite_foreign_keys(self) -> None:
|
||||
"""Enable SQLite FK enforcement for every pooled runtime connection."""
|
||||
engine = self.get_db_engine()
|
||||
if engine.dialect.name != 'sqlite':
|
||||
return
|
||||
if getattr(self, '_sqlite_fk_listener_installed', False):
|
||||
return
|
||||
|
||||
def set_sqlite_pragma(dbapi_connection, _connection_record) -> None:
|
||||
# aiosqlite exposes the normal sqlite cursor API through its
|
||||
# SQLAlchemy adapter. Guard the direct sqlite type too for tests.
|
||||
if isinstance(dbapi_connection, sqlite3.Connection) or hasattr(dbapi_connection, 'cursor'):
|
||||
cursor = dbapi_connection.cursor()
|
||||
cursor.execute('PRAGMA foreign_keys=ON')
|
||||
cursor.close()
|
||||
|
||||
sqlalchemy.event.listen(engine.sync_engine, 'connect', set_sqlite_pragma)
|
||||
self._sqlite_fk_listener_installed = True
|
||||
|
||||
async def _ensure_instance_uuid_metadata(self) -> None:
|
||||
"""Persist the runtime instance identifier before tenant migrations run."""
|
||||
runtime_instance_uuid = constants.instance_id.strip()
|
||||
if not runtime_instance_uuid:
|
||||
raise RuntimeError('LangBot instance UUID is empty before persistence initialization')
|
||||
|
||||
result = await self.execute_async(
|
||||
sqlalchemy.select(metadata.Metadata.value).where(metadata.Metadata.key == 'instance_uuid')
|
||||
)
|
||||
persisted_instance_uuid = result.scalar_one_or_none()
|
||||
|
||||
if persisted_instance_uuid is None:
|
||||
await self.execute_async(
|
||||
sqlalchemy.insert(metadata.Metadata).values(key='instance_uuid', value=runtime_instance_uuid)
|
||||
)
|
||||
return
|
||||
|
||||
if persisted_instance_uuid != runtime_instance_uuid:
|
||||
raise RuntimeError(
|
||||
'LangBot instance UUID does not match the value bound to this database: '
|
||||
f'{runtime_instance_uuid!r} != {persisted_instance_uuid!r}'
|
||||
)
|
||||
|
||||
async def write_space_model_providers(self):
|
||||
if constants.edition != 'community':
|
||||
# SaaS Workspace/provider linkage is explicit control-plane state;
|
||||
# a process-level compatibility provider must never be projected
|
||||
# into an arbitrary cloud Workspace.
|
||||
return
|
||||
|
||||
space_models_gateway_api_url = self.ap.instance_config.data.get('space', {}).get(
|
||||
'models_gateway_api_url', 'https://api.langbot.cloud/v1'
|
||||
)
|
||||
|
||||
# write space model providers
|
||||
workspace_result = await self.execute_async(
|
||||
sqlalchemy.select(persistence_workspace.Workspace.uuid).where(
|
||||
persistence_workspace.Workspace.instance_uuid == constants.instance_id,
|
||||
persistence_workspace.Workspace.source == persistence_workspace.WorkspaceSource.LOCAL.value,
|
||||
)
|
||||
)
|
||||
workspace_uuids = workspace_result.scalars().all()
|
||||
if len(workspace_uuids) != 1:
|
||||
raise RuntimeError(
|
||||
f'The fixed LangBot Models provider requires exactly one local Workspace; found {len(workspace_uuids)}'
|
||||
)
|
||||
workspace_uuid = workspace_uuids[0]
|
||||
|
||||
# The compatibility Space provider belongs to the OSS singleton
|
||||
# Workspace. It must never be discovered or inserted globally.
|
||||
result = await self.execute_async(
|
||||
sqlalchemy.select(persistence_model.ModelProvider).where(
|
||||
persistence_model.ModelProvider.requester == 'space-chat-completions'
|
||||
persistence_model.ModelProvider.workspace_uuid == workspace_uuid,
|
||||
persistence_model.ModelProvider.requester == 'space-chat-completions',
|
||||
)
|
||||
)
|
||||
exists_space_chat_completions_model_provider = result.first()
|
||||
@@ -119,6 +254,7 @@ class PersistenceManager:
|
||||
self.ap.logger.info('Creating space model providers...')
|
||||
space_chat_completions_model_provider = {
|
||||
'uuid': '00000000-0000-0000-0000-000000000000',
|
||||
'workspace_uuid': workspace_uuid,
|
||||
'name': 'LangBot Models',
|
||||
'requester': 'space-chat-completions',
|
||||
'base_url': space_models_gateway_api_url,
|
||||
@@ -132,7 +268,10 @@ class PersistenceManager:
|
||||
if exists_space_chat_completions_model_provider.base_url != space_models_gateway_api_url:
|
||||
await self.execute_async(
|
||||
sqlalchemy.update(persistence_model.ModelProvider)
|
||||
.where(persistence_model.ModelProvider.uuid == exists_space_chat_completions_model_provider.uuid)
|
||||
.where(
|
||||
persistence_model.ModelProvider.workspace_uuid == workspace_uuid,
|
||||
persistence_model.ModelProvider.uuid == exists_space_chat_completions_model_provider.uuid,
|
||||
)
|
||||
.values({'base_url': space_models_gateway_api_url})
|
||||
)
|
||||
|
||||
@@ -153,13 +292,67 @@ class PersistenceManager:
|
||||
await alembic_runner.run_alembic_stamp(engine, '0001_baseline')
|
||||
current_rev = '0001_baseline'
|
||||
|
||||
# Upgrade to head
|
||||
if engine.dialect.name == 'sqlite':
|
||||
if current_rev in _PRE_WORKSPACE_ALEMBIC_REVISIONS:
|
||||
await self._run_verified_sqlite_migration(
|
||||
engine,
|
||||
source_revision=current_rev,
|
||||
target_revision=_WORKSPACE_ALEMBIC_REVISION,
|
||||
)
|
||||
current_rev = await alembic_runner.get_alembic_current(engine)
|
||||
if current_rev == _WORKSPACE_ALEMBIC_REVISION:
|
||||
await self._run_verified_sqlite_migration(
|
||||
engine,
|
||||
source_revision=current_rev,
|
||||
target_revision=_RESOURCE_SCOPE_ALEMBIC_REVISION,
|
||||
)
|
||||
|
||||
# PostgreSQL has transactional DDL. SQLite has already crossed the
|
||||
# two destructive tenancy boundaries under verified backups; this
|
||||
# final call is a no-op today and applies future migrations.
|
||||
await alembic_runner.run_alembic_upgrade(engine, 'head')
|
||||
self.ap.logger.info('Alembic migrations completed.')
|
||||
except Exception as e:
|
||||
self.ap.logger.error(f'Alembic migration failed: {e}', exc_info=True)
|
||||
raise
|
||||
|
||||
async def _run_verified_sqlite_migration(
|
||||
self,
|
||||
engine: sqlalchemy_asyncio.AsyncEngine,
|
||||
*,
|
||||
source_revision: str,
|
||||
target_revision: str,
|
||||
) -> None:
|
||||
from . import alembic_runner
|
||||
|
||||
backup = await sqlite_migration_backup.create_verified_backup(
|
||||
engine,
|
||||
source_revision=source_revision,
|
||||
target_revision=target_revision,
|
||||
)
|
||||
self.ap.logger.info(f'Created verified SQLite migration backup {backup.backup_path} before {target_revision}.')
|
||||
try:
|
||||
await alembic_runner.run_alembic_upgrade(engine, target_revision)
|
||||
completed_revision = await alembic_runner.get_alembic_current(engine)
|
||||
if completed_revision != target_revision:
|
||||
raise RuntimeError(f'Alembic stopped at {completed_revision!r}, expected {target_revision!r}')
|
||||
await sqlite_migration_backup.mark_migration_succeeded(
|
||||
backup,
|
||||
completed_revision=completed_revision,
|
||||
)
|
||||
except BaseException:
|
||||
await sqlite_migration_backup.restore_verified_backup(engine, backup)
|
||||
restored_revision = await alembic_runner.get_alembic_current(engine)
|
||||
if restored_revision != source_revision:
|
||||
raise RuntimeError(
|
||||
f'SQLite migration recovery restored revision {restored_revision!r}, expected {source_revision!r}'
|
||||
)
|
||||
self.ap.logger.error(
|
||||
f'SQLite migration to {target_revision} failed; restored verified backup '
|
||||
f'{backup.backup_path} at revision {source_revision}.'
|
||||
)
|
||||
raise
|
||||
|
||||
async def execute_async(self, *args, **kwargs) -> sqlalchemy.engine.cursor.CursorResult:
|
||||
async with self.get_db_engine().connect() as conn:
|
||||
result = await conn.execute(*args, **kwargs)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
Reference in New Issue
Block a user