core/memory: merge session schema migrations by version

src/plyngent/memory/database/store.py

@@ -15,7 +15,9 @@ from .engine import create_engine
 from .schema import Message, PlyngentBase, Session, User
 
 if TYPE_CHECKING:
-    from collections.abc import Mapping, Sequence
+    from collections.abc import Callable, Mapping, Sequence
+
+    type MigrationStep = Callable[[object, set[str]], None]
 
 
 DEFAULT_USER_NAME = "local"
@@ -63,13 +65,17 @@ class MemoryStore:
         return store
 
     async def create_schema(self) -> None:
-        """Create all tables if they do not exist; apply lightweight migrations."""
+        """Create all tables if they do not exist; apply versioned migrations.
+
+        Schema evolution is grouped by release: ordered steps
+        (:data:`_SCHEMA_MIGRATIONS`) run only when the stored schema version
+        (:attr:`PRAGMA user_version`) is older, then the version is stamped.
+        Each step guards on its own columns, so legacy DBs that were upgraded
+        incrementally (user_version=0 but columns already present) stay safe.
+        """
         async with self._engine.begin() as conn:
             _ = await conn.run_sync(PlyngentBase.metadata.create_all)
-            await conn.run_sync(_migrate_session_workspace)
+            await conn.run_sync(_run_schema_migrations)
-            await conn.run_sync(_migrate_session_llm)
-            await conn.run_sync(_migrate_session_todo_stack)
-            await conn.run_sync(_migrate_session_context_usage)
 
     async def close(self) -> None:
         """Dispose the underlying engine."""
@@ -343,49 +349,45 @@ def _session_columns(sync_conn: object) -> set[str]:
     return {str(row[1]) for row in rows}
 
 
-def _migrate_session_workspace(sync_conn: object) -> None:
+def _schema_version(sync_conn: object) -> int:
-    """Add ``session.workspace`` on existing SQLite DBs created before the column existed."""
+    """Read the SQLite ``user_version`` (0 = never migrated/stamped)."""
+    from sqlalchemy.engine import Connection
+
+    if not isinstance(sync_conn, Connection):
+        return 0
+    return int(sync_conn.scalar(text("PRAGMA user_version")) or 0)
+
+
+def _migrate_session_v1(sync_conn: object, columns: set[str]) -> None:
+    """v0.1.0: add ``workspace`` / ``provider_name`` / ``model``."""
     from sqlalchemy.engine import Connection
 
     if not isinstance(sync_conn, Connection):
         return
-    columns = _session_columns(sync_conn)
+    if "workspace" not in columns:
-    if "workspace" in columns:
+        _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN workspace VARCHAR(1024)"))
-        return
-    _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN workspace VARCHAR(1024)"))
-
-
-def _migrate_session_llm(sync_conn: object) -> None:
-    """Add ``session.provider_name`` / ``session.model`` for remembered LLM selection."""
-    from sqlalchemy.engine import Connection
-
-    if not isinstance(sync_conn, Connection):
-        return
-    columns = _session_columns(sync_conn)
     if "provider_name" not in columns:
         _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN provider_name VARCHAR(128)"))
     if "model" not in columns:
         _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN model VARCHAR(256)"))
 
 
-def _migrate_session_todo_stack(sync_conn: object) -> None:
+def _migrate_session_v2(sync_conn: object, columns: set[str]) -> None:
-    """Add ``session.todo_stack`` JSON for the todo/task sub-task stack."""
+    """v0.1.2: add ``todo_stack`` JSON for the todo/task sub-task stack."""
     from sqlalchemy.engine import Connection
 
     if not isinstance(sync_conn, Connection):
         return
-    columns = _session_columns(sync_conn)
     if "todo_stack" not in columns:
         _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN todo_stack JSON"))
 
 
-def _migrate_session_context_usage(sync_conn: object) -> None:
+def _migrate_session_v3(sync_conn: object, columns: set[str]) -> None:
-    """Add session context usage + directive reminder band columns."""
+    """Post-v0.2.0: add session context usage + directive reminder band columns."""
     from sqlalchemy.engine import Connection
 
     if not isinstance(sync_conn, Connection):
         return
-    columns = _session_columns(sync_conn)
     if "last_prompt_tokens" not in columns:
         _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN last_prompt_tokens INTEGER"))
     if "peak_prompt_tokens" not in columns:
@@ -396,3 +398,34 @@ def _migrate_session_context_usage(sync_conn: object) -> None:
         _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN usage_source VARCHAR(16)"))
     if "reminder_last_band" not in columns:
         _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN reminder_last_band INTEGER"))
+
+
+# Latest schema version this build knows. Bump (and append a step to
+# _SCHEMA_MIGRATIONS) when a new release adds or changes columns.
+_SCHEMA_VERSION = 3
+
+# Ordered (version, step) pairs. Each step receives the current ``session``
+# column set (read once) and adds only its own missing columns, so it is
+# idempotent for DBs that were upgraded incrementally without a version stamp.
+# Version 1 shipped in v0.1.0, version 2 in v0.1.2, version 3 post-v0.2.0.
+_SCHEMA_MIGRATIONS: tuple[tuple[int, MigrationStep], ...] = (
+    (1, _migrate_session_v1),
+    (2, _migrate_session_v2),
+    (3, _migrate_session_v3),
+)
+
+
+def _run_schema_migrations(sync_conn: object) -> None:
+    """Run pending versioned ``session`` migrations, then stamp the version."""
+    from sqlalchemy.engine import Connection
+
+    if not isinstance(sync_conn, Connection):
+        return
+    version = _schema_version(sync_conn)
+    if version >= _SCHEMA_VERSION:
+        return
+    columns = _session_columns(sync_conn)
+    for step_version, step in _SCHEMA_MIGRATIONS:
+        if step_version > version:
+            step(sync_conn, columns)
+    _ = sync_conn.execute(text(f"PRAGMA user_version = {_SCHEMA_VERSION}"))

tests/test_memory/test_store.py

@@ -224,3 +224,134 @@ async def test_workspace_column_migration(tmp_path: object) -> None:
 
     assert session.workspace == normalize_workspace(tmp_path)
     await store.close()
+
+
+async def _user_version(db_path: object) -> int:
+    from sqlalchemy import text
+    from sqlalchemy.ext.asyncio import create_async_engine
+
+    engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}")
+    try:
+        async with engine.connect() as conn:
+            result = await conn.execute(text("PRAGMA user_version"))
+            return int(result.scalar() or 0)
+    finally:
+        await engine.dispose()
+
+
+async def _session_columns(db_path: object) -> set[str]:
+    from sqlalchemy import text
+    from sqlalchemy.ext.asyncio import create_async_engine
+
+    engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}")
+    try:
+        async with engine.connect() as conn:
+            result = await conn.execute(text("PRAGMA table_info(session)"))
+            return {str(row[1]) for row in result.fetchall()}
+    finally:
+        await engine.dispose()
+
+
+async def _create_legacy_db(db_path: object, *, with_workspace: bool) -> None:
+    """Build a pre-migration DB with a minimal session table (optional workspace)."""
+    from sqlalchemy import text
+    from sqlalchemy.ext.asyncio import create_async_engine
+
+    engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}")
+    async with engine.begin() as conn:
+        _ = await conn.execute(
+            text(
+                "CREATE TABLE user ("
+                "uid INTEGER PRIMARY KEY, name VARCHAR(48) UNIQUE, "
+                "email VARCHAR(255) UNIQUE, password_hash VARCHAR(256), "
+                "created_at DATETIME)"
+            )
+        )
+        workspace_col = ", workspace VARCHAR(1024)" if with_workspace else ""
+        _ = await conn.execute(
+            text(
+                "CREATE TABLE session ("
+                "sid INTEGER PRIMARY KEY, uid INTEGER, name VARCHAR(64), "
+                f"created_at DATETIME, updated_at DATETIME{workspace_col})"
+            )
+        )
+        _ = await conn.execute(
+            text(
+                "CREATE TABLE message ("
+                "mid INTEGER PRIMARY KEY, sid INTEGER, seq INTEGER, "
+                "data JSON, created_at DATETIME, updated_at DATETIME)"
+            )
+        )
+    await engine.dispose()
+
+
+async def test_schema_migrations_stamp_version(tmp_path: object) -> None:
+    """A fresh open creates the current schema and stamps PRAGMA user_version."""
+    from pathlib import Path
+
+    assert isinstance(tmp_path, Path)
+    db_path = tmp_path / "fresh.db"
+    store = await MemoryStore.open(DatabaseConfig(url=str(db_path)))
+    await store.close()
+    assert await _user_version(db_path) == 3
+    columns = await _session_columns(db_path)
+    assert {"workspace", "provider_name", "model", "todo_stack", "last_prompt_tokens"} <= columns
+
+
+async def test_legacy_db_migrated_to_latest(tmp_path: object) -> None:
+    """A pre-v0.1.0 session table runs all versioned steps and gets stamped."""
+    from pathlib import Path
+
+    assert isinstance(tmp_path, Path)
+    db_path = tmp_path / "legacy.db"
+    await _create_legacy_db(db_path, with_workspace=False)
+    assert await _user_version(db_path) == 0
+
+    store = await MemoryStore.open(DatabaseConfig(url=str(db_path)))
+    session = await store.create_session(name="migrated")
+    await store.close()
+
+    assert await _user_version(db_path) == 3
+    columns = await _session_columns(db_path)
+    assert {"workspace", "provider_name", "model", "todo_stack", "last_prompt_tokens"} <= columns
+    assert session.sid is not None
+
+
+async def test_partially_migrated_legacy_db_noops_remaining_steps(tmp_path: object) -> None:
+    """A DB upgraded incrementally (user_version=0, workspace present) is safe.
+
+    Versioned steps guard on their own columns, so an unstamped DB that already
+    has the v0.1.0 columns gets the rest applied and then stamped without
+    re-adding (or erroring on) existing columns.
+    """
+    from pathlib import Path
+
+    assert isinstance(tmp_path, Path)
+    db_path = tmp_path / "partial.db"
+    await _create_legacy_db(db_path, with_workspace=True)
+    assert await _user_version(db_path) == 0
+
+    store = await MemoryStore.open(DatabaseConfig(url=str(db_path)))
+    session = await store.create_session(name="partial", workspace=tmp_path)
+    await store.close()
+
+    assert await _user_version(db_path) == 3
+    columns = await _session_columns(db_path)
+    assert "workspace" in columns
+    assert {"provider_name", "model", "todo_stack", "last_prompt_tokens"} <= columns
+    assert session.workspace is not None
+
+
+async def test_stamped_db_skips_migrations_on_reopen(tmp_path: object) -> None:
+    """Reopening a current-schema DB does not re-run migrations or bump the stamp."""
+    from pathlib import Path
+
+    assert isinstance(tmp_path, Path)
+    db_path = tmp_path / "stamped.db"
+    store = await MemoryStore.open(DatabaseConfig(url=str(db_path)))
+    await store.close()
+    assert await _user_version(db_path) == 3
+
+    store = await MemoryStore.open(DatabaseConfig(url=str(db_path)))
+    await store.close()
+    assert await _user_version(db_path) == 3