83 lines
3.2 KiB
Python
83 lines
3.2 KiB
Python
"""SQLite cache of CRZP searches and records (data/crzp/crzp.sqlite).
|
|
|
|
CRZP is the source of truth; this cache only avoids repeating requests. A
|
|
search result is reused for SEARCH_TTL, a record detail for DETAIL_TTL. The
|
|
interface is small (get/put of searches and details), so another database can
|
|
replace SQLite without touching the callers.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from tools.rag.common import DATA_DIR
|
|
|
|
SEARCH_TTL = float(os.getenv("CRZP_SEARCH_TTL_SECONDS", str(7 * 86400)))
|
|
DETAIL_TTL = float(os.getenv("CRZP_DETAIL_TTL_SECONDS", str(30 * 86400)))
|
|
SCHEMA = 2 # 2: school and workplace split, CRZP work type ids
|
|
|
|
|
|
def default_path() -> Path:
|
|
return Path(os.getenv("CRZP_DB_FILE", str(DATA_DIR / "crzp" / "crzp.sqlite")))
|
|
|
|
|
|
class CrzpStore:
|
|
def __init__(self, path: Path | None = None):
|
|
self.path = Path(path or default_path())
|
|
self._lock = threading.Lock()
|
|
self._ready = False
|
|
|
|
def _connect(self) -> sqlite3.Connection:
|
|
if not self._ready:
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
connection = sqlite3.connect(self.path, timeout=10)
|
|
if not self._ready:
|
|
with self._lock:
|
|
connection.execute("PRAGMA journal_mode=WAL")
|
|
connection.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS searches (
|
|
query TEXT PRIMARY KEY, fetched_at REAL NOT NULL, payload TEXT NOT NULL);
|
|
CREATE TABLE IF NOT EXISTS details (
|
|
sid TEXT PRIMARY KEY, fetched_at REAL NOT NULL, payload TEXT NOT NULL);
|
|
"""
|
|
)
|
|
self._ready = True
|
|
return connection
|
|
|
|
def _get(self, table: str, key_column: str, key: str, ttl: float) -> dict | None:
|
|
with self._connect() as connection:
|
|
row = connection.execute(
|
|
f"SELECT fetched_at, payload FROM {table} WHERE {key_column} = ?", (key,)).fetchone()
|
|
if row is None or time.time() - row[0] > ttl:
|
|
return None
|
|
payload = json.loads(row[1])
|
|
return payload if payload.get("schema") == SCHEMA else None
|
|
|
|
def _put(self, table: str, key: str, payload: dict) -> None:
|
|
with self._connect() as connection:
|
|
connection.execute(f"INSERT OR REPLACE INTO {table} VALUES (?, ?, ?)",
|
|
(key, time.time(), json.dumps({**payload, "schema": SCHEMA}, ensure_ascii=False)))
|
|
|
|
def get_search(self, query: str) -> dict | None:
|
|
return self._get("searches", "query", query, SEARCH_TTL)
|
|
|
|
def put_search(self, query: str, payload: dict) -> None:
|
|
self._put("searches", query, payload)
|
|
|
|
def get_detail(self, sid: str) -> dict | None:
|
|
return self._get("details", "sid", sid, DETAIL_TTL)
|
|
|
|
def put_detail(self, sid: str, payload: dict) -> None:
|
|
self._put("details", sid, payload)
|
|
|
|
def stats(self) -> dict:
|
|
if not self.path.exists():
|
|
return {"searches": 0, "details": 0}
|
|
with self._connect() as connection:
|
|
return {t: connection.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0] for t in ("searches", "details")}
|