webapp/tipitaka-api/app/database/sqlite_db.py

import sqlite3
import threading
from pathlib import Path
from contextlib import contextmanager
from typing import Generator

from app.config import get_settings


class SQLiteDB:
    """SQLite database connection manager with optional in-memory cache."""

    def __init__(self, db_path: str) -> None:
        self.db_path = str(Path(db_path).resolve())
        self._mem_conn: sqlite3.Connection | None = None
        self._lock = threading.RLock()

    @contextmanager
    def get_connection(self) -> Generator[sqlite3.Connection, None, None]:
        """Get a connection — uses in-memory DB if loaded, else disk."""
        if self._mem_conn is not None:
            with self._lock:
                yield self._mem_conn
            return

        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
        finally:
            conn.close()

    def load_to_memory(self) -> None:
        """
        Load entire DB to :memory: for fast queries.
        Uses sqlite3.backup() — native C-level copy, faster than iterdump().
        """
        mem = sqlite3.connect(":memory:", check_same_thread=False)
        with sqlite3.connect(self.db_path) as disk:
            disk.backup(mem)

        mem.row_factory = sqlite3.Row
        # Performance PRAGMAs (match v2.1 settings)
        mem.execute("PRAGMA temp_store = MEMORY")
        mem.execute("PRAGMA cache_size = -250000")   # ~300MB query cache
        mem.execute("PRAGMA mmap_size = 0")          # mmap irrelevant for :memory:

        # Create FTS5 virtual table — name: pages_fts (consistent with search_service)
        try:
            mem.execute("""
                CREATE VIRTUAL TABLE IF NOT EXISTS pages_fts
                USING fts5(content_text, content='pages', content_rowid='id',
                           tokenize='unicode61')
            """)
            mem.execute("INSERT INTO pages_fts(pages_fts) VALUES('rebuild')")
        except sqlite3.OperationalError as e:
            import logging
            logging.getLogger(__name__).warning(f"FTS5 build warning: {e}")

        self._mem_conn = mem

    @contextmanager
    def get_disk_connection(self) -> Generator[sqlite3.Connection, None, None]:
        """Always connects to disk — used for persistent writes (e.g. search_log)."""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
            conn.commit()
        finally:
            conn.close()

    def ensure_search_log_table(self) -> None:
        """Create search_log table on disk if not exists."""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS search_log (
                    id      INTEGER PRIMARY KEY AUTOINCREMENT,
                    query   TEXT NOT NULL COLLATE NOCASE,
                    count   INTEGER NOT NULL DEFAULT 1,
                    last_at TEXT NOT NULL DEFAULT (datetime('now','localtime')),
                    UNIQUE(query)
                )
            """)
            conn.execute("CREATE INDEX IF NOT EXISTS idx_search_log_query ON search_log(query)")
            conn.commit()

    def ensure_reference_tables(self) -> None:
        """Create reference_markers table on disk if not exists."""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS reference_markers (
                    id          INTEGER PRIMARY KEY AUTOINCREMENT,
                    volume_num  INTEGER NOT NULL,
                    page_num    INTEGER NOT NULL,
                    marker_id   TEXT NOT NULL,
                    type        TEXT NOT NULL, -- 'footnote' | 'abbrev'
                    content     TEXT NOT NULL,
                    UNIQUE(volume_num, page_num, marker_id, type)
                )
            """)
            conn.execute("CREATE INDEX IF NOT EXISTS idx_ref_vol_page ON reference_markers(volume_num, page_num)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_ref_marker ON reference_markers(marker_id)")
            conn.commit()

    @property
    def is_in_memory(self) -> bool:
        return self._mem_conn is not None

    def ensure_query_cache_table(self) -> None:
        """Create query_cache table on disk if not exists."""
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS query_cache (
                    cache_key   TEXT PRIMARY KEY,
                    result_json TEXT NOT NULL,
                    hit_count   INTEGER DEFAULT 1,
                    created_at  DATETIME DEFAULT CURRENT_TIMESTAMP,
                    last_used   DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            """)
            conn.execute("CREATE INDEX IF NOT EXISTS idx_query_cache_last_used ON query_cache(last_used)")
            conn.commit()

    def get_query_cache(self, key: str) -> dict | None:
        """Retrieve query cache from disk and increment hit count."""
        import json
        try:
            with self.get_disk_connection() as conn:
                row = conn.execute(
                    "SELECT result_json FROM query_cache WHERE cache_key = ?", (key,)
                ).fetchone()
                if row:
                    conn.execute(
                        """UPDATE query_cache
                           SET hit_count = hit_count + 1, last_used = datetime('now','localtime')
                           WHERE cache_key = ?""", (key,)
                    )
                    return json.loads(row["result_json"])
        except Exception:
            pass
        return None

    def set_query_cache(self, key: str, result: dict) -> None:
        """Store query cache to disk and evict oldest if too large."""
        import json
        try:
            self.evict_query_cache(max_entries=10000)
            with self.get_disk_connection() as conn:
                conn.execute(
                    """INSERT INTO query_cache (cache_key, result_json) VALUES (?, ?)
                       ON CONFLICT(cache_key) DO UPDATE SET
                           result_json = excluded.result_json,
                           last_used = datetime('now','localtime')""",
                    (key, json.dumps(result, ensure_ascii=False))
                )
        except Exception:
            pass

    def evict_query_cache(self, max_entries: int = 10000) -> None:
        """Evict oldest cache entries on disk if count exceeds max_entries."""
        try:
            with self.get_disk_connection() as conn:
                count = conn.execute("SELECT COUNT(*) FROM query_cache").fetchone()[0]
                if count > max_entries:
                    conn.execute("""
                        DELETE FROM query_cache WHERE cache_key IN (
                            SELECT cache_key FROM query_cache
                            ORDER BY last_used ASC LIMIT ?
                        )
                    """, (count - max_entries,))
        except Exception:
            pass


# Singleton database instance
_db: SQLiteDB | None = None


def get_db() -> SQLiteDB:
    """Get or create the singleton database instance."""
    global _db
    if _db is None:
        settings = get_settings()
        _db = SQLiteDB(settings.DATABASE_PATH)
        _db.ensure_search_log_table()
        _db.ensure_reference_tables()
        _db.ensure_query_cache_table()
    return _db