brain-of-reese/tests/integration/test_chunks_hnsw_index.py

"""Integration: migration 0029 (chunks.embedding HNSW index) — phase 188, task 01.

Drives the **real Alembic engine** against the live dev database
(``podman compose up -d db``), mirroring the house pattern of
``test_migration_0028.py`` (fixture upgrades to head at setup and
teardown so the dev DB is never left below head; tests target the
0028 → 0029 step explicitly so later migrations cannot break them):

* upgrade 0028 → 0029 → ``ix_chunks_embedding_hnsw`` exists,
  ``pg_index.indisvalid`` is TRUE, and ``pg_get_indexdef`` names
  ``USING hnsw`` + ``vector_cosine_ops`` + ``m=16`` + ``ef_construction=64``
  (the D1 shape the recall pins of task 03 depend on);
* downgrade to 0028 → the index is GONE (A13 — reversible, ``CONCURRENTLY``),
  the ``chunks`` table and its PK survive;
* upgrade back to head → the index is back and valid (round-trip);
* the D2 self-heal: with an INVALID index of that name deliberately
  present, ``upgrade()`` drops it and rebuilds it VALID — no
  ``relation already exists`` error and no leftover INVALID index.
  The INVALID state is forced with the documented superuser catalog
  trick ``UPDATE pg_index SET indisvalid = false WHERE
  indexrelid = CAST(:name AS regclass)`` — the same failure a killed
  ``CREATE INDEX CONCURRENTLY`` leaves behind (pgvector cannot fail a
  build on data, and the dev/test DB role is a superuser; a role
  without catalog UPDATE rights gets a clean skip).

Pure-helper matrix tests (no DB): ``should_drop_invalid`` and
``is_valid`` over ``None`` / ``True`` / ``False``.

**Why every ``command.*`` call here is preceded by ``db.rollback()``
when the session was used.** psycopg3 keeps its transaction snapshot
alive while the session sits idle-in-transaction (measured: ``pg_stat
_activity.backend_xmin`` stays set where psql's clears), and ``CREATE
INDEX CONCURRENTLY`` phase 3 waits for every older-snapshot
transaction to finish — an unreleased session parks the build FOREVER
(observed on the dev box). The test session must therefore hold no
open transaction while the concurrent DDL runs; the rollback is what
makes the file deterministic rather than flaky.
**Task 02 adds the ``ef_search`` half:** the knob the index needs is
applied with ``SET LOCAL`` inside the turn's own transaction, so the file
pins the SCOPING against the real server — after a ``SET LOCAL`` the same
transaction reports the (clamped) value and a NEW transaction reports the
server's own default again (LOCKED A4: no global GUC, nothing surviving
the request or a pool recycle).
"""
from __future__ import annotations

import importlib.util
import logging
import types
from collections.abc import Iterator
from pathlib import Path

import pytest
from alembic.config import Config
from sqlalchemy import text
from sqlalchemy.exc import ProgrammingError
from sqlalchemy.orm import Session

from alembic import command
from app.config import Settings
from app.db import db_available
from app.rag import retriever
from app.rag.retriever import _vector_candidates

MIGRATION_PATH = Path("alembic/versions/0029_chunks_embedding_hnsw.py")
INDEX_NAME = "ix_chunks_embedding_hnsw"


def _load_migration() -> types.ModuleType:
    """Import the 0029 module by file path (``alembic/versions`` is not a
    package — the house way to reach a migration's own helpers)."""
    spec = importlib.util.spec_from_file_location("m0029_chunks_embedding_hnsw", MIGRATION_PATH)
    assert spec is not None and spec.loader is not None, f"missing {MIGRATION_PATH}"
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


@pytest.fixture()
def alembic(db: Session) -> Iterator[Config]:
    """Real Alembic config bound to the dev DB (URL from app settings).

    Starts at head (repairs an interrupted earlier run — 0029's own
    INVALID self-heal makes that repair safe when the interrupted
    migration was 0029 itself); teardown upgrades to head no matter
    what happened, so the dev DB is never left below head.
    """
    if not db_available():
        pytest.skip("Postgres not reachable — run `podman compose up -d db` first")
    cfg = Config()  # no alembic.ini file — env.py gets the URL from app config
    cfg.set_main_option("script_location", "alembic")
    command.upgrade(cfg, "head")
    try:
        yield cfg
    finally:
        # Release the test session's open transaction BEFORE the repair
        # DDL (the 0018 house pattern): a still-running snapshot would
        # park the CONCURRENTLY build, and teardown is exactly where a
        # failed test leaves one open.
        db.rollback()
        command.upgrade(cfg, "head")


def _version(db: Session) -> str | None:
    return db.execute(text("SELECT version_num FROM alembic_version")).scalar()


def _index_present(db: Session) -> bool:
    return (
        db.execute(
            text("SELECT 1 FROM pg_indexes WHERE indexname = :n"), {"n": INDEX_NAME}
        ).scalar()
        is not None
    )


def _index_valid(db: Session) -> bool | None:
    """``pg_index.indisvalid`` for the index, ``None`` when it is absent."""
    return db.execute(
        text(
            "SELECT i.indisvalid FROM pg_index i"
            " JOIN pg_class c ON c.oid = i.indexrelid"
            " WHERE c.relname = :n"
        ),
        {"n": INDEX_NAME},
    ).scalar()


def _index_def(db: Session) -> str | None:
    return db.execute(
        text("SELECT pg_get_indexdef(c.oid) FROM pg_class c WHERE c.relname = :n"),
        {"n": INDEX_NAME},
    ).scalar()


# ---------------------------------------------------------------------------
# Pure helpers (no DB, no I/O)
# ---------------------------------------------------------------------------


def test_should_drop_invalid_matrix() -> None:
    """``should_drop_invalid``: absent → False, valid → False,
    INVALID → True (the only state the pre-drop fires for)."""
    mod = _load_migration()
    assert mod.should_drop_invalid(None) is False, "absent index needs no drop"
    assert mod.should_drop_invalid(True) is False, "a valid index must be left alone"
    assert mod.should_drop_invalid(False) is True, "an INVALID remnant must be dropped"


def test_is_valid_matrix() -> None:
    """``is_valid``: only a present, validated index counts as valid."""
    mod = _load_migration()
    assert mod.is_valid(None) is False, "absent is not valid"
    assert mod.is_valid(True) is True, "validated is valid"
    assert mod.is_valid(False) is False, "INVALID is not valid"


def test_migration_module_contract() -> None:
    """The module's identity: revision 0029 on top of 0028, and the one
    index name the whole phase pins on."""
    mod = _load_migration()
    assert mod.revision == "0029"
    assert mod.down_revision == "0028"
    assert mod.INDEX_NAME == INDEX_NAME


# ---------------------------------------------------------------------------
# Migration round-trip against the real Alembic engine
# ---------------------------------------------------------------------------


def test_upgrade_creates_a_valid_hnsw_index(db: Session, alembic: Config) -> None:
    """Upgrade 0028 → 0029: the index exists, VALIDATED, with the D1
    shape — HNSW, ``vector_cosine_ops`` (the operator class A7's ``<=>``
    needs), ``m=16``, ``ef_construction=64`` — and the ``chunks`` table
    itself survives the additive step untouched."""
    command.downgrade(alembic, "0028")  # start from the pre-0029 state
    assert _version(db) == "0028"
    assert not _index_present(db), f"{INDEX_NAME} must be absent at 0028"

    db.rollback()  # no open snapshot while the CONCURRENTLY build runs
    command.upgrade(alembic, "0029")
    assert _version(db) == "0029", "alembic_version must be at 0029"
    assert _index_present(db), "the index must exist at 0029"
    assert _index_valid(db) is True, "a fresh build must be VALID (D2 post-check)"

    raw = _index_def(db)
    assert raw is not None
    definition = raw.replace("'", "")  # pg renders WITH options quoted
    assert "USING hnsw" in definition, f"not an HNSW index: {definition}"
    assert "vector_cosine_ops" in definition, (
        f"cosine operator class required for <=>: {definition}"
    )
    assert "m=16" in definition, f"m=16 expected: {definition}"
    assert "ef_construction=64" in definition, f"ef_construction=64 expected: {definition}"


def test_downgrade_removes_it(db: Session, alembic: Config) -> None:
    """Downgrade 0029 → 0028: the index is gone (A13 — reversible), and
    the ``chunks`` relation with its PK survives (the drop touches no
    data)."""
    command.upgrade(alembic, "0029")
    assert _index_present(db)

    db.rollback()  # no open snapshot while the CONCURRENTLY drop runs
    command.downgrade(alembic, "0028")
    assert _version(db) == "0028"
    assert not _index_present(db), f"{INDEX_NAME} must be dropped by the downgrade"
    assert _index_valid(db) is None, "no pg_index row may survive the downgrade"

    pk = db.execute(
        text("SELECT 1 FROM pg_indexes WHERE indexname = 'chunks_pkey'")
    ).scalar()
    assert pk is not None, "chunks_pkey must survive the downgrade"


def test_upgrade_head_again(db: Session, alembic: Config) -> None:
    """A13 round-trip: downgrade to 0028, upgrade back to HEAD — the
    index returns, valid, with the same definition."""
    command.downgrade(alembic, "0028")
    command.upgrade(alembic, "head")
    assert _index_present(db), "the index must be back after re-upgrade"
    assert _index_valid(db) is True, "the rebuilt index must be VALID"
    raw = _index_def(db)
    assert raw is not None and "USING hnsw" in raw


def test_invalid_index_is_dropped_and_rebuilt(db: Session, alembic: Config) -> None:
    """D2 self-heal: an INVALID index of this name (what a killed
    ``CREATE INDEX CONCURRENTLY`` leaves behind) is dropped and rebuilt
    VALID on the next ``upgrade()``.

    The INVALID state is forced with the documented superuser catalog
    update ``UPDATE pg_index SET indisvalid = false WHERE indexrelid =
    CAST(:name AS regclass)`` — pgvector cannot fail an HNSW build on
    data, and timing a real interruption is nondeterministic; the dev/
    test DB role is a superuser, and a role without catalog UPDATE
    rights skips cleanly instead of leaving the DB inconsistent (the
    stamp that follows the forced state, never before it)."""
    command.upgrade(alembic, "head")
    assert _index_valid(db) is True, "precondition: a valid index to corrupt"

    try:
        db.execute(
            text(
                "UPDATE pg_index SET indisvalid = false"
                " WHERE indexrelid = CAST(:n AS regclass)"
            ),
            {"n": INDEX_NAME},
        )
        db.commit()
    except ProgrammingError as exc:  # non-superuser role
        db.rollback()
        pytest.skip(f"forcing INVALID needs catalog UPDATE rights: {exc}")

    # The stamp moves the version table WITHOUT touching the index, so
    # the next upgrade() re-runs 0029 against the corrupted state.
    command.stamp(alembic, "0028")
    assert _version(db) == "0028"
    assert _index_valid(db) is False, "the forced INVALID state must be visible"

    db.rollback()  # no open snapshot of ours while the CONCURRENTLY run rebuilds
    command.upgrade(alembic, "0029")
    assert _index_present(db), "no 'relation already exists' — it was pre-dropped"
    assert _index_valid(db) is True, "the rebuilt index must be VALID (no leftover INVALID)"
    raw = _index_def(db)
    assert raw is not None and "USING hnsw" in raw, "the rebuild is the HNSW index itself"


# ---------------------------------------------------------------------------
# Task 02: SET LOCAL hnsw.ef_search — applied, clamped, transaction-scoped
# ---------------------------------------------------------------------------


def _show_ef_search(db: Session) -> str:
    return str(db.execute(text("SHOW hnsw.ef_search")).scalar())


def _server_default_ef_search(db: Session) -> str:
    """The server's own reset value for the GUC (the value a NEW
    transaction must fall back to). Read from ``pg_settings`` rather than
    hardcoded: pgvector renders its default, and the pin is about the
    FALLBACK, not the number."""
    value = db.execute(
        text("SELECT reset_val FROM pg_settings WHERE name = 'hnsw.ef_search'")
    ).scalar()
    assert value is not None, "pgvector must register hnsw.ef_search in this session"
    return str(value)


def _stub_settings(**kwargs: object) -> Settings:
    return Settings(_env_file=None, **kwargs)  # pyright: ignore[reportCallIssue]


def test_set_local_ef_search_reaches_the_server_in_the_same_transaction(
    db: Session, monkeypatch: pytest.MonkeyPatch
) -> None:
    """The D3 rule against a real server: the value the tier sends is the
    value Postgres runs the search with. The setting here (10) sits BELOW
    the limit (250), so the server must report the CLAMPED value — the
    clamp is not a Python-side fiction."""
    monkeypatch.setattr(
        retriever, "get_settings", lambda: _stub_settings(retrieval_ef_search=10)
    )
    db.rollback()  # a clean transaction boundary for the SET LOCAL
    _vector_candidates(db, [0.0] * 768, limit=250)  # pyright: ignore[reportArgumentType]
    assert _show_ef_search(db) == "250", "the clamped value must be live in this tx"
    db.rollback()


def test_set_local_ef_search_does_not_survive_the_transaction(
    db: Session, monkeypatch: pytest.MonkeyPatch
) -> None:
    """LOCKED A4 pinned on the wire: once the turn's transaction ends, the
    server is back on ITS default — no global GUC mutation, no state a
    pooled connection carries into the next request."""
    monkeypatch.setattr(
        retriever, "get_settings", lambda: _stub_settings(retrieval_ef_search=317)
    )
    db.rollback()
    _vector_candidates(db, [0.0] * 768, limit=100)  # pyright: ignore[reportArgumentType]
    assert _show_ef_search(db) == "317"

    db.rollback()  # the end of the turn's transaction
    after = _show_ef_search(db)
    default = _server_default_ef_search(db)
    assert after == default, f"a new tx must see the server default, got {after}"
    assert after != "317"
    db.rollback()


def test_ef_search_clamp_logs_one_warning_line(
    db: Session, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
    """The clamp leaves exactly one operator-visible line in the server log
    (AGENTS rule 10), and the pass-through case leaves none."""
    monkeypatch.setattr(
        retriever, "get_settings", lambda: _stub_settings(retrieval_ef_search=5)
    )
    db.rollback()
    with caplog.at_level(logging.WARNING, logger="app.rag.retriever"):
        _vector_candidates(db, [0.0] * 768, limit=100)  # pyright: ignore[reportArgumentType]
    db.rollback()
    lines = [r.getMessage() for r in caplog.records if "ef_search" in r.getMessage()]
    assert lines == ["retrieval: ef_search=5 raised to 100 (A7 candidate count)"]

    monkeypatch.setattr(
        retriever, "get_settings", lambda: _stub_settings(retrieval_ef_search=100)
    )
    caplog.clear()
    with caplog.at_level(logging.WARNING, logger="app.rag.retriever"):
        _vector_candidates(db, [0.0] * 768, limit=100)  # pyright: ignore[reportArgumentType]
    db.rollback()
    assert [r.getMessage() for r in caplog.records if "ef_search" in r.getMessage()] == []