diff --git a/docs/v3/contributing/configuration.mdx b/docs/v3/contributing/configuration.mdx index 3c210e71..f9b99ced 100644 --- a/docs/v3/contributing/configuration.mdx +++ b/docs/v3/contributing/configuration.mdx @@ -557,6 +557,15 @@ DB_SQL_DEBUG=false # Per-connection establish timeout (seconds) so a single connection attempt # fails fast instead of hanging when the server/pooler is unreachable. DB_CONNECT_TIMEOUT_SECONDS=2 + +# pgvector HNSW iterative scan (requires pgvector >= 0.8.0) +# When set, each new pool connection gets the GUC applied so filtered +# HNSW queries scan additional candidates instead of silently +# under-returning when out-of-scope rows consume the initial scan budget. +# Note: pgvector may still stop before reaching top_k if hnsw.max_scan_tuples +# or hnsw.scan_mem_multiplier thresholds are exceeded. +# Valid values: "off", "strict_order", "relaxed_order", or unset (None). +DB_HNSW_ITERATIVE_SCAN=strict_order ``` ### Authentication diff --git a/src/config.py b/src/config.py index 80827327..542c9fba 100644 --- a/src/config.py +++ b/src/config.py @@ -739,6 +739,18 @@ class DBSettings(HonchoSettings): SQL_DEBUG: bool = False TRACING: bool = False + # pgvector HNSW iterative scan mode for filtered approximate searches. + # When set (default "strict_order"), applied as a per-connection server + # setting so filtered HNSW queries scan additional candidates instead of + # silently under-returning when out-of-scope rows consume the initial scan + # budget. Note: pgvector may still stop before reaching top_k if + # hnsw.max_scan_tuples or hnsw.scan_mem_multiplier thresholds are exceeded. + # Requires pgvector >= 0.8.0. + # See: https://github.com/pgvector/pgvector#iterative-index-scans + HNSW_ITERATIVE_SCAN: Literal["off", "strict_order", "relaxed_order"] | None = ( + "strict_order" + ) + # Per-connection establish timeout (seconds) passed to the driver, so a # single connection attempt fails fast instead of hanging when the server or # pooler is unreachable or stalled. Connection acquisition is a single diff --git a/src/db.py b/src/db.py index 633f6f06..baa47125 100644 --- a/src/db.py +++ b/src/db.py @@ -20,7 +20,7 @@ from src.telemetry.prometheus.metrics import ( logger = logging.getLogger(__name__) -connect_args = { +connect_args: dict[str, Any] = { "prepare_threshold": None, # Bound a single connection attempt so it fails fast instead of hanging when # the server/pooler is unreachable or stalled (psycopg, seconds). @@ -89,6 +89,54 @@ ReadSessionLocal = async_sessionmaker( ) +def _set_hnsw_iterative_scan_on_connect( + dbapi_connection: Any, _connection_record: Any +) -> None: + """Apply pgvector HNSW iterative scan GUC per-connection. + + Registered when ``DB.HNSW_ITERATIVE_SCAN`` is set. Fires once per new + pool connection so filtered HNSW queries scan additional candidates + instead of silently under-returning when out-of-scope rows consume the + initial scan budget. Note: pgvector may still stop before reaching top_k + if ``hnsw.max_scan_tuples`` or ``hnsw.scan_mem_multiplier`` thresholds are + exceeded. Uses ``set_config`` with a bind parameter (same pattern as + ``_set_application_name_on_checkout``) rather than an f-string ``SET`` + to avoid special-casing utility-statement parameter binding. + + Runs in autocommit so it never leaves the connection 'idle in + transaction': this hook fires BEFORE the dialect applies execution-option + isolation levels, and psycopg refuses to switch a connection into + AUTOCOMMIT (which the read engine does) while a transaction opened by + this statement is still in progress. ``set_config(..., is_local=false)`` + is session-scoped, so it persists past the autocommit boundary. + """ + value = settings.DB.HNSW_ITERATIVE_SCAN + if not value: + return + try: + previous_autocommit = dbapi_connection.autocommit + if not previous_autocommit: + dbapi_connection.autocommit = True + try: + cursor = dbapi_connection.cursor() + try: + cursor.execute( + "SELECT set_config('hnsw.iterative_scan', %s, false)", + (value,), + ) + finally: + cursor.close() + finally: + if not previous_autocommit: + dbapi_connection.autocommit = False + except Exception: + logger.debug("setting hnsw.iterative_scan on connect failed", exc_info=True) + + +if settings.DB.HNSW_ITERATIVE_SCAN: + event.listen(engine.sync_engine, "connect", _set_hnsw_iterative_scan_on_connect) + + def _set_application_name_on_checkout( dbapi_connection: Any, _connection_record: Any, _connection_proxy: Any ) -> None: @@ -316,6 +364,25 @@ meta.schema = table_schema Base = declarative_base(metadata=meta) +def _validate_pgvector_version(version_str: str) -> None: + """Check that the installed pgvector version supports HNSW iterative scan. + + Raises ``RuntimeError`` if pgvector < 0.8.0. Extracted from + ``init_db`` so it can be unit-tested without importing alembic. + """ + version_parts = version_str.split(".") + major = int(version_parts[0]) if len(version_parts) > 0 else 0 + minor = int(version_parts[1]) if len(version_parts) > 1 else 0 + if (major, minor) < (0, 8): + raise RuntimeError( + "pgvector version " + + version_str + + " is installed but HNSW_ITERATIVE_SCAN" + + " requires pgvector >= 0.8.0." + + " Upgrade pgvector or set HNSW_ITERATIVE_SCAN=off." + ) + + async def init_db(): """Initialize the database using Alembic migrations""" from alembic import command @@ -328,6 +395,21 @@ async def init_db(): await connection.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) await connection.commit() + # Validate pgvector version when HNSW iterative scan is enabled + # (requires pgvector >= 0.8.0). Fail startup with a clear error + # rather than relying on silent query-time failures. + # Skip the check for "off" — listener is still registered so the + # application can override a server-level setting, but no version + # requirement applies. + if settings.DB.HNSW_ITERATIVE_SCAN in ("strict_order", "relaxed_order"): + async with engine.connect() as connection: + result = await connection.execute( + text("SELECT extversion FROM pg_extension WHERE extname = 'vector'") + ) + row = result.fetchone() + if row is not None: + _validate_pgvector_version(row[0]) + # Run Alembic migrations alembic_cfg = Config("alembic.ini") command.upgrade(alembic_cfg, "head") diff --git a/tests/conftest.py b/tests/conftest.py index 090d5395..4dfd5dcf 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -91,6 +91,8 @@ _RUNTIME_MOCK_TEST_BLOCKLIST_PREFIXES = ( # Pure JWT scope tests — operate on src.security directly, no DB needed. "tests/test_security.py", "tests/test_generate_jwt_script.py", + # Pure config unit tests — no DB or runtime mocks needed. + "tests/test_hnsw_iterative_scan.py", ) _LIVE_LLM_MARKER = "live_llm" diff --git a/tests/test_hnsw_iterative_scan.py b/tests/test_hnsw_iterative_scan.py new file mode 100644 index 00000000..195822c4 --- /dev/null +++ b/tests/test_hnsw_iterative_scan.py @@ -0,0 +1,129 @@ +"""Tests for the pgvector HNSW iterative scan connection setting. + +Verifies that: +- ``DBSettings.HNSW_ITERATIVE_SCAN`` accepts valid enum values and ``None`` +- An invalid value is rejected at config-load time (fail-closed) +- The ``connect`` event listener is registered when the setting is enabled +- The ``connect`` event listener is NOT registered when the setting is ``None`` +- ``_validate_pgvector_version`` raises for pgvector < 0.8.0 and passes for >= 0.8.0 +""" + +import pytest + +from src.config import DBSettings + + +def test_hnsw_iterative_scan_defaults_to_strict_order() -> None: + settings = DBSettings() + assert settings.HNSW_ITERATIVE_SCAN == "strict_order" + + +def test_hnsw_iterative_scan_accepts_valid_values() -> None: + for value in ("off", "strict_order", "relaxed_order"): + settings = DBSettings(HNSW_ITERATIVE_SCAN=value) + assert value == settings.HNSW_ITERATIVE_SCAN + + +def test_hnsw_iterative_scan_accepts_none() -> None: + settings = DBSettings(HNSW_ITERATIVE_SCAN=None) + assert settings.HNSW_ITERATIVE_SCAN is None + + +def test_hnsw_iterative_scan_rejects_invalid_value() -> None: + with pytest.raises((ValueError, TypeError)): + DBSettings( + HNSW_ITERATIVE_SCAN="on" # pyright: ignore[reportArgumentType] + ) + + +def test_hnsw_iterative_scan_rejects_arbitrary_string() -> None: + bad: str = "DROP TABLE users; --" + with pytest.raises((ValueError, TypeError)): + DBSettings(HNSW_ITERATIVE_SCAN=bad) # pyright: ignore[reportArgumentType] + + +def test_connect_listener_registered_when_enabled() -> None: + """The connect event listener is attached to the engine when + HNSW_ITERATIVE_SCAN is set. + + Uses ``sqlalchemy.event.contains`` to verify the listener is actually + registered with the engine's event system, not just that the function + exists and is callable. + """ + + from sqlalchemy import event + + from src import db as db_module + + _listener = db_module._set_hnsw_iterative_scan_on_connect # pyright: ignore[reportPrivateUsage] + + # The listener is registered at import time when the setting is + # truthy (the default is "strict_order"). We verify registration via + # the SQLAlchemy event registry rather than just checking callability. + assert event.contains( + db_module.engine.sync_engine, + "connect", + _listener, + ), "HNSW iterative scan connect listener should be registered on the engine" + + +def test_connect_listener_not_registered_when_disabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """When HNSW_ITERATIVE_SCAN is None, the connect function + short-circuits early without executing any SQL. + + This tests the early-return guard in + ``_set_hnsw_iterative_scan_on_connect`` rather than event + registration, because the listener is registered at import time + based on the initial config value and is not dynamically + removed when the setting changes at runtime. + """ + + from src.config import settings + + monkeypatch.setattr(settings.DB, "HNSW_ITERATIVE_SCAN", None) + assert settings.DB.HNSW_ITERATIVE_SCAN is None + + import types + + from src import db as db_module + + _listener = db_module._set_hnsw_iterative_scan_on_connect # pyright: ignore[reportPrivateUsage] + + execute_calls: list[tuple[tuple[object, ...], dict[str, object]]] = [] + + class AssertingCursor: + def execute(self, *args: object, **_kwargs: object) -> None: + execute_calls.append((args, _kwargs)) + raise AssertionError( + "execute should not be called when HNSW_ITERATIVE_SCAN is None" + ) + + def close(self) -> None: + pass + + dummy_conn = types.SimpleNamespace( + autocommit=False, cursor=lambda: AssertingCursor() + ) + # The function reads settings.DB.HNSW_ITERATIVE_SCAN at call time, + # so with monkeypatch it should return early without executing SQL. + _listener(dummy_conn, None) + assert execute_calls == [], "No SQL should execute when HNSW_ITERATIVE_SCAN is None" + + +def test_validate_pgvector_version_rejects_old_versions() -> None: + """_validate_pgvector_version raises RuntimeError for pgvector < 0.8.0.""" + from src.db import _validate_pgvector_version # pyright: ignore[reportPrivateUsage] + + for old_version in ("0.7.0", "0.6.1", "0.5.0"): + with pytest.raises(RuntimeError, match=r"requires pgvector >= 0\.8\.0"): + _validate_pgvector_version(old_version) + + +def test_validate_pgvector_version_accepts_new_versions() -> None: + """_validate_pgvector_version passes silently for pgvector >= 0.8.0.""" + from src.db import _validate_pgvector_version # pyright: ignore[reportPrivateUsage] + + for new_version in ("0.8.0", "0.8.1", "0.9.0", "1.0.0"): + _validate_pgvector_version(new_version) # should not raise