feat(backends): observable maintenance hooks + pgvector indexed-build path (RFC 001, #1725)
Adds the maintenance contract RFC 001 specifies but #1679 deferred, and gives pgvector an opt-in HNSW index path with concurrency-safe builds. - base.py: MaintenanceResult (status ran/already_running/noop + free-form stats), UnsupportedMaintenanceKindError, BaseBackend.maintenance_kinds ClassVar (reserved: analyze/compact/reindex; a backend with no analogue MUST omit, not no-op), and BaseCollection.maintenance_state()/run_maintenance() defaults. EmbeddingCollection forwards both (BaseCollection methods shadow __getattr__). - sqlite_exact: analyze (ANALYZE) + compact (VACUUM, autocommit + page stats); omits reindex (exact scan, no ANN index). maintenance_state reports row/page counts. - pgvector: reindex builds the optional HNSW index, serialized by a session-level pg_advisory_lock so concurrent daemon writers learn "already_running" instead of each stacking an ACCESS EXCLUSIVE build (the production wedge). It is opt-in: the default exact `<=>` scan is the 100%-recall path; an HNSW index trades exact recall for scale, so an operator invokes it deliberately. Also analyze; omits compact (autovacuum). Advertises supports_server_side_indexes. maintenance_state reports index presence. - qdrant/chroma: empty maintenance_kinds (qdrant self-optimizes; chroma maintenance is the separate repair CLI) — the faithful "omit" default. Tests: contract + sqlite (real, CI-runnable) + pgvector advisory-lock flow via a fake client (ran/noop/already_running, no live Postgres). Full suite green: 2488 passed, 82.47% coverage. Benchmark three-phase wiring is deferred — the existing benchmarks/ are task-benchmarks, not backend-comparison harnesses, so there is nothing to wire into yet. Closes #1725. Refs #743.
This commit is contained in:
Igor Lins e Silva
parent
13de7a6743
commit
fdbafe5420
|
|
@ -27,11 +27,13 @@ from .base import (
|
|||
HealthStatus,
|
||||
LexicalHit,
|
||||
LexicalResult,
|
||||
MaintenanceResult,
|
||||
PalaceNotFoundError,
|
||||
PalaceRef,
|
||||
QueryResult,
|
||||
UnsupportedCapabilityError,
|
||||
UnsupportedFilterError,
|
||||
UnsupportedMaintenanceKindError,
|
||||
)
|
||||
from .chroma import ChromaBackend, ChromaCollection
|
||||
from .pgvector import PgVectorBackend, PgVectorCollection
|
||||
|
|
@ -64,6 +66,7 @@ __all__ = [
|
|||
"HealthStatus",
|
||||
"LexicalHit",
|
||||
"LexicalResult",
|
||||
"MaintenanceResult",
|
||||
"PalaceNotFoundError",
|
||||
"PalaceRef",
|
||||
"PgVectorBackend",
|
||||
|
|
@ -75,6 +78,7 @@ __all__ = [
|
|||
"SQLiteExactCollection",
|
||||
"UnsupportedCapabilityError",
|
||||
"UnsupportedFilterError",
|
||||
"UnsupportedMaintenanceKindError",
|
||||
"available_backends",
|
||||
"detect_backend_for_path",
|
||||
"detect_backends_for_path",
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ conformance suite land in follow-up PRs.
|
|||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from typing import ClassVar, Optional, Protocol, runtime_checkable
|
||||
|
||||
|
||||
|
|
@ -62,6 +62,15 @@ class UnsupportedCapabilityError(BackendError):
|
|||
"""Raised when a backend does not implement an optional capability."""
|
||||
|
||||
|
||||
class UnsupportedMaintenanceKindError(BackendError):
|
||||
"""Raised when ``run_maintenance(kind)`` is called with an unadvertised kind.
|
||||
|
||||
A backend MUST advertise a kind in ``maintenance_kinds`` before it accepts
|
||||
it (RFC 001). Advertising a kind it does not implement is a conformance
|
||||
failure; a kind it has no analogue for MUST be omitted, not no-op'd.
|
||||
"""
|
||||
|
||||
|
||||
class BackendMismatchError(BackendError):
|
||||
"""Raised when a selected backend does not match existing palace artifacts."""
|
||||
|
||||
|
|
@ -140,6 +149,29 @@ class EmbedderIdentity:
|
|||
dimension: int = 0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MaintenanceResult:
|
||||
"""Observable outcome of ``run_maintenance(kind)`` (RFC 001).
|
||||
|
||||
Maintenance is *not* fire-and-forget: a backend MUST serialize concurrent
|
||||
same-kind runs and report the outcome so a caller can learn it must not
|
||||
re-trigger. ``status`` is one of:
|
||||
|
||||
* ``"ran"`` — this call performed the maintenance.
|
||||
* ``"already_running"`` — another caller holds the work; this call did
|
||||
nothing and the caller MUST NOT re-trigger (the production index-build
|
||||
wedge: concurrent writers each issuing the build stacked exclusive locks).
|
||||
* ``"noop"`` — nothing needed doing (e.g. the index already exists).
|
||||
|
||||
``stats`` is free-form per kind (rows analyzed, bytes reclaimed, index
|
||||
build time) for benchmark/operator reporting.
|
||||
"""
|
||||
|
||||
kind: str
|
||||
status: str
|
||||
stats: dict = field(default_factory=dict)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class Embedder(Protocol):
|
||||
"""Minimal embedder contract (RFC 001, normative for identity checking).
|
||||
|
|
@ -436,6 +468,27 @@ class BaseCollection(ABC):
|
|||
"""
|
||||
return None
|
||||
|
||||
def maintenance_state(self) -> dict:
|
||||
"""Return a structured snapshot of this collection's maintenance state.
|
||||
|
||||
Free-form per backend (e.g. row count, whether a vector index exists,
|
||||
last-analyze age). Used by benchmark harnesses to record state
|
||||
alongside each latency/recall measurement so an un-analyzed store is
|
||||
not compared against a settled one (RFC 001). Defaults to empty.
|
||||
"""
|
||||
return {}
|
||||
|
||||
def run_maintenance(self, kind: str) -> "MaintenanceResult":
|
||||
"""Run a maintenance ``kind`` and return an observable result (RFC 001).
|
||||
|
||||
Backends advertise supported kinds in ``BaseBackend.maintenance_kinds``
|
||||
and override this. The default supports nothing, so every kind raises
|
||||
:class:`UnsupportedMaintenanceKindError`. Implementations MUST serialize
|
||||
concurrent same-kind runs and report ``already_running`` rather than
|
||||
stacking the work.
|
||||
"""
|
||||
raise UnsupportedMaintenanceKindError(f"backend does not support maintenance kind {kind!r}")
|
||||
|
||||
def lexical_search(
|
||||
self,
|
||||
*,
|
||||
|
|
@ -522,6 +575,14 @@ class BaseBackend(ABC):
|
|||
#: search converts distance→similarity off this declaration rather than
|
||||
#: assuming cosine. All in-tree backends are cosine today.
|
||||
distance_metric: ClassVar[str] = "cosine"
|
||||
#: Maintenance kinds this backend implements (RFC 001). Reserved names:
|
||||
#: ``"analyze"`` (refresh planner/query statistics), ``"compact"`` (reclaim
|
||||
#: space, rewrite storage), ``"reindex"`` (build/rebuild secondary indexes).
|
||||
#: A backend with no analogue for a kind MUST omit it rather than declare a
|
||||
#: no-op, so a benchmark harness can trust the set. Backends MAY add their
|
||||
#: own kinds. ``run_maintenance`` raises ``UnsupportedMaintenanceKindError``
|
||||
#: for anything not listed here.
|
||||
maintenance_kinds: ClassVar[frozenset[str]] = frozenset()
|
||||
|
||||
@abstractmethod
|
||||
def get_collection(
|
||||
|
|
|
|||
|
|
@ -69,6 +69,12 @@ class EmbeddingCollection(BaseCollection):
|
|||
def effective_embedder_identity(self):
|
||||
return self._inner.effective_embedder_identity()
|
||||
|
||||
def maintenance_state(self) -> dict:
|
||||
return self._inner.maintenance_state()
|
||||
|
||||
def run_maintenance(self, kind: str):
|
||||
return self._inner.run_maintenance(kind)
|
||||
|
||||
def add(self, *, documents, ids, metadatas=None, embeddings=None):
|
||||
documents = _as_list(documents)
|
||||
ids = _as_list(ids)
|
||||
|
|
|
|||
|
|
@ -352,6 +352,30 @@ def _quote_identifier(name: str) -> str:
|
|||
return '"' + name.replace('"', '""') + '"'
|
||||
|
||||
|
||||
# Session-level advisory-lock namespace for serializing HNSW index builds
|
||||
# across daemon writers (RFC 001). classid is a fixed mempalace constant
|
||||
# ("MEMP" in ASCII); objid is a stable per-table key. Both must fit a signed
|
||||
# int4, which ``pg_advisory_lock(int4, int4)`` requires.
|
||||
_MAINTENANCE_LOCK_CLASSID = 0x4D454D50 # "MEMP" — a positive, valid int4
|
||||
|
||||
|
||||
def _advisory_objid(table: str) -> int:
|
||||
"""Stable signed-int4 advisory key derived from the table name."""
|
||||
raw = int(sha256(table.encode("utf-8")).hexdigest()[:8], 16) # 0 .. 2**32-1
|
||||
return raw - 2**32 if raw >= 2**31 else raw
|
||||
|
||||
|
||||
def _hnsw_index_name(table: str) -> str:
|
||||
"""Deterministic, collision-safe index name for ``table``.
|
||||
|
||||
Routes through :func:`_pg_identifier`, which hashes the overflow when the
|
||||
name exceeds Postgres' 63-byte limit. A naive ``[:63]`` truncation could
|
||||
return the table name verbatim (tables and indexes share the ``pg_class``
|
||||
namespace), which would fail with "relation already exists".
|
||||
"""
|
||||
return _pg_identifier(f"{table}_hnsw_idx")
|
||||
|
||||
|
||||
def _field_sql(field: str, expression: Any, params: list) -> str:
|
||||
"""Translate one field predicate to a JSONB containment expression."""
|
||||
if isinstance(expression, dict):
|
||||
|
|
@ -625,6 +649,39 @@ class _PgVectorClient:
|
|||
def drop_table(self, table: str) -> None:
|
||||
self._execute(f"DROP TABLE IF EXISTS {_quote_identifier(table)}")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Maintenance (RFC 001)
|
||||
# ------------------------------------------------------------------
|
||||
def has_vector_index(self, table: str) -> bool:
|
||||
rows = self._execute(
|
||||
"SELECT 1 FROM pg_indexes WHERE schemaname = current_schema() "
|
||||
"AND tablename = %s AND indexdef ILIKE %s",
|
||||
[table, "%using hnsw%"],
|
||||
fetch=True,
|
||||
)
|
||||
return bool(rows)
|
||||
|
||||
def try_advisory_lock(self, classid: int, objid: int) -> bool:
|
||||
rows = self._execute("SELECT pg_try_advisory_lock(%s, %s)", [classid, objid], fetch=True)
|
||||
return bool(rows and rows[0] and rows[0][0])
|
||||
|
||||
def advisory_unlock(self, classid: int, objid: int) -> None:
|
||||
self._execute("SELECT pg_advisory_unlock(%s, %s)", [classid, objid], fetch=True)
|
||||
|
||||
def create_hnsw_index(self, table: str) -> None:
|
||||
qi = _quote_identifier(table)
|
||||
idx = _quote_identifier(_hnsw_index_name(table))
|
||||
# Non-concurrent build takes ACCESS EXCLUSIVE for the build duration;
|
||||
# the advisory lock in the caller ensures only one session builds, so
|
||||
# writes are blocked once rather than by every writer that crossed the
|
||||
# threshold (the production wedge this serialization fixes).
|
||||
self._execute(
|
||||
f"CREATE INDEX IF NOT EXISTS {idx} ON {qi} USING hnsw (embedding vector_cosine_ops)"
|
||||
)
|
||||
|
||||
def analyze_table(self, table: str) -> None:
|
||||
self._execute(f"ANALYZE {_quote_identifier(table)}")
|
||||
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
if self._conn is not None:
|
||||
|
|
@ -1014,6 +1071,60 @@ class PgVectorCollection(BaseCollection):
|
|||
return HealthStatus.unhealthy(str(exc))
|
||||
return HealthStatus.healthy()
|
||||
|
||||
def maintenance_state(self) -> dict:
|
||||
empty = {"row_count": 0, "vector_index": None, "index_build_complete": False}
|
||||
self._ensure_open()
|
||||
try:
|
||||
if not self._table_exists():
|
||||
return empty
|
||||
rows = self._client.count_rows(self._table)
|
||||
has_index = self._client.has_vector_index(self._table)
|
||||
except Exception: # noqa: BLE001 - state report must not raise
|
||||
logger.debug("pgvector maintenance state probe failed", exc_info=True)
|
||||
return empty
|
||||
return {
|
||||
"row_count": rows,
|
||||
"vector_index": "hnsw" if has_index else None,
|
||||
"index_build_complete": has_index,
|
||||
}
|
||||
|
||||
def run_maintenance(self, kind: str):
|
||||
from .base import MaintenanceResult, UnsupportedMaintenanceKindError
|
||||
|
||||
if kind not in PgVectorBackend.maintenance_kinds:
|
||||
raise UnsupportedMaintenanceKindError(
|
||||
f"pgvector does not support maintenance kind {kind!r}"
|
||||
)
|
||||
self._ensure_open()
|
||||
# Nothing to maintain on a not-yet-materialized table (collection opened
|
||||
# create=True but never written) — return noop rather than letting a
|
||||
# raw "relation does not exist" error escape.
|
||||
if not self._table_exists():
|
||||
return MaintenanceResult(kind=kind, status="noop", stats={"reason": "no table"})
|
||||
if kind == "analyze":
|
||||
self._client.analyze_table(self._table)
|
||||
return MaintenanceResult(kind="analyze", status="ran")
|
||||
|
||||
# reindex → build the optional HNSW index. Opt-in: it makes search
|
||||
# approximate, trading the exact-scan 100%-recall default for scale.
|
||||
# Serialized with a session advisory lock so concurrent daemon writers
|
||||
# learn "already_running" instead of each stacking an ACCESS EXCLUSIVE
|
||||
# index build.
|
||||
if self._client.has_vector_index(self._table):
|
||||
return MaintenanceResult(kind="reindex", status="noop", stats={"vector_index": "hnsw"})
|
||||
classid, objid = _MAINTENANCE_LOCK_CLASSID, _advisory_objid(self._table)
|
||||
if not self._client.try_advisory_lock(classid, objid):
|
||||
return MaintenanceResult(kind="reindex", status="already_running")
|
||||
try:
|
||||
if self._client.has_vector_index(self._table): # re-check under lock
|
||||
return MaintenanceResult(
|
||||
kind="reindex", status="noop", stats={"vector_index": "hnsw"}
|
||||
)
|
||||
self._client.create_hnsw_index(self._table)
|
||||
return MaintenanceResult(kind="reindex", status="ran", stats={"vector_index": "hnsw"})
|
||||
finally:
|
||||
self._client.advisory_unlock(classid, objid)
|
||||
|
||||
|
||||
class PgVectorBackend(BaseBackend):
|
||||
name = "pgvector"
|
||||
|
|
@ -1026,9 +1137,15 @@ class PgVectorBackend(BaseBackend):
|
|||
"supports_metadata_filters",
|
||||
"supports_lexical_search",
|
||||
"supports_namespace_isolation",
|
||||
"supports_server_side_indexes",
|
||||
"server_mode",
|
||||
}
|
||||
)
|
||||
# "compact" is omitted: Postgres autovacuum reclaims space automatically,
|
||||
# so a manual VACUUM kind would be redundant. "reindex" builds the optional
|
||||
# HNSW index — an opt-in scale lever, NOT on by default, because it makes
|
||||
# vector search approximate (the exact ``<=>`` scan is the 100%-recall path).
|
||||
maintenance_kinds = frozenset({"analyze", "reindex"})
|
||||
|
||||
def __init__(self):
|
||||
self._clients: dict[_PgVectorConfig, _PgVectorClient] = {}
|
||||
|
|
|
|||
|
|
@ -740,6 +740,63 @@ class SQLiteExactCollection(BaseCollection):
|
|||
return HealthStatus.unhealthy("collection closed")
|
||||
return HealthStatus.healthy()
|
||||
|
||||
def maintenance_state(self) -> dict:
|
||||
try:
|
||||
rows = self.count()
|
||||
except Exception:
|
||||
rows = 0
|
||||
# vector_index is null by design — exact cosine over every row, no ANN.
|
||||
state = {"row_count": rows, "vector_index": None}
|
||||
try:
|
||||
with self._cursor() as cur:
|
||||
page_count = cur.execute("PRAGMA page_count").fetchone()
|
||||
freelist = cur.execute("PRAGMA freelist_count").fetchone()
|
||||
state["page_count"] = int(page_count[0]) if page_count else 0
|
||||
state["freelist_pages"] = int(freelist[0]) if freelist else 0
|
||||
except Exception:
|
||||
pass
|
||||
return state
|
||||
|
||||
def run_maintenance(self, kind: str):
|
||||
from .base import MaintenanceResult, UnsupportedMaintenanceKindError
|
||||
|
||||
if kind not in SQLiteExactBackend.maintenance_kinds:
|
||||
raise UnsupportedMaintenanceKindError(
|
||||
f"sqlite_exact does not support maintenance kind {kind!r}"
|
||||
)
|
||||
if kind == "analyze":
|
||||
# Refresh planner stats. Concurrent runs serialize on the handle lock.
|
||||
with self._cursor() as cur:
|
||||
cur.execute("ANALYZE")
|
||||
return MaintenanceResult(kind="analyze", status="ran")
|
||||
|
||||
# compact → VACUUM. It cannot run inside a transaction, so flip the
|
||||
# connection to autocommit for the duration. The handle lock serializes
|
||||
# concurrent runs in-process; SQLite's own write lock serializes across
|
||||
# processes.
|
||||
before = self.maintenance_state()
|
||||
with self._handle.lock:
|
||||
self._ensure_open()
|
||||
conn = self._handle.conn
|
||||
prev_isolation = conn.isolation_level
|
||||
try:
|
||||
conn.commit()
|
||||
conn.isolation_level = None
|
||||
conn.execute("VACUUM")
|
||||
finally:
|
||||
conn.isolation_level = prev_isolation
|
||||
after = self.maintenance_state()
|
||||
reclaimed = max(0, before.get("page_count", 0) - after.get("page_count", 0))
|
||||
return MaintenanceResult(
|
||||
kind="compact",
|
||||
status="ran",
|
||||
stats={
|
||||
"pages_before": before.get("page_count", 0),
|
||||
"pages_after": after.get("page_count", 0),
|
||||
"pages_reclaimed": reclaimed,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class SQLiteExactBackend(BaseBackend):
|
||||
name = "sqlite_exact"
|
||||
|
|
@ -754,6 +811,9 @@ class SQLiteExactBackend(BaseBackend):
|
|||
"local_mode",
|
||||
}
|
||||
)
|
||||
# "reindex" is intentionally omitted: sqlite_exact does exact cosine over
|
||||
# every row (no ANN index to build), so it has no analogue for it.
|
||||
maintenance_kinds = frozenset({"analyze", "compact"})
|
||||
|
||||
def __init__(self):
|
||||
self._clients: dict[str, _SQLiteExactHandle] = {}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,275 @@
|
|||
"""Backend maintenance hooks (RFC 001).
|
||||
|
||||
Maintenance is observable, not fire-and-forget: ``run_maintenance(kind)``
|
||||
returns a ``MaintenanceResult`` and MUST serialize concurrent same-kind runs.
|
||||
The pgvector ``reindex`` path (the opt-in HNSW build) is exercised here with a
|
||||
fake client so the advisory-lock flow is tested without a live Postgres.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from mempalace.backends.base import (
|
||||
BaseCollection,
|
||||
MaintenanceResult,
|
||||
PalaceRef,
|
||||
UnsupportedMaintenanceKindError,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Contract surface
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_maintenance_result_shape():
|
||||
r = MaintenanceResult(kind="reindex", status="ran", stats={"ms": 12})
|
||||
assert r.kind == "reindex" and r.status == "ran" and r.stats["ms"] == 12
|
||||
assert MaintenanceResult(kind="analyze", status="noop").stats == {}
|
||||
|
||||
|
||||
def test_default_collection_rejects_all_kinds():
|
||||
class _Col(BaseCollection):
|
||||
def add(self, **k): ...
|
||||
def upsert(self, **k): ...
|
||||
def query(self, **k): ...
|
||||
def get(self, **k): ...
|
||||
def delete(self, **k): ...
|
||||
def count(self):
|
||||
return 0
|
||||
|
||||
col = _Col()
|
||||
assert col.maintenance_state() == {}
|
||||
with pytest.raises(UnsupportedMaintenanceKindError):
|
||||
col.run_maintenance("analyze")
|
||||
|
||||
|
||||
def test_backend_maintenance_kinds_declared():
|
||||
from mempalace.backends.chroma import ChromaBackend
|
||||
from mempalace.backends.pgvector import PgVectorBackend
|
||||
from mempalace.backends.qdrant import QdrantBackend
|
||||
from mempalace.backends.sqlite_exact import SQLiteExactBackend
|
||||
|
||||
assert SQLiteExactBackend.maintenance_kinds == frozenset({"analyze", "compact"})
|
||||
assert PgVectorBackend.maintenance_kinds == frozenset({"analyze", "reindex"})
|
||||
# qdrant self-optimizes; chroma maintenance is the separate repair CLI.
|
||||
assert QdrantBackend.maintenance_kinds == frozenset()
|
||||
assert ChromaBackend.maintenance_kinds == frozenset()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# sqlite_exact (CI-runnable, real backend)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _sqlite_collection(tmp_path, rows=20):
|
||||
from mempalace.backends.sqlite_exact import SQLiteExactBackend
|
||||
|
||||
col = SQLiteExactBackend().get_collection(
|
||||
palace=PalaceRef(id=str(tmp_path), local_path=str(tmp_path)),
|
||||
collection_name="mempalace_drawers",
|
||||
create=True,
|
||||
)
|
||||
for i in range(rows):
|
||||
col.add(
|
||||
documents=[f"doc {i}"],
|
||||
ids=[f"id{i}"],
|
||||
metadatas=[{}],
|
||||
embeddings=[[0.1, 0.2, 0.3, 0.4]],
|
||||
)
|
||||
return col
|
||||
|
||||
|
||||
def test_sqlite_maintenance_state(tmp_path):
|
||||
col = _sqlite_collection(tmp_path, rows=5)
|
||||
state = col.maintenance_state()
|
||||
assert state["row_count"] == 5
|
||||
assert state["vector_index"] is None # exact scan — no ANN index
|
||||
assert "page_count" in state and "freelist_pages" in state
|
||||
|
||||
|
||||
def test_sqlite_analyze_runs(tmp_path):
|
||||
col = _sqlite_collection(tmp_path, rows=5)
|
||||
r = col.run_maintenance("analyze")
|
||||
assert r.kind == "analyze" and r.status == "ran"
|
||||
|
||||
|
||||
def test_sqlite_compact_runs_and_reports_pages(tmp_path):
|
||||
col = _sqlite_collection(tmp_path, rows=30)
|
||||
col.delete(ids=[f"id{i}" for i in range(20)])
|
||||
r = col.run_maintenance("compact")
|
||||
assert r.kind == "compact" and r.status == "ran"
|
||||
assert "pages_reclaimed" in r.stats
|
||||
|
||||
|
||||
def test_sqlite_omits_reindex(tmp_path):
|
||||
# sqlite_exact has no ANN index, so reindex is omitted, not no-op'd.
|
||||
col = _sqlite_collection(tmp_path, rows=2)
|
||||
with pytest.raises(UnsupportedMaintenanceKindError):
|
||||
col.run_maintenance("reindex")
|
||||
|
||||
|
||||
def test_sqlite_unknown_kind_raises(tmp_path):
|
||||
col = _sqlite_collection(tmp_path, rows=2)
|
||||
with pytest.raises(UnsupportedMaintenanceKindError):
|
||||
col.run_maintenance("bogus")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# pgvector advisory-lock reindex flow (fake client, no live Postgres)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeClient:
|
||||
def __init__(self, has_index=False):
|
||||
self.has_index = has_index
|
||||
self.locked = False
|
||||
self.created = 0
|
||||
self.analyzed = 0
|
||||
|
||||
def table_exists(self, table):
|
||||
return True
|
||||
|
||||
def count_rows(self, table):
|
||||
return 7
|
||||
|
||||
def has_vector_index(self, table):
|
||||
return self.has_index
|
||||
|
||||
def try_advisory_lock(self, classid, objid):
|
||||
if self.locked:
|
||||
return False
|
||||
self.locked = True
|
||||
return True
|
||||
|
||||
def advisory_unlock(self, classid, objid):
|
||||
self.locked = False
|
||||
|
||||
def create_hnsw_index(self, table):
|
||||
self.has_index = True
|
||||
self.created += 1
|
||||
|
||||
def analyze_table(self, table):
|
||||
self.analyzed += 1
|
||||
|
||||
|
||||
class _FakeBackend:
|
||||
_closed = False
|
||||
|
||||
|
||||
def _pg_collection(client):
|
||||
from mempalace.backends.pgvector import PgVectorCollection, _PgVectorConfig
|
||||
|
||||
return PgVectorCollection(
|
||||
backend=_FakeBackend(),
|
||||
client=client,
|
||||
config=_PgVectorConfig(dsn="postgresql://example", namespace=None),
|
||||
palace=PalaceRef(id="/tmp/p", local_path="/tmp/p"),
|
||||
collection_name="mempalace_drawers",
|
||||
table="mp_drawers_t",
|
||||
)
|
||||
|
||||
|
||||
def test_pgvector_reindex_builds_index_under_lock():
|
||||
client = _FakeClient(has_index=False)
|
||||
col = _pg_collection(client)
|
||||
r = col.run_maintenance("reindex")
|
||||
assert r.status == "ran" and r.stats.get("vector_index") == "hnsw"
|
||||
assert client.created == 1
|
||||
assert client.locked is False # lock released in finally
|
||||
|
||||
|
||||
def test_pgvector_reindex_noop_when_index_exists():
|
||||
client = _FakeClient(has_index=True)
|
||||
col = _pg_collection(client)
|
||||
r = col.run_maintenance("reindex")
|
||||
assert r.status == "noop"
|
||||
assert client.created == 0 # never attempted a build
|
||||
|
||||
|
||||
def test_pgvector_reindex_already_running_when_lock_held():
|
||||
client = _FakeClient(has_index=False)
|
||||
client.locked = True # another session is building
|
||||
col = _pg_collection(client)
|
||||
r = col.run_maintenance("reindex")
|
||||
assert r.status == "already_running"
|
||||
assert client.created == 0 # did not re-trigger the build
|
||||
|
||||
|
||||
def test_pgvector_analyze_runs():
|
||||
client = _FakeClient()
|
||||
col = _pg_collection(client)
|
||||
r = col.run_maintenance("analyze")
|
||||
assert r.status == "ran" and client.analyzed == 1
|
||||
|
||||
|
||||
def test_pgvector_unknown_kind_raises():
|
||||
col = _pg_collection(_FakeClient())
|
||||
with pytest.raises(UnsupportedMaintenanceKindError):
|
||||
col.run_maintenance("compact") # pgvector omits compact (autovacuum)
|
||||
|
||||
|
||||
def test_pgvector_maintenance_state_reports_index():
|
||||
col = _pg_collection(_FakeClient(has_index=True))
|
||||
state = col.maintenance_state()
|
||||
assert state["row_count"] == 7
|
||||
assert state["vector_index"] == "hnsw" and state["index_build_complete"] is True
|
||||
|
||||
|
||||
def test_pgvector_maintenance_noop_when_table_missing():
|
||||
# Collection opened create=True but never written: no table yet. Maintenance
|
||||
# must noop, not let a raw "relation does not exist" error escape.
|
||||
client = _FakeClient()
|
||||
client.table_exists = lambda table: False
|
||||
col = _pg_collection(client)
|
||||
assert col.run_maintenance("reindex").status == "noop"
|
||||
assert col.run_maintenance("analyze").status == "noop"
|
||||
assert col.maintenance_state()["row_count"] == 0
|
||||
|
||||
|
||||
def test_hnsw_index_name_never_collides_with_table_name():
|
||||
# A naive [:63] truncation would return a 63-char table name verbatim,
|
||||
# colliding in pg_class. _pg_identifier hashes the overflow instead.
|
||||
from mempalace.backends.pgvector import _hnsw_index_name
|
||||
|
||||
for table in ("t", "mp_drawers", "x" * 63, "y" * 200):
|
||||
name = _hnsw_index_name(table)
|
||||
assert name != table
|
||||
assert len(name.encode("utf-8")) <= 63
|
||||
|
||||
|
||||
def test_pgvector_advisory_key_is_signed_int4_and_stable():
|
||||
from mempalace.backends.pgvector import _MAINTENANCE_LOCK_CLASSID, _advisory_objid
|
||||
|
||||
for table in ("a", "mempalace_drawers_xyz", "x" * 80):
|
||||
objid = _advisory_objid(table)
|
||||
assert -(2**31) <= objid < 2**31
|
||||
assert _advisory_objid(table) == objid # stable
|
||||
assert -(2**31) <= _MAINTENANCE_LOCK_CLASSID < 2**31
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# EmbeddingCollection delegation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_embeddingcollection_delegates_maintenance():
|
||||
from mempalace.backends.embedding_wrapper import EmbeddingCollection
|
||||
|
||||
class _Inner(BaseCollection):
|
||||
def add(self, **k): ...
|
||||
def upsert(self, **k): ...
|
||||
def query(self, **k): ...
|
||||
def get(self, **k): ...
|
||||
def delete(self, **k): ...
|
||||
def count(self):
|
||||
return 0
|
||||
|
||||
def maintenance_state(self):
|
||||
return {"row_count": 3}
|
||||
|
||||
def run_maintenance(self, kind):
|
||||
return MaintenanceResult(kind=kind, status="ran")
|
||||
|
||||
wrapped = EmbeddingCollection(_Inner())
|
||||
assert wrapped.maintenance_state() == {"row_count": 3}
|
||||
assert wrapped.run_maintenance("analyze").status == "ran"
|
||||
Loading…
Reference in New Issue