import pytest from psycopg2.errors import InsufficientPrivilege from sqlalchemy import create_engine, text from sqlalchemy.exc import ProgrammingError from testcontainers.postgres import PostgresContainer from tht.adapters.vector.thoth_http import ThothHttpVectorStore from tht.config import DatabaseConfig from tht.ports.vector import ( VectorReadUnavailable, VectorRecord, VectorStoreError, VectorWriteRecord, VectorWriteUnavailable, ) def _record(content_hash: str, embedding: list[float], *, kind: str = "memory"): return VectorWriteRecord( record=VectorRecord( id=f"record:{content_hash}", kind=kind, ref="session:test", title=content_hash, content=f"content {content_hash}", metadata={"content_hash": content_hash}, ), embedding=embedding, content_hash=content_hash, ) @pytest.fixture(scope="module") def vector_configs(): with PostgresContainer("pgvector/pgvector:pg16") as pg: host = pg.get_container_host_ip() port = int(pg.get_exposed_port(5432)) admin_config = DatabaseConfig( host=host, port=port, database=pg.dbname, schema="vectors", user=pg.username, password=pg.password, ) engine = create_engine(pg.get_connection_url()) with engine.begin() as connection: connection.exec_driver_sql("CREATE SCHEMA vectors") # Match the co-located Supabase deployment: tables are in vectors, extension in public. connection.exec_driver_sql("CREATE EXTENSION vector WITH SCHEMA public") for table in ("schema_records", "evidence", "memory"): connection.exec_driver_sql(f""" CREATE TABLE vectors.{table} ( id bigserial PRIMARY KEY, record_key text UNIQUE NOT NULL, kind text NOT NULL, content_hash text NOT NULL, metadata jsonb NOT NULL, embedding public.vector(2) NOT NULL, indexed_at timestamptz NOT NULL DEFAULT now() ) """) connection.exec_driver_sql("CREATE ROLE vector_l0_reader LOGIN PASSWORD 'reader'") connection.exec_driver_sql("CREATE ROLE vector_l0_writer LOGIN PASSWORD 'writer'") connection.exec_driver_sql( "CREATE ROLE vector_l0_no_sequence LOGIN PASSWORD 'no_sequence'" ) connection.exec_driver_sql( "GRANT USAGE ON SCHEMA vectors TO vector_l0_reader, vector_l0_writer, " "vector_l0_no_sequence" ) connection.exec_driver_sql( "GRANT SELECT ON ALL TABLES IN SCHEMA vectors TO vector_l0_reader" ) connection.exec_driver_sql( "GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA vectors TO vector_l0_writer" ) for table in ("schema_records", "evidence", "memory"): connection.exec_driver_sql( f"GRANT INSERT, UPDATE ON vectors.{table} " "TO vector_l0_writer, vector_l0_no_sequence" ) if table == "evidence": connection.exec_driver_sql( "GRANT DELETE ON vectors.evidence TO vector_l0_writer" ) connection.exec_driver_sql( "GRANT SELECT (kind, metadata) ON vectors.evidence TO vector_l0_writer" ) connection.exec_driver_sql( f"GRANT SELECT (record_key, kind, content_hash) " f"ON vectors.{table} TO vector_l0_writer, vector_l0_no_sequence" ) engine.dispose() reader_config = admin_config.model_copy( update={"user": "vector_l0_reader", "password": "reader"} ) writer_config = admin_config.model_copy( update={"user": "vector_l0_writer", "password": "writer"} ) no_sequence_config = admin_config.model_copy( update={"user": "vector_l0_no_sequence", "password": "no_sequence"} ) yield admin_config, reader_config, writer_config, no_sequence_config @pytest.fixture def store(vector_configs): from tht.adapters.vector.pgvector import PgVectorStore _, reader_config, writer_config, _ = vector_configs store = PgVectorStore(reader_config, writer_config, expected_dimension=2) store.upsert("memory", [_record("reset", [0.0, 1.0])]) yield store def test_pgvector_round_trip_hash_and_upsert(store): assert store.upsert("memory", [_record("a", [1.0, 0.0])]) == 1 assert store.existing_hashes("memory", ["memory"])["record:a"] == "a" hits = store.search(["memory"], [1.0, 0.0], limit=5, kinds=["memory"]) assert hits[0].metadata["content_hash"] == "a" assert hits[0].id == "record:a" assert store.upsert("memory", [_record("a", [0.8, 0.2])]) == 1 assert store.search(["memory"], [0.8, 0.2], limit=1)[0].id == "record:a" def test_pgvector_lists_and_deletes_exact_evidence_generation(store): generation = "gen:" + "a" * 32 value = VectorWriteRecord( record=VectorRecord( id="evidence-generation-a", kind="evidence", ref="doc:a", title="a", content="content", metadata={"vector_generation": generation, "workspace_id": "default"}, ), embedding=[1.0, 0.0], content_hash="sha256:" + "a" * 64, ) store.upsert("evidence", [value]) assert generation in store.list_evidence_generations("evidence", "default") assert store.delete_generation("evidence", generation, "default") == 1 assert generation not in store.list_evidence_generations("evidence", "default") def test_pgvector_generation_cleanup_isolated_between_workspaces(store): generation = "gen:" + "b" * 32 records = [VectorWriteRecord( record=VectorRecord( id=f"evidence-{workspace}", kind="evidence", ref=f"doc:{workspace}", title=workspace, content=workspace, metadata={"vector_generation": generation, "workspace_id": workspace}, ), embedding=[1.0, 0.0], content_hash="sha256:" + key * 64, ) for workspace, key in (("workspace-a", "b"), ("workspace-b", "c"))] store.upsert("evidence", records) assert store.delete_generation("evidence", generation, "workspace-a") == 1 assert generation not in store.list_evidence_generations("evidence", "workspace-a") assert generation in store.list_evidence_generations("evidence", "workspace-b") def test_pgvector_search_filters_kinds_before_limit(store): store.upsert("memory", [_record("solved", [1.0, 0.0], kind="solved_question")]) hits = store.search("memory".split(), [1.0, 0.0], limit=1, kinds=["memory"]) assert len(hits) == 1 assert hits[0].kind == "memory" def test_pgvector_multi_collection_search_skips_collections_unrelated_to_kinds(store): store.upsert("evidence", [_record("evidence", [1.0, 0.0], kind="evidence")]) hits = store.search(["evidence", "memory"], [1.0, 0.0], limit=3, kinds=["memory"]) assert hits assert {hit.kind for hit in hits} == {"memory"} def test_pgvector_multi_collection_kind_filter_matches_http_adapter(store): class Reader: def search_similar(self, collection, embedding, limit, kinds=None): if collection != "memory" or "memory" not in (kinds or []): return [] return [ { "similarity": 1.0, "metadata": { "record_key": "record:a", "kind": "memory", "ref": "session:test", "title": "a", "content": "content a", "content_hash": "a", }, } ] direct = store.search(["evidence", "memory"], [1.0, 0.0], limit=1, kinds=["memory"]) http = ThothHttpVectorStore(Reader(), None).search( ["evidence", "memory"], [1.0, 0.0], limit=1, kinds=["memory"] ) assert [(hit.id, hit.kind) for hit in direct] == [(hit.id, hit.kind) for hit in http] def test_pgvector_search_rejects_unknown_kind_globally(store): with pytest.raises(VectorStoreError, match="Kind not allowed"): store.search(["memory"], [1.0, 0.0], limit=1, kinds=["unknown"]) @pytest.mark.parametrize("limit", [True, False, 1.0, 0, -1]) def test_pgvector_search_requires_strict_positive_limit(store, limit): with pytest.raises(ValueError, match="positive integer"): store.search(["memory"], [1.0, 0.0], limit=limit) def test_pgvector_allowlists_collections(store): with pytest.raises(VectorStoreError, match="Collection not allowed"): store.search(["memory; DROP SCHEMA vectors"], [1.0, 0.0], limit=1) with pytest.raises(VectorStoreError, match="Collection not allowed"): store.upsert("unknown", []) def test_pgvector_rejects_kinds_not_belonging_to_collection(store): with pytest.raises(VectorStoreError, match="Kind not allowed"): store.existing_hashes("evidence", ["memory"]) with pytest.raises(VectorStoreError, match="Kind not allowed"): store.upsert("evidence", [_record("wrong", [1.0, 0.0])]) def test_pgvector_separates_read_and_write_credentials(vector_configs): from tht.adapters.vector.pgvector import PgVectorStore _, reader_config, writer_config, _ = vector_configs reader = PgVectorStore(reader_config, expected_dimension=2) assert reader.capabilities.search is True assert reader.capabilities.upsert is False with pytest.raises(VectorWriteUnavailable): reader.upsert("memory", []) writer = PgVectorStore(None, writer_config, expected_dimension=2) assert writer.capabilities.search is False assert writer.capabilities.upsert is True with pytest.raises(VectorReadUnavailable): writer.search(["memory"], [1.0, 0.0], limit=1) def test_pgvector_database_roles_are_least_privilege(vector_configs): _, reader_config, writer_config, _ = vector_configs reader_engine = create_engine( f"postgresql+psycopg2://{reader_config.user}:{reader_config.password}" f"@{reader_config.host}:{reader_config.port}/{reader_config.database}" ) writer_engine = create_engine( f"postgresql+psycopg2://{writer_config.user}:{writer_config.password}" f"@{writer_config.host}:{writer_config.port}/{writer_config.database}" ) with pytest.raises(ProgrammingError): with reader_engine.begin() as connection: connection.execute( text( "INSERT INTO vectors.memory " "(record_key, kind, content_hash, metadata, embedding) " "VALUES ('forbidden', 'memory', 'x', '{}', '[1,0]')" ) ) with pytest.raises(ProgrammingError): with writer_engine.connect() as connection: connection.execute( text( "SELECT metadata, 1 - (embedding <=> '[1,0]'::vector) AS similarity " "FROM vectors.memory ORDER BY embedding <=> '[1,0]'::vector LIMIT 1" ) ) reader_engine.dispose() writer_engine.dispose() def test_pgvector_writer_health_requires_sequence_usage(vector_configs): from tht.adapters.vector.pgvector import PgVectorStore admin_config, _, _, no_sequence_config = vector_configs store = PgVectorStore(None, no_sequence_config, expected_dimension=2) health = store.health() assert health.ok is False assert health.write_reachable is False assert health.write_detail == ( "vector schema incomplete: missing sequence privileges evidence, memory, schema_records" ) with pytest.raises(VectorWriteUnavailable) as error: store.upsert("memory", [_record("needs-sequence", [1.0, 0.0])]) assert isinstance(error.value.__cause__, InsufficientPrivilege) admin_engine = create_engine( f"postgresql+psycopg2://{admin_config.user}:{admin_config.password}" f"@{admin_config.host}:{admin_config.port}/{admin_config.database}" ) with admin_engine.begin() as connection: connection.exec_driver_sql( "GRANT USAGE ON ALL SEQUENCES IN SCHEMA vectors TO vector_l0_no_sequence" ) admin_engine.dispose() assert store.health().ok is True assert store.upsert("memory", [_record("has-sequence", [1.0, 0.0])]) == 1 def test_pgvector_health_requires_schema_usage_for_reader_and_writer(vector_configs): from tht.adapters.vector.pgvector import PgVectorStore admin_config, reader_config, writer_config, _ = vector_configs admin_engine = create_engine( f"postgresql+psycopg2://{admin_config.user}:{admin_config.password}" f"@{admin_config.host}:{admin_config.port}/{admin_config.database}" ) store = PgVectorStore(reader_config, writer_config, expected_dimension=2) with admin_engine.begin() as connection: connection.exec_driver_sql( f"REVOKE USAGE ON SCHEMA vectors FROM {reader_config.user}, {writer_config.user}" ) health = store.health() assert health.read_reachable is False and health.write_reachable is False assert "missing schema usage" in health.read_detail assert "missing schema usage" in health.write_detail with pytest.raises(VectorReadUnavailable, match="Vector read operation unavailable"): store.search(["memory"], [1.0, 0.0], limit=1) with pytest.raises(VectorWriteUnavailable, match="Vector write operation unavailable"): store.upsert("memory", [_record("blocked", [1.0, 0.0])]) with admin_engine.begin() as connection: connection.exec_driver_sql( f"GRANT USAGE ON SCHEMA vectors TO {reader_config.user}, {writer_config.user}" ) admin_engine.dispose() assert store.health().ok is True def test_pgvector_maps_unavailable_connections_without_leaking_password(vector_configs): from tht.adapters.vector.pgvector import PgVectorStore _, reader_config, writer_config, _ = vector_configs password = "never-leak-this" reader = reader_config.model_copy(update={"port": 1, "password": password}) writer = writer_config.model_copy(update={"port": 1, "password": password}) with pytest.raises(VectorReadUnavailable) as read_error: PgVectorStore(reader, None).search(["memory"], [1.0, 0.0], limit=1) with pytest.raises(VectorWriteUnavailable) as hash_error: PgVectorStore(None, writer).existing_hashes("memory", ["memory"]) with pytest.raises(VectorWriteUnavailable) as write_error: PgVectorStore(None, writer).upsert("memory", [_record("x", [1.0, 0.0])]) assert password not in str(read_error.value) assert password not in str(hash_error.value) assert password not in str(write_error.value) def test_pgvector_health_reports_dimension_and_each_connection(vector_configs): from tht.adapters.vector.pgvector import PgVectorStore _, reader_config, writer_config, _ = vector_configs health = PgVectorStore(reader_config, writer_config, expected_dimension=2).health() assert health.ok is True assert health.read_reachable is True assert health.write_reachable is True assert health.observed_dimensions == (2,) assert health.dimension_compatible is True mismatch = PgVectorStore(reader_config, None, expected_dimension=3).health() assert mismatch.ok is False assert mismatch.read_reachable is False assert mismatch.read_detail == ( "embedding dimension mismatch: evidence=2, memory=2, schema_records=2" ) assert mismatch.dimension_compatible is False def test_pgvector_health_rejects_clean_and_partial_schemas(vector_configs): from tht.adapters.vector.pgvector import PgVectorStore admin_config, _, _, _ = vector_configs engine = create_engine( f"postgresql+psycopg2://{admin_config.user}:{admin_config.password}" f"@{admin_config.host}:{admin_config.port}/{admin_config.database}" ) with engine.begin() as connection: connection.exec_driver_sql("CREATE SCHEMA clean_vectors") connection.exec_driver_sql("CREATE SCHEMA partial_vectors") connection.exec_driver_sql( "CREATE TABLE partial_vectors.memory " "(record_key text, kind text, content_hash text, metadata jsonb)" ) engine.dispose() clean = PgVectorStore( admin_config.model_copy(update={"db_schema": "clean_vectors"}), expected_dimension=2, ).health() assert clean.ok is False assert clean.read_reachable is False assert clean.read_detail == ( "vector schema incomplete: missing tables evidence, memory, schema_records" ) partial = PgVectorStore( admin_config.model_copy(update={"db_schema": "partial_vectors"}), expected_dimension=2, ).health() assert partial.ok is False assert partial.read_reachable is False assert partial.read_detail == ( "vector schema incomplete: missing tables evidence, schema_records; " "missing embedding columns memory" )