54 lines
1.9 KiB
Python
54 lines
1.9 KiB
Python
"""Small SQLite store for OpenViking user keys."""
|
|
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
|
|
class OpenVikingUserKeyStore:
|
|
def __init__(self, sqlite_path: str) -> None:
|
|
self.sqlite_path = sqlite_path
|
|
self._ensure_table()
|
|
|
|
def get_user_key(self, user_id: str) -> str | None:
|
|
with self._connect() as conn:
|
|
row = conn.execute(
|
|
"SELECT user_key FROM memory_system_openviking_users WHERE user_id = ?",
|
|
(user_id,),
|
|
).fetchone()
|
|
return str(row[0]) if row else None
|
|
|
|
def save_user_key(self, user_id: str, user_key: str) -> None:
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
with self._connect() as conn:
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO memory_system_openviking_users (user_id, account_id, user_key, created_at, updated_at)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
ON CONFLICT(user_id) DO UPDATE SET
|
|
user_key = excluded.user_key,
|
|
updated_at = excluded.updated_at
|
|
""",
|
|
(user_id, user_id, user_key, now, now),
|
|
)
|
|
|
|
def _ensure_table(self) -> None:
|
|
path = Path(self.sqlite_path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with self._connect() as conn:
|
|
conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS memory_system_openviking_users (
|
|
user_id TEXT PRIMARY KEY,
|
|
account_id TEXT NOT NULL,
|
|
user_key TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
)
|
|
"""
|
|
)
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
return sqlite3.connect(self.sqlite_path)
|