moving and cleaning things

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

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.",

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():

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.

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