brain-of-reese/alembic/versions/0029_chunks_embedding_hnsw.py

"""chunks.embedding HNSW index: cosine search stops full-scanning (phase 188, task 01)

Revision ID: 0029
Revises: 0028
Create Date: 2026-10-15

The 2026-10-15 incident investigation found ``pg_indexes`` for ``chunks``
listing only the PK, ``ix_chunks_document_id`` and ``ix_chunks_tsv`` —
**no index on ``embedding``** — so every chat turn's A7 cosine tier
(``ORDER BY embedding <=> :vec LIMIT 100``) runs as a sequential scan of
a ~3 GB, 414,220-row relation: per-turn CPU on the interactive path plus
page-cache churn on a 3,904 MB box the sync also needs.

* ``ix_chunks_embedding_hnsw`` — HNSW over the existing ``vector(768)``
  column (A6 unchanged) with ``vector_cosine_ops``, so the ``<=>``
  operator in ``_vector_candidates`` uses it with NO query rewrite (A3's
  one-system rule: no new service, no new dependency). ``m = 16`` and
  ``ef_construction = 64`` are the pgvector defaults' neighbourhood,
  which keeps the graph buildable on a 4 GB box. ``halfvec(768)`` was
  considered and REJECTED for this phase (it needs a statement cast plus
  its own quantisation-recall study — two quality risks in one phase).

**Why ``CONCURRENTLY`` inside an autocommit block.** The index serves a
live interactive path; a plain ``CREATE INDEX`` would hold a lock that
blocks writes to ``chunks`` for the whole (long) build. ``CREATE INDEX
CONCURRENTLY`` cannot run inside a transaction, so the DDL executes in
``op.get_context().autocommit_block()``.

**The failure mode this migration guards (D2).** A ``CONCURRENTLY``
build that dies mid-way — killed by the OOM killer on a 3,904 MB box,
cancelled, or disconnected — leaves the index behind marked INVALID in
``pg_index``: it costs insert time on every sync and serves no queries,
*SILENTLY*. On this hardware that is the expected failure, not the
exotic one. Therefore ``upgrade()`` (a) pre-drops any existing INVALID
index of this name before building (otherwise the rebuild dies on
``relation "ix_chunks_embedding_hnsw" already exists``), and (b)
re-checks ``pg_index.indisvalid`` after the build and raises if the new
index did not validate — a migration that reports success while leaving
an INVALID index would paper over exactly the silent failure above.

Pure helpers (unit-testable, no I/O): :func:`should_drop_invalid` and
:func:`is_valid` are total functions over the ``indisvalid`` lookup
(``None`` = index absent).

**Reversibility (A13).** The downgrade is
``DROP INDEX CONCURRENTLY IF EXISTS`` — an index is a derived structure,
so the drop touches no data and 0028's schema is byte-identical after
it. No ``chunks``/``documents`` column changes and no data rewrite: the
migration is cheap on a healthy box and interruptible (the INVALID
self-heal makes re-running it the recovery path). The production build
(414k × 768) is an operator action with a pre-flight runbook — see
README "Production Deployment".
"""
from __future__ import annotations

from sqlalchemy import text
from sqlalchemy.engine import Connection

from alembic import op

revision = "0029"
down_revision = "0028"
branch_labels = None
depends_on = None

INDEX_NAME = "ix_chunks_embedding_hnsw"


def should_drop_invalid(indisvalid: bool | None) -> bool:
    """True only when the index EXISTS but is INVALID (phase 188 D2).

    ``None`` (absent) needs no drop; ``True`` (valid) must be left
    alone; ``False`` is an interrupted ``CREATE INDEX CONCURRENTLY``
    remnant that would otherwise block the rebuild with
    ``relation already exists``.
    """
    return indisvalid is False


def is_valid(indisvalid: bool | None) -> bool:
    """True only when the index EXISTS and validated (phase 188 D2)."""
    return indisvalid is True


def _lookup_index_validity(bind: Connection) -> bool | None:
    """``pg_index.indisvalid`` for :data:`INDEX_NAME`, ``None`` if absent."""
    row = bind.execute(
        text(
            "SELECT i.indisvalid"
            " FROM pg_index i"
            " JOIN pg_class c ON c.oid = i.indexrelid"
            " WHERE c.relname = :name"
        ),
        {"name": INDEX_NAME},
    ).scalar()
    return None if row is None else bool(row)


def upgrade() -> None:
    # CONCURRENTLY DDL cannot run inside a transaction — the autocommit
    # block commits alembic's surrounding transaction, runs the DDL
    # statement-by-statement, and reopens one afterwards.
    with op.get_context().autocommit_block():
        # Self-heal an interrupted build (D2): drop the INVALID remnant
        # BEFORE creating, or the CREATE dies on "already exists".
        if should_drop_invalid(_lookup_index_validity(op.get_bind())):
            op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {INDEX_NAME}")
        op.execute(
            f"CREATE INDEX CONCURRENTLY {INDEX_NAME} ON chunks "
            "USING hnsw (embedding vector_cosine_ops) "
            "WITH (m = 16, ef_construction = 64)"
        )
        # Post-validate: a CONCURRENTLY build that dies leaves an
        # INVALID index silently — fail the migration loudly instead
        # (the INVALID self-heal above makes re-running the recovery).
        if not is_valid(_lookup_index_validity(op.get_bind())):
            raise RuntimeError(f"{INDEX_NAME} did not validate — drop it and re-run")


def downgrade() -> None:
    # Same autocommit requirement as the create. IF EXISTS keeps the
    # downgrade idempotent; no data is touched (A13 — an index is a
    # derived structure, 0028's schema is byte-identical after the drop).
    with op.get_context().autocommit_block():
        op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {INDEX_NAME}")