mirror of
https://github.com/VectifyAI/PageIndex.git
synced 2026-07-24 21:41:04 +02:00
11ba63debe
- Remove stream_metadata from chat/completions payloads — the server's ChatCompletionRequest model has no such field; it was silently ignored. - Local list_documents tie-breaker: rowid DESC → doc_id ASC, matching the server's "createdAt" DESC, "id" ASC ordering. - Cloud get_document: skip page_count when pageNum is 0 (server returns 0 for NULL, meaning "not yet computed", not "zero pages").
309 lines
14 KiB
Python
309 lines
14 KiB
Python
import json
|
|
import sqlite3
|
|
import threading
|
|
from pathlib import Path
|
|
|
|
from ..errors import CollectionAlreadyExistsError
|
|
from .._validation import validate_collection_name
|
|
|
|
|
|
def _validate_collection_name(name: str) -> None:
|
|
# SQLiteStorage is a public StorageEngine and may be used without
|
|
# LocalBackend, so enforce the shared contract at this boundary too.
|
|
validate_collection_name(name)
|
|
|
|
|
|
class SQLiteStorage:
|
|
def __init__(self, db_path: str):
|
|
self._db_path = Path(db_path).expanduser()
|
|
self._db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._local = threading.local()
|
|
self._connections: list[sqlite3.Connection] = []
|
|
self._conn_lock = threading.Lock()
|
|
# Bumped by close(). A thread caches its connection in thread-local
|
|
# storage, so after close() every OTHER thread's thread-local still
|
|
# points at a now-closed connection. Comparing the cached generation
|
|
# against this counter lets _get_conn detect that and reconnect, instead
|
|
# of handing back a closed connection (sqlite3.ProgrammingError). close()
|
|
# can only touch its OWN thread-local, so this is the only way to
|
|
# invalidate the others consistently.
|
|
self._generation = 0
|
|
# Serializes the (fast) write operations within this process so
|
|
# concurrent indexing threads don't collide on WAL's single writer
|
|
# ("database is locked"). Reads stay concurrent; the expensive LLM
|
|
# indexing runs outside this lock. busy_timeout above covers the
|
|
# cross-process case.
|
|
self._write_lock = threading.Lock()
|
|
self._init_schema()
|
|
|
|
def _get_conn(self) -> sqlite3.Connection:
|
|
"""Return a thread-local SQLite connection.
|
|
|
|
Reconnects if this thread has no connection yet OR its cached connection
|
|
was invalidated by a close() on another thread (generation mismatch).
|
|
"""
|
|
if (not hasattr(self._local, "conn")
|
|
or getattr(self._local, "generation", None) != self._generation):
|
|
# Each thread gets its own connection (threading.local), so
|
|
# statements never race. check_same_thread=False exists solely so
|
|
# close() can close every tracked connection from whichever thread
|
|
# calls it — with the default True those closes raise
|
|
# ProgrammingError and the connections leak.
|
|
# isolation_level=None -> autocommit: a plain SELECT (e.g. the
|
|
# dedup hash lookup) never leaves a lingering read snapshot that a
|
|
# later write on the same connection would conflict with
|
|
# (SQLITE_BUSY_SNAPSHOT, which busy_timeout can't retry). Each
|
|
# statement is its own transaction, so busy_timeout can actually
|
|
# wait for the WAL single-writer lock under concurrency.
|
|
conn = sqlite3.connect(str(self._db_path), check_same_thread=False,
|
|
isolation_level=None)
|
|
conn.execute("PRAGMA journal_mode=WAL")
|
|
conn.execute("PRAGMA foreign_keys=ON")
|
|
conn.execute("PRAGMA busy_timeout=10000")
|
|
self._local.conn = conn
|
|
self._local.generation = self._generation
|
|
with self._conn_lock:
|
|
self._connections.append(conn)
|
|
return self._local.conn
|
|
|
|
def _init_schema(self):
|
|
conn = self._get_conn()
|
|
# SQLite cannot ALTER a CHECK constraint. Rebuild the small parent table
|
|
# transactionally when opening a v1 database; documents keep referring to
|
|
# the same table name and are verified after foreign keys are re-enabled.
|
|
conn.execute("PRAGMA foreign_keys=OFF")
|
|
try:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
schema_version = conn.execute("PRAGMA user_version").fetchone()[0]
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS collections (
|
|
name TEXT PRIMARY KEY NOT NULL
|
|
CHECK(
|
|
length(name) BETWEEN 1 AND 255
|
|
-- SQLite length(TEXT) stops at the first NUL, while
|
|
-- Python and the cloud API count it as a character.
|
|
-- The Python boundary still enforces the 255 limit.
|
|
OR (instr(name, char(0)) > 0 AND name <> '')
|
|
),
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
schema_row = conn.execute(
|
|
"SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'collections'"
|
|
).fetchone()
|
|
schema_sql = schema_row[0] if schema_row else ""
|
|
schema_upper = schema_sql.upper()
|
|
if "BETWEEN 1 AND 128" in schema_upper or "NAME GLOB" in schema_upper:
|
|
conn.execute("""
|
|
CREATE TABLE _pageindex_collections_v2 (
|
|
name TEXT PRIMARY KEY NOT NULL
|
|
CHECK(
|
|
length(name) BETWEEN 1 AND 255
|
|
OR (instr(name, char(0)) > 0 AND name <> '')
|
|
),
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
conn.execute("""
|
|
INSERT INTO _pageindex_collections_v2 (name, created_at)
|
|
SELECT name, created_at FROM collections
|
|
""")
|
|
conn.execute("DROP TABLE collections")
|
|
conn.execute("ALTER TABLE _pageindex_collections_v2 RENAME TO collections")
|
|
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS documents (
|
|
doc_id TEXT PRIMARY KEY,
|
|
collection_name TEXT NOT NULL REFERENCES collections(name) ON DELETE CASCADE,
|
|
doc_name TEXT,
|
|
doc_description TEXT,
|
|
file_path TEXT,
|
|
file_hash TEXT,
|
|
doc_type TEXT NOT NULL,
|
|
status TEXT NOT NULL,
|
|
page_count INTEGER,
|
|
line_count INTEGER,
|
|
structure JSON,
|
|
pages JSON,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
UNIQUE(collection_name, file_hash)
|
|
)
|
|
""")
|
|
conn.execute(
|
|
"CREATE INDEX IF NOT EXISTS idx_docs_collection ON documents(collection_name)"
|
|
)
|
|
conn.execute(
|
|
"CREATE INDEX IF NOT EXISTS idx_docs_hash ON documents(collection_name, file_hash)"
|
|
)
|
|
if schema_version < 2:
|
|
conn.execute("PRAGMA user_version = 2")
|
|
conn.commit()
|
|
except Exception:
|
|
conn.rollback()
|
|
raise
|
|
finally:
|
|
conn.execute("PRAGMA foreign_keys=ON")
|
|
|
|
violations = conn.execute("PRAGMA foreign_key_check").fetchall()
|
|
if violations:
|
|
raise sqlite3.IntegrityError(
|
|
f"Foreign-key violations after SQLite schema migration: {violations!r}"
|
|
)
|
|
|
|
# DBs created before these columns existed: add them in place.
|
|
cols = {r[1] for r in conn.execute("PRAGMA table_info(documents)")}
|
|
for col_name, col_def in (
|
|
("page_count", "INTEGER"),
|
|
("line_count", "INTEGER"),
|
|
("status", "TEXT NOT NULL DEFAULT 'completed'"),
|
|
):
|
|
if col_name not in cols:
|
|
try:
|
|
conn.execute(f"ALTER TABLE documents ADD COLUMN {col_name} {col_def}")
|
|
except sqlite3.OperationalError:
|
|
pass # concurrent open of the same legacy DB already added it
|
|
conn.commit()
|
|
|
|
def create_collection(self, name: str) -> None:
|
|
_validate_collection_name(name)
|
|
with self._write_lock:
|
|
conn = self._get_conn()
|
|
try:
|
|
conn.execute("INSERT INTO collections (name) VALUES (?)", (name,))
|
|
except sqlite3.IntegrityError as e:
|
|
raise CollectionAlreadyExistsError(f"Collection '{name}' already exists") from e
|
|
conn.commit()
|
|
|
|
def get_or_create_collection(self, name: str) -> None:
|
|
_validate_collection_name(name)
|
|
with self._write_lock:
|
|
conn = self._get_conn()
|
|
# INSERT OR IGNORE works with older SQLite versions, but it can also
|
|
# suppress CHECK failures. Verify the row so ignored non-duplicate
|
|
# constraints never falsely report success.
|
|
conn.execute("INSERT OR IGNORE INTO collections (name) VALUES (?)", (name,))
|
|
if conn.execute(
|
|
"SELECT 1 FROM collections WHERE name = ?", (name,)
|
|
).fetchone() is None:
|
|
raise sqlite3.IntegrityError(
|
|
f"Collection {name!r} was rejected by the SQLite schema"
|
|
)
|
|
conn.commit()
|
|
|
|
def list_collections(self) -> list[str]:
|
|
conn = self._get_conn()
|
|
rows = conn.execute("SELECT name FROM collections ORDER BY name").fetchall()
|
|
return [r[0] for r in rows]
|
|
|
|
def delete_collection(self, name: str) -> None:
|
|
with self._write_lock:
|
|
conn = self._get_conn()
|
|
conn.execute("DELETE FROM collections WHERE name = ?", (name,))
|
|
conn.commit()
|
|
|
|
def save_document(self, collection: str, doc_id: str, doc: dict) -> None:
|
|
# Plain INSERT (doc_id is a fresh uuid, never pre-existing). A duplicate
|
|
# (collection_name, file_hash) raises sqlite3.IntegrityError, which the
|
|
# caller uses to resolve a concurrent add-of-same-file race.
|
|
with self._write_lock:
|
|
conn = self._get_conn()
|
|
conn.execute(
|
|
"""INSERT INTO documents
|
|
(doc_id, collection_name, doc_name, doc_description, file_path, file_hash, doc_type, status, page_count, line_count, structure, pages)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
|
(doc_id, collection, doc.get("doc_name"), doc.get("doc_description"),
|
|
doc.get("file_path"), doc.get("file_hash"), doc["doc_type"],
|
|
doc["status"],
|
|
doc.get("page_count"), doc.get("line_count"),
|
|
json.dumps(doc.get("structure", [])),
|
|
json.dumps(doc.get("pages")) if doc.get("pages") else None),
|
|
)
|
|
conn.commit()
|
|
|
|
def find_document_by_hash(self, collection: str, file_hash: str) -> str | None:
|
|
conn = self._get_conn()
|
|
row = conn.execute(
|
|
"SELECT doc_id FROM documents WHERE collection_name = ? AND file_hash = ?",
|
|
(collection, file_hash),
|
|
).fetchone()
|
|
return row[0] if row else None
|
|
|
|
def get_document(self, collection: str, doc_id: str) -> dict:
|
|
conn = self._get_conn()
|
|
row = conn.execute(
|
|
"SELECT doc_id, doc_name, doc_description, file_path, doc_type, status, page_count, line_count FROM documents WHERE doc_id = ? AND collection_name = ?",
|
|
(doc_id, collection),
|
|
).fetchone()
|
|
if not row:
|
|
return {}
|
|
doc = {"doc_id": row[0], "doc_name": row[1], "doc_description": row[2],
|
|
"file_path": row[3], "doc_type": row[4], "status": row[5]}
|
|
doc.update({k: v for k, v in (("page_count", row[6]), ("line_count", row[7])) if v is not None})
|
|
return doc
|
|
|
|
def get_document_structure(self, collection: str, doc_id: str) -> list:
|
|
conn = self._get_conn()
|
|
row = conn.execute(
|
|
"SELECT structure FROM documents WHERE doc_id = ? AND collection_name = ?",
|
|
(doc_id, collection),
|
|
).fetchone()
|
|
if not row:
|
|
return []
|
|
return json.loads(row[0])
|
|
|
|
def get_pages(self, collection: str, doc_id: str) -> list | None:
|
|
"""Return cached page content, or None if not cached."""
|
|
conn = self._get_conn()
|
|
row = conn.execute(
|
|
"SELECT pages FROM documents WHERE doc_id = ? AND collection_name = ?",
|
|
(doc_id, collection),
|
|
).fetchone()
|
|
if not row or not row[0]:
|
|
return None
|
|
return json.loads(row[0])
|
|
|
|
def list_documents(self, collection: str) -> list[dict]:
|
|
conn = self._get_conn()
|
|
rows = conn.execute(
|
|
"SELECT doc_id, doc_name, doc_description, doc_type FROM documents WHERE collection_name = ? ORDER BY created_at DESC, doc_id ASC",
|
|
(collection,),
|
|
).fetchall()
|
|
return [{"doc_id": r[0], "doc_name": r[1], "doc_description": r[2] or "", "doc_type": r[3]} for r in rows]
|
|
|
|
def delete_document(self, collection: str, doc_id: str) -> None:
|
|
with self._write_lock:
|
|
conn = self._get_conn()
|
|
conn.execute(
|
|
"DELETE FROM documents WHERE doc_id = ? AND collection_name = ?",
|
|
(doc_id, collection),
|
|
)
|
|
conn.commit()
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
self.close()
|
|
return False
|
|
|
|
def close(self) -> None:
|
|
"""Close all tracked SQLite connections across all threads."""
|
|
with self._conn_lock:
|
|
for conn in self._connections:
|
|
try:
|
|
conn.close()
|
|
except Exception:
|
|
pass
|
|
self._connections.clear()
|
|
# Invalidate every thread's cached connection. close() can only
|
|
# del its OWN thread-local, so the bump is what makes _get_conn on
|
|
# any other thread reconnect instead of reusing a closed handle.
|
|
self._generation += 1
|
|
if hasattr(self._local, "conn"):
|
|
del self._local.conn
|
|
|
|
def __del__(self):
|
|
try:
|
|
self.close()
|
|
except Exception:
|
|
pass
|