""" Unit tests for PGKVStorage.upsert batch optimization (PR #2742 fixes). Verifies: 1. Each namespace builds correct tuple ordering matching SQL positional params. 2. _run_with_retry is used (not the removed PostgreSQLDB.executemany wrapper). 3. Sub-batching splits data when len(data) > _max_batch_size. 4. Unknown namespace raises ValueError. 5. Empty data returns without any DB call. """ import json import pytest import numpy as np from unittest.mock import AsyncMock, MagicMock from lightrag.kg.postgres_impl import PGDocStatusStorage, PGKVStorage, PGVectorStorage from lightrag.namespace import NameSpace from lightrag.utils import EmbeddingFunc # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- GLOBAL_CONFIG = {"embedding_batch_num": 10} def make_storage(namespace: str) -> PGKVStorage: """Construct a PGKVStorage instance with a mocked db.""" db = MagicMock() captured: list[tuple] = [] retry_kwargs: list[dict] = [] async def fake_run_with_retry(operation, **kwargs): """Call the closure with a mock connection to capture executemany args.""" retry_kwargs.append(kwargs) mock_conn = AsyncMock() await operation(mock_conn) # Store (sql, data) from each executemany call for call in mock_conn.executemany.call_args_list: captured.append((call.args[0], call.args[1])) db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry) db.workspace = "test_ws" storage = PGKVStorage.__new__(PGKVStorage) storage.namespace = namespace storage.workspace = "test_ws" storage.global_config = GLOBAL_CONFIG storage.db = db storage.__post_init__() storage._captured = captured storage._retry_kwargs = retry_kwargs return storage def make_doc_status_storage() -> PGDocStatusStorage: """Construct a PGDocStatusStorage instance with a mocked db.""" db = MagicMock() captured: list[tuple] = [] retry_kwargs: list[dict] = [] async def fake_run_with_retry(operation, **kwargs): retry_kwargs.append(kwargs) mock_conn = AsyncMock() tx = AsyncMock() tx.__aenter__.return_value = tx tx.__aexit__.return_value = False mock_conn.transaction = MagicMock(return_value=tx) await operation(mock_conn) for call in mock_conn.executemany.call_args_list: captured.append((call.args[0], call.args[1])) db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry) db.workspace = "test_ws" storage = PGDocStatusStorage.__new__(PGDocStatusStorage) storage.namespace = NameSpace.DOC_STATUS storage.workspace = "test_ws" storage.global_config = GLOBAL_CONFIG storage.db = db storage._captured = captured storage._retry_kwargs = retry_kwargs return storage def make_vector_storage(namespace: str) -> PGVectorStorage: """Construct a PGVectorStorage instance with a mocked db and embedding func.""" db = MagicMock() captured: list[tuple] = [] retry_kwargs: list[dict] = [] async def fake_run_with_retry(operation, **kwargs): retry_kwargs.append(kwargs) mock_conn = AsyncMock() await operation(mock_conn) for call in mock_conn.executemany.call_args_list: captured.append((call.args[0], call.args[1])) db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry) db.workspace = "test_ws" async def embed_func(texts, **kwargs): return np.array([[0.1, 0.2, 0.3] for _ in texts], dtype=np.float32) embedding = EmbeddingFunc( embedding_dim=3, func=embed_func, model_name="test_model", ) storage = PGVectorStorage( namespace=namespace, workspace="test_ws", global_config={ "embedding_batch_num": 10, "vector_db_storage_cls_kwargs": { "cosine_better_than_threshold": 0.5, }, }, embedding_func=embedding, ) storage.db = db storage._captured = captured storage._retry_kwargs = retry_kwargs return storage # --------------------------------------------------------------------------- # 1. _max_batch_size is always 200 (not embedding_batch_num) # --------------------------------------------------------------------------- def test_max_batch_size_is_constant(): storage = make_storage(NameSpace.KV_STORE_TEXT_CHUNKS) assert storage._max_batch_size == 200 # --------------------------------------------------------------------------- # 2. Namespace: TEXT_CHUNKS # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_upsert_text_chunks_tuple_order(): storage = make_storage(NameSpace.KV_STORE_TEXT_CHUNKS) data = { "chunk-1": { "tokens": 42, "chunk_order_index": 0, "full_doc_id": "doc-1", "content": "hello world", "file_path": "/a/b.txt", "llm_cache_list": ["cache-key"], } } await storage.upsert(data) assert len(storage._captured) == 1 sql, rows = storage._captured[0] assert "LIGHTRAG_DOC_CHUNKS" in sql assert len(rows) == 1 row = rows[0] # SQL: (workspace, id, tokens, chunk_order_index, full_doc_id, # content, file_path, llm_cache_list, create_time, update_time) assert row[0] == "test_ws" # workspace assert row[1] == "chunk-1" # id assert row[2] == 42 # tokens assert row[3] == 0 # chunk_order_index assert row[4] == "doc-1" # full_doc_id assert row[5] == "hello world" # content assert row[6] == "/a/b.txt" # file_path assert json.loads(row[7]) == ["cache-key"] # llm_cache_list # --------------------------------------------------------------------------- # 3. Namespace: FULL_DOCS # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_upsert_full_docs_tuple_order(): storage = make_storage(NameSpace.KV_STORE_FULL_DOCS) data = {"doc-1": {"content": "full text", "file_path": "/path/doc.pdf"}} await storage.upsert(data) assert len(storage._captured) == 1 _, rows = storage._captured[0] row = rows[0] # SQL: (id, content, doc_name, workspace) assert row[0] == "doc-1" assert row[1] == "full text" assert row[2] == "/path/doc.pdf" assert row[3] == "test_ws" # --------------------------------------------------------------------------- # 4. Namespace: LLM_RESPONSE_CACHE # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_upsert_llm_cache_tuple_order(): storage = make_storage(NameSpace.KV_STORE_LLM_RESPONSE_CACHE) data = { "key-1": { "original_prompt": "what is X?", "return": "X is Y", "chunk_id": "chunk-1", "cache_type": "query", "queryparam": {"mode": "hybrid"}, } } await storage.upsert(data) assert len(storage._captured) == 1 _, rows = storage._captured[0] row = rows[0] # SQL: (workspace, id, original_prompt, return_value, chunk_id, cache_type, queryparam) assert row[0] == "test_ws" assert row[1] == "key-1" assert row[2] == "what is X?" assert row[3] == "X is Y" assert row[4] == "chunk-1" assert row[5] == "query" assert json.loads(row[6]) == {"mode": "hybrid"} @pytest.mark.asyncio async def test_upsert_llm_cache_null_queryparam(): storage = make_storage(NameSpace.KV_STORE_LLM_RESPONSE_CACHE) data = { "key-2": { "original_prompt": "prompt", "return": "answer", "cache_type": "extract", } } await storage.upsert(data) _, rows = storage._captured[0] assert rows[0][6] is None # queryparam should be None # --------------------------------------------------------------------------- # 5. Namespace: FULL_ENTITIES # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_upsert_full_entities_tuple_order(): storage = make_storage(NameSpace.KV_STORE_FULL_ENTITIES) data = {"ent-1": {"entity_names": ["EntityA", "EntityB"], "count": 2}} await storage.upsert(data) _, rows = storage._captured[0] row = rows[0] # SQL: (workspace, id, entity_names, count, create_time, update_time) assert row[0] == "test_ws" assert row[1] == "ent-1" assert json.loads(row[2]) == ["EntityA", "EntityB"] assert row[3] == 2 # --------------------------------------------------------------------------- # 6. Namespace: FULL_RELATIONS # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_upsert_full_relations_tuple_order(): storage = make_storage(NameSpace.KV_STORE_FULL_RELATIONS) data = {"rel-1": {"relation_pairs": [["A", "B"]], "count": 1}} await storage.upsert(data) _, rows = storage._captured[0] row = rows[0] # SQL: (workspace, id, relation_pairs, count, create_time, update_time) assert row[0] == "test_ws" assert row[1] == "rel-1" assert json.loads(row[2]) == [["A", "B"]] assert row[3] == 1 # --------------------------------------------------------------------------- # 7. Namespace: ENTITY_CHUNKS / RELATION_CHUNKS # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_upsert_entity_chunks_tuple_order(): storage = make_storage(NameSpace.KV_STORE_ENTITY_CHUNKS) data = {"ec-1": {"chunk_ids": ["c1", "c2"], "count": 2}} await storage.upsert(data) _, rows = storage._captured[0] row = rows[0] # SQL: (workspace, id, chunk_ids, count, create_time, update_time) assert row[0] == "test_ws" assert row[1] == "ec-1" assert json.loads(row[2]) == ["c1", "c2"] assert row[3] == 2 @pytest.mark.asyncio async def test_upsert_relation_chunks_tuple_order(): storage = make_storage(NameSpace.KV_STORE_RELATION_CHUNKS) data = {"rc-1": {"chunk_ids": ["c3"], "count": 1}} await storage.upsert(data) _, rows = storage._captured[0] row = rows[0] assert row[0] == "test_ws" assert row[1] == "rc-1" assert json.loads(row[2]) == ["c3"] assert row[3] == 1 # --------------------------------------------------------------------------- # 8. Sub-batching: data > _max_batch_size splits into multiple _run_with_retry calls # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_sub_batching_splits_correctly(): storage = make_storage(NameSpace.KV_STORE_FULL_DOCS) storage._max_batch_size = 3 # Override to small value for testing data = {f"doc-{i}": {"content": f"text {i}", "file_path": ""} for i in range(7)} await storage.upsert(data) # 7 records / batch_size 3 => 3 batches (3 + 3 + 1) assert len(storage._captured) == 3 assert len(storage._captured[0][1]) == 3 assert len(storage._captured[1][1]) == 3 assert len(storage._captured[2][1]) == 1 @pytest.mark.asyncio async def test_sub_batching_exact_multiple(): storage = make_storage(NameSpace.KV_STORE_FULL_DOCS) storage._max_batch_size = 3 data = {f"doc-{i}": {"content": f"text {i}", "file_path": ""} for i in range(6)} await storage.upsert(data) # 6 / 3 => exactly 2 batches assert len(storage._captured) == 2 assert len(storage._captured[0][1]) == 3 assert len(storage._captured[1][1]) == 3 # --------------------------------------------------------------------------- # 9. Empty data: no DB call # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_upsert_empty_data_no_db_call(): storage = make_storage(NameSpace.KV_STORE_FULL_DOCS) await storage.upsert({}) assert len(storage._captured) == 0 storage.db._run_with_retry.assert_not_called() # --------------------------------------------------------------------------- # 10. Unknown namespace raises ValueError # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_upsert_unknown_namespace_raises(): storage = make_storage("unknown_namespace") with pytest.raises(ValueError, match="Unknown namespace"): await storage.upsert({"k": {"v": 1}}) # --------------------------------------------------------------------------- # 11. Multiple records go into one batch when within limit # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_multiple_records_single_batch(): storage = make_storage(NameSpace.KV_STORE_FULL_DOCS) data = { "doc-1": {"content": "text 1", "file_path": "/a"}, "doc-2": {"content": "text 2", "file_path": "/b"}, "doc-3": {"content": "text 3", "file_path": "/c"}, } await storage.upsert(data) # All 3 fit within default batch size of 200 assert len(storage._captured) == 1 _, rows = storage._captured[0] assert len(rows) == 3 ids = {row[0] for row in rows} # id is $1 for FULL_DOCS assert ids == {"doc-1", "doc-2", "doc-3"} @pytest.mark.asyncio async def test_kv_upsert_passes_timing_label(): storage = make_storage(NameSpace.KV_STORE_FULL_DOCS) await storage.upsert({"doc-1": {"content": "text 1", "file_path": "/a"}}) assert storage._retry_kwargs[0]["timing_label"] == ( f"test_ws PGKVStorage.upsert[{NameSpace.KV_STORE_FULL_DOCS}]" ) @pytest.mark.asyncio async def test_doc_status_upsert_passes_timing_label(): storage = make_doc_status_storage() await storage.upsert( { "doc-1": { "content_summary": "summary", "content_length": 12, "chunks_count": 1, "status": "processed", "file_path": "/a.txt", "chunks_list": ["chunk-1"], "metadata": {"source": "test"}, "created_at": "2024-01-01T00:00:00+00:00", "updated_at": "2024-01-01T00:00:00+00:00", } } ) assert storage._retry_kwargs[0]["timing_label"] == ( "test_ws PGDocStatusStorage.upsert" ) @pytest.mark.asyncio async def test_vector_upsert_passes_timing_label(): storage = make_vector_storage(NameSpace.VECTOR_STORE_CHUNKS) await storage.upsert( { "chunk-1": { "tokens": 42, "chunk_order_index": 0, "full_doc_id": "doc-1", "content": "hello world", "file_path": "/a/b.txt", } } ) assert storage._retry_kwargs[0]["timing_label"] == ( f"test_ws PGVectorStorage.upsert[{NameSpace.VECTOR_STORE_CHUNKS}]" )