diff --git a/src/langbot/pkg/persistence/alembic/versions/8f24d6c9b1a0_ensure_mcp_readme_column.py b/src/langbot/pkg/persistence/alembic/versions/8f24d6c9b1a0_ensure_mcp_readme_column.py new file mode 100644 index 00000000..e859c74a --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/versions/8f24d6c9b1a0_ensure_mcp_readme_column.py @@ -0,0 +1,37 @@ +"""ensure mcp_servers readme column exists + +Revision ID: 8f24d6c9b1a0 +Revises: 7b2c1d9e4f30 +Create Date: 2026-06-13 +""" + +from alembic import op +import sqlalchemy as sa + + +revision = '8f24d6c9b1a0' +down_revision = '7b2c1d9e4f30' +branch_labels = None +depends_on = None + + +def _table_exists(table_name: str) -> bool: + return table_name in sa.inspect(op.get_bind()).get_table_names() + + +def _column_exists(table_name: str, column_name: str) -> bool: + return column_name in {column['name'] for column in sa.inspect(op.get_bind()).get_columns(table_name)} + + +def upgrade() -> None: + if not _table_exists('mcp_servers') or _column_exists('mcp_servers', 'readme'): + return + with op.batch_alter_table('mcp_servers', schema=None) as batch_op: + batch_op.add_column(sa.Column('readme', sa.Text(), nullable=False, server_default='')) + + +def downgrade() -> None: + if not _table_exists('mcp_servers') or not _column_exists('mcp_servers', 'readme'): + return + with op.batch_alter_table('mcp_servers', schema=None) as batch_op: + batch_op.drop_column('readme') diff --git a/tests/integration/persistence/test_migrations.py b/tests/integration/persistence/test_migrations.py index 1eadfcc3..e576042c 100644 --- a/tests/integration/persistence/test_migrations.py +++ b/tests/integration/persistence/test_migrations.py @@ -11,6 +11,7 @@ from __future__ import annotations import pytest from alembic.script import ScriptDirectory +from sqlalchemy import inspect, text from sqlalchemy.ext.asyncio import create_async_engine from langbot.pkg.entity.persistence.base import Base @@ -147,6 +148,41 @@ class TestSQLiteMigrationUpgrade: rev2 = await get_alembic_current(sqlite_engine) assert rev2 == rev1, f"Expected {rev1}, got {rev2}" + @pytest.mark.asyncio + async def test_upgrade_repairs_head_stamped_mcp_readme_column(self, sqlite_engine): + """ + A database may already be stamped at the previous head while missing a + column added by an earlier guarded migration. Upgrade should still + repair mcp_servers.readme so startup ORM queries do not fail. + """ + async with sqlite_engine.begin() as conn: + await conn.execute( + text( + """ + CREATE TABLE mcp_servers ( + uuid VARCHAR(255) NOT NULL PRIMARY KEY, + name VARCHAR(255) NOT NULL, + enable BOOLEAN NOT NULL, + mode VARCHAR(255) NOT NULL, + extra_args JSON NOT NULL, + created_at DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL, + updated_at DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL + ) + """ + ) + ) + + await run_alembic_stamp(sqlite_engine, '7b2c1d9e4f30') + await run_alembic_upgrade(sqlite_engine, 'head') + + async with sqlite_engine.connect() as conn: + columns = await conn.run_sync( + lambda sync_conn: {column['name'] for column in inspect(sync_conn).get_columns('mcp_servers')} + ) + + assert 'readme' in columns + assert await get_alembic_current(sqlite_engine) == alembic_head_revision() + class TestSQLiteMigrationFreshDatabase: """Tests for fresh database workflow.""" diff --git a/tests/integration/persistence/test_migrations_postgres.py b/tests/integration/persistence/test_migrations_postgres.py index 7867d4af..3b0abc6b 100644 --- a/tests/integration/persistence/test_migrations_postgres.py +++ b/tests/integration/persistence/test_migrations_postgres.py @@ -15,11 +15,14 @@ from __future__ import annotations import os import pytest +from alembic.config import Config +from alembic.script import ScriptDirectory from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy import text from langbot.pkg.entity.persistence.base import Base from langbot.pkg.persistence.alembic_runner import ( + _ALEMBIC_DIR, run_alembic_upgrade, run_alembic_stamp, get_alembic_current, @@ -29,6 +32,13 @@ from langbot.pkg.persistence.alembic_runner import ( pytestmark = [pytest.mark.integration, pytest.mark.slow] +def alembic_head_revision() -> str: + """Return the repository's current Alembic head revision.""" + cfg = Config() + cfg.set_main_option('script_location', _ALEMBIC_DIR) + return ScriptDirectory.from_config(cfg).get_current_head() + + @pytest.fixture def postgres_url(): """Get PostgreSQL URL from environment.""" @@ -150,8 +160,7 @@ class TestPostgreSQLMigrationUpgrade: # Verify revision rev = await get_alembic_current(postgres_engine) assert rev is not None, "Expected a revision after upgrade" - # Head should be the latest migration (0004 for current state) - assert rev.startswith('0004'), f"Expected head to be 0004_*, got {rev}" + assert rev == alembic_head_revision() @pytest.mark.asyncio async def test_postgres_upgrade_idempotent( @@ -214,4 +223,4 @@ class TestPostgreSQLMigrationGetCurrent: await run_alembic_stamp(postgres_engine, '0001_baseline') rev = await get_alembic_current(postgres_engine) - assert rev == '0001_baseline' \ No newline at end of file + assert rev == '0001_baseline'