diff --git a/src/ell/stores/migrations/__init__.py b/src/ell/stores/migrations/__init__.py new file mode 100644 index 000000000..14602592c --- /dev/null +++ b/src/ell/stores/migrations/__init__.py @@ -0,0 +1,80 @@ + +from alembic import command +from alembic.config import Config +from sqlalchemy import inspect, text +from pathlib import Path + +from sqlmodel import Session, SQLModel, create_engine, select +import logging + +logger = logging.getLogger(__name__) + +def get_alembic_config(engine_url: str) -> Config: + """Create Alembic config programmatically""" + alembic_cfg = Config() + migrations_dir = Path(__file__).parent + + alembic_cfg.set_main_option("script_location", str(migrations_dir)) + alembic_cfg.set_main_option("sqlalchemy.url", str(engine_url)) + alembic_cfg.set_main_option("version_table", "ell_alembic_version") + alembic_cfg.set_main_option("timezone", "UTC") + + return alembic_cfg + +def init_or_migrate_database(engine) -> None: + """Initialize or migrate database with ELL schema + + Handles three cases: + 1. Existing database with our tables but no Alembic -> stamp with initial migration + 2. Database with Alembic -> upgrade to head + 3. New/empty database or database without our tables -> create tables and stamp with head + + Args: + engine_or_url: SQLAlchemy engine or database URL string + """ + inspector = inspect(engine) + + # Check database state + our_tables_v1 = {'serializedlmp', 'invocation', 'invocationcontents', + 'invocationtrace', 'serializedlmpuses'} + existing_tables = set(inspector.get_table_names()) + has_our_tables = bool(our_tables_v1 & existing_tables) # Intersection + has_alembic = 'ell_alembic_version' in existing_tables + + alembic_cfg = get_alembic_config(engine.url) + try: + if has_our_tables and not has_alembic: + # Case 1: Existing database with our tables but no Alembic + # This is likely a database from version <= 0.14 + logger.debug("Found existing tables but no Alembic - stamping with initial migration") + + command.stamp(alembic_cfg, "4524fb60d23e") + # Verify table was created + after_tables = set(inspect(engine).get_table_names()) + logger.debug(f"Tables after stamp: {after_tables}") + + # Check if version table has our stamp + with engine.connect() as connection: + version_result = connection.execute(text("SELECT version_num FROM ell_alembic_version")).first() + if not version_result or version_result[0] != "4524fb60d23e": + raise RuntimeError("Failed to stamp database - version table empty or incorrect version") + logger.debug(f"Successfully stamped database with version {version_result[0]}") + + has_alembic = True + + if has_alembic: + # Case 2: Database has Alembic - run any pending migrations + logger.debug("Running any pending Alembic migrations") + command.upgrade(alembic_cfg, "head") + + else: + # Case 3: New database or database without our tables + logger.debug("New database detected - creating schema and stamping with latest migration") + # Create all tables according to current schema + SQLModel.metadata.create_all(engine) + # Stamp with latest migration + command.stamp(alembic_cfg, "head") + + except Exception as e: + logger.error(f"Failed to initialize/migrate database: {e}") + raise diff --git a/src/ell/stores/migrations/env.py b/src/ell/stores/migrations/env.py index aa284ef50..bf2dc98e2 100644 --- a/src/ell/stores/migrations/env.py +++ b/src/ell/stores/migrations/env.py @@ -42,6 +42,7 @@ def run_migrations_offline() -> None: script output. """ + #XXX: These are currently untested and unused. url = config.get_main_option("sqlalchemy.url") version_table = config.get_main_option("version_table", "ell_alembic_version") @@ -68,6 +69,8 @@ def run_migrations_online() -> None: and associate a connection with the context. """ + + #XXX: These are currently untested and unused. connectable = engine_from_config( config.get_section(config.config_ini_section, {}), prefix="sqlalchemy.", diff --git a/src/ell/stores/migrations/make.py b/src/ell/stores/migrations/make.py index 6193aea60..fec1aaac3 100644 --- a/src/ell/stores/migrations/make.py +++ b/src/ell/stores/migrations/make.py @@ -1,6 +1,6 @@ import argparse from sqlalchemy import create_engine -from ell.stores.sql import get_alembic_config +from ell.stores.migrations import get_alembic_config from alembic import command def main(): diff --git a/src/ell/stores/sql.py b/src/ell/stores/sql.py index 22f3dd451..ebfbb9ab5 100644 --- a/src/ell/stores/sql.py +++ b/src/ell/stores/sql.py @@ -3,6 +3,7 @@ from pathlib import Path from typing import Any, Optional, Dict, List, Set from sqlmodel import Session, SQLModel, create_engine, select +from ell.stores.migrations import init_or_migrate_database import ell.stores.store from sqlalchemy.sql import text from ell.stores.studio import InvocationTrace, SerializedLMP, Invocation @@ -11,84 +12,10 @@ import gzip import json -from alembic import command -from alembic.config import Config -from sqlalchemy import inspect - import logging logger = logging.getLogger(__name__) -def get_alembic_config(engine_url: str) -> Config: - """Create Alembic config programmatically""" - alembic_cfg = Config() - migrations_dir = Path(__file__).parent / "migrations" - - alembic_cfg.set_main_option("script_location", str(migrations_dir)) - alembic_cfg.set_main_option("sqlalchemy.url", str(engine_url)) - alembic_cfg.set_main_option("version_table", "ell_alembic_version") - alembic_cfg.set_main_option("timezone", "UTC") - - return alembic_cfg - -def init_or_migrate_database(engine) -> None: - """Initialize or migrate database with ELL schema - - Handles three cases: - 1. Existing database with our tables but no Alembic -> stamp with initial migration - 2. Database with Alembic -> upgrade to head - 3. New/empty database or database without our tables -> create tables and stamp with head - - Args: - engine_or_url: SQLAlchemy engine or database URL string - """ - inspector = inspect(engine) - - # Check database state - our_tables_v1 = {'serializedlmp', 'invocation', 'invocationcontents', - 'invocationtrace', 'serializedlmpuses'} - existing_tables = set(inspector.get_table_names()) - has_our_tables = bool(our_tables_v1 & existing_tables) # Intersection - has_alembic = 'ell_alembic_version' in existing_tables - - alembic_cfg = get_alembic_config(engine.url) - try: - if has_our_tables and not has_alembic: - # Case 1: Existing database with our tables but no Alembic - # This is likely a database from version <= 0.14 - logger.debug("Found existing tables but no Alembic - stamping with initial migration") - - command.stamp(alembic_cfg, "4524fb60d23e") - # Verify table was created - after_tables = set(inspect(engine).get_table_names()) - logger.debug(f"Tables after stamp: {after_tables}") - - # Check if version table has our stamp - with engine.connect() as connection: - version_result = connection.execute(text("SELECT version_num FROM ell_alembic_version")).first() - if not version_result or version_result[0] != "4524fb60d23e": - raise RuntimeError("Failed to stamp database - version table empty or incorrect version") - logger.debug(f"Successfully stamped database with version {version_result[0]}") - - has_alembic = True - - if has_alembic: - # Case 2: Database has Alembic - run any pending migrations - logger.debug("Running any pending Alembic migrations") - command.upgrade(alembic_cfg, "head") - - else: - # Case 3: New database or database without our tables - logger.debug("New database detected - creating schema and stamping with latest migration") - # Create all tables according to current schema - SQLModel.metadata.create_all(engine) - # Stamp with latest migration - command.stamp(alembic_cfg, "head") - - except Exception as e: - logger.error(f"Failed to initialize/migrate database: {e}") - raise - class SQLStore(ell.stores.store.Store): def __init__(self, db_uri: str, blob_store: Optional[ell.stores.store.BlobStore] = None): # XXX: Use Serialization serialzie_object in incoming PR. diff --git a/tests/test_migrations.py b/tests/test_migrations.py index a50f85766..ebb0cf190 100644 --- a/tests/test_migrations.py +++ b/tests/test_migrations.py @@ -5,7 +5,7 @@ from sqlalchemy import inspect, MetaData, create_engine, text from sqlmodel import SQLModel -from ell.stores.sql import init_or_migrate_database, get_alembic_config +from ell.stores.migrations import init_or_migrate_database, get_alembic_config from alembic import command from alembic.config import Config