arsvendg/Stirling-PDF
1
0
Fork 0
mirror of https://github.com/arsvendg/Stirling-PDF.git synced 2026-07-16 03:20:46 +02:00
Code Issues Actions 3 Packages Projects Releases 1 Wiki Activity

Improvements to Stirling Engine to prepare for SaaS release (#6603)

Browse Source 
# Description of Changes
- Use pool for postgres connections
- Add ability to require user ID to be set on API calls to the engine
- Add process-wide concurrency cap on AI access (in addition to existing
user caps)
- Allow number of workers (threads) to be specified for stirling engine
- Update env var names to reflect that the DB is not just for RAG
This commit is contained in:
James Brunton
2026-06-11 16:31:35 +01:00
committed by GitHub
parent 606964ee52
commit d52c7ced7c
17 changed files with 315 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/src/stirling/documents/pgvector_store.py
+59 -36
View File
@@ -1,10 +1,12 @@
from __future__ import annotations
import asyncio
import json
from datetime import datetime
import psycopg
from pgvector.psycopg import register_vector_async
from psycopg_pool import AsyncConnectionPool
from stirling.contracts.documents import Page, PageRange
from stirling.documents.store import Document, DocumentStore, SearchResult, StoredPage
@@ -13,12 +15,20 @@ from stirling.models import OwnerId, PrincipalId
_READ_PERMISSION = "read"
async def _register_vector(conn: psycopg.AsyncConnection) -> None:
await register_vector_async(conn)
class PgVectorStore(DocumentStore):
"""PostgreSQL + pgvector backed store, scoped by owner with ACL-gated reads.
Connects to an external Postgres instance (DSN provided via config) and uses the
`vector` extension for similarity search. The schema is created on first use.
Queries run on a shared async connection pool. The pool is opened lazily on
first use, after the schema bootstrap, because each pooled connection's
configure hook registers the ``vector`` type, which must already exist.
Owned tables:
* ``documents_meta`` - parent row, one per ``(collection, owner_id)``.
@@ -31,21 +41,36 @@ class PgVectorStore(DocumentStore):
Writes are owner-scoped; reads are ACL-scoped.
"""
def __init__(self, dsn: str) -> None:
def __init__(self, dsn: str, pool_min_size: int, pool_max_size: int) -> None:
if not dsn:
raise ValueError("pgvector backend requires a non-empty DSN (STIRLING_RAG_PGVECTOR_DSN)")
raise ValueError("pgvector backend requires a non-empty DSN (STIRLING_DOCUMENTS_PGVECTOR_DSN)")
self._dsn = dsn
self._pool = AsyncConnectionPool(
dsn,
min_size=pool_min_size,
max_size=pool_max_size,
open=False,
configure=_register_vector,
check=AsyncConnectionPool.check_connection,
)
self._initialized = False
self._init_lock = asyncio.Lock()
async def _connect(self) -> psycopg.AsyncConnection:
conn = await psycopg.AsyncConnection.connect(self._dsn)
await register_vector_async(conn)
return conn
async def _ensure_schema(self) -> None:
async def _ensure_ready(self) -> None:
"""Create the schema and open the pool, exactly once across concurrent callers."""
if self._initialized:
return
async with await self._connect() as conn:
async with self._init_lock:
if self._initialized:
return
await self._bootstrap_schema()
await self._pool.open()
self._initialized = True
async def _bootstrap_schema(self) -> None:
# Uses a dedicated non-pool connection: pooled connections register the
# vector type on checkout, which fails until the extension exists.
async with await psycopg.AsyncConnection.connect(self._dsn) as conn:
async with conn.cursor() as cur:
await cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
await cur.execute(
@@ -118,7 +143,6 @@ class PgVectorStore(DocumentStore):
"CREATE INDEX IF NOT EXISTS idx_acl_principal_permission ON document_acl(principal_id, permission)"
)
await conn.commit()
self._initialized = True
# ── lifecycle of the (collection, owner_id) row ────────────────────────
@@ -129,8 +153,8 @@ class PgVectorStore(DocumentStore):
owner_id: OwnerId,
expires_at: datetime | None,
) -> None:
await self._ensure_schema()
async with await self._connect() as conn:
await self._ensure_ready()
async with self._pool.connection() as conn:
async with conn.cursor() as cur:
await cur.execute(
"""
@@ -145,8 +169,8 @@ class PgVectorStore(DocumentStore):
await conn.commit()
async def purge_owner(self, owner_id: OwnerId) -> int:
await self._ensure_schema()
async with await self._connect() as conn:
await self._ensure_ready()
async with self._pool.connection() as conn:
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM documents_meta WHERE owner_id = %s",
@@ -157,8 +181,8 @@ class PgVectorStore(DocumentStore):
return deleted
async def reap_expired(self) -> int:
await self._ensure_schema()
async with await self._connect() as conn:
await self._ensure_ready()
async with self._pool.connection() as conn:
async with conn.cursor() as cur:
await cur.execute("DELETE FROM documents_meta WHERE expires_at IS NOT NULL AND expires_at < NOW()")
deleted = cur.rowcount
@@ -166,8 +190,8 @@ class PgVectorStore(DocumentStore):
return deleted
async def delete_collection(self, collection: str, owner_id: OwnerId) -> bool:
await self._ensure_schema()
async with await self._connect() as conn:
await self._ensure_ready()
async with self._pool.connection() as conn:
async with conn.cursor() as cur:
# Cascade FKs handle rag_documents, document_pages, document_acl.
await cur.execute(
@@ -192,8 +216,8 @@ class PgVectorStore(DocumentStore):
if not documents:
return
await self._ensure_schema()
async with await self._connect() as conn:
await self._ensure_ready()
async with self._pool.connection() as conn:
async with conn.cursor() as cur:
for doc, emb in zip(documents, embeddings):
await cur.execute(
@@ -211,8 +235,8 @@ class PgVectorStore(DocumentStore):
await conn.commit()
async def add_pages(self, collection: str, pages: list[StoredPage], owner_id: OwnerId) -> None:
await self._ensure_schema()
async with await self._connect() as conn:
await self._ensure_ready()
async with self._pool.connection() as conn:
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM document_pages WHERE collection = %s AND owner_id = %s",
@@ -238,8 +262,8 @@ class PgVectorStore(DocumentStore):
) -> None:
if not principals:
return
await self._ensure_schema()
async with await self._connect() as conn:
await self._ensure_ready()
async with self._pool.connection() as conn:
async with conn.cursor() as cur:
await cur.executemany(
"""
@@ -257,8 +281,8 @@ class PgVectorStore(DocumentStore):
owner_id: OwnerId,
principal: PrincipalId,
) -> None:
await self._ensure_schema()
async with await self._connect() as conn:
await self._ensure_ready()
async with self._pool.connection() as conn:
async with conn.cursor() as cur:
await cur.execute(
"DELETE FROM document_acl WHERE collection = %s AND owner_id = %s AND principal_id = %s",
@@ -277,8 +301,8 @@ class PgVectorStore(DocumentStore):
) -> list[SearchResult]:
if not principals:
return []
await self._ensure_schema()
async with await self._connect() as conn:
await self._ensure_ready()
async with self._pool.connection() as conn:
async with conn.cursor() as cur:
owner_id = await self._readable_owner_for(cur, collection, principals)
if owner_id is None:
@@ -332,8 +356,8 @@ class PgVectorStore(DocumentStore):
) -> list[Page]:
if not principals:
return []
await self._ensure_schema()
async with await self._connect() as conn:
await self._ensure_ready()
async with self._pool.connection() as conn:
async with conn.cursor() as cur:
owner_id = await self._readable_owner_for(cur, collection, principals)
if owner_id is None:
@@ -358,8 +382,8 @@ class PgVectorStore(DocumentStore):
async def has_collection(self, collection: str, principals: list[PrincipalId]) -> bool:
if not principals:
return False
await self._ensure_schema()
async with await self._connect() as conn:
await self._ensure_ready()
async with self._pool.connection() as conn:
async with conn.cursor() as cur:
owner_id = await self._readable_owner_for(cur, collection, principals)
return owner_id is not None
@@ -367,8 +391,8 @@ class PgVectorStore(DocumentStore):
async def list_collections(self, principals: list[PrincipalId]) -> list[str]:
if not principals:
return []
await self._ensure_schema()
async with await self._connect() as conn:
await self._ensure_ready()
async with self._pool.connection() as conn:
async with conn.cursor() as cur:
await cur.execute(
"""
@@ -383,5 +407,4 @@ class PgVectorStore(DocumentStore):
return [r[0] for r in rows]
async def close(self) -> None:
# Connections are opened and closed per call, so nothing persistent to release.
return None
await self._pool.close()