Spaces:
Sleeping
Sleeping
| """ | |
| storage/cache.py — SQLite-backed search history and financial data cache. | |
| """ | |
| import sqlite3 | |
| import json | |
| import datetime | |
| from pathlib import Path | |
| from typing import Optional, List, Dict, Any | |
| import sys | |
| sys.path.insert(0, str(Path(__file__).parent.parent)) | |
| from config import DB_PATH, CACHE_TTL_DAYS | |
| # ── Database setup ──────────────────────────────────────────────────────────── | |
| def _connect() -> sqlite3.Connection: | |
| conn = sqlite3.connect(DB_PATH) | |
| conn.row_factory = sqlite3.Row | |
| return conn | |
| def init_db() -> None: | |
| """Create tables if they don't exist.""" | |
| with _connect() as conn: | |
| conn.executescript(""" | |
| CREATE TABLE IF NOT EXISTS searches ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| company_name TEXT NOT NULL, | |
| ticker TEXT, | |
| exchange TEXT, | |
| searched_at TEXT NOT NULL, | |
| models_used TEXT, | |
| output_docx TEXT, | |
| output_xlsx TEXT, | |
| status TEXT DEFAULT 'completed' | |
| ); | |
| CREATE TABLE IF NOT EXISTS financial_cache ( | |
| ticker TEXT NOT NULL, | |
| data_type TEXT NOT NULL, | |
| fetched_at TEXT NOT NULL, | |
| expires_at TEXT NOT NULL, | |
| raw_json TEXT NOT NULL, | |
| PRIMARY KEY (ticker, data_type) | |
| ); | |
| CREATE INDEX IF NOT EXISTS idx_searches_ticker ON searches(ticker); | |
| CREATE INDEX IF NOT EXISTS idx_cache_ticker ON financial_cache(ticker); | |
| """) | |
| # ── Search history ──────────────────────────────────────────────────────────── | |
| def add_search( | |
| company_name: str, | |
| ticker: str, | |
| exchange: str, | |
| models_used: List[str], | |
| output_docx: str = "", | |
| output_xlsx: str = "", | |
| status: str = "completed", | |
| ) -> int: | |
| """Record a completed analysis in search history. Returns row id.""" | |
| now = datetime.datetime.now().isoformat(timespec="seconds") | |
| with _connect() as conn: | |
| cur = conn.execute( | |
| """INSERT INTO searches | |
| (company_name, ticker, exchange, searched_at, models_used, output_docx, output_xlsx, status) | |
| VALUES (?, ?, ?, ?, ?, ?, ?, ?)""", | |
| (company_name, ticker, exchange, now, json.dumps(models_used), output_docx, output_xlsx, status), | |
| ) | |
| return cur.lastrowid | |
| def update_search_outputs(row_id: int, output_docx: str, output_xlsx: str, status: str = "completed") -> None: | |
| with _connect() as conn: | |
| conn.execute( | |
| "UPDATE searches SET output_docx=?, output_xlsx=?, status=? WHERE id=?", | |
| (output_docx, output_xlsx, status, row_id), | |
| ) | |
| def get_history(limit: int = 50) -> List[Dict]: | |
| """Return recent searches, newest first.""" | |
| with _connect() as conn: | |
| rows = conn.execute( | |
| "SELECT * FROM searches ORDER BY searched_at DESC LIMIT ?", (limit,) | |
| ).fetchall() | |
| return [dict(r) for r in rows] | |
| def delete_history_entry(row_id: int) -> None: | |
| with _connect() as conn: | |
| conn.execute("DELETE FROM searches WHERE id=?", (row_id,)) | |
| # ── Financial data cache ────────────────────────────────────────────────────── | |
| def get_cached(ticker: str, data_type: str) -> Optional[Any]: | |
| """Return cached JSON data if still valid, else None.""" | |
| now = datetime.datetime.now().isoformat() | |
| with _connect() as conn: | |
| row = conn.execute( | |
| "SELECT raw_json, expires_at FROM financial_cache WHERE ticker=? AND data_type=?", | |
| (ticker.upper(), data_type), | |
| ).fetchone() | |
| if row is None: | |
| return None | |
| if row["expires_at"] < now: | |
| # Expired | |
| return None | |
| return json.loads(row["raw_json"]) | |
| def set_cache(ticker: str, data_type: str, data: Any) -> None: | |
| """Store data in cache with TTL.""" | |
| now = datetime.datetime.now() | |
| expires = (now + datetime.timedelta(days=CACHE_TTL_DAYS)).isoformat() | |
| with _connect() as conn: | |
| conn.execute( | |
| """INSERT OR REPLACE INTO financial_cache | |
| (ticker, data_type, fetched_at, expires_at, raw_json) | |
| VALUES (?, ?, ?, ?, ?)""", | |
| (ticker.upper(), data_type, now.isoformat(), expires, json.dumps(data)), | |
| ) | |
| def clear_cache(ticker: str) -> None: | |
| """Remove all cached data for a ticker (force refresh).""" | |
| with _connect() as conn: | |
| conn.execute("DELETE FROM financial_cache WHERE ticker=?", (ticker.upper(),)) | |
| def cache_stats() -> Dict: | |
| """Return some stats for diagnostics.""" | |
| with _connect() as conn: | |
| total_searches = conn.execute("SELECT COUNT(*) FROM searches").fetchone()[0] | |
| cached_tickers = conn.execute( | |
| "SELECT COUNT(DISTINCT ticker) FROM financial_cache" | |
| ).fetchone()[0] | |
| return {"total_searches": total_searches, "cached_tickers": cached_tickers} | |
| # Initialise on import | |
| init_db() | |