| """ |
| SQLite-backed session store with TTL for RAG sessions. |
| Replaces the in-memory dict to persist sessions across server restarts |
| and prevent unbounded memory growth. |
| """ |
| import sqlite3 |
| import json |
| import time |
| from pathlib import Path |
| from typing import Optional, Dict, Any |
|
|
| SESSION_TTL_SECONDS = 24 * 60 * 60 |
|
|
|
|
| class SessionStore: |
| def __init__(self, db_path: Optional[str] = None): |
| if db_path is None: |
| db_path = str(Path(__file__).parent.parent / "sessions.db") |
| self.db_path = db_path |
| self._init_db() |
| print(f"[SessionStore] SQLite store at: {self.db_path}") |
|
|
| def _get_conn(self) -> sqlite3.Connection: |
| conn = sqlite3.connect(self.db_path, check_same_thread=False) |
| conn.execute("PRAGMA journal_mode=WAL") |
| conn.row_factory = sqlite3.Row |
| return conn |
|
|
| def _init_db(self): |
| with self._get_conn() as conn: |
| conn.execute( |
| """ |
| CREATE TABLE IF NOT EXISTS sessions ( |
| session_id TEXT PRIMARY KEY, |
| data TEXT NOT NULL, |
| created_at REAL NOT NULL, |
| updated_at REAL NOT NULL |
| ) |
| """ |
| ) |
| conn.commit() |
|
|
| def get(self, session_id: str) -> Optional[Dict[str, Any]]: |
| """Return session data or None if not found / expired.""" |
| with self._get_conn() as conn: |
| row = conn.execute( |
| "SELECT data, updated_at FROM sessions WHERE session_id = ?", |
| (session_id,), |
| ).fetchone() |
|
|
| if row is None: |
| return None |
|
|
| if time.time() - row["updated_at"] > SESSION_TTL_SECONDS: |
| self.delete(session_id) |
| return None |
|
|
| return json.loads(row["data"]) |
|
|
| def set(self, session_id: str, data: Dict[str, Any]): |
| """Persist or update a session (all values must be JSON-serialisable).""" |
| now = time.time() |
| data_json = json.dumps(data, ensure_ascii=False) |
| with self._get_conn() as conn: |
| conn.execute( |
| """ |
| INSERT INTO sessions (session_id, data, created_at, updated_at) |
| VALUES (?, ?, ?, ?) |
| ON CONFLICT(session_id) DO UPDATE |
| SET data = excluded.data, updated_at = excluded.updated_at |
| """, |
| (session_id, data_json, now, now), |
| ) |
| conn.commit() |
|
|
| def delete(self, session_id: str): |
| with self._get_conn() as conn: |
| conn.execute( |
| "DELETE FROM sessions WHERE session_id = ?", (session_id,) |
| ) |
| conn.commit() |
|
|
| def cleanup_expired(self) -> int: |
| """Delete all sessions older than SESSION_TTL_SECONDS. Returns count deleted.""" |
| cutoff = time.time() - SESSION_TTL_SECONDS |
| with self._get_conn() as conn: |
| cur = conn.execute( |
| "DELETE FROM sessions WHERE updated_at < ?", (cutoff,) |
| ) |
| conn.commit() |
| deleted = cur.rowcount |
| if deleted: |
| print(f"[SessionStore] Cleaned up {deleted} expired session(s)") |
| return deleted |
|
|
| def cleanup_expired(self) -> int: |
| """Remove sessions older than TTL. Returns number of rows deleted.""" |
| cutoff = time.time() - SESSION_TTL_SECONDS |
| with self._get_conn() as conn: |
| deleted = conn.execute( |
| "DELETE FROM sessions WHERE updated_at < ?", (cutoff,) |
| ).rowcount |
| conn.commit() |
| if deleted: |
| print(f"[SessionStore] Cleaned up {deleted} expired session(s)") |
| return deleted |
|
|
