Coverage for node / src / stigmem_node / storage / sqlite_backend.py: 59%
105 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-06-18 05:34 +0000
1"""SQLite implementation of StorageBackend — the default backend.
3When *encryption_key* is provided (32 bytes), the backend uses SQLCipher via
4the ``sqlcipher3`` package instead of stdlib ``sqlite3``. Install the extra
5before enabling encryption::
7 pip install 'stigmem-node[sqlcipher]'
9When *embed_enabled* is True the backend loads the ``sqlite-vec`` extension on
10every new connection and creates the ``vec_facts`` virtual table if absent::
12 pip install 'stigmem-node[sqlite-vec]'
13"""
15from __future__ import annotations
17import logging
18import re
19import sqlite3
20from collections.abc import Generator
21from contextlib import contextmanager
22from datetime import UTC, datetime
23from pathlib import Path
24from typing import Any
26from .base import StorageBackend
28logger = logging.getLogger("stigmem.storage.sqlite")
30# A migration that opens its own top-level transaction (e.g. 019, 029) must not
31# be wrapped again — `BEGIN` inside an open transaction errors. Trigger bodies
32# use a bare `BEGIN` (no semicolon), so this only matches transaction control.
33_MANAGES_OWN_TXN = re.compile(r"^\s*BEGIN\s*;", re.IGNORECASE | re.MULTILINE)
36class SQLiteBackend(StorageBackend):
37 """Default SQLite backend.
39 Behaviour is identical to the pre-trait implementation in ``db.py``.
40 Uses WAL journal mode and enforces foreign-key constraints on every
41 connection. When *encryption_key* is provided, SQLCipher is used
42 transparently — the key is set via ``PRAGMA key`` immediately after open.
43 When *embed_enabled* is True, sqlite-vec is loaded and ``vec_facts`` is
44 created with the given *embed_dimension*.
45 """
47 def __init__(
48 self,
49 db_path: str,
50 encryption_key: bytes | None = None,
51 embed_enabled: bool = False,
52 embed_dimension: int = 768,
53 ) -> None:
54 self._db_path = db_path
55 self._encryption_key = encryption_key
56 self._embed_enabled = embed_enabled
57 self._embed_dimension = embed_dimension
59 @property
60 def backend_name(self) -> str:
61 return "sqlite"
63 def _open_conn(self) -> Any:
64 """Open and return a raw (un-transacted) connection, WAL + FK enabled."""
65 if self._encryption_key is not None: 65 ↛ 66line 65 didn't jump to line 66 because the condition on line 65 was never true
66 try:
67 import sqlcipher3 as _sc
68 except ImportError as exc:
69 raise RuntimeError(
70 "sqlcipher3 is required for SQLite encryption-at-rest. "
71 "Install it with: pip install 'stigmem-node[sqlcipher]'"
72 ) from exc
73 conn = _sc.connect(self._db_path)
74 hex_key = self._encryption_key.hex()
75 conn.execute(f"PRAGMA key = \"x'{hex_key}'\"") # noqa: S608
76 conn.row_factory = _sc.Row
77 else:
78 conn = sqlite3.connect(self._db_path)
79 conn.row_factory = sqlite3.Row
80 conn.execute("PRAGMA journal_mode=WAL")
81 conn.execute("PRAGMA foreign_keys=ON")
82 if self._embed_enabled: 82 ↛ 83line 82 didn't jump to line 83 because the condition on line 82 was never true
83 self._load_sqlite_vec(conn)
84 return conn
86 def _load_sqlite_vec(self, conn: Any) -> None:
87 """Load the sqlite-vec extension and ensure vec_facts virtual table exists."""
88 try:
89 import sqlite_vec
90 except ImportError as exc:
91 raise RuntimeError(
92 "sqlite-vec is required when embed_enabled=true. "
93 "Install it with: pip install 'stigmem-node[sqlite-vec]'"
94 ) from exc
96 try:
97 conn.enable_load_extension(True)
98 sqlite_vec.load(conn)
99 conn.enable_load_extension(False)
100 except Exception as exc:
101 raise RuntimeError(f"Failed to load sqlite-vec extension: {exc}") from exc
103 from stigmem_node.vector_search import ensure_vec_table
105 ensure_vec_table(conn, self._embed_dimension)
107 @contextmanager
108 def connection(self) -> Generator[Any, None, None]:
109 conn = self._open_conn()
110 try:
111 yield conn
112 conn.commit()
113 except Exception:
114 conn.rollback()
115 raise
116 finally:
117 conn.close()
119 def apply_migrations(self, migrations_dir: Path) -> None:
120 conn = self._open_conn()
121 try:
122 conn.execute(
123 """CREATE TABLE IF NOT EXISTS schema_migrations (
124 id INTEGER PRIMARY KEY AUTOINCREMENT,
125 version TEXT NOT NULL UNIQUE,
126 applied_at TEXT NOT NULL
127 )"""
128 )
129 conn.commit()
131 applied = {r["version"] for r in conn.execute("SELECT version FROM schema_migrations")}
133 for f in sorted(migrations_dir.glob("*.sql")):
134 version = f.stem
135 if version in applied:
136 continue
137 # Apply each migration atomically (audit F-MIG-TXN). executescript()
138 # runs in autocommit and disregards isolation_level, so a multi-
139 # statement rebuild (CREATE/DROP/RENAME) that failed partway left a
140 # half-applied, committed schema. Wrap the body in an explicit
141 # transaction unless the migration already manages its own (e.g.
142 # 019, 029) — double-BEGIN would error.
143 sql = f.read_text()
144 script = sql if _MANAGES_OWN_TXN.search(sql) else f"BEGIN;\n{sql}\nCOMMIT;"
145 try:
146 conn.executescript(script)
147 except Exception:
148 conn.rollback()
149 raise
150 conn.execute(
151 "INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)",
152 (version, datetime.now(UTC).isoformat()),
153 )
154 conn.commit()
155 finally:
156 conn.close()
158 def export_snapshot(self, dest: Path) -> None:
159 """Online backup via ``sqlite3.Connection.backup()``.
161 Encrypted databases produce encrypted snapshots (same key).
162 """
163 if self._encryption_key is not None:
164 try:
165 import sqlcipher3 as _sc
166 except ImportError as exc:
167 raise RuntimeError(
168 "sqlcipher3 is required to snapshot an encrypted SQLite database."
169 ) from exc
170 hex_key = self._encryption_key.hex()
171 src_conn = _sc.connect(self._db_path)
172 src_conn.execute(f"PRAGMA key = \"x'{hex_key}'\"") # noqa: S608
173 dst_conn = _sc.connect(str(dest))
174 dst_conn.execute(f"PRAGMA key = \"x'{hex_key}'\"") # noqa: S608
175 try:
176 src_conn.backup(dst_conn)
177 finally:
178 dst_conn.close()
179 src_conn.close()
180 else:
181 src_conn = sqlite3.connect(self._db_path)
182 dst_conn = sqlite3.connect(str(dest))
183 try:
184 src_conn.backup(dst_conn)
185 finally:
186 dst_conn.close()
187 src_conn.close()
189 def import_snapshot(self, src: Path) -> None:
190 """Restore by replacing the current database file."""
191 import shutil
193 shutil.copy2(str(src), self._db_path)