stack/tests/llm/test_migrate.py

"""llm.migrate — idempotent DDL for the llm database."""

from unittest.mock import MagicMock

from llm import migrate


class TestDdl:
    def test_index_state_ddl_is_idempotent_sql(self):
        assert "CREATE TABLE IF NOT EXISTS index_state" in migrate.INDEX_STATE_DDL
        assert "PRIMARY KEY (item_key, collection)" in migrate.INDEX_STATE_DDL

    def test_migrate_executes_ddl(self):
        engine = MagicMock()
        conn = engine.begin.return_value.__enter__.return_value
        migrate.migrate(engine)
        executed = " ".join(str(call.args[0]) for call in conn.execute.call_args_list)
        assert "index_state" in executed

    def test_ensure_hnsw_types_column_then_indexes(self):
        engine = MagicMock()
        conn = engine.begin.return_value.__enter__.return_value
        migrate.ensure_hnsw(engine, 768)
        executed = " ".join(str(call.args[0]) for call in conn.execute.call_args_list)
        assert "vector(768)" in executed
        assert "USING hnsw" in executed
        assert (
            "ix_embedding_collection" in executed
            and "langchain_pg_embedding (collection_id)" in executed
        )
        # F6: the metadata GIN indexes are ensure_metadata_indexes' job now,
        # not ensure_hnsw's — this must never take the ACCESS EXCLUSIVE lock.
        assert "ix_embedding_codes" not in executed
        assert "ix_embedding_families" not in executed
        assert "ix_embedding_elements" not in executed

    def test_ensure_metadata_indexes_issues_the_filter_indexes(self):
        engine = MagicMock()
        conn = engine.begin.return_value.__enter__.return_value
        migrate.ensure_metadata_indexes(engine)
        calls = [str(call.args[0]) for call in conn.execute.call_args_list]
        assert len(calls) == 4
        executed = " ".join(calls)
        # code-cited lookups scan cmetadata->>'codes' inside one collection
        assert (
            "ix_embedding_codes" in executed
            and "USING gin (string_to_array" in executed
        )
        assert (
            "ix_embedding_families" in executed and "cmetadata->>'families'" in executed
        )
        assert (
            "ix_embedding_elements" in executed and "cmetadata->>'elements'" in executed
        )
        # comment tagging looks up one item_key; `->>` cannot use the jsonb GIN
        assert (
            "ix_embedding_item_key" in executed
            and "(cmetadata->>'item_key')" in executed
        )
        # No vector-column ALTER, no HNSW build — this must be safe on a
        # live table.
        assert "ALTER TABLE" not in executed
        assert "hnsw" not in executed.lower()

    def test_fingerprint_column_and_docket_state(self):
        assert "fingerprint" in migrate.INDEX_STATE_DDL
        assert "ADD COLUMN IF NOT EXISTS fingerprint" in migrate.INDEX_STATE_ALTER
        assert (
            "CREATE TABLE IF NOT EXISTS index_docket_state"
            in migrate.INDEX_DOCKET_STATE_DDL
        )
        assert "PRIMARY KEY (collection, docket)" in migrate.INDEX_DOCKET_STATE_DDL

    def test_migrate_runs_alter_and_docket_state(self):
        engine = MagicMock()
        conn = engine.begin.return_value.__enter__.return_value
        migrate.migrate(engine)
        executed = " ".join(str(call.args[0]) for call in conn.execute.call_args_list)
        assert "ADD COLUMN IF NOT EXISTS fingerprint" in executed
        assert "index_docket_state" in executed


class TestIterativeScan:
    def test_migrate_sets_iterative_scan_on_the_database(self):
        from unittest.mock import MagicMock

        from llm.migrate import migrate

        engine = MagicMock()
        conn = engine.begin.return_value.__enter__.return_value

        def answer(stmt, *a, **k):
            r = MagicMock()
            r.scalar.return_value = (
                False if "pg_db_role_setting" in str(stmt) else "llm"
            )
            return r

        conn.execute.side_effect = answer
        migrate(engine)
        sql = [str(c.args[0]) for c in conn.execute.call_args_list]
        assert any(
            "ALTER DATABASE llm SET hnsw.iterative_scan = 'relaxed_order'" in q
            for q in sql
        )

    def test_already_set_database_wide_issues_no_alter(self):
        # The llm role may not ALTER DATABASE; re-issuing it on every index
        # run only logged two permission ERRORs in postgres per run.
        from unittest.mock import MagicMock

        from llm.migrate import ensure_iterative_scan

        engine = MagicMock()
        conn = engine.begin.return_value.__enter__.return_value

        def answer(stmt, *a, **k):
            r = MagicMock()
            r.scalar.return_value = True if "pg_db_role_setting" in str(stmt) else "llm"
            return r

        conn.execute.side_effect = answer
        ensure_iterative_scan(engine)
        sql = [str(c.args[0]) for c in conn.execute.call_args_list]
        assert not any("ALTER DATABASE" in q for q in sql)

    def test_failure_is_a_warning_not_an_error(self, caplog):
        from unittest.mock import MagicMock

        from llm.migrate import ensure_iterative_scan

        engine = MagicMock()
        engine.begin.return_value.__enter__.return_value.execute.side_effect = (
            RuntimeError("permission denied")
        )
        import logging

        caplog.set_level(logging.INFO, logger="llm.migrate")
        ensure_iterative_scan(engine)  # no raise
        assert "hnsw.iterative_scan not set" in caplog.text