Set up document management for Stirling Engine (#6476)
# Description of Changes Change Stirling Engine to support deleting documents automatically. This happens both on user logout and after an amount of time specified by the Java when ingesting a document (allowing for personal documents to have short lifetimes but org documents to be left in the db with no expiry date). Also sets up an [ACL policy](https://en.wikipedia.org/wiki/Access-control_list) for the documents so the database knows which users have access to which documents. This is not fully implemented in the Java, so currently all docs are treated as having a single owner, the uploader, but theoretically when we need to support org storage, we shouldn't need to change the db schema.
This commit is contained in:
James Brunton
@@ -5,19 +5,35 @@ 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 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.
|
||||
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:
|
||||
@@ -51,17 +67,30 @@ class SqliteVecStore(DocumentStore):
|
||||
self._conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS documents_meta (
|
||||
collection TEXT PRIMARY KEY,
|
||||
source TEXT NOT NULL
|
||||
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 (
|
||||
name TEXT PRIMARY KEY REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||
collection TEXT NOT NULL,
|
||||
owner_id TEXT NOT NULL,
|
||||
dim INTEGER NOT NULL,
|
||||
table_name TEXT 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
|
||||
)
|
||||
"""
|
||||
)
|
||||
@@ -69,47 +98,154 @@ class SqliteVecStore(DocumentStore):
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
id TEXT NOT NULL,
|
||||
collection TEXT NOT NULL REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||
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)
|
||||
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 ON documents(collection)")
|
||||
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 REFERENCES documents_meta(collection) ON DELETE CASCADE,
|
||||
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, page_number)
|
||||
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 ON document_pages(collection)")
|
||||
self._conn.commit()
|
||||
|
||||
async def ensure_collection(self, collection: str, source: str) -> None:
|
||||
async with self._lock:
|
||||
await asyncio.to_thread(self._sync_ensure_collection, collection, source)
|
||||
|
||||
def _sync_ensure_collection(self, collection: str, source: str) -> None:
|
||||
self._conn.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_pages_collection_owner ON document_pages(collection, owner_id)"
|
||||
)
|
||||
self._conn.execute(
|
||||
"""
|
||||
INSERT INTO documents_meta(collection, source) VALUES (?, ?)
|
||||
ON CONFLICT(collection) DO UPDATE SET source = excluded.source
|
||||
""",
|
||||
(collection, source),
|
||||
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) -> str:
|
||||
safe = re.sub(r"[^a-zA-Z0-9_]", "_", collection)
|
||||
return f"vec_{safe}"
|
||||
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]:
|
||||
@@ -123,6 +259,7 @@ class SqliteVecStore(DocumentStore):
|
||||
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")
|
||||
@@ -130,34 +267,40 @@ class SqliteVecStore(DocumentStore):
|
||||
return
|
||||
|
||||
async with self._lock:
|
||||
await asyncio.to_thread(self._sync_add, collection, documents, embeddings)
|
||||
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 name = ?", (collection,)).fetchone()
|
||||
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)
|
||||
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(name, dim, table_name) VALUES (?, ?, ?)",
|
||||
(collection, dim, table_name),
|
||||
"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} has dim {existing_dim}, got embedding of 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 id IN ({placeholders})",
|
||||
(collection, *ids),
|
||||
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]
|
||||
@@ -167,8 +310,8 @@ class SqliteVecStore(DocumentStore):
|
||||
vec_rowids,
|
||||
)
|
||||
self._conn.execute(
|
||||
f"DELETE FROM documents WHERE collection = ? AND id IN ({placeholders})",
|
||||
(collection, *ids),
|
||||
f"DELETE FROM documents WHERE collection = ? AND owner_id = ? AND id IN ({placeholders})",
|
||||
(collection, owner_id, *ids),
|
||||
)
|
||||
|
||||
for doc, emb in zip(documents, embeddings):
|
||||
@@ -179,27 +322,101 @@ class SqliteVecStore(DocumentStore):
|
||||
)
|
||||
vec_rowid = cursor.lastrowid
|
||||
self._conn.execute(
|
||||
"INSERT INTO documents(id, collection, text, metadata, vec_rowid) VALUES (?, ?, ?, ?, ?)",
|
||||
(doc.id, collection, doc.text, json.dumps(doc.metadata), vec_rowid),
|
||||
"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 = 5,
|
||||
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)
|
||||
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]:
|
||||
row = self._conn.execute("SELECT table_name, dim FROM collections WHERE name = ?", (collection,)).fetchone()
|
||||
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
|
||||
@@ -213,11 +430,14 @@ class SqliteVecStore(DocumentStore):
|
||||
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 = ?
|
||||
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, query_blob, top_k),
|
||||
(collection, owner_id, query_blob, top_k),
|
||||
).fetchall()
|
||||
|
||||
return [
|
||||
@@ -233,77 +453,85 @@ class SqliteVecStore(DocumentStore):
|
||||
for r in results
|
||||
]
|
||||
|
||||
async def add_pages(self, collection: str, pages: list[StoredPage]) -> None:
|
||||
async with self._lock:
|
||||
await asyncio.to_thread(self._sync_add_pages, collection, pages)
|
||||
|
||||
def _sync_add_pages(self, collection: str, pages: list[StoredPage]) -> None:
|
||||
self._conn.execute("DELETE FROM document_pages WHERE collection = ?", (collection,))
|
||||
if pages:
|
||||
self._conn.executemany(
|
||||
"INSERT INTO document_pages(collection, page_number, text, char_count) VALUES (?, ?, ?, ?)",
|
||||
[(collection, p.page_number, p.text, p.char_count) for p in pages],
|
||||
)
|
||||
self._conn.commit()
|
||||
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 = None,
|
||||
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)
|
||||
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 = ? ORDER BY page_number",
|
||||
(collection,),
|
||||
"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 page_number BETWEEN ? AND ? ORDER BY page_number",
|
||||
(collection, page_range.start, page_range.end),
|
||||
"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 delete_collection(self, collection: str) -> None:
|
||||
async def has_collection(self, collection: str, principals: list[PrincipalId]) -> bool:
|
||||
async with self._lock:
|
||||
await asyncio.to_thread(self._sync_delete_collection, collection)
|
||||
return await asyncio.to_thread(self._sync_has_collection, collection, principals)
|
||||
|
||||
def _sync_delete_collection(self, collection: str) -> None:
|
||||
# Drop the sqlite-vec virtual table first; FK cascade handles the regular tables
|
||||
# (collections, documents, document_pages) when documents_meta is deleted.
|
||||
row = self._conn.execute("SELECT table_name FROM collections WHERE name = ?", (collection,)).fetchone()
|
||||
if row is not None:
|
||||
self._conn.execute(f"DROP TABLE IF EXISTS {row[0]}")
|
||||
self._conn.execute("DELETE FROM documents_meta WHERE collection = ?", (collection,))
|
||||
self._conn.commit()
|
||||
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) -> list[str]:
|
||||
async def list_collections(self, principals: list[PrincipalId]) -> list[str]:
|
||||
async with self._lock:
|
||||
return await asyncio.to_thread(self._sync_list_collections)
|
||||
return await asyncio.to_thread(self._sync_list_collections, principals)
|
||||
|
||||
def _sync_list_collections(self) -> list[str]:
|
||||
rows = self._conn.execute("SELECT collection FROM documents_meta ORDER BY collection").fetchall()
|
||||
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 has_collection(self, collection: str) -> bool:
|
||||
async with self._lock:
|
||||
return await asyncio.to_thread(self._sync_has_collection, collection)
|
||||
|
||||
def _sync_has_collection(self, collection: str) -> bool:
|
||||
row = self._conn.execute(
|
||||
"SELECT 1 FROM documents_meta WHERE collection = ?",
|
||||
(collection,),
|
||||
).fetchone()
|
||||
return row is not None
|
||||
|
||||
async def close(self) -> None:
|
||||
async with self._lock:
|
||||
await asyncio.to_thread(self._sync_close)
|
||||
|
||||
Reference in New Issue
Block a user