b85aea0738
tagger.pg_chunks filters on cmetadata->>'item_key'; the jsonb GIN index cannot serve ->>, so every comment lookup scanned the embedding table.
113 lines
4.4 KiB
Python
113 lines
4.4 KiB
Python
"""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
|
|
conn.execute.return_value.scalar.return_value = "llm"
|
|
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_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
|