from __future__ import annotations import asyncio import json import math import re import sqlite3 from datetime import UTC, datetime from pathlib import Path import sqlite_vec from stirling.contracts.documents import Page, PageRange from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage from stirling.models import OwnerId, PrincipalId _READ_PERMISSION = "read" # sqlite stores TIMESTAMP as TEXT. We normalise to UTC ISO 8601 ``YYYY-MM-DD HH:MM:SS`` # so lexicographic comparison against ``datetime('now')`` matches chronological order. _SQLITE_DATETIME_FMT = "%Y-%m-%d %H:%M:%S" def _to_sqlite_utc(dt: datetime | None) -> str | None: if dt is None: return None if dt.tzinfo is not None: dt = dt.astimezone(UTC).replace(tzinfo=None) return dt.strftime(_SQLITE_DATETIME_FMT) class SqliteVecStore(DocumentStore): """sqlite-vec backed vector store. Single-file SQLite database, embedded, no server. Each ``(collection, owner_id)`` pair gets its own `vec0` virtual table with a fixed embedding dimension (detected on first insert). Document metadata lives in a regular table joined by rowid. """ def __init__(self, db_path: str | Path) -> None: is_memory = str(db_path) == ":memory:" self._db_path: Path | None = None if is_memory else Path(db_path) if self._db_path is not None: self._db_path.parent.mkdir(parents=True, exist_ok=True) conn = sqlite3.connect(str(self._db_path), check_same_thread=False) else: conn = sqlite3.connect(":memory:", check_same_thread=False) conn.enable_load_extension(True) sqlite_vec.load(conn) conn.enable_load_extension(False) # Required so cascade deletes from documents_meta clean up child tables. conn.execute("PRAGMA foreign_keys=ON") if self._db_path is not None: conn.execute("PRAGMA journal_mode=WAL") self._conn = conn self._lock = asyncio.Lock() self._init_schema() @classmethod def ephemeral(cls) -> SqliteVecStore: """In-memory store for testing.""" return cls(":memory:") def _init_schema(self) -> None: self._conn.execute( """ CREATE TABLE IF NOT EXISTS documents_meta ( collection TEXT NOT NULL, owner_id TEXT NOT NULL, source TEXT NOT NULL, expires_at TIMESTAMP, PRIMARY KEY (collection, owner_id) ) """ ) # The reaper filters on ``expires_at IS NOT NULL AND expires_at < now`` so # a partial index over non-null rows keeps the scan tight even when most # rows are persistent (org docs). self._conn.execute( "CREATE INDEX IF NOT EXISTS idx_meta_expires_at ON documents_meta(expires_at) WHERE expires_at IS NOT NULL" ) self._conn.execute( """ CREATE TABLE IF NOT EXISTS collections ( collection TEXT NOT NULL, owner_id TEXT NOT NULL, dim INTEGER NOT NULL, table_name TEXT NOT NULL, PRIMARY KEY (collection, owner_id), FOREIGN KEY (collection, owner_id) REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE ) """ ) self._conn.execute( """ CREATE TABLE IF NOT EXISTS documents ( id TEXT NOT NULL, collection TEXT NOT NULL, owner_id TEXT NOT NULL, text TEXT NOT NULL, metadata TEXT NOT NULL DEFAULT '{}', vec_rowid INTEGER NOT NULL, PRIMARY KEY (id, collection, owner_id), FOREIGN KEY (collection, owner_id) REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE ) """ ) self._conn.execute("CREATE INDEX IF NOT EXISTS idx_doc_collection_owner ON documents(collection, owner_id)") self._conn.execute( """ CREATE TABLE IF NOT EXISTS document_pages ( collection TEXT NOT NULL, owner_id TEXT NOT NULL, page_number INTEGER NOT NULL, text TEXT NOT NULL, char_count INTEGER NOT NULL, PRIMARY KEY (collection, owner_id, page_number), FOREIGN KEY (collection, owner_id) REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE ) """ ) self._conn.execute( "CREATE INDEX IF NOT EXISTS idx_pages_collection_owner ON document_pages(collection, owner_id)" ) self._conn.execute( """ CREATE TABLE IF NOT EXISTS document_acl ( collection TEXT NOT NULL, owner_id TEXT NOT NULL, principal_id TEXT NOT NULL, permission TEXT NOT NULL, PRIMARY KEY (collection, owner_id, principal_id, permission), FOREIGN KEY (collection, owner_id) REFERENCES documents_meta(collection, owner_id) ON DELETE CASCADE ) """ ) # Lookup by principal is the hot path for search/list (every read # joins through this index). Composite ordering matches the WHERE. self._conn.execute( "CREATE INDEX IF NOT EXISTS idx_acl_principal_permission ON document_acl(principal_id, permission)" ) self._conn.commit() # ── lifecycle of the (collection, owner_id) row ──────────────────────── async def ensure_collection( self, collection: str, source: str, owner_id: OwnerId, expires_at: datetime | None, ) -> None: async with self._lock: await asyncio.to_thread(self._sync_ensure_collection, collection, source, owner_id, expires_at) def _sync_ensure_collection( self, collection: str, source: str, owner_id: OwnerId, expires_at: datetime | None, ) -> None: self._conn.execute( """ INSERT INTO documents_meta(collection, owner_id, source, expires_at) VALUES (?, ?, ?, ?) ON CONFLICT(collection, owner_id) DO UPDATE SET source = excluded.source, expires_at = excluded.expires_at """, (collection, owner_id, source, _to_sqlite_utc(expires_at)), ) self._conn.commit() async def delete_collection(self, collection: str, owner_id: OwnerId) -> bool: async with self._lock: return await asyncio.to_thread(self._sync_delete_collection, collection, owner_id) def _sync_delete_collection(self, collection: str, owner_id: OwnerId) -> bool: # Drop the sqlite-vec virtual table first; FK cascade handles the regular tables # (collections, documents, document_pages, document_acl) when documents_meta is deleted. row = self._conn.execute( "SELECT table_name FROM collections WHERE collection = ? AND owner_id = ?", (collection, owner_id), ).fetchone() if row is not None: self._conn.execute(f"DROP TABLE IF EXISTS {row[0]}") cursor = self._conn.execute( "DELETE FROM documents_meta WHERE collection = ? AND owner_id = ?", (collection, owner_id), ) self._conn.commit() return cursor.rowcount > 0 async def purge_owner(self, owner_id: OwnerId) -> int: async with self._lock: return await asyncio.to_thread(self._sync_purge_owner, owner_id) def _sync_purge_owner(self, owner_id: OwnerId) -> int: # Drop all vec0 virtual tables for this owner first (FK cascade can't reach them). vec_tables = [ r[0] for r in self._conn.execute("SELECT table_name FROM collections WHERE owner_id = ?", (owner_id,)).fetchall() ] for name in vec_tables: self._conn.execute(f"DROP TABLE IF EXISTS {name}") cursor = self._conn.execute("DELETE FROM documents_meta WHERE owner_id = ?", (owner_id,)) self._conn.commit() return cursor.rowcount async def reap_expired(self) -> int: async with self._lock: return await asyncio.to_thread(self._sync_reap_expired) def _sync_reap_expired(self) -> int: # Drop vec0 virtual tables for expired collections first (FK cascade can't reach them). vec_tables = [ r[0] for r in self._conn.execute( """ SELECT c.table_name FROM collections c JOIN documents_meta m ON m.collection = c.collection AND m.owner_id = c.owner_id WHERE m.expires_at IS NOT NULL AND m.expires_at < datetime('now') """ ).fetchall() ] for name in vec_tables: self._conn.execute(f"DROP TABLE IF EXISTS {name}") cursor = self._conn.execute( "DELETE FROM documents_meta WHERE expires_at IS NOT NULL AND expires_at < datetime('now')" ) self._conn.commit() return cursor.rowcount # ── write paths ──────────────────────────────────────────────────────── @staticmethod def _sanitize_table_name(collection: str, owner_id: OwnerId) -> str: safe_col = re.sub(r"[^a-zA-Z0-9_]", "_", collection) safe_owner = re.sub(r"[^a-zA-Z0-9_]", "_", owner_id) return f"vec_{safe_owner}_{safe_col}" @staticmethod def _normalize(vector: list[float]) -> list[float]: norm = math.sqrt(sum(x * x for x in vector)) if norm == 0: return list(vector) return [x / norm for x in vector] async def add_documents( self, collection: str, documents: list[Document], embeddings: list[list[float]], owner_id: OwnerId, ) -> None: if len(documents) != len(embeddings): raise ValueError(f"Got {len(documents)} documents but {len(embeddings)} embeddings") if not documents: return async with self._lock: await asyncio.to_thread(self._sync_add, collection, documents, embeddings, owner_id) def _sync_add( self, collection: str, documents: list[Document], embeddings: list[list[float]], owner_id: OwnerId, ) -> None: dim = len(embeddings[0]) row = self._conn.execute( "SELECT dim, table_name FROM collections WHERE collection = ? AND owner_id = ?", (collection, owner_id), ).fetchone() if row is None: table_name = self._sanitize_table_name(collection, owner_id) self._conn.execute(f"CREATE VIRTUAL TABLE IF NOT EXISTS {table_name} USING vec0(embedding float[{dim}])") self._conn.execute( "INSERT INTO collections(collection, owner_id, dim, table_name) VALUES (?, ?, ?, ?)", (collection, owner_id, dim, table_name), ) else: existing_dim, table_name = row if existing_dim != dim: raise ValueError( f"Collection {collection} for owner {owner_id} has dim {existing_dim}, got embedding of dim {dim}" ) # Upsert: delete existing docs with matching IDs first ids = [doc.id for doc in documents] placeholders = ",".join("?" * len(ids)) existing = self._conn.execute( f"SELECT vec_rowid FROM documents WHERE collection = ? AND owner_id = ? AND id IN ({placeholders})", (collection, owner_id, *ids), ).fetchall() if existing: vec_rowids = [r[0] for r in existing] row_placeholders = ",".join("?" * len(vec_rowids)) self._conn.execute( f"DELETE FROM {table_name} WHERE rowid IN ({row_placeholders})", vec_rowids, ) self._conn.execute( f"DELETE FROM documents WHERE collection = ? AND owner_id = ? AND id IN ({placeholders})", (collection, owner_id, *ids), ) for doc, emb in zip(documents, embeddings): normalized = self._normalize(list(emb)) cursor = self._conn.execute( f"INSERT INTO {table_name}(embedding) VALUES (?)", (sqlite_vec.serialize_float32(normalized),), ) vec_rowid = cursor.lastrowid self._conn.execute( "INSERT INTO documents(id, collection, owner_id, text, metadata, vec_rowid) VALUES (?, ?, ?, ?, ?, ?)", (doc.id, collection, owner_id, doc.text, json.dumps(doc.metadata), vec_rowid), ) self._conn.commit() async def add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None: async with self._lock: await asyncio.to_thread(self._sync_add_pages, collection, pages, owner_id) def _sync_add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None: self._conn.execute( "DELETE FROM document_pages WHERE collection = ? AND owner_id = ?", (collection, owner_id), ) if pages: self._conn.executemany( "INSERT INTO document_pages(collection, owner_id, page_number, text, char_count) " "VALUES (?, ?, ?, ?, ?)", [(collection, owner_id, p.page_number, p.text, p.char_count) for p in pages], ) self._conn.commit() # ── ACL management ───────────────────────────────────────────────────── async def grant_read( self, collection: str, owner_id: OwnerId, principals: list[PrincipalId], ) -> None: if not principals: return async with self._lock: await asyncio.to_thread(self._sync_grant_read, collection, owner_id, principals) def _sync_grant_read( self, collection: str, owner_id: OwnerId, principals: list[PrincipalId], ) -> None: self._conn.executemany( """ INSERT INTO document_acl(collection, owner_id, principal_id, permission) VALUES (?, ?, ?, ?) ON CONFLICT(collection, owner_id, principal_id, permission) DO NOTHING """, [(collection, owner_id, p, _READ_PERMISSION) for p in principals], ) self._conn.commit() async def revoke( self, collection: str, owner_id: OwnerId, principal: PrincipalId, ) -> None: async with self._lock: await asyncio.to_thread(self._sync_revoke, collection, owner_id, principal) def _sync_revoke(self, collection: str, owner_id: OwnerId, principal: PrincipalId) -> None: self._conn.execute( "DELETE FROM document_acl WHERE collection = ? AND owner_id = ? AND principal_id = ?", (collection, owner_id, principal), ) self._conn.commit() # ── read paths (ACL-gated) ───────────────────────────────────────────── async def search( self, collection: str, query_embedding: list[float], top_k: int, principals: list[PrincipalId], ) -> list[SearchResult]: async with self._lock: return await asyncio.to_thread(self._sync_search, collection, query_embedding, top_k, principals) def _sync_search( self, collection: str, query_embedding: list[float], top_k: int, principals: list[PrincipalId], ) -> list[SearchResult]: if not principals: return [] owner_id = self._readable_owner_for(collection, principals) if owner_id is None: return [] row = self._conn.execute( "SELECT table_name, dim FROM collections WHERE collection = ? AND owner_id = ?", (collection, owner_id), ).fetchone() if row is None: return [] table_name, dim = row if len(query_embedding) != dim: raise ValueError(f"Query embedding dim {len(query_embedding)} does not match collection dim {dim}") normalized = self._normalize(list(query_embedding)) query_blob = sqlite_vec.serialize_float32(normalized) results = self._conn.execute( f""" SELECT d.id, d.text, d.metadata, v.distance FROM {table_name} v JOIN documents d ON d.vec_rowid = v.rowid AND d.collection = ? AND d.owner_id = ? WHERE v.embedding MATCH ? AND k = ? ORDER BY v.distance """, (collection, owner_id, query_blob, top_k), ).fetchall() return [ SearchResult( document=Document( id=r[0], text=r[1], metadata=json.loads(r[2]) if r[2] else {}, ), # For normalized vectors: cosine_sim = 1 - (L2^2 / 2) score=max(0.0, 1.0 - (r[3] ** 2) / 2.0), ) for r in results ] def _readable_owner_for(self, collection: str, principals: list[PrincipalId]) -> str | None: """Resolve which owner_id this caller is reading. ``None`` means no access.""" placeholders = ",".join("?" * len(principals)) row = self._conn.execute( f""" SELECT owner_id FROM document_acl WHERE collection = ? AND permission = ? AND principal_id IN ({placeholders}) ORDER BY owner_id LIMIT 1 """, (collection, _READ_PERMISSION, *principals), ).fetchone() return row[0] if row else None async def read_pages( self, collection: str, page_range: PageRange | None, principals: list[PrincipalId], ) -> list[Page]: async with self._lock: return await asyncio.to_thread(self._sync_read_pages, collection, page_range, principals) def _sync_read_pages( self, collection: str, page_range: PageRange | None, principals: list[PrincipalId], ) -> list[Page]: if not principals: return [] owner_id = self._readable_owner_for(collection, principals) if owner_id is None: return [] if page_range is None: rows = self._conn.execute( "SELECT page_number, text, char_count FROM document_pages " "WHERE collection = ? AND owner_id = ? ORDER BY page_number", (collection, owner_id), ).fetchall() else: rows = self._conn.execute( "SELECT page_number, text, char_count FROM document_pages " "WHERE collection = ? AND owner_id = ? AND page_number BETWEEN ? AND ? " "ORDER BY page_number", (collection, owner_id, page_range.start, page_range.end), ).fetchall() return [Page(page_number=r[0], text=r[1], char_count=r[2]) for r in rows] async def has_collection(self, collection: str, principals: list[PrincipalId]) -> bool: async with self._lock: return await asyncio.to_thread(self._sync_has_collection, collection, principals) def _sync_has_collection(self, collection: str, principals: list[PrincipalId]) -> bool: if not principals: return False return self._readable_owner_for(collection, principals) is not None async def list_collections(self, principals: list[PrincipalId]) -> list[str]: async with self._lock: return await asyncio.to_thread(self._sync_list_collections, principals) def _sync_list_collections(self, principals: list[PrincipalId]) -> list[str]: if not principals: return [] placeholders = ",".join("?" * len(principals)) rows = self._conn.execute( f""" SELECT DISTINCT collection FROM document_acl WHERE permission = ? AND principal_id IN ({placeholders}) ORDER BY collection """, (_READ_PERMISSION, *principals), ).fetchall() return [r[0] for r in rows] async def close(self) -> None: async with self._lock: await asyncio.to_thread(self._sync_close) def _sync_close(self) -> None: """Checkpoint the WAL into the main database file and close the connection so the .db-shm and .db-wal files are cleaned up on graceful shutdown.""" if self._db_path is not None: try: self._conn.execute("PRAGMA wal_checkpoint(TRUNCATE)") self._conn.commit() except sqlite3.Error: # Best effort: if checkpointing fails we still want to close the connection. pass self._conn.close()