152 lines
4.5 KiB
Python
152 lines
4.5 KiB
Python
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import os
|
||
|
|
import sqlite3
|
||
|
|
import stat
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from datetime import datetime, timezone
|
||
|
|
from pathlib import Path
|
||
|
|
from urllib.parse import urlencode
|
||
|
|
|
||
|
|
DEFAULT_READ_CACHE_PATH = Path(
|
||
|
|
os.environ.get("STATEHUB_READ_CACHE_PATH", "~/.statehub/edge-read-cache.sqlite3")
|
||
|
|
).expanduser()
|
||
|
|
MAX_CACHE_ENTRY_BYTES = 1024 * 1024
|
||
|
|
|
||
|
|
CACHEABLE_GET_PREFIXES = (
|
||
|
|
"/state/",
|
||
|
|
"/workplans/",
|
||
|
|
"/messages/",
|
||
|
|
"/decisions/",
|
||
|
|
"/tasks/",
|
||
|
|
"/progress/",
|
||
|
|
"/sbom/",
|
||
|
|
"/repos",
|
||
|
|
"/legacy-meter/",
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class ReadCacheEntry:
|
||
|
|
cache_key: str
|
||
|
|
status_code: int
|
||
|
|
content: bytes
|
||
|
|
content_type: str | None
|
||
|
|
cached_at: str
|
||
|
|
|
||
|
|
|
||
|
|
def default_read_cache_path() -> Path:
|
||
|
|
return DEFAULT_READ_CACHE_PATH
|
||
|
|
|
||
|
|
|
||
|
|
def utcnow() -> str:
|
||
|
|
return datetime.now(tz=timezone.utc).isoformat()
|
||
|
|
|
||
|
|
|
||
|
|
def build_cache_key(method: str, path: str, query_items: list[tuple[str, str]]) -> str:
|
||
|
|
normalized_path = path if path.startswith("/") else f"/{path}"
|
||
|
|
query = urlencode(sorted(query_items))
|
||
|
|
if query:
|
||
|
|
return f"{method.upper()}:{normalized_path}?{query}"
|
||
|
|
return f"{method.upper()}:{normalized_path}"
|
||
|
|
|
||
|
|
|
||
|
|
def is_cacheable_get(path: str) -> bool:
|
||
|
|
normalized = path if path.startswith("/") else f"/{path}"
|
||
|
|
if normalized.startswith("/edge/"):
|
||
|
|
return False
|
||
|
|
return any(normalized.startswith(prefix) for prefix in CACHEABLE_GET_PREFIXES)
|
||
|
|
|
||
|
|
|
||
|
|
class ReadCacheStore:
|
||
|
|
def __init__(self, path: str | Path | None = None) -> None:
|
||
|
|
self.path = Path(path).expanduser() if path is not None else default_read_cache_path()
|
||
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
self._init_db()
|
||
|
|
self._chmod_private()
|
||
|
|
|
||
|
|
def _connect(self) -> sqlite3.Connection:
|
||
|
|
conn = sqlite3.connect(self.path)
|
||
|
|
conn.row_factory = sqlite3.Row
|
||
|
|
return conn
|
||
|
|
|
||
|
|
def _init_db(self) -> None:
|
||
|
|
with self._connect() as conn:
|
||
|
|
conn.execute(
|
||
|
|
"""
|
||
|
|
CREATE TABLE IF NOT EXISTS read_cache_entries (
|
||
|
|
cache_key TEXT PRIMARY KEY,
|
||
|
|
status_code INTEGER NOT NULL,
|
||
|
|
content BLOB NOT NULL,
|
||
|
|
content_type TEXT,
|
||
|
|
cached_at TEXT NOT NULL
|
||
|
|
)
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
conn.execute(
|
||
|
|
"CREATE INDEX IF NOT EXISTS ix_read_cache_cached_at ON read_cache_entries(cached_at)"
|
||
|
|
)
|
||
|
|
conn.commit()
|
||
|
|
|
||
|
|
def _chmod_private(self) -> None:
|
||
|
|
try:
|
||
|
|
os.chmod(self.path, stat.S_IRUSR | stat.S_IWUSR)
|
||
|
|
except OSError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
def put(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
cache_key: str,
|
||
|
|
status_code: int,
|
||
|
|
content: bytes,
|
||
|
|
content_type: str | None,
|
||
|
|
) -> None:
|
||
|
|
if len(content) > MAX_CACHE_ENTRY_BYTES:
|
||
|
|
return
|
||
|
|
now = utcnow()
|
||
|
|
with self._connect() as conn:
|
||
|
|
conn.execute(
|
||
|
|
"""
|
||
|
|
INSERT INTO read_cache_entries (
|
||
|
|
cache_key, status_code, content, content_type, cached_at
|
||
|
|
) VALUES (?, ?, ?, ?, ?)
|
||
|
|
ON CONFLICT(cache_key) DO UPDATE SET
|
||
|
|
status_code = excluded.status_code,
|
||
|
|
content = excluded.content,
|
||
|
|
content_type = excluded.content_type,
|
||
|
|
cached_at = excluded.cached_at
|
||
|
|
""",
|
||
|
|
(cache_key, status_code, content, content_type, now),
|
||
|
|
)
|
||
|
|
conn.commit()
|
||
|
|
|
||
|
|
def get(self, cache_key: str) -> ReadCacheEntry | None:
|
||
|
|
with self._connect() as conn:
|
||
|
|
row = conn.execute(
|
||
|
|
"SELECT * FROM read_cache_entries WHERE cache_key = ?",
|
||
|
|
(cache_key,),
|
||
|
|
).fetchone()
|
||
|
|
if row is None:
|
||
|
|
return None
|
||
|
|
return ReadCacheEntry(
|
||
|
|
cache_key=row["cache_key"],
|
||
|
|
status_code=row["status_code"],
|
||
|
|
content=row["content"],
|
||
|
|
content_type=row["content_type"],
|
||
|
|
cached_at=row["cached_at"],
|
||
|
|
)
|
||
|
|
|
||
|
|
def summary(self) -> dict[str, int | str | None]:
|
||
|
|
with self._connect() as conn:
|
||
|
|
row = conn.execute(
|
||
|
|
"""
|
||
|
|
SELECT COUNT(*) AS entry_count,
|
||
|
|
MIN(cached_at) AS oldest_cached_at
|
||
|
|
FROM read_cache_entries
|
||
|
|
"""
|
||
|
|
).fetchone()
|
||
|
|
return {
|
||
|
|
"entry_count": int(row["entry_count"] or 0),
|
||
|
|
"oldest_cached_at": row["oldest_cached_at"],
|
||
|
|
}
|