Spaces:
Running
Running
File size: 7,501 Bytes
8052574 de8a899 8052574 de8a899 8052574 de8a899 8052574 9b44273 165565d 8052574 c6f014f 8052574 165565d c6f014f 8052574 | 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 | 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
|