mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-23 00:45:25 +08:00
refactor(persistence): consolidate monitoring index migrations
- Merge migrations 0039 through 0041 into a single 0039 migration - Preserve existing indexes and support partial historical schemas - Update migration tests for the consolidated revision
This commit is contained in:
@@ -1,4 +1,4 @@
|
||||
"""Bound host lease maintenance to active executions within each Workspace."""
|
||||
"""Index execution leases, conversation paging, and historical inline image expiry."""
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
@@ -8,26 +8,40 @@ down_revision = '0038_llm_call_details'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
INDEXES = (
|
||||
('agent_run', 'ix_agent_run_workspace_status_lease', ['workspace_id', 'status', 'execution_lease_expires_at']),
|
||||
(
|
||||
'monitoring_messages',
|
||||
'ix_monitoring_messages_workspace_status_lease',
|
||||
['workspace_uuid', 'status', 'execution_lease_expires_at'],
|
||||
),
|
||||
(
|
||||
'monitoring_messages',
|
||||
'ix_monitoring_messages_conversation_page',
|
||||
['workspace_uuid', 'bot_id', 'session_id', 'timestamp', 'id'],
|
||||
),
|
||||
('monitoring_messages', 'ix_monitoring_messages_media_ttl', ['workspace_uuid', 'timestamp', 'id']),
|
||||
('event_log', 'ix_event_log_media_ttl', ['workspace_id', 'created_at', 'id']),
|
||||
('agent_run', 'ix_agent_run_media_ttl', ['workspace_id', 'created_at', 'id']),
|
||||
('agent_run_event', 'ix_agent_run_event_media_ttl', ['created_at', 'id', 'run_id']),
|
||||
('transcript', 'ix_transcript_media_ttl', ['workspace_id', 'created_at', 'id']),
|
||||
)
|
||||
|
||||
|
||||
def upgrade():
|
||||
inspector = sa.inspect(op.get_bind())
|
||||
for table, workspace_column in (
|
||||
('agent_run', 'workspace_id'),
|
||||
('monitoring_messages', 'workspace_uuid'),
|
||||
):
|
||||
for table, index, columns in INDEXES:
|
||||
if not inspector.has_table(table):
|
||||
continue
|
||||
index = f'ix_{table}_workspace_status_lease'
|
||||
# Fresh installations create these indexes from ORM metadata before migrations.
|
||||
if index not in {item['name'] for item in inspector.get_indexes(table)}:
|
||||
op.create_index(index, table, [workspace_column, 'status', 'execution_lease_expires_at'])
|
||||
op.create_index(index, table, columns)
|
||||
|
||||
|
||||
def downgrade():
|
||||
inspector = sa.inspect(op.get_bind())
|
||||
for table in ('agent_run', 'monitoring_messages'):
|
||||
for table, index, _ in reversed(INDEXES):
|
||||
if not inspector.has_table(table):
|
||||
continue
|
||||
index = f'ix_{table}_workspace_status_lease'
|
||||
if index in {item['name'] for item in inspector.get_indexes(table)}:
|
||||
op.drop_index(index, table_name=table)
|
||||
|
||||
@@ -1,25 +0,0 @@
|
||||
"""Page conversation identities without reading historical message payloads."""
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision = '0040_conversation_page_index'
|
||||
down_revision = '0039_execution_lease_indexes'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
TABLE = 'monitoring_messages'
|
||||
INDEX = 'ix_monitoring_messages_conversation_page'
|
||||
|
||||
|
||||
def upgrade():
|
||||
inspector = sa.inspect(op.get_bind())
|
||||
# Fresh databases already receive this index from ORM metadata.
|
||||
if inspector.has_table(TABLE) and INDEX not in {index['name'] for index in inspector.get_indexes(TABLE)}:
|
||||
op.create_index(INDEX, TABLE, ['workspace_uuid', 'bot_id', 'session_id', 'timestamp', 'id'])
|
||||
|
||||
|
||||
def downgrade():
|
||||
inspector = sa.inspect(op.get_bind())
|
||||
if inspector.has_table(TABLE) and INDEX in {index['name'] for index in inspector.get_indexes(TABLE)}:
|
||||
op.drop_index(INDEX, table_name=TABLE)
|
||||
@@ -1,37 +0,0 @@
|
||||
"""Index incremental expiry of historical inline images."""
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision = '0041_inline_media_ttl_indexes'
|
||||
down_revision = '0040_conversation_page_index'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
INDEXES = (
|
||||
('monitoring_messages', ['workspace_uuid', 'timestamp', 'id']),
|
||||
('event_log', ['workspace_id', 'created_at', 'id']),
|
||||
('agent_run', ['workspace_id', 'created_at', 'id']),
|
||||
('agent_run_event', ['created_at', 'id', 'run_id']),
|
||||
('transcript', ['workspace_id', 'created_at', 'id']),
|
||||
)
|
||||
|
||||
|
||||
def upgrade():
|
||||
inspector = sa.inspect(op.get_bind())
|
||||
for table, columns in INDEXES:
|
||||
if not inspector.has_table(table):
|
||||
continue
|
||||
name = f'ix_{table}_media_ttl'
|
||||
if name not in {index['name'] for index in inspector.get_indexes(table)}:
|
||||
op.create_index(name, table, columns)
|
||||
|
||||
|
||||
def downgrade():
|
||||
inspector = sa.inspect(op.get_bind())
|
||||
for table, _ in reversed(INDEXES):
|
||||
if not inspector.has_table(table):
|
||||
continue
|
||||
name = f'ix_{table}_media_ttl'
|
||||
if name in {index['name'] for index in inspector.get_indexes(table)}:
|
||||
op.drop_index(name, table_name=table)
|
||||
@@ -10,7 +10,8 @@ from langbot.pkg.entity.persistence.monitoring import MonitoringMessage
|
||||
|
||||
|
||||
def test_conversation_index_migration_preserves_history_and_matches_metadata():
|
||||
migration = importlib.import_module('langbot.pkg.persistence.alembic.versions.0040_conversation_page_index')
|
||||
migration = importlib.import_module('langbot.pkg.persistence.alembic.versions.0039_execution_lease_indexes')
|
||||
index_name = 'ix_monitoring_messages_conversation_page'
|
||||
engine = sa.create_engine('sqlite:///:memory:')
|
||||
metadata = sa.MetaData()
|
||||
table = sa.Table(
|
||||
@@ -18,7 +19,16 @@ def test_conversation_index_migration_preserves_history_and_matches_metadata():
|
||||
metadata,
|
||||
*[
|
||||
sa.Column(name, sa.String, primary_key=name == 'id')
|
||||
for name in ('id', 'workspace_uuid', 'bot_id', 'session_id', 'timestamp', 'message_content')
|
||||
for name in (
|
||||
'id',
|
||||
'workspace_uuid',
|
||||
'bot_id',
|
||||
'session_id',
|
||||
'timestamp',
|
||||
'message_content',
|
||||
'status',
|
||||
'execution_lease_expires_at',
|
||||
)
|
||||
],
|
||||
)
|
||||
sa.Index('existing_index', table.c.workspace_uuid)
|
||||
@@ -32,8 +42,8 @@ def test_conversation_index_migration_preserves_history_and_matches_metadata():
|
||||
indexes = {
|
||||
index['name']: index['column_names'] for index in sa.inspect(connection).get_indexes(table.name)
|
||||
}
|
||||
expected = next(index for index in MonitoringMessage.__table__.indexes if index.name == migration.INDEX)
|
||||
assert indexes[migration.INDEX] == [column.name for column in expected.columns]
|
||||
expected = next(index for index in MonitoringMessage.__table__.indexes if index.name == index_name)
|
||||
assert indexes[index_name] == [column.name for column in expected.columns]
|
||||
migration.downgrade()
|
||||
assert {index['name'] for index in sa.inspect(connection).get_indexes(table.name)} == {'existing_index'}
|
||||
assert connection.scalar(sa.select(table.c.message_content)) == 'inline media'
|
||||
|
||||
@@ -53,9 +53,10 @@ def test_lease_index_migration_preserves_data_and_existing_indexes(source):
|
||||
assert original_indexes[table] <= indexes.keys()
|
||||
|
||||
MIGRATION.downgrade()
|
||||
for table, (index, _) in LEASE_INDEXES.items():
|
||||
for table in LEASE_INDEXES:
|
||||
remaining = {item['name'] for item in sa.inspect(connection).get_indexes(table)}
|
||||
assert remaining == original_indexes[table] - {index}
|
||||
removed = {index for index_table, index, _ in MIGRATION.INDEXES if index_table == table}
|
||||
assert remaining == original_indexes[table] - removed
|
||||
assert connection.scalar(sa.select(AgentRun.status)) == 'completed'
|
||||
|
||||
MIGRATION.upgrade()
|
||||
|
||||
@@ -306,7 +306,7 @@ async def test_storage_maintenance_uses_media_ttl_for_inline_records(application
|
||||
|
||||
|
||||
def test_ttl_indexes_migrate_idempotently_and_support_cursor_queries():
|
||||
migration = importlib.import_module('langbot.pkg.persistence.alembic.versions.0041_inline_media_ttl_indexes')
|
||||
migration = importlib.import_module('langbot.pkg.persistence.alembic.versions.0039_execution_lease_indexes')
|
||||
engine = sa.create_engine('sqlite:///:memory:')
|
||||
try:
|
||||
with engine.begin() as connection:
|
||||
@@ -315,14 +315,12 @@ def test_ttl_indexes_migrate_idempotently_and_support_cursor_queries():
|
||||
migration.downgrade()
|
||||
migration.upgrade()
|
||||
migration.upgrade()
|
||||
for table, columns in migration.INDEXES:
|
||||
for table, name, columns in migration.INDEXES:
|
||||
indexes = {
|
||||
index['name']: index['column_names'] for index in sa.inspect(connection).get_indexes(table)
|
||||
}
|
||||
assert indexes[f'ix_{table}_media_ttl'] == columns
|
||||
orm_index = next(
|
||||
index for index in Base.metadata.tables[table].indexes if index.name == f'ix_{table}_media_ttl'
|
||||
)
|
||||
assert indexes[name] == columns
|
||||
orm_index = next(index for index in Base.metadata.tables[table].indexes if index.name == name)
|
||||
assert [column.name for column in orm_index.columns] == columns
|
||||
query = (
|
||||
sa.select(MonitoringMessage.id, MonitoringMessage.message_content)
|
||||
@@ -343,11 +341,8 @@ def test_ttl_indexes_migrate_idempotently_and_support_cursor_queries():
|
||||
engine.dispose()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'revision', ['0039_execution_lease_indexes', '0040_conversation_page_index', '0041_inline_media_ttl_indexes']
|
||||
)
|
||||
def test_index_migrations_allow_historical_schemas_without_monitoring_tables(revision):
|
||||
migration = importlib.import_module(f'langbot.pkg.persistence.alembic.versions.{revision}')
|
||||
def test_index_migration_allows_historical_schemas_without_monitoring_tables():
|
||||
migration = importlib.import_module('langbot.pkg.persistence.alembic.versions.0039_execution_lease_indexes')
|
||||
engine = sa.create_engine('sqlite:///:memory:')
|
||||
try:
|
||||
with engine.begin() as connection:
|
||||
|
||||
Reference in New Issue
Block a user