From a834984bcc4d52ac723fbf9a72886601acb0896e Mon Sep 17 00:00:00 2001 From: worldmozara Date: Sat, 1 Aug 2026 02:47:08 +0800 Subject: [PATCH] core/memory: merge session schema migrations by version --- src/plyngent/memory/database/store.py | 87 +++++++++++------ tests/test_memory/test_store.py | 131 ++++++++++++++++++++++++++ 2 files changed, 191 insertions(+), 27 deletions(-) diff --git a/src/plyngent/memory/database/store.py b/src/plyngent/memory/database/store.py index 75b77e1..ecd6a5b 100644 --- a/src/plyngent/memory/database/store.py +++ b/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(_migrate_session_llm) - await conn.run_sync(_migrate_session_todo_stack) - await conn.run_sync(_migrate_session_context_usage) + await conn.run_sync(_run_schema_migrations) 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: - """Add ``session.workspace`` on existing SQLite DBs created before the column existed.""" +def _schema_version(sync_conn: object) -> int: + """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" in columns: - 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 "workspace" not in columns: + _ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN workspace VARCHAR(1024)")) 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: - """Add ``session.todo_stack`` JSON for the todo/task sub-task stack.""" +def _migrate_session_v2(sync_conn: object, columns: set[str]) -> None: + """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: - """Add session context usage + directive reminder band columns.""" +def _migrate_session_v3(sync_conn: object, columns: set[str]) -> None: + """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}")) diff --git a/tests/test_memory/test_store.py b/tests/test_memory/test_store.py index f0b4502..56afd29 100644 --- a/tests/test_memory/test_store.py +++ b/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