NCBM/plyngent
1
0
Fork 0
mirror of https://github.com/NCBM/plyngent.git synced 2026-11-07 07:17:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

core/memory: merge session schema migrations by version

Browse Source
This commit is contained in:
worldmozara
2026-08-01 02:47:08 +08:00
parent 8954d855d5
commit a834984bcc
2 changed files with 191 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/plyngent/memory/database/store.py
+60 -27
View File
@@ -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}"))
tests/test_memory/test_store.py
+131
View File
@@ -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