From 9cd3544d59600fcb88700d05e4b211f59ac00445 Mon Sep 17 00:00:00 2001 From: RockChinQ Date: Wed, 8 Apr 2026 23:33:13 +0800 Subject: [PATCH] feat: integrate Alembic for database migrations Replace manual if-sqlite/if-postgres branching with Alembic: - Add alembic dependency - Create programmatic alembic env (no CLI/alembic.ini needed) - Support async engines via run_sync passthrough - render_as_batch=True for SQLite ALTER TABLE compatibility - Auto-stamp baseline on first run (existing DB at version 25) - Run alembic upgrade head after legacy migrations - Include sample migration showing schema + data migration patterns - Add alembic dir to package-data for distribution --- pyproject.toml | 3 +- .../pkg/persistence/alembic/__init__.py | 0 src/langbot/pkg/persistence/alembic/env.py | 51 ++++++++++++++ .../pkg/persistence/alembic/script.py.mako | 24 +++++++ .../alembic/versions/0001_baseline.py | 24 +++++++ .../alembic/versions/0002_sample.py | 62 +++++++++++++++++ .../persistence/alembic/versions/__init__.py | 0 src/langbot/pkg/persistence/alembic_runner.py | 68 +++++++++++++++++++ src/langbot/pkg/persistence/mgr.py | 25 +++++++ 9 files changed, 256 insertions(+), 1 deletion(-) create mode 100644 src/langbot/pkg/persistence/alembic/__init__.py create mode 100644 src/langbot/pkg/persistence/alembic/env.py create mode 100644 src/langbot/pkg/persistence/alembic/script.py.mako create mode 100644 src/langbot/pkg/persistence/alembic/versions/0001_baseline.py create mode 100644 src/langbot/pkg/persistence/alembic/versions/0002_sample.py create mode 100644 src/langbot/pkg/persistence/alembic/versions/__init__.py create mode 100644 src/langbot/pkg/persistence/alembic_runner.py diff --git a/pyproject.toml b/pyproject.toml index 9cb7faeb..6ef7ebf3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ dependencies = [ "quart-cors>=0.8.0", "requests>=2.32.3", "slack-sdk>=3.35.0", + "alembic>=1.15.0", "sqlalchemy[asyncio]>=2.0.40", "sqlmodel>=0.0.24", "telegramify-markdown>=0.5.1", @@ -111,7 +112,7 @@ requires = ["setuptools>=61.0", "wheel"] build-backend = "setuptools.build_meta" [tool.setuptools] -package-data = { "langbot" = ["templates/**", "pkg/provider/modelmgr/requesters/*", "pkg/platform/sources/*", "web/dist/**"] } +package-data = { "langbot" = ["templates/**", "pkg/provider/modelmgr/requesters/*", "pkg/platform/sources/*", "web/dist/**", "pkg/persistence/alembic/**"] } [dependency-groups] dev = [ diff --git a/src/langbot/pkg/persistence/alembic/__init__.py b/src/langbot/pkg/persistence/alembic/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/langbot/pkg/persistence/alembic/env.py b/src/langbot/pkg/persistence/alembic/env.py new file mode 100644 index 00000000..40543edd --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/env.py @@ -0,0 +1,51 @@ +"""Alembic environment for LangBot. + +This env.py is designed to be called programmatically (not via CLI). +It supports both SQLite and PostgreSQL. + +The sync connection is passed via config attributes by the runner. +""" + +from __future__ import annotations + +from alembic import context +from sqlalchemy.engine import Connection + +from langbot.pkg.entity.persistence.base import Base + +target_metadata = Base.metadata + + +def run_migrations_offline() -> None: + """Run migrations in 'offline' mode — emit SQL without a live connection.""" + url = context.config.get_main_option('sqlalchemy.url') + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={'paramstyle': 'named'}, + ) + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + """Run migrations with a live sync connection passed via config attributes.""" + connection: Connection = context.config.attributes.get('connection') + if connection is None: + raise RuntimeError('connection not provided in alembic config attributes') + + context.configure( + connection=connection, + target_metadata=target_metadata, + # render_as_batch=True is critical for SQLite ALTER TABLE support + render_as_batch=True, + ) + with context.begin_transaction(): + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/src/langbot/pkg/persistence/alembic/script.py.mako b/src/langbot/pkg/persistence/alembic/script.py.mako new file mode 100644 index 00000000..68feb3ea --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/script.py.mako @@ -0,0 +1,24 @@ +# Alembic script.py.mako — template for auto-generated revisions +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} +""" +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers +revision = ${repr(up_revision)} +down_revision = ${repr(down_revision)} +branch_labels = ${repr(branch_labels)} +depends_on = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/src/langbot/pkg/persistence/alembic/versions/0001_baseline.py b/src/langbot/pkg/persistence/alembic/versions/0001_baseline.py new file mode 100644 index 00000000..929356e0 --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/versions/0001_baseline.py @@ -0,0 +1,24 @@ +"""baseline: stamp existing schema (db version 25) + +This is a no-op migration that marks the starting point for Alembic. +All tables already exist via create_all() + legacy DBMigration system. + +Revision ID: 0001_baseline +Revises: None +Create Date: 2026-04-08 +""" + +revision = '0001_baseline' +down_revision = None +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # No-op: existing schema is already at database_version=25 + # This revision serves as the Alembic baseline. + pass + + +def downgrade() -> None: + pass diff --git a/src/langbot/pkg/persistence/alembic/versions/0002_sample.py b/src/langbot/pkg/persistence/alembic/versions/0002_sample.py new file mode 100644 index 00000000..9d869c12 --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/versions/0002_sample.py @@ -0,0 +1,62 @@ +"""example: sample migration demonstrating Alembic patterns + +This is a SAMPLE showing how to write migrations that work +seamlessly across SQLite and PostgreSQL. Delete or adapt as needed. + +Revision ID: 0002_sample +Revises: 0001_baseline +Create Date: 2026-04-08 + +Patterns demonstrated: + 1. Schema change (add column) — works on both DBs via render_as_batch + 2. Data migration (read + modify JSON) — pure SQLAlchemy, no dialect branching +""" + +revision = '0002_sample' +down_revision = '0001_baseline' +branch_labels = None +depends_on = None + + +def upgrade() -> None: + """ + EXAMPLE: Uncomment to use. This shows the patterns. + + # --- Pattern 1: Schema change (add/drop column) --- + # render_as_batch=True in env.py makes this work on SQLite too. + # + # op.add_column('pipelines', sa.Column('description', sa.String(512), server_default='')) + + # --- Pattern 2: Data migration (read + modify JSON field) --- + # No if/else for sqlite vs postgres needed! + # + # conn = op.get_bind() + # rows = conn.execute(sa.text("SELECT uuid, config FROM pipelines")).fetchall() + # for row in rows: + # config = json.loads(row[1]) if isinstance(row[1], str) else row[1] + # # Modify the config + # config.setdefault('ai', {}).setdefault('some_new_key', 'default_value') + # conn.execute( + # sa.text("UPDATE pipelines SET config = :cfg WHERE uuid = :uuid"), + # {"cfg": json.dumps(config), "uuid": row[0]} + # ) + + # --- Pattern 3: Create a new table --- + # + # op.create_table( + # 'audit_log', + # sa.Column('id', sa.Integer, primary_key=True, autoincrement=True), + # sa.Column('action', sa.String(255), nullable=False), + # sa.Column('detail', sa.Text), + # sa.Column('created_at', sa.DateTime, server_default=sa.func.now()), + # ) + """ + pass + + +def downgrade() -> None: + """ + # op.drop_column('pipelines', 'description') + # op.drop_table('audit_log') + """ + pass diff --git a/src/langbot/pkg/persistence/alembic/versions/__init__.py b/src/langbot/pkg/persistence/alembic/versions/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/langbot/pkg/persistence/alembic_runner.py b/src/langbot/pkg/persistence/alembic_runner.py new file mode 100644 index 00000000..535931f0 --- /dev/null +++ b/src/langbot/pkg/persistence/alembic_runner.py @@ -0,0 +1,68 @@ +"""Programmatic Alembic runner for LangBot. + +Usage from async code: + from langbot.pkg.persistence.alembic_runner import run_alembic_upgrade + await run_alembic_upgrade(async_engine) +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +from alembic.config import Config +from alembic import command +from alembic.runtime.migration import MigrationContext + +if TYPE_CHECKING: + from sqlalchemy.ext.asyncio import AsyncEngine + from sqlalchemy.engine import Connection + + +_ALEMBIC_DIR = os.path.join(os.path.dirname(__file__), 'alembic') + + +def _build_config(connection: Connection) -> Config: + """Build an Alembic Config with sync connection attached.""" + cfg = Config() + cfg.set_main_option('script_location', _ALEMBIC_DIR) + cfg.attributes['connection'] = connection + return cfg + + +def _do_upgrade(connection: Connection, revision: str = 'head') -> None: + """Synchronous upgrade — runs inside run_sync.""" + cfg = _build_config(connection) + command.upgrade(cfg, revision) + + +def _do_stamp(connection: Connection, revision: str = 'head') -> None: + """Synchronous stamp — runs inside run_sync.""" + cfg = _build_config(connection) + command.stamp(cfg, revision) + + +def _do_get_current(connection: Connection) -> str | None: + """Get current alembic revision synchronously.""" + ctx = MigrationContext.configure(connection) + return ctx.get_current_revision() + + +async def run_alembic_upgrade(async_engine: AsyncEngine, revision: str = 'head') -> None: + """Run Alembic upgrade to the given revision.""" + async with async_engine.connect() as conn: + await conn.run_sync(_do_upgrade, revision) + await conn.commit() + + +async def run_alembic_stamp(async_engine: AsyncEngine, revision: str = 'head') -> None: + """Stamp the database with a revision without running migrations.""" + async with async_engine.connect() as conn: + await conn.run_sync(_do_stamp, 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: + return await conn.run_sync(_do_get_current) diff --git a/src/langbot/pkg/persistence/mgr.py b/src/langbot/pkg/persistence/mgr.py index ead20a8b..7ad7b968 100644 --- a/src/langbot/pkg/persistence/mgr.py +++ b/src/langbot/pkg/persistence/mgr.py @@ -76,6 +76,9 @@ class PersistenceManager: self.ap.logger.info(f'Successfully upgraded database to version {last_migration_number}.') + # Run Alembic migrations (new migration system) + await self._run_alembic_migrations() + await self.write_space_model_providers() async def create_tables(self): @@ -135,6 +138,28 @@ class PersistenceManager: # ================================= + async def _run_alembic_migrations(self): + """Run Alembic-based migrations after legacy migrations complete.""" + from . import alembic_runner + + engine = self.get_db_engine() + + try: + current_rev = await alembic_runner.get_alembic_current(engine) + + if current_rev is None: + # First time: stamp baseline so Alembic knows existing schema is up-to-date + self.ap.logger.info('Alembic: no revision found, stamping baseline...') + await alembic_runner.run_alembic_stamp(engine, '0001_baseline') + current_rev = '0001_baseline' + + # Upgrade to head + 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 execute_async(self, *args, **kwargs) -> sqlalchemy.engine.cursor.CursorResult: async with self.get_db_engine().connect() as conn: result = await conn.execute(*args, **kwargs)