mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
core/memory: bind sessions to workspace path with SQLite migrate
Add session.workspace column, create/list filters, and ALTER TABLE migration for existing chat databases.
This commit is contained in:
@@ -27,6 +27,8 @@ class Session(PlyngentBase):
|
|||||||
sid: Mapped[int] = mapped_column(primary_key=True)
|
sid: Mapped[int] = mapped_column(primary_key=True)
|
||||||
uid: Mapped[int] = mapped_column(ForeignKey(User.uid, ondelete="CASCADE"), index=True)
|
uid: Mapped[int] = mapped_column(ForeignKey(User.uid, ondelete="CASCADE"), index=True)
|
||||||
name: Mapped[str] = mapped_column(String(64))
|
name: Mapped[str] = mapped_column(String(64))
|
||||||
|
# Absolute workspace path this chat is bound to (tools root); null = legacy unbound.
|
||||||
|
workspace: Mapped[str | None] = mapped_column(String(1024), nullable=True, index=True)
|
||||||
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
|
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
|
||||||
updated_at: Mapped[datetime] = mapped_column(
|
updated_at: Mapped[datetime] = mapped_column(
|
||||||
DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
|
DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
|
||||||
|
|||||||
@@ -1,9 +1,10 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING, Self
|
from typing import TYPE_CHECKING, Self
|
||||||
|
|
||||||
import msgspec
|
import msgspec
|
||||||
from sqlalchemy import select
|
from sqlalchemy import select, text
|
||||||
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
|
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
|
||||||
|
|
||||||
from plyngent.config.models import DatabaseConfig
|
from plyngent.config.models import DatabaseConfig
|
||||||
@@ -21,6 +22,16 @@ DEFAULT_USER_EMAIL = "local@localhost"
|
|||||||
DEFAULT_USER_PASSWORD_HASH = ""
|
DEFAULT_USER_PASSWORD_HASH = ""
|
||||||
|
|
||||||
|
|
||||||
|
def normalize_workspace(path: str | Path | None) -> str | None:
|
||||||
|
"""Return a stable absolute workspace path string, or None if unset."""
|
||||||
|
if path is None:
|
||||||
|
return None
|
||||||
|
text_path = str(path).strip()
|
||||||
|
if not text_path:
|
||||||
|
return None
|
||||||
|
return str(Path(text_path).expanduser().resolve())
|
||||||
|
|
||||||
|
|
||||||
class MemoryStore:
|
class MemoryStore:
|
||||||
"""Async persistence for users, chat sessions, and messages."""
|
"""Async persistence for users, chat sessions, and messages."""
|
||||||
|
|
||||||
@@ -50,9 +61,10 @@ class MemoryStore:
|
|||||||
return store
|
return store
|
||||||
|
|
||||||
async def create_schema(self) -> None:
|
async def create_schema(self) -> None:
|
||||||
"""Create all tables if they do not exist."""
|
"""Create all tables if they do not exist; apply lightweight migrations."""
|
||||||
async with self._engine.begin() as conn:
|
async with self._engine.begin() as conn:
|
||||||
_ = await conn.run_sync(PlyngentBase.metadata.create_all)
|
_ = await conn.run_sync(PlyngentBase.metadata.create_all)
|
||||||
|
await conn.run_sync(_migrate_session_workspace)
|
||||||
|
|
||||||
async def close(self) -> None:
|
async def close(self) -> None:
|
||||||
"""Dispose the underlying engine."""
|
"""Dispose the underlying engine."""
|
||||||
@@ -80,13 +92,23 @@ class MemoryStore:
|
|||||||
result = await session.execute(select(User).where(User.name == name))
|
result = await session.execute(select(User).where(User.name == name))
|
||||||
return result.scalar_one_or_none()
|
return result.scalar_one_or_none()
|
||||||
|
|
||||||
async def create_session(self, *, uid: int | None = None, name: str = "default") -> Session:
|
async def create_session(
|
||||||
"""Create a chat session for ``uid`` (default local user when omitted)."""
|
self,
|
||||||
|
*,
|
||||||
|
uid: int | None = None,
|
||||||
|
name: str = "default",
|
||||||
|
workspace: str | Path | None = None,
|
||||||
|
) -> Session:
|
||||||
|
"""Create a chat session for ``uid`` (default local user when omitted).
|
||||||
|
|
||||||
|
``workspace`` is stored as a resolved absolute path when provided.
|
||||||
|
"""
|
||||||
if uid is None:
|
if uid is None:
|
||||||
user = await self.ensure_default_user()
|
user = await self.ensure_default_user()
|
||||||
uid = user.uid
|
uid = user.uid
|
||||||
|
ws = normalize_workspace(workspace)
|
||||||
async with self._session_factory() as session:
|
async with self._session_factory() as session:
|
||||||
row = Session(uid=uid, name=name)
|
row = Session(uid=uid, name=name, workspace=ws)
|
||||||
session.add(row)
|
session.add(row)
|
||||||
await session.commit()
|
await session.commit()
|
||||||
await session.refresh(row)
|
await session.refresh(row)
|
||||||
@@ -96,14 +118,22 @@ class MemoryStore:
|
|||||||
async with self._session_factory() as session:
|
async with self._session_factory() as session:
|
||||||
return await session.get(Session, sid)
|
return await session.get(Session, sid)
|
||||||
|
|
||||||
async def list_sessions(self, *, uid: int | None = None) -> Sequence[Session]:
|
async def list_sessions(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
uid: int | None = None,
|
||||||
|
workspace: str | Path | None = None,
|
||||||
|
) -> Sequence[Session]:
|
||||||
|
"""List sessions; optionally filter by user and/or bound workspace path."""
|
||||||
|
ws = normalize_workspace(workspace)
|
||||||
async with self._session_factory() as session:
|
async with self._session_factory() as session:
|
||||||
stmt = select(Session).order_by(Session.sid)
|
stmt = select(Session).order_by(Session.sid)
|
||||||
if uid is not None:
|
if uid is not None:
|
||||||
stmt = stmt.where(Session.uid == uid)
|
stmt = stmt.where(Session.uid == uid)
|
||||||
|
if ws is not None:
|
||||||
|
stmt = stmt.where(Session.workspace == ws)
|
||||||
result = await session.execute(stmt)
|
result = await session.execute(stmt)
|
||||||
return result.scalars().all()
|
return result.scalars().all()
|
||||||
|
|
||||||
async def append_message(self, sid: int, message: AnyChatMessage) -> Message:
|
async def append_message(self, sid: int, message: AnyChatMessage) -> Message:
|
||||||
"""Append a chat message to a session with the next sequence number."""
|
"""Append a chat message to a session with the next sequence number."""
|
||||||
data = msgspec.to_builtins(message)
|
data = msgspec.to_builtins(message)
|
||||||
@@ -135,3 +165,16 @@ class MemoryStore:
|
|||||||
async with self._session_factory() as session:
|
async with self._session_factory() as session:
|
||||||
result = await session.execute(select(Message).where(Message.sid == sid).order_by(Message.seq))
|
result = await session.execute(select(Message).where(Message.sid == sid).order_by(Message.seq))
|
||||||
return result.scalars().all()
|
return result.scalars().all()
|
||||||
|
|
||||||
|
|
||||||
|
def _migrate_session_workspace(sync_conn: object) -> None:
|
||||||
|
"""Add ``session.workspace`` on existing SQLite DBs created before the column existed."""
|
||||||
|
from sqlalchemy.engine import Connection
|
||||||
|
|
||||||
|
if not isinstance(sync_conn, Connection):
|
||||||
|
return
|
||||||
|
rows = sync_conn.execute(text("PRAGMA table_info(session)")).fetchall()
|
||||||
|
columns = {str(row[1]) for row in rows}
|
||||||
|
if "workspace" in columns:
|
||||||
|
return
|
||||||
|
_ = sync_conn.execute(text("ALTER TABLE session ADD COLUMN workspace VARCHAR(1024)"))
|
||||||
|
|||||||
@@ -59,3 +59,63 @@ async def test_list_sessions(store: MemoryStore) -> None:
|
|||||||
ids = {s.sid for s in sessions}
|
ids = {s.sid for s in sessions}
|
||||||
assert s1.sid in ids
|
assert s1.sid in ids
|
||||||
assert s2.sid in ids
|
assert s2.sid in ids
|
||||||
|
|
||||||
|
|
||||||
|
async def test_session_workspace_binding(store: MemoryStore, tmp_path: object) -> None:
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
assert isinstance(tmp_path, Path)
|
||||||
|
a = tmp_path / "proj-a"
|
||||||
|
b = tmp_path / "proj-b"
|
||||||
|
a.mkdir()
|
||||||
|
b.mkdir()
|
||||||
|
sa = await store.create_session(name="a", workspace=a)
|
||||||
|
sb = await store.create_session(name="b", workspace=b)
|
||||||
|
assert sa.workspace == str(a.resolve())
|
||||||
|
listed_a = await store.list_sessions(workspace=a)
|
||||||
|
assert {s.sid for s in listed_a} == {sa.sid}
|
||||||
|
listed_b = await store.list_sessions(workspace=b)
|
||||||
|
assert {s.sid for s in listed_b} == {sb.sid}
|
||||||
|
|
||||||
|
|
||||||
|
async def test_workspace_column_migration(tmp_path: object) -> None:
|
||||||
|
"""Existing DBs without session.workspace get the column via ALTER."""
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from sqlalchemy import text
|
||||||
|
from sqlalchemy.ext.asyncio import create_async_engine
|
||||||
|
|
||||||
|
assert isinstance(tmp_path, Path)
|
||||||
|
db_path = tmp_path / "legacy.db"
|
||||||
|
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)"
|
||||||
|
)
|
||||||
|
)
|
||||||
|
_ = await conn.execute(
|
||||||
|
text(
|
||||||
|
"CREATE TABLE session ("
|
||||||
|
"sid INTEGER PRIMARY KEY, uid INTEGER, name VARCHAR(64), "
|
||||||
|
"created_at DATETIME, updated_at DATETIME)"
|
||||||
|
)
|
||||||
|
)
|
||||||
|
_ = 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()
|
||||||
|
|
||||||
|
store = await MemoryStore.open(DatabaseConfig(url=str(db_path)))
|
||||||
|
session = await store.create_session(name="migrated", workspace=tmp_path)
|
||||||
|
from plyngent.memory.database.store import normalize_workspace
|
||||||
|
|
||||||
|
assert session.workspace == normalize_workspace(tmp_path)
|
||||||
|
await store.close()
|
||||||
|
|||||||
Reference in New Issue
Block a user