Spaces:
Sleeping
Sleeping
File size: 11,322 Bytes
901e5fa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 | # utils/cache_db.py
"""Thread-safe SQLite-backed cache with TTL and LRU eviction.
Replaces the in-memory dict/OrderedDict caches used by preview, image-metadata,
and repo-info modules. A single database file holds all namespaces in one table
so there is exactly one connection to manage.
Usage::
from ai_agent.utils.cache_db import get_cache_db
db = get_cache_db()
db.set("meta", key, value, max_entries=128)
hit = db.get("meta", key)
"""
from __future__ import annotations
import logging
import os
import sqlite3
import tempfile
import threading
import time
from pathlib import Path
from typing import Optional
log = logging.getLogger("cache_db")
def _default_db_path() -> str:
"""Return a per-user cache DB path so different OS users cannot share data.
On POSIX systems the current UID is embedded in the filename, ensuring
each user gets an isolated file under the system temp directory. The
caller can override the path entirely via ``CACHE_DB_PATH``.
"""
env = os.getenv("CACHE_DB_PATH")
if env:
return env
uid_suffix = ""
try:
uid_suffix = f"_{os.getuid()}"
except AttributeError:
pass # Windows — no getuid(); fall back to a single shared filename
return str(Path(tempfile.gettempdir()) / f"ai_agent_cache{uid_suffix}.db")
_DEFAULT_DB_PATH = _default_db_path()
class CacheDB:
"""Thread-safe SQLite-backed key/value cache.
* **Namespaces** isolate unrelated caches inside one DB.
* **TTL** – entries expire after *ttl_seconds* (wallclock).
``ttl_seconds <= 0`` means the entry never expires.
* **Capacity** – when *max_entries > 0*, the oldest-accessed rows in
that namespace are evicted after each write.
* All values are stored as plain text; callers handle serialisation.
"""
@staticmethod
def _ensure_private_db_file(path: str) -> None:
"""Ensure an on-disk SQLite DB file is owner-readable/writable only.
This is a best-effort hardening step for local cache data. For new
files, create the file first so we can set restrictive permissions
before SQLite opens it. For existing files, tighten permissions.
"""
try:
fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
except FileExistsError:
os.chmod(path, 0o600)
else:
os.close(fd)
os.chmod(path, 0o600)
def __init__(self, db_path: str | Path | None = None) -> None:
path = str(db_path or _DEFAULT_DB_PATH)
if path != ":memory:":
Path(path).parent.mkdir(parents=True, exist_ok=True)
self._ensure_private_db_file(path)
self._conn = sqlite3.connect(path, check_same_thread=False)
self._lock = threading.Lock()
# Restrict the DB file to the owner only so other OS users cannot
# read cached data. The chmod is best-effort (no-op on Windows).
if path != ":memory:":
try:
os.chmod(path, 0o600)
except Exception:
log.warning("Could not set owner-only permissions on cache DB: %s", path)
self._conn.execute("PRAGMA journal_mode=WAL")
self._conn.execute("PRAGMA synchronous=NORMAL")
self._conn.executescript(
"""
CREATE TABLE IF NOT EXISTS cache (
namespace TEXT NOT NULL,
key TEXT NOT NULL,
value TEXT NOT NULL,
expires_at REAL, -- NULL → never expires
accessed_at REAL NOT NULL,
PRIMARY KEY (namespace, key)
);
CREATE INDEX IF NOT EXISTS idx_cache_ns_expires
ON cache (namespace, expires_at);
CREATE INDEX IF NOT EXISTS idx_cache_ns_accessed
ON cache (namespace, accessed_at);
"""
)
self._conn.commit()
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
def get(
self, namespace: str, key: str, *, track_lru: bool = True
) -> Optional[str]:
"""Return the cached value or *None* if missing / expired.
*track_lru* (default ``True``) updates *accessed_at* on a hit so that
LRU eviction in :meth:`set` works correctly. Pass ``False`` when the
namespace has no capacity limit and you want to avoid the extra write.
"""
now = time.time()
with self._lock:
row = self._conn.execute(
"SELECT value, expires_at FROM cache"
" WHERE namespace = ? AND key = ?",
(namespace, key),
).fetchone()
if row is None:
return None
value, expires_at = row
if expires_at is not None and expires_at <= now:
self._conn.execute(
"DELETE FROM cache WHERE namespace = ? AND key = ?",
(namespace, key),
)
self._conn.commit()
return None
if track_lru:
# Touch accessed_at so LRU eviction in set() stays accurate.
self._conn.execute(
"UPDATE cache SET accessed_at = ? WHERE namespace = ? AND key = ?",
(now, namespace, key),
)
self._conn.commit()
return value
def set(
self,
namespace: str,
key: str,
value: str,
*,
ttl_seconds: float = 0,
max_entries: int = 0,
) -> None:
"""Store *value* under *(namespace, key)*.
*ttl_seconds ≤ 0* → no expiration.
*max_entries ≤ 0* → no capacity limit.
"""
now = time.time()
expires_at = (now + ttl_seconds) if ttl_seconds > 0 else None
with self._lock:
# Lazy sweep of expired rows in this namespace
self._conn.execute(
"DELETE FROM cache"
" WHERE namespace = ? AND expires_at IS NOT NULL AND expires_at <= ?",
(namespace, now),
)
self._conn.execute(
"""
INSERT INTO cache (namespace, key, value, expires_at, accessed_at)
VALUES (?, ?, ?, ?, ?)
ON CONFLICT (namespace, key)
DO UPDATE SET value = excluded.value,
expires_at = excluded.expires_at,
accessed_at = excluded.accessed_at
""",
(namespace, key, value, expires_at, now),
)
# Enforce capacity by evicting least-recently-accessed rows
if max_entries > 0:
count = self._conn.execute(
"SELECT COUNT(*) FROM cache WHERE namespace = ?",
(namespace,),
).fetchone()[0]
overage = count - max_entries
if overage > 0:
self._conn.execute(
"""
DELETE FROM cache
WHERE rowid IN (
SELECT rowid FROM cache
WHERE namespace = ?
ORDER BY accessed_at ASC
LIMIT ?
)
""",
(namespace, overage),
)
self._conn.commit()
def delete(self, namespace: str, key: str) -> None:
"""Remove a single entry."""
with self._lock:
self._conn.execute(
"DELETE FROM cache WHERE namespace = ? AND key = ?",
(namespace, key),
)
self._conn.commit()
def clear(self, namespace: Optional[str] = None) -> None:
"""Clear entries. *None* → all namespaces."""
with self._lock:
if namespace is not None:
self._conn.execute(
"DELETE FROM cache WHERE namespace = ?", (namespace,)
)
else:
self._conn.execute("DELETE FROM cache")
self._conn.commit()
def sweep_expired(self) -> int:
"""Delete all expired rows across every namespace.
Returns the number of rows removed. Designed to be called from a
background thread or an atexit hook without touching private state.
"""
now = time.time()
with self._lock:
deleted = self._conn.execute(
"DELETE FROM cache WHERE expires_at IS NOT NULL AND expires_at <= ?",
(now,),
).rowcount
self._conn.commit()
return deleted
def vacuum_and_close(self) -> None:
"""Sweep expired rows, VACUUM the database, then close the connection.
VACUUM must run outside any open transaction; this method handles the
``isolation_level`` toggling safely with a ``finally`` guard.
"""
try:
self.sweep_expired()
with self._lock:
previous_isolation_level = self._conn.isolation_level
try:
self._conn.isolation_level = None
self._conn.execute("VACUUM")
finally:
self._conn.isolation_level = previous_isolation_level
finally:
self.close()
def close(self) -> None:
"""Close the underlying database connection."""
with self._lock:
try:
self._conn.close()
except Exception:
pass
# ------------------------------------------------------------------
# Module-level singleton
# ------------------------------------------------------------------
_db: Optional[CacheDB] = None
_db_lock = threading.Lock()
def get_cache_db() -> CacheDB:
"""Return (or create) the process-wide :class:`CacheDB` singleton."""
global _db
if _db is None:
with _db_lock:
if _db is None:
_db = CacheDB()
log.info("Opened cache DB at %s", _DEFAULT_DB_PATH)
return _db
def get_cache_db_or_none() -> Optional[CacheDB]:
"""Return the singleton if it has already been initialised, else ``None``.
Unlike :func:`get_cache_db` this never creates the database — useful for
background housekeeping that should be a no-op before the first real cache
access.
"""
return _db
def reset_cache_db(db: Optional[CacheDB] = None) -> None:
"""Replace the singleton — mainly useful for tests.
If *db* is ``None`` a fresh in-memory database is **not** created
automatically; the next call to :func:`get_cache_db` will open the
default path again.
"""
global _db
with _db_lock:
if _db is not None:
_db.close()
_db = db
|