chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,91 @@
|
||||
"""LIGHTRAG_DOC_FULL.parse_engine must be TEXT (Phase 2).
|
||||
|
||||
The parse_engine field may carry an encoded engine-parameter directive
|
||||
(``mineru(page_range=1-3,language=en)``) longer than the original VARCHAR(32),
|
||||
so the CREATE DDL uses TEXT and the migration widens an existing VARCHAR(32)
|
||||
column. (A real overflow round-trip against live Postgres is the separate
|
||||
@pytest.mark.integration test; these mocked assertions cannot run real SQL.)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from lightrag.kg.postgres_impl import TABLES, PostgreSQLDB
|
||||
|
||||
|
||||
def test_create_ddl_uses_text_for_parse_engine():
|
||||
ddl = TABLES["LIGHTRAG_DOC_FULL"]["ddl"]
|
||||
assert "parse_engine TEXT" in ddl
|
||||
assert "parse_engine VARCHAR(32)" not in ddl
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_migration_widens_existing_varchar_column_to_text():
|
||||
db = PostgreSQLDB.__new__(PostgreSQLDB) # skip __init__/connection setup
|
||||
|
||||
executed: list[str] = []
|
||||
|
||||
async def fake_query(sql, *args, **kwargs):
|
||||
# All columns already present -> ADD COLUMN loop is a no-op; the
|
||||
# parse_engine column reports the legacy VARCHAR type so the widening
|
||||
# ALTER fires.
|
||||
if "information_schema.columns" in sql and "column_name = ANY" in sql:
|
||||
return [
|
||||
{"column_name": c}
|
||||
for c in (
|
||||
"sidecar_location",
|
||||
"parse_format",
|
||||
"content_hash",
|
||||
"process_options",
|
||||
"chunk_options",
|
||||
"parse_engine",
|
||||
)
|
||||
]
|
||||
if "information_schema.columns" in sql: # the parse_engine type probe
|
||||
return {"data_type": "character varying"}
|
||||
return None
|
||||
|
||||
async def fake_execute(sql, *args, **kwargs):
|
||||
executed.append(sql)
|
||||
|
||||
db.query = AsyncMock(side_effect=fake_query)
|
||||
db.execute = AsyncMock(side_effect=fake_execute)
|
||||
|
||||
await db._migrate_doc_full_add_pipeline_fields()
|
||||
|
||||
assert any("ALTER COLUMN parse_engine TYPE TEXT" in sql for sql in executed), (
|
||||
executed
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_migration_skips_alter_when_already_text():
|
||||
db = PostgreSQLDB.__new__(PostgreSQLDB)
|
||||
executed: list[str] = []
|
||||
|
||||
async def fake_query(sql, *args, **kwargs):
|
||||
if "information_schema.columns" in sql and "column_name = ANY" in sql:
|
||||
return [
|
||||
{"column_name": c}
|
||||
for c in (
|
||||
"sidecar_location",
|
||||
"parse_format",
|
||||
"content_hash",
|
||||
"process_options",
|
||||
"chunk_options",
|
||||
"parse_engine",
|
||||
)
|
||||
]
|
||||
if "information_schema.columns" in sql:
|
||||
return {"data_type": "text"}
|
||||
return None
|
||||
|
||||
db.query = AsyncMock(side_effect=fake_query)
|
||||
db.execute = AsyncMock(side_effect=lambda sql, *a, **k: executed.append(sql))
|
||||
|
||||
await db._migrate_doc_full_add_pipeline_fields()
|
||||
|
||||
assert not any("ALTER COLUMN parse_engine" in sql for sql in executed)
|
||||
@@ -0,0 +1,252 @@
|
||||
"""
|
||||
Unit tests for PGGraphStorage.get_nodes_edges_batch and get_node_edges
|
||||
with special characters in entity names.
|
||||
|
||||
Verifies the fix for KeyError when entity names contain double quotes (PR #2872)
|
||||
and the follow-up Option C refactor to parameterized Cypher queries.
|
||||
|
||||
The root cause: AGE returns the original un-escaped entity_id, but the edges_norm
|
||||
dict was previously keyed with the normalized (escaped) ID, causing a KeyError on lookup.
|
||||
|
||||
The Option C fix: use $node_ids / $entity_id parameters instead of string interpolation,
|
||||
eliminating the need for _normalize_node_id in these read paths entirely.
|
||||
"""
|
||||
|
||||
import json
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from lightrag.kg.postgres_impl import PGGraphStorage
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def make_graph_storage() -> PGGraphStorage:
|
||||
"""Construct a PGGraphStorage instance with a mocked _query method."""
|
||||
storage = PGGraphStorage.__new__(PGGraphStorage)
|
||||
storage.workspace = "test_ws"
|
||||
storage.namespace = "test_graph"
|
||||
storage.graph_name = "test_graph"
|
||||
storage.db = MagicMock()
|
||||
return storage
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _normalize_node_id (still used by write paths: remove_nodes, upsert_node, etc.)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_normalize_plain_id():
|
||||
assert PGGraphStorage._normalize_node_id("Alice") == "Alice"
|
||||
|
||||
|
||||
def test_normalize_double_quote():
|
||||
assert PGGraphStorage._normalize_node_id('John "Smith"') == 'John \\"Smith\\"'
|
||||
|
||||
|
||||
def test_normalize_backslash():
|
||||
assert PGGraphStorage._normalize_node_id("C:\\path") == "C:\\\\path"
|
||||
|
||||
|
||||
def test_normalize_both_special_chars():
|
||||
assert (
|
||||
PGGraphStorage._normalize_node_id('say \\"hello\\"')
|
||||
== 'say \\\\\\"hello\\\\\\"'
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_node_edges — parameterized query (Option C)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_node_edges_passes_original_id_as_parameter():
|
||||
"""entity_id must be passed as a JSON parameter, not interpolated into Cypher."""
|
||||
storage = make_graph_storage()
|
||||
entity = 'John "Smith"'
|
||||
captured_params: list[dict] = []
|
||||
|
||||
async def fake_query(sql, **kwargs):
|
||||
if kwargs.get("params"):
|
||||
captured_params.append(json.loads(list(kwargs["params"].values())[0]))
|
||||
return []
|
||||
|
||||
with patch.object(storage, "_query", side_effect=fake_query):
|
||||
await storage.get_node_edges(entity)
|
||||
|
||||
assert len(captured_params) == 1
|
||||
assert captured_params[0]["entity_id"] == entity
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_node_edges_cypher_uses_parameter_syntax():
|
||||
"""The SQL sent to _query must use $1::agtype, not a hardcoded escaped string."""
|
||||
storage = make_graph_storage()
|
||||
entity = 'John "Smith"'
|
||||
captured_sql: list[str] = []
|
||||
|
||||
async def fake_query(sql, **kwargs):
|
||||
captured_sql.append(sql)
|
||||
return []
|
||||
|
||||
with patch.object(storage, "_query", side_effect=fake_query):
|
||||
await storage.get_node_edges(entity)
|
||||
|
||||
assert len(captured_sql) == 1
|
||||
assert "$1::agtype" in captured_sql[0]
|
||||
# Entity name must NOT appear literally in the SQL string
|
||||
assert entity not in captured_sql[0]
|
||||
assert '\\"' not in captured_sql[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_node_edges_returns_edges():
|
||||
storage = make_graph_storage()
|
||||
|
||||
async def fake_query(_sql, **_kwargs):
|
||||
return [
|
||||
{"source_id": "Alice", "connected_id": "Bob"},
|
||||
{"source_id": "Alice", "connected_id": None},
|
||||
]
|
||||
|
||||
with patch.object(storage, "_query", side_effect=fake_query):
|
||||
result = await storage.get_node_edges("Alice")
|
||||
|
||||
assert result == [("Alice", "Bob")]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_nodes_edges_batch — parameterized query (Option C)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_nodes_edges_batch_passes_original_ids_as_parameter():
|
||||
"""node_ids batch must be passed as a JSON parameter, not interpolated."""
|
||||
storage = make_graph_storage()
|
||||
entities = ['John "Smith"', "Alice", "O\\Brien"]
|
||||
captured_params: list[dict] = []
|
||||
|
||||
async def fake_query(_sql, **kwargs):
|
||||
if kwargs.get("params"):
|
||||
captured_params.append(json.loads(list(kwargs["params"].values())[0]))
|
||||
return []
|
||||
|
||||
with patch.object(storage, "_query", side_effect=fake_query):
|
||||
await storage.get_nodes_edges_batch(entities)
|
||||
|
||||
assert len(captured_params) == 2 # outgoing + incoming
|
||||
assert captured_params[0]["node_ids"] == entities
|
||||
assert captured_params[1]["node_ids"] == entities
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_nodes_edges_batch_cypher_uses_parameter_syntax():
|
||||
"""The SQL must use $1::agtype, not hardcoded escaped entity names."""
|
||||
storage = make_graph_storage()
|
||||
entity = 'John "Smith"'
|
||||
captured_sql: list[str] = []
|
||||
|
||||
async def fake_query(sql, **_kwargs):
|
||||
captured_sql.append(sql)
|
||||
return []
|
||||
|
||||
with patch.object(storage, "_query", side_effect=fake_query):
|
||||
await storage.get_nodes_edges_batch([entity])
|
||||
|
||||
assert len(captured_sql) == 2
|
||||
for sql in captured_sql:
|
||||
assert "$1::agtype" in sql
|
||||
assert entity not in sql
|
||||
assert '\\"' not in sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_nodes_edges_batch_with_quoted_entity():
|
||||
"""
|
||||
AGE returns the original un-escaped node_id in query results.
|
||||
The result dict must be keyed by the original ID.
|
||||
"""
|
||||
storage = make_graph_storage()
|
||||
entity = 'John "Smith"'
|
||||
|
||||
async def fake_query(sql, **_kwargs):
|
||||
if "OPTIONAL MATCH (n:base)-[]->" in sql:
|
||||
return [{"node_id": entity, "connected_id": "Alice"}]
|
||||
if "OPTIONAL MATCH (n:base)<-[]-" in sql:
|
||||
return [{"node_id": entity, "connected_id": "Bob"}]
|
||||
return []
|
||||
|
||||
with patch.object(storage, "_query", side_effect=fake_query):
|
||||
result = await storage.get_nodes_edges_batch([entity])
|
||||
|
||||
assert entity in result
|
||||
assert (entity, "Alice") in result[entity]
|
||||
assert ("Bob", entity) in result[entity]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_nodes_edges_batch_plain_entity():
|
||||
"""Entity names without special chars still work correctly."""
|
||||
storage = make_graph_storage()
|
||||
entity = "Alice"
|
||||
|
||||
async def fake_query(sql, **_kwargs):
|
||||
if "OPTIONAL MATCH (n:base)-[]->" in sql:
|
||||
return [{"node_id": entity, "connected_id": "Bob"}]
|
||||
return []
|
||||
|
||||
with patch.object(storage, "_query", side_effect=fake_query):
|
||||
result = await storage.get_nodes_edges_batch([entity])
|
||||
|
||||
assert entity in result
|
||||
assert (entity, "Bob") in result[entity]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_nodes_edges_batch_no_results():
|
||||
"""Nodes with no edges return an empty list, not a KeyError."""
|
||||
storage = make_graph_storage()
|
||||
entity = 'Entity "X"'
|
||||
|
||||
async def fake_query(_sql, **_kwargs):
|
||||
return []
|
||||
|
||||
with patch.object(storage, "_query", side_effect=fake_query):
|
||||
result = await storage.get_nodes_edges_batch([entity])
|
||||
|
||||
assert entity in result
|
||||
assert result[entity] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_nodes_edges_batch_deduplication():
|
||||
"""Duplicate input IDs are deduplicated; each maps to the same edge list."""
|
||||
storage = make_graph_storage()
|
||||
entity = 'Dup "Entity"'
|
||||
|
||||
async def fake_query(sql, **_kwargs):
|
||||
if "OPTIONAL MATCH (n:base)-[]->" in sql:
|
||||
return [{"node_id": entity, "connected_id": "Other"}]
|
||||
return []
|
||||
|
||||
with patch.object(storage, "_query", side_effect=fake_query):
|
||||
result = await storage.get_nodes_edges_batch([entity, entity])
|
||||
|
||||
assert result[entity] == [(entity, "Other")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_nodes_edges_batch_empty_input():
|
||||
"""Empty input returns empty dict without calling _query."""
|
||||
storage = make_graph_storage()
|
||||
|
||||
with patch.object(storage, "_query") as mock_q:
|
||||
result = await storage.get_nodes_edges_batch([])
|
||||
|
||||
assert result == {}
|
||||
mock_q.assert_not_called()
|
||||
@@ -0,0 +1,160 @@
|
||||
"""Unit tests for the PostgreSQL batch-limit helpers.
|
||||
|
||||
Covers the module-level batching primitives shared by the PG non-graph write
|
||||
paths (mirrors the MONGO_* / OPENSEARCH_* contract):
|
||||
|
||||
- ``_resolve_pg_batch_limits``: env parsing, defaults, non-positive disable.
|
||||
- ``_estimate_record_bytes``: per-field byte estimation for asyncpg tuples.
|
||||
- ``_chunk_by_budget``: record-count + payload-byte splitting.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
|
||||
from lightrag.kg.postgres_impl import (
|
||||
DEFAULT_PG_DELETE_MAX_RECORDS_PER_BATCH,
|
||||
DEFAULT_PG_UPSERT_MAX_PAYLOAD_BYTES,
|
||||
DEFAULT_PG_UPSERT_MAX_RECORDS_PER_BATCH,
|
||||
_chunk_by_budget,
|
||||
_estimate_record_bytes,
|
||||
_resolve_pg_batch_limits,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _resolve_pg_batch_limits
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestResolvePgBatchLimits:
|
||||
def test_defaults(self):
|
||||
with patch.dict("os.environ", {}, clear=True):
|
||||
payload, upserts, deletes = _resolve_pg_batch_limits()
|
||||
assert payload == DEFAULT_PG_UPSERT_MAX_PAYLOAD_BYTES
|
||||
assert upserts == DEFAULT_PG_UPSERT_MAX_RECORDS_PER_BATCH
|
||||
assert deletes == DEFAULT_PG_DELETE_MAX_RECORDS_PER_BATCH
|
||||
|
||||
def test_default_record_cap_keeps_kv_historical_200(self):
|
||||
# The PG-wide upsert record default deliberately stays at 200 (KV's
|
||||
# historical batch size), not Mongo/OS's 128.
|
||||
assert DEFAULT_PG_UPSERT_MAX_RECORDS_PER_BATCH == 200
|
||||
|
||||
def test_env_override(self):
|
||||
env = {
|
||||
"POSTGRES_UPSERT_MAX_PAYLOAD_BYTES": "12345",
|
||||
"POSTGRES_UPSERT_MAX_RECORDS_PER_BATCH": "7",
|
||||
"POSTGRES_DELETE_MAX_RECORDS_PER_BATCH": "9",
|
||||
}
|
||||
with patch.dict("os.environ", env, clear=True):
|
||||
assert _resolve_pg_batch_limits() == (12345, 7, 9)
|
||||
|
||||
def test_non_positive_disables_and_warns(self):
|
||||
env = {
|
||||
"POSTGRES_UPSERT_MAX_PAYLOAD_BYTES": "0",
|
||||
"POSTGRES_UPSERT_MAX_RECORDS_PER_BATCH": "-1",
|
||||
"POSTGRES_DELETE_MAX_RECORDS_PER_BATCH": "0",
|
||||
}
|
||||
with patch.dict("os.environ", env, clear=True):
|
||||
with patch("lightrag.kg.postgres_impl.logger") as mock_logger:
|
||||
payload, upserts, deletes = _resolve_pg_batch_limits()
|
||||
assert (payload, upserts, deletes) == (0, -1, 0)
|
||||
warnings = [c.args[0] for c in mock_logger.warning.call_args_list]
|
||||
assert sum("non-positive" in msg for msg in warnings) == 3
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _estimate_record_bytes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEstimateRecordBytes:
|
||||
def test_ndarray_uses_nbytes(self):
|
||||
vec = np.zeros(128, dtype=np.float32)
|
||||
# 128 * 4 bytes (nbytes) plus the 2-byte "id" str.
|
||||
assert _estimate_record_bytes(("id", vec)) == vec.nbytes + 2
|
||||
|
||||
def test_str_uses_utf8_length(self):
|
||||
assert _estimate_record_bytes(("abc",)) == 3
|
||||
|
||||
def test_multibyte_str(self):
|
||||
# Each CJK char is 3 bytes in UTF-8.
|
||||
assert _estimate_record_bytes(("中文",)) == 6
|
||||
|
||||
def test_bytes_and_none(self):
|
||||
assert _estimate_record_bytes((b"abcd", None)) == 4
|
||||
|
||||
def test_dict_and_list_use_compact_json(self):
|
||||
payload = {"a": 1, "b": [1, 2, 3]}
|
||||
expected = len(
|
||||
json.dumps(
|
||||
payload, ensure_ascii=False, separators=(",", ":"), default=str
|
||||
).encode("utf-8")
|
||||
)
|
||||
assert _estimate_record_bytes((payload,)) == expected
|
||||
|
||||
def test_scalar_constant(self):
|
||||
# int + float + None -> 16 + 16 + 0
|
||||
assert _estimate_record_bytes((1, 2.0, None)) == 32
|
||||
|
||||
def test_mixed_tuple(self):
|
||||
vec = np.ones(4, dtype=np.float32) # 16 bytes
|
||||
record = ("ws", "id", "hello", vec, None, 42)
|
||||
# 2 + 2 + 5 + 16 + 0 + 16
|
||||
assert _estimate_record_bytes(record) == 41
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _chunk_by_budget
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestChunkByBudget:
|
||||
def test_empty(self):
|
||||
assert _chunk_by_budget([], lambda x: 1, 100, 100) == []
|
||||
|
||||
def test_split_by_record_count(self):
|
||||
items = list(range(7))
|
||||
batches = _chunk_by_budget(
|
||||
items, lambda x: 1, max_payload_bytes=0, max_records_per_batch=3
|
||||
)
|
||||
sizes = [len(b) for b, _ in batches]
|
||||
assert sizes == [3, 3, 1]
|
||||
|
||||
def test_split_by_payload_bytes(self):
|
||||
# Each item ~10 bytes; a 25-byte budget fits 2 per batch (2 overhead +
|
||||
# 10 + 1 + 10 = 23 <= 25; adding a third would exceed).
|
||||
items = ["x" * 10, "y" * 10, "z" * 10, "w" * 10]
|
||||
batches = _chunk_by_budget(
|
||||
items, lambda s: len(s), max_payload_bytes=25, max_records_per_batch=0
|
||||
)
|
||||
sizes = [len(b) for b, _ in batches]
|
||||
assert sizes == [2, 2]
|
||||
|
||||
def test_both_dimensions_record_cap_wins(self):
|
||||
items = ["x" * 5] * 10
|
||||
# record cap 4 binds before the generous byte budget.
|
||||
batches = _chunk_by_budget(
|
||||
items, lambda s: len(s), max_payload_bytes=10_000, max_records_per_batch=4
|
||||
)
|
||||
sizes = [len(b) for b, _ in batches]
|
||||
assert sizes == [4, 4, 2]
|
||||
|
||||
def test_oversized_single_item_becomes_own_batch(self):
|
||||
items = ["small", "X" * 1000, "small2"]
|
||||
batches = _chunk_by_budget(
|
||||
items, lambda s: len(s), max_payload_bytes=50, max_records_per_batch=0
|
||||
)
|
||||
# The 1000-byte item is emitted alone rather than raising.
|
||||
sizes = [len(b) for b, _ in batches]
|
||||
assert sizes == [1, 1, 1]
|
||||
assert batches[1][0] == ["X" * 1000]
|
||||
|
||||
def test_non_positive_disables_both_dimensions(self):
|
||||
items = list(range(100))
|
||||
batches = _chunk_by_budget(
|
||||
items, lambda x: 999, max_payload_bytes=0, max_records_per_batch=0
|
||||
)
|
||||
assert len(batches) == 1
|
||||
assert len(batches[0][0]) == 100
|
||||
@@ -0,0 +1,73 @@
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from lightrag.kg.postgres_impl import ClientManager
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_client_manager_state() -> None:
|
||||
ClientManager._instances = {"db": None, "ref_count": 0, "vector_signature": None}
|
||||
|
||||
|
||||
def test_pg_vector_storage_enables_vector() -> None:
|
||||
config = ClientManager.get_config("PGVectorStorage")
|
||||
assert config["enable_vector"] is True
|
||||
|
||||
|
||||
def test_non_pg_vector_storage_disables_vector() -> None:
|
||||
config = ClientManager.get_config("NanoVectorDBStorage")
|
||||
assert config["enable_vector"] is False
|
||||
|
||||
|
||||
def test_milvus_storage_disables_vector() -> None:
|
||||
config = ClientManager.get_config("MilvusVectorDBStorage")
|
||||
assert config["enable_vector"] is False
|
||||
|
||||
|
||||
def test_qdrant_storage_disables_vector() -> None:
|
||||
config = ClientManager.get_config("QdrantVectorDBStorage")
|
||||
assert config["enable_vector"] is False
|
||||
|
||||
|
||||
def test_none_vector_storage_defaults_to_true() -> None:
|
||||
# Backward compatibility: when vector_storage is unknown (None), default to True.
|
||||
config = ClientManager.get_config(None)
|
||||
assert config["enable_vector"] is True
|
||||
|
||||
|
||||
def test_no_args_defaults_to_true() -> None:
|
||||
# Backward compatibility: calling without arguments preserves prior behavior.
|
||||
config = ClientManager.get_config()
|
||||
assert config["enable_vector"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_client_reuses_shared_pool_for_same_vector_settings() -> None:
|
||||
db = MagicMock()
|
||||
db.initdb = AsyncMock()
|
||||
db.check_tables = AsyncMock()
|
||||
|
||||
with patch("lightrag.kg.postgres_impl.PostgreSQLDB", return_value=db) as db_cls:
|
||||
first = await ClientManager.get_client("PGVectorStorage")
|
||||
second = await ClientManager.get_client("PGVectorStorage")
|
||||
|
||||
assert first is db
|
||||
assert second is db
|
||||
assert ClientManager._instances["ref_count"] == 2
|
||||
db_cls.assert_called_once()
|
||||
db.initdb.assert_awaited_once()
|
||||
db.check_tables.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_client_rejects_conflicting_vector_storage_settings() -> None:
|
||||
db = MagicMock()
|
||||
db.initdb = AsyncMock()
|
||||
db.check_tables = AsyncMock()
|
||||
|
||||
with patch("lightrag.kg.postgres_impl.PostgreSQLDB", return_value=db):
|
||||
await ClientManager.get_client("NanoVectorDBStorage")
|
||||
|
||||
with pytest.raises(RuntimeError, match="process-wide"):
|
||||
await ClientManager.get_client("PGVectorStorage")
|
||||
@@ -0,0 +1,362 @@
|
||||
"""
|
||||
Unit tests for Cypher injection prevention in PGGraphStorage write paths.
|
||||
|
||||
Verifies that upsert_node and upsert_edge keep entity IDs parameterized while
|
||||
rendering property maps as safely escaped Cypher literals, which is required by
|
||||
Apache AGE because ``SET ... += $props`` is not supported.
|
||||
"""
|
||||
|
||||
import json
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from lightrag.kg.postgres_impl import PGGraphStorage
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def make_graph_storage() -> PGGraphStorage:
|
||||
"""Construct a PGGraphStorage instance with a mocked db."""
|
||||
storage = PGGraphStorage.__new__(PGGraphStorage)
|
||||
storage.workspace = "test_ws"
|
||||
storage.namespace = "test_graph"
|
||||
storage.graph_name = "test_graph"
|
||||
storage.db = MagicMock()
|
||||
return storage
|
||||
|
||||
|
||||
class _FakeConnection:
|
||||
"""Captures statements + args passed to a fake asyncpg connection."""
|
||||
|
||||
def __init__(self):
|
||||
self.calls: list[dict] = []
|
||||
|
||||
def transaction(self):
|
||||
return _FakeTransaction()
|
||||
|
||||
async def execute(self, sql, *args):
|
||||
self.calls.append({"sql": sql, "args": args})
|
||||
return ""
|
||||
|
||||
|
||||
class _FakeTransaction:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
|
||||
async def _capture_upsert_edge(storage: PGGraphStorage, src: str, tgt: str, edge_data):
|
||||
"""Invoke upsert_edge against a fake connection and return the captured calls."""
|
||||
conn = _FakeConnection()
|
||||
|
||||
async def fake_run_with_retry(operation, **_kwargs):
|
||||
return await operation(conn)
|
||||
|
||||
storage.db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry)
|
||||
await storage.upsert_edge(src, tgt, edge_data)
|
||||
return conn.calls
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# upsert_node — parameterized Cypher
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_node_uses_parameterized_cypher():
|
||||
"""upsert_node must pass entity_id as a Cypher parameter, not interpolate it."""
|
||||
storage = make_graph_storage()
|
||||
captured_calls: list[dict] = []
|
||||
|
||||
async def fake_query(sql, **kwargs):
|
||||
captured_calls.append({"sql": sql, **kwargs})
|
||||
return []
|
||||
|
||||
with patch.object(storage, "_query", side_effect=fake_query):
|
||||
await storage.upsert_node(
|
||||
"Alice", {"entity_id": "Alice", "description": "A person"}
|
||||
)
|
||||
|
||||
assert len(captured_calls) == 1
|
||||
call = captured_calls[0]
|
||||
assert "$1::agtype" in call["sql"]
|
||||
assert '"Alice"' not in call["sql"].replace("$1::agtype", "")
|
||||
assert "params" in call
|
||||
params = json.loads(call["params"]["params"])
|
||||
assert params["entity_id"] == "Alice"
|
||||
assert "props" not in params
|
||||
assert '`description`: "A person"' in call["sql"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_node_injection_payload_in_entity_id():
|
||||
"""A Cypher injection payload in entity_id must be treated as data, not code."""
|
||||
storage = make_graph_storage()
|
||||
injection = 'test"}) RETURN n; MATCH (m) DETACH DELETE m; //'
|
||||
captured_calls: list[dict] = []
|
||||
|
||||
async def fake_query(sql, **kwargs):
|
||||
captured_calls.append({"sql": sql, **kwargs})
|
||||
return []
|
||||
|
||||
with patch.object(storage, "_query", side_effect=fake_query):
|
||||
await storage.upsert_node(
|
||||
injection, {"entity_id": injection, "description": "malicious"}
|
||||
)
|
||||
|
||||
call = captured_calls[0]
|
||||
# The injection payload must NOT appear in the SQL string
|
||||
assert "DETACH DELETE" not in call["sql"]
|
||||
assert injection not in call["sql"]
|
||||
# It must be safely contained in the JSON parameter
|
||||
params = json.loads(call["params"]["params"])
|
||||
assert params["entity_id"] == injection
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_node_special_chars_in_properties():
|
||||
"""Property values with special characters are safely escaped in Cypher."""
|
||||
storage = make_graph_storage()
|
||||
captured_calls: list[dict] = []
|
||||
|
||||
async def fake_query(sql, **kwargs):
|
||||
captured_calls.append({"sql": sql, **kwargs})
|
||||
return []
|
||||
|
||||
node_data = {
|
||||
"entity_id": "test_node",
|
||||
"description": 'He said "hello" and used a backslash \\',
|
||||
"notes": "Line1\nLine2\tTabbed",
|
||||
"formula": "x < 5 && y > 3",
|
||||
}
|
||||
|
||||
with patch.object(storage, "_query", side_effect=fake_query):
|
||||
await storage.upsert_node("test_node", node_data)
|
||||
|
||||
call = captured_calls[0]
|
||||
assert (
|
||||
'`description`: "He said \\"hello\\" and used a backslash \\\\"' in call["sql"]
|
||||
)
|
||||
assert '`notes`: "Line1\\nLine2\\tTabbed"' in call["sql"]
|
||||
assert '`formula`: "x < 5 && y > 3"' in call["sql"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_node_unicode_entity_id():
|
||||
"""Unicode entity names are safely parameterized."""
|
||||
storage = make_graph_storage()
|
||||
captured_calls: list[dict] = []
|
||||
|
||||
async def fake_query(sql, **kwargs):
|
||||
captured_calls.append({"sql": sql, **kwargs})
|
||||
return []
|
||||
|
||||
unicode_id = "\u4e2d\u6587\u5b9e\u4f53" # Chinese characters
|
||||
with patch.object(storage, "_query", side_effect=fake_query):
|
||||
await storage.upsert_node(
|
||||
unicode_id, {"entity_id": unicode_id, "description": "\u63cf\u8ff0"}
|
||||
)
|
||||
|
||||
call = captured_calls[0]
|
||||
params = json.loads(call["params"]["params"])
|
||||
assert params["entity_id"] == unicode_id
|
||||
assert '`description`: "描述"' in call["sql"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_node_dollar_signs_in_entity_id():
|
||||
"""Dollar signs in entity_id don't break dollar-quoting of the Cypher template."""
|
||||
storage = make_graph_storage()
|
||||
captured_calls: list[dict] = []
|
||||
|
||||
async def fake_query(sql, **kwargs):
|
||||
captured_calls.append({"sql": sql, **kwargs})
|
||||
return []
|
||||
|
||||
dollar_id = "price is $100 or $$200$$"
|
||||
with patch.object(storage, "_query", side_effect=fake_query):
|
||||
await storage.upsert_node(
|
||||
dollar_id, {"entity_id": dollar_id, "description": "has dollars"}
|
||||
)
|
||||
|
||||
call = captured_calls[0]
|
||||
# The dollar signs are in the params, not the SQL template
|
||||
params = json.loads(call["params"]["params"])
|
||||
assert params["entity_id"] == dollar_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_node_escapes_backticks_in_property_keys():
|
||||
"""Backticks in property keys must be escaped before inlining the map."""
|
||||
storage = make_graph_storage()
|
||||
captured_calls: list[dict] = []
|
||||
|
||||
async def fake_query(sql, **kwargs):
|
||||
captured_calls.append({"sql": sql, **kwargs})
|
||||
return []
|
||||
|
||||
with patch.object(storage, "_query", side_effect=fake_query):
|
||||
await storage.upsert_node(
|
||||
"node",
|
||||
{"entity_id": "node", "danger`key": 'value "quoted"'},
|
||||
)
|
||||
|
||||
assert '`danger``key`: "value \\"quoted\\""' in captured_calls[0]["sql"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_node_requires_entity_id():
|
||||
"""upsert_node still raises ValueError when entity_id is missing."""
|
||||
storage = make_graph_storage()
|
||||
with pytest.raises(ValueError, match="entity_id"):
|
||||
await storage.upsert_node("test", {"description": "no entity_id"})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# upsert_edge — parameterized Cypher
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edge_uses_parameterized_cypher():
|
||||
"""upsert_edge must pass entity IDs as Cypher parameters."""
|
||||
storage = make_graph_storage()
|
||||
calls = await _capture_upsert_edge(
|
||||
storage, "Alice", "Bob", {"weight": "1.0", "description": "knows"}
|
||||
)
|
||||
|
||||
# Three statements: per-edge lock, graph-wide shared lock, then cypher.
|
||||
assert len(calls) == 3
|
||||
|
||||
lock_sql = calls[0]["sql"]
|
||||
# Raw node IDs are positional params on the lock, never interpolated.
|
||||
assert "Alice" not in lock_sql
|
||||
assert "Bob" not in lock_sql
|
||||
# graph_name flows as $1, the endpoint pair as $2/$3.
|
||||
assert calls[0]["args"] == ("test_graph", "Alice", "Bob")
|
||||
|
||||
cypher_call = calls[2]
|
||||
cypher_sql = cypher_call["sql"]
|
||||
assert "$1::agtype" in cypher_sql
|
||||
assert '"Alice"' not in cypher_sql.replace("$1::agtype", "")
|
||||
assert '"Bob"' not in cypher_sql.replace("$1::agtype", "")
|
||||
# Cypher params arrive as a single positional agtype JSON arg.
|
||||
params = json.loads(cypher_call["args"][0])
|
||||
assert params["src_id"] == "Alice"
|
||||
assert params["tgt_id"] == "Bob"
|
||||
assert "props" not in params
|
||||
assert '`weight`: "1.0"' in cypher_sql
|
||||
assert '`description`: "knows"' in cypher_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edge_injection_payload():
|
||||
"""Injection payloads in edge entity IDs are safely parameterized."""
|
||||
storage = make_graph_storage()
|
||||
injection_src = 'src"}) MATCH (x) DETACH DELETE x; //'
|
||||
injection_tgt = 'tgt"})-[r]-() DELETE r; //'
|
||||
calls = await _capture_upsert_edge(
|
||||
storage, injection_src, injection_tgt, {"description": "edge"}
|
||||
)
|
||||
|
||||
# Injection payloads must never appear in either SQL template — they only
|
||||
# flow through positional params.
|
||||
for call in calls:
|
||||
assert "DETACH DELETE" not in call["sql"]
|
||||
assert "DELETE r" not in call["sql"]
|
||||
assert injection_src not in call["sql"]
|
||||
assert injection_tgt not in call["sql"]
|
||||
|
||||
# Lock statement passes graph_name + raw IDs as positional params.
|
||||
assert calls[0]["args"] == ("test_graph", injection_src, injection_tgt)
|
||||
|
||||
# Cypher params arrive as a single positional agtype JSON arg (3rd statement,
|
||||
# after the per-edge and graph-wide-shared locks).
|
||||
params = json.loads(calls[2]["args"][0])
|
||||
assert params["src_id"] == injection_src
|
||||
assert params["tgt_id"] == injection_tgt
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edge_unicode_entity_ids():
|
||||
"""Unicode entity IDs in edges are safely parameterized."""
|
||||
storage = make_graph_storage()
|
||||
src = "\u5317\u4eac"
|
||||
tgt = "\u4e0a\u6d77"
|
||||
calls = await _capture_upsert_edge(
|
||||
storage, src, tgt, {"description": "\u8def\u7ebf"}
|
||||
)
|
||||
|
||||
# Lock statement carries graph_name + raw IDs as positional params, not
|
||||
# interpolated.
|
||||
assert calls[0]["args"] == ("test_graph", src, tgt)
|
||||
assert src not in calls[0]["sql"]
|
||||
assert tgt not in calls[0]["sql"]
|
||||
|
||||
# Cypher params parsed from the positional agtype JSON arg (3rd statement).
|
||||
cypher_sql = calls[2]["sql"]
|
||||
params = json.loads(calls[2]["args"][0])
|
||||
assert params["src_id"] == src
|
||||
assert params["tgt_id"] == tgt
|
||||
assert '`description`: "路线"' in cypher_sql
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _normalize_node_id — defence-in-depth for remaining interpolation paths
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_normalize_node_id_strips_null_bytes():
|
||||
"""Null bytes are stripped to prevent string truncation."""
|
||||
assert PGGraphStorage._normalize_node_id("before\x00after") == "beforeafter"
|
||||
|
||||
|
||||
def test_normalize_node_id_escapes_backslash_and_quote():
|
||||
"""Backslashes and double quotes are escaped."""
|
||||
assert PGGraphStorage._normalize_node_id('a\\"b') == 'a\\\\\\"b'
|
||||
|
||||
|
||||
def test_normalize_node_id_injection_payload():
|
||||
"""Injection payload is escaped so it cannot break out of Cypher string."""
|
||||
payload = 'test"}) RETURN n; MATCH (m) DETACH DELETE m; //'
|
||||
normalized = PGGraphStorage._normalize_node_id(payload)
|
||||
# The double quote must be escaped
|
||||
assert '\\"' in normalized
|
||||
# The escaped string must not contain an unescaped double quote
|
||||
# (remove all escaped quotes and check no raw ones remain)
|
||||
unescaped = normalized.replace('\\"', "")
|
||||
assert '"' not in unescaped
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _query write path passes params to db.execute
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_write_path_passes_params():
|
||||
"""When readonly=False, _query must forward params to db.execute."""
|
||||
storage = make_graph_storage()
|
||||
captured_execute_kwargs: list[dict] = []
|
||||
|
||||
async def fake_execute(sql, **kwargs):
|
||||
captured_execute_kwargs.append(kwargs)
|
||||
return None
|
||||
|
||||
storage.db.execute = fake_execute
|
||||
|
||||
test_params = {"params": json.dumps({"entity_id": "test"})}
|
||||
await storage._query(
|
||||
"SELECT 1",
|
||||
readonly=False,
|
||||
upsert=True,
|
||||
params=test_params,
|
||||
)
|
||||
|
||||
assert len(captured_execute_kwargs) == 1
|
||||
assert captured_execute_kwargs[0]["data"] == test_params
|
||||
@@ -0,0 +1,98 @@
|
||||
"""Regression test for Issue #3286.
|
||||
|
||||
``adelete_by_doc_id`` / ``_purge_doc_chunks_and_kg`` historically passed
|
||||
``chunk_ids`` as a ``set`` to ``PGKVStorage.delete`` / ``PGDocStatusStorage.delete``.
|
||||
After delete-batching was introduced, the chunked path slices the id collection
|
||||
(``ids[i : i + chunk]``); slicing a ``set`` raises ``'set' object is not
|
||||
subscriptable``, which the broad ``except Exception`` in ``delete`` then swallows
|
||||
as a log line -- so the records were silently NOT deleted.
|
||||
|
||||
These tests feed a ``set`` straight into ``delete`` with a positive batch cap
|
||||
(forcing the chunked path) and assert that every id is actually handed to the
|
||||
SQL layer. Asserting "no exception" alone is insufficient because the bug is
|
||||
swallowed; we must assert the records really got deleted.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("asyncpg")
|
||||
|
||||
from lightrag.kg.postgres_impl import ( # noqa: E402
|
||||
PGDocStatusStorage,
|
||||
PGKVStorage,
|
||||
)
|
||||
from lightrag.namespace import NameSpace # noqa: E402
|
||||
|
||||
|
||||
class _FakeTransaction:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc):
|
||||
return False
|
||||
|
||||
|
||||
class _RecordingConnection:
|
||||
"""Captures the id-array argument of each ``execute`` call."""
|
||||
|
||||
def __init__(self):
|
||||
self.deleted_id_batches: list[list[str]] = []
|
||||
|
||||
def transaction(self):
|
||||
return _FakeTransaction()
|
||||
|
||||
async def execute(self, sql, workspace, id_slice):
|
||||
# The chunked path slices ids; a slice of a list is a list. Record it.
|
||||
assert isinstance(id_slice, list)
|
||||
self.deleted_id_batches.append(id_slice)
|
||||
|
||||
|
||||
class _FakeDB:
|
||||
def __init__(self, connection):
|
||||
self._connection = connection
|
||||
|
||||
async def _run_with_retry(self, operation):
|
||||
# Mirror the real helper: invoke the closure against a live connection.
|
||||
return await operation(self._connection)
|
||||
|
||||
|
||||
def _make_storage(cls, namespace, batch_cap):
|
||||
"""Build a storage instance exercising only the attributes ``delete`` reads."""
|
||||
storage = object.__new__(cls)
|
||||
storage.namespace = namespace
|
||||
storage.workspace = "test_ws"
|
||||
storage._max_delete_records_per_batch = batch_cap
|
||||
connection = _RecordingConnection()
|
||||
storage.db = _FakeDB(connection)
|
||||
return storage, connection
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pgkv_delete_accepts_set_and_chunks_all_ids():
|
||||
ids = {f"chunk-{i}" for i in range(5)}
|
||||
# Positive cap < len(ids) forces the chunked slicing path that broke on sets.
|
||||
storage, connection = _make_storage(
|
||||
PGKVStorage, NameSpace.KV_STORE_TEXT_CHUNKS, batch_cap=2
|
||||
)
|
||||
|
||||
await storage.delete(ids)
|
||||
|
||||
flattened = [cid for batch in connection.deleted_id_batches for cid in batch]
|
||||
assert set(flattened) == ids
|
||||
assert len(flattened) == len(ids) # no duplicates / drops
|
||||
assert all(len(batch) <= 2 for batch in connection.deleted_id_batches)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pgdocstatus_delete_accepts_set_and_chunks_all_ids():
|
||||
ids = {f"doc-{i}" for i in range(5)}
|
||||
storage, connection = _make_storage(
|
||||
PGDocStatusStorage, NameSpace.DOC_STATUS, batch_cap=2
|
||||
)
|
||||
|
||||
await storage.delete(ids)
|
||||
|
||||
flattened = [did for batch in connection.deleted_id_batches for did in batch]
|
||||
assert set(flattened) == ids
|
||||
assert len(flattened) == len(ids)
|
||||
assert all(len(batch) <= 2 for batch in connection.deleted_id_batches)
|
||||
@@ -0,0 +1,166 @@
|
||||
"""Unit tests for PGDocStatusStorage database-native overrides.
|
||||
|
||||
Covers the PG-specific implementations of:
|
||||
* get_doc_by_file_basename
|
||||
* get_doc_by_content_hash
|
||||
|
||||
Both override the base-class full-table scan with indexed SQL queries.
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from lightrag.kg.postgres_impl import PGDocStatusStorage
|
||||
from lightrag.namespace import NameSpace
|
||||
|
||||
|
||||
def _make_storage():
|
||||
storage = PGDocStatusStorage.__new__(PGDocStatusStorage)
|
||||
storage.namespace = NameSpace.DOC_STATUS
|
||||
storage.workspace = "test_ws"
|
||||
storage.global_config = {"embedding_batch_num": 10}
|
||||
storage.db = MagicMock()
|
||||
storage.db.query = AsyncMock()
|
||||
return storage
|
||||
|
||||
|
||||
def _row(**overrides):
|
||||
base = {
|
||||
"id": "doc-1",
|
||||
"content_summary": "summary",
|
||||
"content_length": 12,
|
||||
"chunks_count": 1,
|
||||
"status": "processed",
|
||||
"file_path": "report.pdf",
|
||||
"chunks_list": "[]",
|
||||
"metadata": "{}",
|
||||
"error_msg": None,
|
||||
"track_id": None,
|
||||
"content_hash": "abc123",
|
||||
"created_at": datetime(2024, 1, 1, 0, 0, 0),
|
||||
"updated_at": datetime(2024, 1, 1, 0, 0, 0),
|
||||
}
|
||||
base.update(overrides)
|
||||
return base
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_doc_by_file_basename
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_doc_by_file_basename_empty_returns_none():
|
||||
storage = _make_storage()
|
||||
assert await storage.get_doc_by_file_basename("") is None
|
||||
storage.db.query.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_doc_by_file_basename_unknown_source_returns_none():
|
||||
storage = _make_storage()
|
||||
# normalize_document_file_path returns "unknown_source" for None-ish inputs
|
||||
assert await storage.get_doc_by_file_basename("unknown_source") is None
|
||||
storage.db.query.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_doc_by_file_basename_exact_match():
|
||||
storage = _make_storage()
|
||||
storage.db.query.return_value = [_row(file_path="report.pdf")]
|
||||
|
||||
result = await storage.get_doc_by_file_basename("report.pdf")
|
||||
|
||||
assert result is not None
|
||||
doc_id, doc = result
|
||||
assert doc_id == "doc-1"
|
||||
assert doc["file_path"] == "report.pdf"
|
||||
assert doc["content_hash"] == "abc123"
|
||||
|
||||
call = storage.db.query.call_args
|
||||
sql = call.args[0]
|
||||
params = call.args[1]
|
||||
assert "LIGHTRAG_DOC_STATUS" in sql
|
||||
assert "workspace=$1" in sql
|
||||
assert params[0] == "test_ws"
|
||||
assert params[1] == "report.pdf"
|
||||
assert params == ["test_ws", "report.pdf"]
|
||||
assert "LIKE" not in sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_doc_by_file_basename_orders_stably_for_canonical_rows():
|
||||
storage = _make_storage()
|
||||
storage.db.query.return_value = [_row(id="doc-exact", file_path="report.pdf")]
|
||||
|
||||
result = await storage.get_doc_by_file_basename("report.pdf")
|
||||
|
||||
assert result is not None
|
||||
assert result[0] == "doc-exact"
|
||||
|
||||
sql = storage.db.query.call_args.args[0]
|
||||
assert "file_path = $2" in sql
|
||||
assert "created_at ASC" in sql
|
||||
assert "id ASC" in sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_doc_by_file_basename_uses_exact_match_for_like_metacharacters():
|
||||
storage = _make_storage()
|
||||
storage.db.query.return_value = []
|
||||
|
||||
await storage.get_doc_by_file_basename("100%_off.pdf")
|
||||
|
||||
sql = storage.db.query.call_args.args[0]
|
||||
params = storage.db.query.call_args.args[1]
|
||||
assert "LIKE" not in sql
|
||||
assert params == ["test_ws", "100%_off.pdf"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_doc_by_file_basename_no_match_returns_none():
|
||||
storage = _make_storage()
|
||||
storage.db.query.return_value = []
|
||||
assert await storage.get_doc_by_file_basename("missing.pdf") is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_doc_by_content_hash
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_doc_by_content_hash_empty_returns_none():
|
||||
storage = _make_storage()
|
||||
assert await storage.get_doc_by_content_hash("") is None
|
||||
storage.db.query.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_doc_by_content_hash_match():
|
||||
storage = _make_storage()
|
||||
storage.db.query.return_value = [_row(content_hash="hash-abc")]
|
||||
|
||||
result = await storage.get_doc_by_content_hash("hash-abc")
|
||||
|
||||
assert result is not None
|
||||
doc_id, doc = result
|
||||
assert doc_id == "doc-1"
|
||||
assert doc["content_hash"] == "hash-abc"
|
||||
|
||||
call = storage.db.query.call_args
|
||||
sql = call.args[0]
|
||||
params = call.args[1]
|
||||
assert "content_hash=$2" in sql
|
||||
# Stable ordering for repeatability across re-runs / replicas
|
||||
assert "ORDER BY created_at ASC, id ASC" in sql
|
||||
assert "LIMIT 1" in sql
|
||||
assert params == ["test_ws", "hash-abc"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_doc_by_content_hash_no_match_returns_none():
|
||||
storage = _make_storage()
|
||||
storage.db.query.return_value = []
|
||||
assert await storage.get_doc_by_content_hash("nope") is None
|
||||
@@ -0,0 +1,189 @@
|
||||
"""
|
||||
Unit tests for PGGraphStorage.get_knowledge_graph("*") — the full-graph view.
|
||||
|
||||
These cover the undirected fix for the degree-based node selection used when the
|
||||
graph is truncated to ``max_nodes``: the previous implementation counted only
|
||||
outgoing edges (``OPTIONAL MATCH (n)-[r]->()``), so a node that is mostly an edge
|
||||
*target* was under-ranked and could be dropped on truncation. The fix selects
|
||||
top-degree nodes via native SQL that counts both ``start_id`` and ``end_id``
|
||||
(undirected degree), LEFT JOIN-ing the base vertex table so isolated nodes are
|
||||
preserved when the graph is not truncated.
|
||||
|
||||
The subgraph edge read intentionally stays directed: ``a`` iterates over every
|
||||
selected node, so ``(a)-[r]->(b)`` already captures every edge whose endpoints
|
||||
are both selected. These tests guard against an accidental regression there.
|
||||
|
||||
All tests mock ``PGGraphStorage._query`` and inspect the SQL it receives.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from lightrag.kg.postgres_impl import PGGraphStorage
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def make_graph_storage() -> PGGraphStorage:
|
||||
"""Construct a PGGraphStorage instance with a mocked _query method."""
|
||||
storage = PGGraphStorage.__new__(PGGraphStorage)
|
||||
storage.workspace = "test_ws"
|
||||
storage.namespace = "test_graph"
|
||||
storage.graph_name = "test_graph"
|
||||
storage.global_config = {"max_graph_nodes": 1000}
|
||||
storage.db = MagicMock()
|
||||
return storage
|
||||
|
||||
|
||||
class _QueryCapture:
|
||||
"""Dispatch the three _query calls of the '*' branch by SQL content."""
|
||||
|
||||
def __init__(self, *, total_nodes, degree_rows, subgraph_rows):
|
||||
self._total_nodes = total_nodes
|
||||
self._degree_rows = degree_rows
|
||||
self._subgraph_rows = subgraph_rows
|
||||
self.count_sql = None
|
||||
self.degree_sql = None
|
||||
self.degree_params = None
|
||||
self.subgraph_sql = None
|
||||
|
||||
def as_side_effect(self):
|
||||
"""Return an ``async def`` so AsyncMock awaits it (a callable instance is not)."""
|
||||
|
||||
async def fake_query(query, **kwargs):
|
||||
if "count(distinct n)" in query:
|
||||
self.count_sql = query
|
||||
return [{"total_nodes": self._total_nodes}]
|
||||
if "node_degrees" in query:
|
||||
self.degree_sql = query
|
||||
self.degree_params = kwargs.get("params")
|
||||
return self._degree_rows
|
||||
# subgraph read
|
||||
self.subgraph_sql = query
|
||||
return self._subgraph_rows
|
||||
|
||||
return fake_query
|
||||
|
||||
|
||||
def _node(node_id, entity_id):
|
||||
return {"id": node_id, "properties": {"entity_id": entity_id}}
|
||||
|
||||
|
||||
def _edge(edge_id, start_id, end_id, weight="1"):
|
||||
return {
|
||||
"id": edge_id,
|
||||
"label": "DIRECTED",
|
||||
"start_id": start_id,
|
||||
"end_id": end_id,
|
||||
"properties": {"weight": weight},
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# degree node selection SQL
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_degree_selection_sql_is_undirected_and_preserves_isolated():
|
||||
"""The degree query must count start_id + end_id and keep degree-0 nodes."""
|
||||
capture = _QueryCapture(
|
||||
total_nodes=2,
|
||||
degree_rows=[{"node_id": 1, "degree": 3}, {"node_id": 2, "degree": 0}],
|
||||
subgraph_rows=[{"a": _node(1, "Alice"), "r": None, "b": None}],
|
||||
)
|
||||
storage = make_graph_storage()
|
||||
|
||||
with patch.object(storage, "_query", side_effect=capture.as_side_effect()):
|
||||
await storage.get_knowledge_graph("*", max_nodes=50)
|
||||
|
||||
sql = capture.degree_sql
|
||||
assert sql is not None, "degree selection query was never issued"
|
||||
# Undirected: both edge endpoints are counted.
|
||||
assert "start_id" in sql
|
||||
assert "end_id" in sql
|
||||
assert "UNION ALL" in sql
|
||||
# Isolated nodes preserved.
|
||||
assert "LEFT JOIN" in sql
|
||||
assert "COALESCE" in sql
|
||||
# Stable ordering with id tie-break.
|
||||
assert "ORDER BY degree DESC" in sql
|
||||
assert "v.id ASC" in sql
|
||||
# The old outgoing-only Cypher must be gone.
|
||||
assert "-[r]->()" not in sql
|
||||
assert "OPTIONAL MATCH (n)-[r]->()" not in sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_degree_selection_limit_is_parameterized():
|
||||
"""max_nodes must be passed via params, not interpolated into the SQL."""
|
||||
capture = _QueryCapture(
|
||||
total_nodes=2,
|
||||
degree_rows=[{"node_id": 1, "degree": 3}],
|
||||
subgraph_rows=[{"a": _node(1, "Alice"), "r": None, "b": None}],
|
||||
)
|
||||
storage = make_graph_storage()
|
||||
|
||||
with patch.object(storage, "_query", side_effect=capture.as_side_effect()):
|
||||
await storage.get_knowledge_graph("*", max_nodes=37)
|
||||
|
||||
assert "LIMIT $1" in capture.degree_sql
|
||||
assert "37" not in capture.degree_sql
|
||||
assert capture.degree_params == {"limit": 37}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_isolated_node_survives_end_to_end():
|
||||
"""A degree-0 node selected by the degree query reaches the final KnowledgeGraph."""
|
||||
capture = _QueryCapture(
|
||||
total_nodes=2,
|
||||
degree_rows=[{"node_id": 1, "degree": 2}, {"node_id": 2, "degree": 0}],
|
||||
subgraph_rows=[
|
||||
{"a": _node(1, "Alice"), "r": None, "b": None},
|
||||
# MATCH (a) still returns the isolated node with null r/b.
|
||||
{"a": _node(2, "Bob"), "r": None, "b": None},
|
||||
],
|
||||
)
|
||||
storage = make_graph_storage()
|
||||
|
||||
with patch.object(storage, "_query", side_effect=capture.as_side_effect()):
|
||||
kg = await storage.get_knowledge_graph("*", max_nodes=50)
|
||||
|
||||
# The isolated node id is formatted into the subgraph id list...
|
||||
assert "2" in capture.subgraph_sql
|
||||
# ...and present in the returned graph.
|
||||
assert {node.id for node in kg.nodes} == {"1", "2"}
|
||||
assert kg.is_truncated is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# subgraph edge read — regression guard
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subgraph_read_stays_directed_and_dedupes_edges():
|
||||
"""Subgraph read keeps the directed (a)-[r]->(b) match and dedupes by edge id."""
|
||||
duplicate_edge = _edge(10, 1, 2)
|
||||
capture = _QueryCapture(
|
||||
total_nodes=2,
|
||||
degree_rows=[{"node_id": 1, "degree": 1}, {"node_id": 2, "degree": 1}],
|
||||
subgraph_rows=[
|
||||
{"a": _node(1, "Alice"), "r": duplicate_edge, "b": _node(2, "Bob")},
|
||||
# Same AGE edge id appearing twice must collapse to one edge.
|
||||
{"a": _node(1, "Alice"), "r": duplicate_edge, "b": _node(2, "Bob")},
|
||||
],
|
||||
)
|
||||
storage = make_graph_storage()
|
||||
|
||||
with patch.object(storage, "_query", side_effect=capture.as_side_effect()):
|
||||
kg = await storage.get_knowledge_graph("*", max_nodes=50)
|
||||
|
||||
assert "(a)-[r]->(b)" in capture.subgraph_sql
|
||||
assert len(kg.edges) == 1
|
||||
edge = kg.edges[0]
|
||||
assert edge.source == "1"
|
||||
assert edge.target == "2"
|
||||
@@ -0,0 +1,236 @@
|
||||
"""Unit tests for PGGraphStorage chunk-level transaction batch paths.
|
||||
|
||||
Covers the PR-2 chunk-level fallback: ``upsert_nodes_batch`` /
|
||||
``upsert_edges_batch`` group per-row Cypher into payload/record-bounded chunks
|
||||
run in a single transaction; ``remove_nodes`` runs all chunks in ONE
|
||||
transaction (all-or-nothing); ``remove_edges`` runs one transaction per chunk.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from lightrag.kg.postgres_impl import PGGraphStorage
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Capture harness
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _Capture:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[dict] = [] # every connection.execute(sql, *args)
|
||||
self.tx_count = 0 # connection.transaction() opens
|
||||
self.run_count = 0 # _run_with_retry invocations (= chunks issued)
|
||||
|
||||
|
||||
class _FakeTransaction:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
|
||||
class _FakeConnection:
|
||||
def __init__(self, capture: _Capture):
|
||||
self._capture = capture
|
||||
|
||||
def transaction(self):
|
||||
self._capture.tx_count += 1
|
||||
return _FakeTransaction()
|
||||
|
||||
async def execute(self, sql, *args):
|
||||
self._capture.calls.append({"sql": sql, "args": args})
|
||||
return ""
|
||||
|
||||
|
||||
def make_graph_storage(
|
||||
*,
|
||||
max_upsert_records: int | None = None,
|
||||
max_upsert_payload: int | None = None,
|
||||
max_delete_records: int | None = None,
|
||||
) -> tuple[PGGraphStorage, _Capture]:
|
||||
storage = PGGraphStorage.__new__(PGGraphStorage)
|
||||
storage.workspace = "test_ws"
|
||||
storage.namespace = "test_graph"
|
||||
storage.graph_name = "test_graph"
|
||||
storage.__post_init__() # resolves the chunk-level batch limits
|
||||
if max_upsert_records is not None:
|
||||
storage._max_upsert_records_per_batch = max_upsert_records
|
||||
if max_upsert_payload is not None:
|
||||
storage._max_upsert_payload_bytes = max_upsert_payload
|
||||
if max_delete_records is not None:
|
||||
storage._max_delete_records_per_batch = max_delete_records
|
||||
|
||||
capture = _Capture()
|
||||
conn = _FakeConnection(capture)
|
||||
|
||||
async def fake_run_with_retry(operation, **_kwargs):
|
||||
capture.run_count += 1
|
||||
return await operation(conn)
|
||||
|
||||
storage.db = AsyncMock()
|
||||
storage.db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry)
|
||||
return storage, capture
|
||||
|
||||
|
||||
def _sql_count(capture: _Capture, needle: str) -> int:
|
||||
return sum(1 for c in capture.calls if needle in c["sql"])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# upsert_nodes_batch
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_nodes_batch_splits_into_chunk_transactions():
|
||||
storage, cap = make_graph_storage(max_upsert_records=2)
|
||||
nodes = [(f"n{i}", {"entity_id": f"n{i}", "v": str(i)}) for i in range(5)]
|
||||
|
||||
await storage.upsert_nodes_batch(nodes)
|
||||
|
||||
# 5 nodes / cap 2 => 3 chunks, each its own _run_with_retry + transaction.
|
||||
assert cap.run_count == 3
|
||||
assert cap.tx_count == 3
|
||||
# One MERGE per node.
|
||||
assert _sql_count(cap, "MERGE (n:base") == 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_nodes_batch_dedupes_last_write_wins():
|
||||
storage, cap = make_graph_storage()
|
||||
await storage.upsert_nodes_batch(
|
||||
[
|
||||
("A", {"entity_id": "A", "weight": "first"}),
|
||||
("A", {"entity_id": "A", "weight": "second"}),
|
||||
]
|
||||
)
|
||||
|
||||
merge_calls = [c for c in cap.calls if "MERGE (n:base" in c["sql"]]
|
||||
assert len(merge_calls) == 1
|
||||
assert '"second"' in merge_calls[0]["sql"]
|
||||
assert '"first"' not in merge_calls[0]["sql"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_nodes_batch_splits_by_payload_bytes():
|
||||
# Disable the record cap, force a small byte budget so each big node lands
|
||||
# in its own chunk.
|
||||
storage, cap = make_graph_storage(max_upsert_records=0, max_upsert_payload=200)
|
||||
nodes = [(f"n{i}", {"entity_id": f"n{i}", "blob": "X" * 150}) for i in range(3)]
|
||||
|
||||
await storage.upsert_nodes_batch(nodes)
|
||||
|
||||
assert cap.run_count == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_nodes_batch_missing_entity_id_raises_before_tx():
|
||||
storage, cap = make_graph_storage()
|
||||
with pytest.raises(ValueError, match="entity_id"):
|
||||
await storage.upsert_nodes_batch([("A", {"foo": "bar"})])
|
||||
# The builder runs before the transaction opens, so nothing was issued.
|
||||
assert cap.run_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_nodes_batch_empty_noop():
|
||||
storage, cap = make_graph_storage()
|
||||
await storage.upsert_nodes_batch([])
|
||||
assert cap.run_count == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# upsert_edges_batch
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edges_batch_splits_into_chunk_transactions():
|
||||
storage, cap = make_graph_storage(max_upsert_records=2)
|
||||
edges = [(f"s{i}", f"t{i}", {"w": str(i)}) for i in range(5)]
|
||||
|
||||
await storage.upsert_edges_batch(edges)
|
||||
|
||||
# 5 edges / cap 2 => 3 chunks; each chunk = one transaction.
|
||||
assert cap.run_count == 3
|
||||
assert cap.tx_count == 3
|
||||
# One CREATE cypher per edge, and exactly one graph-wide lock per chunk.
|
||||
assert _sql_count(cap, "CREATE (source)-[r:DIRECTED") == 5
|
||||
assert _sql_count(cap, "pg_advisory_xact_lock") == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edges_batch_takes_one_graph_lock_per_chunk():
|
||||
"""The batch chunk takes a single graph-wide advisory lock (keyed on
|
||||
graph_name only), not one lock per edge -- bounded regardless of edge count."""
|
||||
storage, cap = make_graph_storage()
|
||||
await storage.upsert_edges_batch(
|
||||
[(f"s{i}", f"t{i}", {"w": str(i)}) for i in range(4)]
|
||||
)
|
||||
|
||||
lock_calls = [c for c in cap.calls if "pg_advisory_xact_lock" in c["sql"]]
|
||||
assert len(lock_calls) == 1 # one chunk, one lock (not 4)
|
||||
assert lock_calls[0]["args"] == ("test_graph",) # graph-keyed, no endpoints
|
||||
assert _sql_count(cap, "CREATE (source)-[r:DIRECTED") == 4
|
||||
# The lock is acquired before any edge cypher in the chunk.
|
||||
assert "pg_advisory_xact_lock" in cap.calls[0]["sql"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edges_batch_empty_noop():
|
||||
storage, cap = make_graph_storage()
|
||||
await storage.upsert_edges_batch([])
|
||||
assert cap.run_count == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# remove_nodes — all chunks in ONE transaction (all-or-nothing)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remove_nodes_chunks_share_one_transaction():
|
||||
storage, cap = make_graph_storage(max_delete_records=2)
|
||||
await storage.remove_nodes([f"n{i}" for i in range(5)])
|
||||
|
||||
# One transaction wraps all chunks (preserves single-statement atomicity).
|
||||
assert cap.run_count == 1
|
||||
assert cap.tx_count == 1
|
||||
# 5 ids / cap 2 => 3 bounded IN [...] DETACH DELETE statements.
|
||||
assert _sql_count(cap, "DETACH DELETE n") == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remove_nodes_empty_noop():
|
||||
storage, cap = make_graph_storage()
|
||||
await storage.remove_nodes([])
|
||||
assert cap.run_count == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# remove_edges — one transaction per chunk
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remove_edges_one_transaction_per_chunk():
|
||||
storage, cap = make_graph_storage(max_delete_records=2)
|
||||
edges = [(f"s{i}", f"t{i}") for i in range(5)]
|
||||
|
||||
await storage.remove_edges(edges)
|
||||
|
||||
# 5 edges / cap 2 => 3 chunks, each its own transaction.
|
||||
assert cap.run_count == 3
|
||||
assert cap.tx_count == 3
|
||||
# One DELETE r statement per edge.
|
||||
assert _sql_count(cap, "DELETE r") == 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_remove_edges_empty_noop():
|
||||
storage, cap = make_graph_storage()
|
||||
await storage.remove_edges([])
|
||||
assert cap.run_count == 0
|
||||
@@ -0,0 +1,374 @@
|
||||
import pytest
|
||||
import numpy as np
|
||||
from unittest.mock import patch, AsyncMock
|
||||
from lightrag.utils import EmbeddingFunc
|
||||
from lightrag.kg.postgres_impl import (
|
||||
PGVectorStorage,
|
||||
PostgreSQLDB,
|
||||
_safe_index_name,
|
||||
)
|
||||
from lightrag.exceptions import DataMigrationError
|
||||
from lightrag.namespace import NameSpace
|
||||
|
||||
|
||||
# Mock PostgreSQLDB
|
||||
@pytest.fixture
|
||||
def mock_pg_db():
|
||||
"""Mock PostgreSQL database connection"""
|
||||
db = AsyncMock()
|
||||
db.workspace = "test_workspace"
|
||||
db.vector_index_type = None
|
||||
|
||||
# Mock query responses: list for search queries (multirows=True), dict for DDL checks
|
||||
async def mock_query(sql, params=None, multirows=False, **kwargs):
|
||||
if multirows:
|
||||
return []
|
||||
return {"exists": False, "count": 0}
|
||||
|
||||
# Mock for execute that mimics PostgreSQLDB.execute() behavior
|
||||
async def mock_execute(sql, data=None, **kwargs):
|
||||
return None
|
||||
|
||||
db.query = AsyncMock(side_effect=mock_query)
|
||||
db.execute = AsyncMock(side_effect=mock_execute)
|
||||
|
||||
return db
|
||||
|
||||
|
||||
# Mock get_data_init_lock to avoid async lock issues in tests
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_data_init_lock():
|
||||
with patch("lightrag.kg.postgres_impl.get_data_init_lock") as mock_lock:
|
||||
mock_lock_ctx = AsyncMock()
|
||||
mock_lock.return_value = mock_lock_ctx
|
||||
yield mock_lock
|
||||
|
||||
|
||||
# Mock ClientManager
|
||||
@pytest.fixture
|
||||
def mock_client_manager(mock_pg_db):
|
||||
with patch("lightrag.kg.postgres_impl.ClientManager") as mock_manager:
|
||||
mock_manager.get_client = AsyncMock(return_value=mock_pg_db)
|
||||
mock_manager.release_client = AsyncMock()
|
||||
yield mock_manager
|
||||
|
||||
|
||||
# Mock Embedding function
|
||||
@pytest.fixture
|
||||
def mock_embedding_func():
|
||||
async def embed_func(texts, **kwargs):
|
||||
return np.array([[0.1] * 768 for _ in texts])
|
||||
|
||||
# Note: EmbeddingFunc in this version of lightrag supports model_name
|
||||
func = EmbeddingFunc(embedding_dim=768, func=embed_func, model_name="test_model")
|
||||
return func
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_postgres_halfvec_table_creation(
|
||||
mock_client_manager, mock_pg_db, mock_embedding_func
|
||||
):
|
||||
"""Test if table is created with HALFVEC type when HNSW_HALFVEC is selected"""
|
||||
# Set index type to HNSW_HALFVEC
|
||||
mock_pg_db.vector_index_type = "HNSW_HALFVEC"
|
||||
|
||||
config = {
|
||||
"embedding_batch_num": 10,
|
||||
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
||||
}
|
||||
|
||||
storage = PGVectorStorage(
|
||||
namespace=NameSpace.VECTOR_STORE_CHUNKS,
|
||||
global_config=config,
|
||||
embedding_func=mock_embedding_func,
|
||||
workspace="test_ws",
|
||||
)
|
||||
|
||||
# Mock table doesn't exist
|
||||
mock_pg_db.check_table_exists = AsyncMock(return_value=False)
|
||||
|
||||
# Initialize storage (should trigger table creation)
|
||||
await storage.initialize()
|
||||
|
||||
# Verify table creation SQL contains HALFVEC(768)
|
||||
create_table_calls = [
|
||||
call
|
||||
for call in mock_pg_db.execute.call_args_list
|
||||
if "CREATE TABLE" in call[0][0]
|
||||
]
|
||||
|
||||
assert len(create_table_calls) > 0
|
||||
create_sql = create_table_calls[0][0][0]
|
||||
assert "HALFVEC(768)" in create_sql
|
||||
assert "VECTOR(768)" not in create_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_postgres_vector_table_creation_default(
|
||||
mock_client_manager, mock_pg_db, mock_embedding_func
|
||||
):
|
||||
"""Test if table is created with default VECTOR type when other index type is selected"""
|
||||
# Set index type to HNSW (default)
|
||||
mock_pg_db.vector_index_type = "HNSW"
|
||||
|
||||
config = {
|
||||
"embedding_batch_num": 10,
|
||||
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
||||
}
|
||||
|
||||
storage = PGVectorStorage(
|
||||
namespace=NameSpace.VECTOR_STORE_CHUNKS,
|
||||
global_config=config,
|
||||
embedding_func=mock_embedding_func,
|
||||
workspace="test_ws",
|
||||
)
|
||||
|
||||
# Mock table doesn't exist
|
||||
mock_pg_db.check_table_exists = AsyncMock(return_value=False)
|
||||
|
||||
# Initialize storage (should trigger table creation)
|
||||
await storage.initialize()
|
||||
|
||||
# Verify table creation SQL contains VECTOR(768)
|
||||
create_table_calls = [
|
||||
call
|
||||
for call in mock_pg_db.execute.call_args_list
|
||||
if "CREATE TABLE" in call[0][0]
|
||||
]
|
||||
|
||||
assert len(create_table_calls) > 0
|
||||
create_sql = create_table_calls[0][0][0]
|
||||
assert "VECTOR(768)" in create_sql
|
||||
assert "HALFVEC(768)" not in create_sql
|
||||
|
||||
|
||||
# Namespaces that use vector search SQL templates (query path)
|
||||
QUERY_NAMESPACES = [
|
||||
NameSpace.VECTOR_STORE_CHUNKS,
|
||||
NameSpace.VECTOR_STORE_ENTITIES,
|
||||
NameSpace.VECTOR_STORE_RELATIONSHIPS,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("namespace", QUERY_NAMESPACES)
|
||||
async def test_query_uses_halfvec_cast_when_hnsw_halfvec(
|
||||
mock_client_manager, mock_pg_db, mock_embedding_func, namespace
|
||||
):
|
||||
"""When HNSW_HALFVEC is set, generated search SQL uses ::halfvec (not ::vector)."""
|
||||
mock_pg_db.vector_index_type = "HNSW_HALFVEC"
|
||||
mock_pg_db.check_table_exists = AsyncMock(return_value=True)
|
||||
|
||||
config = {
|
||||
"embedding_batch_num": 10,
|
||||
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
||||
}
|
||||
storage = PGVectorStorage(
|
||||
namespace=namespace,
|
||||
global_config=config,
|
||||
embedding_func=mock_embedding_func,
|
||||
workspace="test_ws",
|
||||
)
|
||||
await storage.initialize()
|
||||
|
||||
query_embedding = [0.1] * 768
|
||||
await storage.query("test query", top_k=5, query_embedding=query_embedding)
|
||||
|
||||
assert mock_pg_db.query.called
|
||||
call_args = mock_pg_db.query.call_args
|
||||
sql = call_args[0][0]
|
||||
assert "::halfvec" in sql
|
||||
assert "::vector" not in sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("namespace", QUERY_NAMESPACES)
|
||||
async def test_query_uses_vector_cast_when_hnsw_default(
|
||||
mock_client_manager, mock_pg_db, mock_embedding_func, namespace
|
||||
):
|
||||
"""When HNSW (default) is set, generated search SQL uses ::vector (not ::halfvec)."""
|
||||
mock_pg_db.vector_index_type = "HNSW"
|
||||
mock_pg_db.check_table_exists = AsyncMock(return_value=True)
|
||||
|
||||
config = {
|
||||
"embedding_batch_num": 10,
|
||||
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
||||
}
|
||||
storage = PGVectorStorage(
|
||||
namespace=namespace,
|
||||
global_config=config,
|
||||
embedding_func=mock_embedding_func,
|
||||
workspace="test_ws",
|
||||
)
|
||||
await storage.initialize()
|
||||
|
||||
query_embedding = [0.1] * 768
|
||||
await storage.query("test query", top_k=5, query_embedding=query_embedding)
|
||||
|
||||
assert mock_pg_db.query.called
|
||||
call_args = mock_pg_db.query.call_args
|
||||
sql = call_args[0][0]
|
||||
assert "::vector" in sql
|
||||
assert "::halfvec" not in sql
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Index switching: old conflicting indexes are dropped
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_vector_index_drops_old_indexes_when_switching(mock_pg_db):
|
||||
"""Switching from HNSW to HNSW_HALFVEC drops the old hnsw_cosine index."""
|
||||
mock_pg_db.vector_index_type = "HNSW_HALFVEC"
|
||||
mock_pg_db.hnsw_m = 16
|
||||
mock_pg_db.hnsw_ef = 64
|
||||
mock_pg_db.ivfflat_lists = 100
|
||||
mock_pg_db.vchordrq_build_options = ""
|
||||
|
||||
table_name = "lightrag_vdb_chunks_test"
|
||||
|
||||
async def mock_query(sql, params=None, multirows=False, **kwargs):
|
||||
if "pg_indexes" in sql:
|
||||
return None
|
||||
return None
|
||||
|
||||
mock_pg_db.query = AsyncMock(side_effect=mock_query)
|
||||
mock_pg_db.execute = AsyncMock()
|
||||
|
||||
# Call the real method with mock_pg_db as self
|
||||
await PostgreSQLDB._create_vector_index(mock_pg_db, table_name, 3072)
|
||||
|
||||
execute_calls = [call[0][0] for call in mock_pg_db.execute.call_args_list]
|
||||
|
||||
old_hnsw_name = _safe_index_name(table_name, "hnsw_cosine")
|
||||
old_ivfflat_name = _safe_index_name(table_name, "ivfflat_cosine")
|
||||
old_vchordrq_name = _safe_index_name(table_name, "vchordrq_cosine")
|
||||
|
||||
drop_calls = [c for c in execute_calls if "DROP INDEX IF EXISTS" in c]
|
||||
dropped_names = {c.split("DROP INDEX IF EXISTS ")[1].strip() for c in drop_calls}
|
||||
assert old_hnsw_name in dropped_names
|
||||
assert old_ivfflat_name in dropped_names
|
||||
assert old_vchordrq_name in dropped_names
|
||||
|
||||
new_index_name = _safe_index_name(table_name, "hnsw_halfvec_cosine")
|
||||
assert new_index_name not in dropped_names
|
||||
|
||||
alter_calls = [c for c in execute_calls if "ALTER TABLE" in c]
|
||||
assert any("HALFVEC(3072)" in c for c in alter_calls)
|
||||
|
||||
create_calls = [c for c in execute_calls if "CREATE INDEX" in c]
|
||||
assert any("halfvec_cosine_ops" in c for c in create_calls)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_vector_index_no_drop_when_index_exists(mock_pg_db):
|
||||
"""If the target index already exists, no DROP or CREATE is issued."""
|
||||
mock_pg_db.vector_index_type = "HNSW_HALFVEC"
|
||||
mock_pg_db.hnsw_m = 16
|
||||
mock_pg_db.hnsw_ef = 64
|
||||
mock_pg_db.ivfflat_lists = 100
|
||||
mock_pg_db.vchordrq_build_options = ""
|
||||
|
||||
table_name = "lightrag_vdb_chunks_test"
|
||||
|
||||
async def mock_query(sql, params=None, multirows=False, **kwargs):
|
||||
if "pg_indexes" in sql:
|
||||
return {"?column?": 1}
|
||||
return None
|
||||
|
||||
mock_pg_db.query = AsyncMock(side_effect=mock_query)
|
||||
mock_pg_db.execute = AsyncMock()
|
||||
|
||||
await PostgreSQLDB._create_vector_index(mock_pg_db, table_name, 3072)
|
||||
|
||||
execute_calls = [call[0][0] for call in mock_pg_db.execute.call_args_list]
|
||||
assert not any("DROP INDEX" in c for c in execute_calls)
|
||||
assert not any("CREATE INDEX" in c for c in execute_calls)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HalfVector dimension detection in setup_table
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _MockHalfVector:
|
||||
"""Mimics pgvector.halfvec.HalfVector for testing dimension detection."""
|
||||
|
||||
def __init__(self, dim: int):
|
||||
self._dim = dim
|
||||
|
||||
def dimensions(self) -> int:
|
||||
return self._dim
|
||||
|
||||
def to_list(self):
|
||||
return [0.0] * self._dim
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_setup_table_detects_halfvector_dimension_mismatch(mock_pg_db):
|
||||
"""DataMigrationError is raised when a HalfVector column has a different dimension."""
|
||||
table_name = "lightrag_vdb_chunks_new"
|
||||
legacy_table = "lightrag_vdb_chunks"
|
||||
|
||||
mock_pg_db.check_table_exists = AsyncMock(
|
||||
side_effect=lambda t: t.lower() == legacy_table.lower()
|
||||
)
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def mock_query(sql, params=None, multirows=False, **kwargs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if "COUNT(*)" in sql:
|
||||
return {"count": 5}
|
||||
if "content_vector" in sql:
|
||||
return {"content_vector": _MockHalfVector(1024)}
|
||||
return None
|
||||
|
||||
mock_pg_db.query = AsyncMock(side_effect=mock_query)
|
||||
mock_pg_db.execute = AsyncMock()
|
||||
|
||||
with pytest.raises(DataMigrationError, match="Dimension mismatch"):
|
||||
await PGVectorStorage.setup_table(
|
||||
db=mock_pg_db,
|
||||
table_name=table_name,
|
||||
workspace="test_ws",
|
||||
embedding_dim=768,
|
||||
legacy_table_name=legacy_table,
|
||||
base_table=legacy_table,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_setup_table_accepts_matching_halfvector_dimension(mock_pg_db):
|
||||
"""No error when HalfVector dimension matches the expected embedding_dim."""
|
||||
table_name = "lightrag_vdb_chunks_new"
|
||||
legacy_table = "lightrag_vdb_chunks"
|
||||
|
||||
mock_pg_db.check_table_exists = AsyncMock(
|
||||
side_effect=lambda t: t.lower() == legacy_table.lower()
|
||||
)
|
||||
mock_pg_db.vector_index_type = "HNSW_HALFVEC"
|
||||
|
||||
async def mock_query(sql, params=None, multirows=False, **kwargs):
|
||||
if "COUNT(*)" in sql:
|
||||
return {"count": 5}
|
||||
if "content_vector" in sql:
|
||||
return {"content_vector": _MockHalfVector(768)}
|
||||
if multirows:
|
||||
return []
|
||||
return None
|
||||
|
||||
mock_pg_db.query = AsyncMock(side_effect=mock_query)
|
||||
mock_pg_db.execute = AsyncMock()
|
||||
|
||||
with patch.object(PGVectorStorage, "_pg_create_table", new_callable=AsyncMock):
|
||||
await PGVectorStorage.setup_table(
|
||||
db=mock_pg_db,
|
||||
table_name=table_name,
|
||||
workspace="test_ws",
|
||||
embedding_dim=768,
|
||||
legacy_table_name=legacy_table,
|
||||
base_table=legacy_table,
|
||||
)
|
||||
@@ -0,0 +1,210 @@
|
||||
"""
|
||||
Unit tests for PostgreSQL safe index name generation.
|
||||
|
||||
This module tests the _safe_index_name helper function which prevents
|
||||
PostgreSQL's silent 63-byte identifier truncation from causing index
|
||||
lookup failures.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
# Mark all tests as offline (no external dependencies)
|
||||
pytestmark = pytest.mark.offline
|
||||
|
||||
|
||||
class TestSafeIndexName:
|
||||
"""Test suite for _safe_index_name function."""
|
||||
|
||||
def test_short_name_unchanged(self):
|
||||
"""Short index names should remain unchanged."""
|
||||
from lightrag.kg.postgres_impl import _safe_index_name
|
||||
|
||||
# Short table name - should return unchanged
|
||||
result = _safe_index_name("lightrag_vdb_entity", "hnsw_cosine")
|
||||
assert result == "idx_lightrag_vdb_entity_hnsw_cosine"
|
||||
assert len(result.encode("utf-8")) <= 63
|
||||
|
||||
def test_long_name_gets_hashed(self):
|
||||
"""Long table names exceeding 63 bytes should get hashed."""
|
||||
from lightrag.kg.postgres_impl import _safe_index_name
|
||||
|
||||
# Long table name that would exceed 63 bytes
|
||||
long_table_name = "LIGHTRAG_VDB_ENTITY_text_embedding_3_large_3072d"
|
||||
result = _safe_index_name(long_table_name, "hnsw_cosine")
|
||||
|
||||
# Should be within 63 bytes
|
||||
assert len(result.encode("utf-8")) <= 63
|
||||
|
||||
# Should start with idx_ prefix
|
||||
assert result.startswith("idx_")
|
||||
|
||||
# Should contain the suffix
|
||||
assert result.endswith("_hnsw_cosine")
|
||||
|
||||
# Should NOT be the naive concatenation (which would be truncated)
|
||||
naive_name = f"idx_{long_table_name.lower()}_hnsw_cosine"
|
||||
assert result != naive_name
|
||||
|
||||
def test_deterministic_output(self):
|
||||
"""Same input should always produce same output (deterministic)."""
|
||||
from lightrag.kg.postgres_impl import _safe_index_name
|
||||
|
||||
table_name = "LIGHTRAG_VDB_CHUNKS_text_embedding_3_large_3072d"
|
||||
suffix = "hnsw_cosine"
|
||||
|
||||
result1 = _safe_index_name(table_name, suffix)
|
||||
result2 = _safe_index_name(table_name, suffix)
|
||||
|
||||
assert result1 == result2
|
||||
|
||||
def test_different_suffixes_different_results(self):
|
||||
"""Different suffixes should produce different index names."""
|
||||
from lightrag.kg.postgres_impl import _safe_index_name
|
||||
|
||||
table_name = "LIGHTRAG_VDB_ENTITY_text_embedding_3_large_3072d"
|
||||
|
||||
result1 = _safe_index_name(table_name, "hnsw_cosine")
|
||||
result2 = _safe_index_name(table_name, "ivfflat_cosine")
|
||||
|
||||
assert result1 != result2
|
||||
|
||||
def test_case_insensitive(self):
|
||||
"""Table names should be normalized to lowercase."""
|
||||
from lightrag.kg.postgres_impl import _safe_index_name
|
||||
|
||||
result_upper = _safe_index_name("LIGHTRAG_VDB_ENTITY", "hnsw_cosine")
|
||||
result_lower = _safe_index_name("lightrag_vdb_entity", "hnsw_cosine")
|
||||
|
||||
assert result_upper == result_lower
|
||||
|
||||
def test_boundary_case_exactly_63_bytes(self):
|
||||
"""Test boundary case where name is exactly at 63-byte limit."""
|
||||
from lightrag.kg.postgres_impl import _safe_index_name
|
||||
|
||||
# Create a table name that results in exactly 63 bytes
|
||||
# idx_ (4) + table_name + _ (1) + suffix = 63
|
||||
# So table_name + suffix = 58
|
||||
|
||||
# Test a name that's just under the limit (should remain unchanged)
|
||||
short_suffix = "id"
|
||||
# idx_ (4) + 56 chars + _ (1) + id (2) = 63
|
||||
table_56 = "a" * 56
|
||||
result = _safe_index_name(table_56, short_suffix)
|
||||
expected = f"idx_{table_56}_{short_suffix}"
|
||||
assert result == expected
|
||||
assert len(result.encode("utf-8")) == 63
|
||||
|
||||
def test_unicode_handling(self):
|
||||
"""Unicode characters should be properly handled (bytes, not chars)."""
|
||||
from lightrag.kg.postgres_impl import _safe_index_name
|
||||
|
||||
# Unicode characters can take more bytes than visible chars
|
||||
# Chinese characters are 3 bytes each in UTF-8
|
||||
table_name = "lightrag_测试_table" # Contains Chinese chars
|
||||
result = _safe_index_name(table_name, "hnsw_cosine")
|
||||
|
||||
# Should always be within 63 bytes
|
||||
assert len(result.encode("utf-8")) <= 63
|
||||
|
||||
def test_real_world_model_names(self):
|
||||
"""Test with real-world embedding model names that cause issues."""
|
||||
from lightrag.kg.postgres_impl import _safe_index_name
|
||||
|
||||
# These are actual model names that have caused issues
|
||||
test_cases = [
|
||||
("LIGHTRAG_VDB_CHUNKS_text_embedding_3_large_3072d", "hnsw_cosine"),
|
||||
("LIGHTRAG_VDB_ENTITY_text_embedding_3_large_3072d", "hnsw_cosine"),
|
||||
("LIGHTRAG_VDB_RELATION_text_embedding_3_large_3072d", "hnsw_cosine"),
|
||||
(
|
||||
"LIGHTRAG_VDB_ENTITY_bge_m3_1024d",
|
||||
"hnsw_cosine",
|
||||
), # Shorter model name
|
||||
(
|
||||
"LIGHTRAG_VDB_CHUNKS_nomic_embed_text_v1_768d",
|
||||
"ivfflat_cosine",
|
||||
), # Different index type
|
||||
]
|
||||
|
||||
for table_name, suffix in test_cases:
|
||||
result = _safe_index_name(table_name, suffix)
|
||||
|
||||
# Critical: must be within PostgreSQL's 63-byte limit
|
||||
assert len(result.encode("utf-8")) <= 63, (
|
||||
f"Index name too long: {result} for table {table_name}"
|
||||
)
|
||||
|
||||
# Must have consistent format
|
||||
assert result.startswith("idx_"), f"Missing idx_ prefix: {result}"
|
||||
assert result.endswith(f"_{suffix}"), f"Missing suffix {suffix}: {result}"
|
||||
|
||||
def test_hash_uniqueness_for_similar_tables(self):
|
||||
"""Similar but different table names should produce different hashes."""
|
||||
from lightrag.kg.postgres_impl import _safe_index_name
|
||||
|
||||
# These tables have similar names but should have different hashes
|
||||
tables = [
|
||||
"LIGHTRAG_VDB_CHUNKS_model_a_1024d",
|
||||
"LIGHTRAG_VDB_CHUNKS_model_b_1024d",
|
||||
"LIGHTRAG_VDB_ENTITY_model_a_1024d",
|
||||
]
|
||||
|
||||
results = [_safe_index_name(t, "hnsw_cosine") for t in tables]
|
||||
|
||||
# All results should be unique
|
||||
assert len(set(results)) == len(results), "Hash collision detected!"
|
||||
|
||||
|
||||
class TestIndexNameIntegration:
|
||||
"""Integration-style tests for index name usage patterns."""
|
||||
|
||||
def test_pg_indexes_lookup_compatibility(self):
|
||||
"""
|
||||
Test that the generated index name will work with pg_indexes lookup.
|
||||
|
||||
This is the core problem: PostgreSQL stores the truncated name,
|
||||
but we were looking up the untruncated name. Our fix ensures we
|
||||
always use a name that fits within 63 bytes.
|
||||
"""
|
||||
from lightrag.kg.postgres_impl import _safe_index_name
|
||||
|
||||
table_name = "LIGHTRAG_VDB_CHUNKS_text_embedding_3_large_3072d"
|
||||
suffix = "hnsw_cosine"
|
||||
|
||||
# Generate the index name
|
||||
index_name = _safe_index_name(table_name, suffix)
|
||||
|
||||
# Simulate what PostgreSQL would store (truncate at 63 bytes)
|
||||
stored_name = index_name.encode("utf-8")[:63].decode("utf-8", errors="ignore")
|
||||
|
||||
# The key fix: our generated name should equal the stored name
|
||||
# because it's already within the 63-byte limit
|
||||
assert index_name == stored_name, (
|
||||
"Index name would be truncated by PostgreSQL, causing lookup failures!"
|
||||
)
|
||||
|
||||
def test_backward_compatibility_short_names(self):
|
||||
"""
|
||||
Ensure backward compatibility with existing short index names.
|
||||
|
||||
For tables that have existing indexes with short names (pre-model-suffix era),
|
||||
the function should not change their names.
|
||||
"""
|
||||
from lightrag.kg.postgres_impl import _safe_index_name
|
||||
|
||||
# Legacy table names without model suffix
|
||||
legacy_tables = [
|
||||
"LIGHTRAG_VDB_ENTITY",
|
||||
"LIGHTRAG_VDB_RELATION",
|
||||
"LIGHTRAG_VDB_CHUNKS",
|
||||
]
|
||||
|
||||
for table in legacy_tables:
|
||||
for suffix in ["hnsw_cosine", "ivfflat_cosine", "id"]:
|
||||
result = _safe_index_name(table, suffix)
|
||||
expected = f"idx_{table.lower()}_{suffix}"
|
||||
|
||||
# Short names should remain unchanged for backward compatibility
|
||||
if len(expected.encode("utf-8")) <= 63:
|
||||
assert result == expected, (
|
||||
f"Short name changed unexpectedly: {result} != {expected}"
|
||||
)
|
||||
@@ -0,0 +1,856 @@
|
||||
import pytest
|
||||
from unittest.mock import patch, AsyncMock
|
||||
import numpy as np
|
||||
from lightrag.utils import EmbeddingFunc
|
||||
from lightrag.kg.postgres_impl import (
|
||||
PGVectorStorage,
|
||||
)
|
||||
from lightrag.namespace import NameSpace
|
||||
|
||||
|
||||
# Mock PostgreSQLDB
|
||||
@pytest.fixture
|
||||
def mock_pg_db():
|
||||
"""Mock PostgreSQL database connection"""
|
||||
db = AsyncMock()
|
||||
db.workspace = "test_workspace"
|
||||
|
||||
# Mock query responses with multirows support
|
||||
async def mock_query(sql, params=None, multirows=False, **kwargs):
|
||||
# Default return value
|
||||
if multirows:
|
||||
return [] # Return empty list for multirows
|
||||
return {"exists": False, "count": 0}
|
||||
|
||||
# Mock for execute that mimics PostgreSQLDB.execute() behavior
|
||||
async def mock_execute(sql, data=None, **kwargs):
|
||||
"""
|
||||
Mock that mimics PostgreSQLDB.execute() behavior:
|
||||
- Accepts data as dict[str, Any] | None (second parameter)
|
||||
- Internally converts dict.values() to tuple for AsyncPG
|
||||
"""
|
||||
# Mimic real execute() which accepts dict and converts to tuple
|
||||
if data is not None and not isinstance(data, dict):
|
||||
raise TypeError(
|
||||
f"PostgreSQLDB.execute() expects data as dict, got {type(data).__name__}"
|
||||
)
|
||||
return None
|
||||
|
||||
db.query = AsyncMock(side_effect=mock_query)
|
||||
db.execute = AsyncMock(side_effect=mock_execute)
|
||||
|
||||
return db
|
||||
|
||||
|
||||
# Mock get_data_init_lock to avoid async lock issues in tests
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_data_init_lock():
|
||||
with patch("lightrag.kg.postgres_impl.get_data_init_lock") as mock_lock:
|
||||
mock_lock_ctx = AsyncMock()
|
||||
mock_lock.return_value = mock_lock_ctx
|
||||
yield mock_lock
|
||||
|
||||
|
||||
# Mock ClientManager
|
||||
@pytest.fixture
|
||||
def mock_client_manager(mock_pg_db):
|
||||
with patch("lightrag.kg.postgres_impl.ClientManager") as mock_manager:
|
||||
mock_manager.get_client = AsyncMock(return_value=mock_pg_db)
|
||||
mock_manager.release_client = AsyncMock()
|
||||
yield mock_manager
|
||||
|
||||
|
||||
# Mock Embedding function
|
||||
@pytest.fixture
|
||||
def mock_embedding_func():
|
||||
async def embed_func(texts, **kwargs):
|
||||
return np.array([[0.1] * 768 for _ in texts])
|
||||
|
||||
func = EmbeddingFunc(embedding_dim=768, func=embed_func, model_name="test_model")
|
||||
return func
|
||||
|
||||
|
||||
async def test_postgres_table_naming(
|
||||
mock_client_manager, mock_pg_db, mock_embedding_func
|
||||
):
|
||||
"""Test if table name is correctly generated with model suffix"""
|
||||
config = {
|
||||
"embedding_batch_num": 10,
|
||||
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
||||
}
|
||||
|
||||
storage = PGVectorStorage(
|
||||
namespace=NameSpace.VECTOR_STORE_CHUNKS,
|
||||
global_config=config,
|
||||
embedding_func=mock_embedding_func,
|
||||
workspace="test_ws",
|
||||
)
|
||||
|
||||
# Verify table name contains model suffix
|
||||
expected_suffix = "test_model_768d"
|
||||
assert expected_suffix in storage.table_name
|
||||
assert storage.table_name == f"LIGHTRAG_VDB_CHUNKS_{expected_suffix}"
|
||||
|
||||
# Verify legacy table name
|
||||
assert storage.legacy_table_name == "LIGHTRAG_VDB_CHUNKS"
|
||||
|
||||
|
||||
async def test_postgres_migration_trigger(
|
||||
mock_client_manager, mock_pg_db, mock_embedding_func
|
||||
):
|
||||
"""Test if migration logic is triggered correctly"""
|
||||
config = {
|
||||
"embedding_batch_num": 10,
|
||||
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
||||
}
|
||||
|
||||
storage = PGVectorStorage(
|
||||
namespace=NameSpace.VECTOR_STORE_CHUNKS,
|
||||
global_config=config,
|
||||
embedding_func=mock_embedding_func,
|
||||
workspace="test_ws",
|
||||
)
|
||||
|
||||
# Setup mocks for migration scenario
|
||||
# 1. New table does not exist, legacy table exists
|
||||
async def mock_check_table_exists(table_name):
|
||||
return table_name == storage.legacy_table_name
|
||||
|
||||
mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists)
|
||||
|
||||
# 2. Legacy table has 100 records
|
||||
mock_rows = [
|
||||
{"id": f"test_id_{i}", "content": f"content_{i}", "workspace": "test_ws"}
|
||||
for i in range(100)
|
||||
]
|
||||
migration_state = {"new_table_count": 0}
|
||||
|
||||
async def mock_query(sql, params=None, multirows=False, **kwargs):
|
||||
if "COUNT(*)" in sql:
|
||||
sql_upper = sql.upper()
|
||||
legacy_table = storage.legacy_table_name.upper()
|
||||
new_table = storage.table_name.upper()
|
||||
is_new_table = new_table in sql_upper
|
||||
is_legacy_table = legacy_table in sql_upper and not is_new_table
|
||||
|
||||
if is_new_table:
|
||||
return {"count": migration_state["new_table_count"]}
|
||||
if is_legacy_table:
|
||||
return {"count": 100}
|
||||
return {"count": 0}
|
||||
elif multirows and "SELECT *" in sql:
|
||||
# Mock batch fetch for migration using keyset pagination
|
||||
# New pattern: WHERE workspace = $1 AND id > $2 ORDER BY id LIMIT $3
|
||||
# or first batch: WHERE workspace = $1 ORDER BY id LIMIT $2
|
||||
if "WHERE workspace" in sql:
|
||||
if "id >" in sql:
|
||||
# Keyset pagination: params = [workspace, last_id, limit]
|
||||
last_id = params[1] if len(params) > 1 else None
|
||||
# Find rows after last_id
|
||||
start_idx = 0
|
||||
for i, row in enumerate(mock_rows):
|
||||
if row["id"] == last_id:
|
||||
start_idx = i + 1
|
||||
break
|
||||
limit = params[2] if len(params) > 2 else 500
|
||||
else:
|
||||
# First batch (no last_id): params = [workspace, limit]
|
||||
start_idx = 0
|
||||
limit = params[1] if len(params) > 1 else 500
|
||||
else:
|
||||
# No workspace filter with keyset
|
||||
if "id >" in sql:
|
||||
last_id = params[0] if params else None
|
||||
start_idx = 0
|
||||
for i, row in enumerate(mock_rows):
|
||||
if row["id"] == last_id:
|
||||
start_idx = i + 1
|
||||
break
|
||||
limit = params[1] if len(params) > 1 else 500
|
||||
else:
|
||||
start_idx = 0
|
||||
limit = params[0] if params else 500
|
||||
end = min(start_idx + limit, len(mock_rows))
|
||||
return mock_rows[start_idx:end]
|
||||
return {}
|
||||
|
||||
mock_pg_db.query = AsyncMock(side_effect=mock_query)
|
||||
|
||||
# Track migration through _run_with_retry calls
|
||||
migration_executed = []
|
||||
|
||||
async def mock_run_with_retry(operation, **kwargs):
|
||||
# Track that migration batch operation was called
|
||||
migration_executed.append(True)
|
||||
migration_state["new_table_count"] = 100
|
||||
return None
|
||||
|
||||
mock_pg_db._run_with_retry = AsyncMock(side_effect=mock_run_with_retry)
|
||||
|
||||
with patch(
|
||||
"lightrag.kg.postgres_impl.PGVectorStorage._pg_create_table", AsyncMock()
|
||||
):
|
||||
# Initialize storage (should trigger migration)
|
||||
await storage.initialize()
|
||||
|
||||
# Verify migration was executed by checking _run_with_retry was called
|
||||
# (batch migration uses _run_with_retry with executemany)
|
||||
assert len(migration_executed) > 0, "Migration should have been executed"
|
||||
|
||||
|
||||
async def test_postgres_no_migration_needed(
|
||||
mock_client_manager, mock_pg_db, mock_embedding_func
|
||||
):
|
||||
"""Test scenario where new table already exists (no migration needed)"""
|
||||
config = {
|
||||
"embedding_batch_num": 10,
|
||||
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
||||
}
|
||||
|
||||
storage = PGVectorStorage(
|
||||
namespace=NameSpace.VECTOR_STORE_CHUNKS,
|
||||
global_config=config,
|
||||
embedding_func=mock_embedding_func,
|
||||
workspace="test_ws",
|
||||
)
|
||||
|
||||
# Mock: new table already exists
|
||||
async def mock_check_table_exists(table_name):
|
||||
return table_name == storage.table_name
|
||||
|
||||
mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists)
|
||||
|
||||
with patch(
|
||||
"lightrag.kg.postgres_impl.PGVectorStorage._pg_create_table", AsyncMock()
|
||||
) as mock_create:
|
||||
await storage.initialize()
|
||||
|
||||
# Verify no table creation was attempted
|
||||
mock_create.assert_not_called()
|
||||
|
||||
|
||||
async def test_scenario_1_new_workspace_creation(
|
||||
mock_client_manager, mock_pg_db, mock_embedding_func
|
||||
):
|
||||
"""
|
||||
Scenario 1: New workspace creation
|
||||
|
||||
Expected behavior:
|
||||
- No legacy table exists
|
||||
- Directly create new table with model suffix
|
||||
- No migration needed
|
||||
"""
|
||||
config = {
|
||||
"embedding_batch_num": 10,
|
||||
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
||||
}
|
||||
|
||||
embedding_func = EmbeddingFunc(
|
||||
embedding_dim=3072,
|
||||
func=mock_embedding_func.func,
|
||||
model_name="text-embedding-3-large",
|
||||
)
|
||||
|
||||
storage = PGVectorStorage(
|
||||
namespace=NameSpace.VECTOR_STORE_CHUNKS,
|
||||
global_config=config,
|
||||
embedding_func=embedding_func,
|
||||
workspace="new_workspace",
|
||||
)
|
||||
|
||||
# Mock: neither table exists
|
||||
async def mock_check_table_exists(table_name):
|
||||
return False
|
||||
|
||||
mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists)
|
||||
|
||||
with patch(
|
||||
"lightrag.kg.postgres_impl.PGVectorStorage._pg_create_table", AsyncMock()
|
||||
) as mock_create:
|
||||
await storage.initialize()
|
||||
|
||||
# Verify table name format
|
||||
assert "text_embedding_3_large_3072d" in storage.table_name
|
||||
|
||||
# Verify new table creation was called
|
||||
mock_create.assert_called_once()
|
||||
call_args = mock_create.call_args
|
||||
assert (
|
||||
call_args[0][1] == storage.table_name
|
||||
) # table_name is second positional arg
|
||||
|
||||
|
||||
async def test_scenario_2_legacy_upgrade_migration(
|
||||
mock_client_manager, mock_pg_db, mock_embedding_func
|
||||
):
|
||||
"""
|
||||
Scenario 2: Upgrade from legacy version
|
||||
|
||||
Expected behavior:
|
||||
- Legacy table exists (without model suffix)
|
||||
- New table doesn't exist
|
||||
- Automatically migrate data to new table with suffix
|
||||
"""
|
||||
config = {
|
||||
"embedding_batch_num": 10,
|
||||
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
||||
}
|
||||
|
||||
embedding_func = EmbeddingFunc(
|
||||
embedding_dim=1536,
|
||||
func=mock_embedding_func.func,
|
||||
model_name="text-embedding-ada-002",
|
||||
)
|
||||
|
||||
storage = PGVectorStorage(
|
||||
namespace=NameSpace.VECTOR_STORE_CHUNKS,
|
||||
global_config=config,
|
||||
embedding_func=embedding_func,
|
||||
workspace="legacy_workspace",
|
||||
)
|
||||
|
||||
# Mock: only legacy table exists
|
||||
async def mock_check_table_exists(table_name):
|
||||
return table_name == storage.legacy_table_name
|
||||
|
||||
mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists)
|
||||
|
||||
# Mock: legacy table has 50 records
|
||||
mock_rows = [
|
||||
{
|
||||
"id": f"legacy_id_{i}",
|
||||
"content": f"legacy_content_{i}",
|
||||
"workspace": "legacy_workspace",
|
||||
}
|
||||
for i in range(50)
|
||||
]
|
||||
|
||||
# Track which queries have been made for proper response
|
||||
query_history = []
|
||||
migration_state = {"new_table_count": 0}
|
||||
|
||||
async def mock_query(sql, params=None, multirows=False, **kwargs):
|
||||
query_history.append(sql)
|
||||
|
||||
if "COUNT(*)" in sql:
|
||||
# Determine table type:
|
||||
# - Legacy: contains base name but NOT model suffix
|
||||
# - New: contains model suffix (e.g., text_embedding_ada_002_1536d)
|
||||
sql_upper = sql.upper()
|
||||
base_name = storage.legacy_table_name.upper()
|
||||
|
||||
# Check if this is querying the new table (has model suffix)
|
||||
has_model_suffix = storage.table_name.upper() in sql_upper
|
||||
|
||||
is_legacy_table = base_name in sql_upper and not has_model_suffix
|
||||
has_workspace_filter = "WHERE workspace" in sql
|
||||
|
||||
if is_legacy_table and has_workspace_filter:
|
||||
# Count for legacy table with workspace filter (before migration)
|
||||
return {"count": 50}
|
||||
elif is_legacy_table and not has_workspace_filter:
|
||||
# Total count for legacy table
|
||||
return {"count": 50}
|
||||
else:
|
||||
# New table count (before/after migration)
|
||||
return {"count": migration_state["new_table_count"]}
|
||||
elif multirows and "SELECT *" in sql:
|
||||
# Mock batch fetch for migration using keyset pagination
|
||||
# New pattern: WHERE workspace = $1 AND id > $2 ORDER BY id LIMIT $3
|
||||
# or first batch: WHERE workspace = $1 ORDER BY id LIMIT $2
|
||||
if "WHERE workspace" in sql:
|
||||
if "id >" in sql:
|
||||
# Keyset pagination: params = [workspace, last_id, limit]
|
||||
last_id = params[1] if len(params) > 1 else None
|
||||
# Find rows after last_id
|
||||
start_idx = 0
|
||||
for i, row in enumerate(mock_rows):
|
||||
if row["id"] == last_id:
|
||||
start_idx = i + 1
|
||||
break
|
||||
limit = params[2] if len(params) > 2 else 500
|
||||
else:
|
||||
# First batch (no last_id): params = [workspace, limit]
|
||||
start_idx = 0
|
||||
limit = params[1] if len(params) > 1 else 500
|
||||
else:
|
||||
# No workspace filter with keyset
|
||||
if "id >" in sql:
|
||||
last_id = params[0] if params else None
|
||||
start_idx = 0
|
||||
for i, row in enumerate(mock_rows):
|
||||
if row["id"] == last_id:
|
||||
start_idx = i + 1
|
||||
break
|
||||
limit = params[1] if len(params) > 1 else 500
|
||||
else:
|
||||
start_idx = 0
|
||||
limit = params[0] if params else 500
|
||||
end = min(start_idx + limit, len(mock_rows))
|
||||
return mock_rows[start_idx:end]
|
||||
return {}
|
||||
|
||||
mock_pg_db.query = AsyncMock(side_effect=mock_query)
|
||||
|
||||
# Track migration through _run_with_retry calls
|
||||
migration_executed = []
|
||||
|
||||
async def mock_run_with_retry(operation, **kwargs):
|
||||
# Track that migration batch operation was called
|
||||
migration_executed.append(True)
|
||||
migration_state["new_table_count"] = 50
|
||||
return None
|
||||
|
||||
mock_pg_db._run_with_retry = AsyncMock(side_effect=mock_run_with_retry)
|
||||
|
||||
with patch(
|
||||
"lightrag.kg.postgres_impl.PGVectorStorage._pg_create_table", AsyncMock()
|
||||
) as mock_create:
|
||||
await storage.initialize()
|
||||
|
||||
# Verify table name contains ada-002
|
||||
assert "text_embedding_ada_002_1536d" in storage.table_name
|
||||
|
||||
# Verify migration was executed (batch migration uses _run_with_retry)
|
||||
assert len(migration_executed) > 0, "Migration should have been executed"
|
||||
mock_create.assert_called_once()
|
||||
|
||||
|
||||
async def test_scenario_3_multi_model_coexistence(
|
||||
mock_client_manager, mock_pg_db, mock_embedding_func
|
||||
):
|
||||
"""
|
||||
Scenario 3: Multiple embedding models coexist
|
||||
|
||||
Expected behavior:
|
||||
- Different embedding models create separate tables
|
||||
- Tables are isolated by model suffix
|
||||
- No interference between different models
|
||||
"""
|
||||
config = {
|
||||
"embedding_batch_num": 10,
|
||||
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
||||
}
|
||||
|
||||
# Workspace A: uses bge-small (768d)
|
||||
embedding_func_a = EmbeddingFunc(
|
||||
embedding_dim=768, func=mock_embedding_func.func, model_name="bge-small"
|
||||
)
|
||||
|
||||
storage_a = PGVectorStorage(
|
||||
namespace=NameSpace.VECTOR_STORE_CHUNKS,
|
||||
global_config=config,
|
||||
embedding_func=embedding_func_a,
|
||||
workspace="workspace_a",
|
||||
)
|
||||
|
||||
# Workspace B: uses bge-large (1024d)
|
||||
async def embed_func_b(texts, **kwargs):
|
||||
return np.array([[0.1] * 1024 for _ in texts])
|
||||
|
||||
embedding_func_b = EmbeddingFunc(
|
||||
embedding_dim=1024, func=embed_func_b, model_name="bge-large"
|
||||
)
|
||||
|
||||
storage_b = PGVectorStorage(
|
||||
namespace=NameSpace.VECTOR_STORE_CHUNKS,
|
||||
global_config=config,
|
||||
embedding_func=embedding_func_b,
|
||||
workspace="workspace_b",
|
||||
)
|
||||
|
||||
# Verify different table names
|
||||
assert storage_a.table_name != storage_b.table_name
|
||||
assert "bge_small_768d" in storage_a.table_name
|
||||
assert "bge_large_1024d" in storage_b.table_name
|
||||
|
||||
# Mock: both tables don't exist yet
|
||||
async def mock_check_table_exists(table_name):
|
||||
return False
|
||||
|
||||
mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists)
|
||||
|
||||
with patch(
|
||||
"lightrag.kg.postgres_impl.PGVectorStorage._pg_create_table", AsyncMock()
|
||||
) as mock_create:
|
||||
# Initialize both storages
|
||||
await storage_a.initialize()
|
||||
await storage_b.initialize()
|
||||
|
||||
# Verify two separate tables were created
|
||||
assert mock_create.call_count == 2
|
||||
|
||||
# Verify table names are different
|
||||
call_args_list = mock_create.call_args_list
|
||||
table_names = [call[0][1] for call in call_args_list] # Second positional arg
|
||||
assert len(set(table_names)) == 2 # Two unique table names
|
||||
assert storage_a.table_name in table_names
|
||||
assert storage_b.table_name in table_names
|
||||
|
||||
|
||||
async def test_case1_empty_legacy_auto_cleanup(
|
||||
mock_client_manager, mock_pg_db, mock_embedding_func
|
||||
):
|
||||
"""
|
||||
Case 1a: Both new and legacy tables exist, but legacy is EMPTY
|
||||
Expected: Automatically delete empty legacy table (safe cleanup)
|
||||
"""
|
||||
config = {
|
||||
"embedding_batch_num": 10,
|
||||
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
||||
}
|
||||
|
||||
embedding_func = EmbeddingFunc(
|
||||
embedding_dim=1536,
|
||||
func=mock_embedding_func.func,
|
||||
model_name="test-model",
|
||||
)
|
||||
|
||||
storage = PGVectorStorage(
|
||||
namespace=NameSpace.VECTOR_STORE_CHUNKS,
|
||||
global_config=config,
|
||||
embedding_func=embedding_func,
|
||||
workspace="test_ws",
|
||||
)
|
||||
|
||||
# Mock: Both tables exist
|
||||
async def mock_check_table_exists(table_name):
|
||||
return True # Both new and legacy exist
|
||||
|
||||
mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists)
|
||||
|
||||
# Mock: Legacy table is empty (0 records)
|
||||
async def mock_query(sql, params=None, multirows=False, **kwargs):
|
||||
if "COUNT(*)" in sql:
|
||||
if storage.legacy_table_name in sql:
|
||||
return {"count": 0} # Empty legacy table
|
||||
else:
|
||||
return {"count": 100} # New table has data
|
||||
return {}
|
||||
|
||||
mock_pg_db.query = AsyncMock(side_effect=mock_query)
|
||||
|
||||
with patch("lightrag.kg.postgres_impl.logger"):
|
||||
await storage.initialize()
|
||||
|
||||
# Verify: Empty legacy table should be automatically cleaned up
|
||||
# Empty tables are safe to delete without data loss risk
|
||||
delete_calls = [
|
||||
call
|
||||
for call in mock_pg_db.execute.call_args_list
|
||||
if call[0][0] and "DROP TABLE" in call[0][0]
|
||||
]
|
||||
assert len(delete_calls) >= 1, "Empty legacy table should be auto-deleted"
|
||||
# Check if legacy table was dropped
|
||||
dropped_table = storage.legacy_table_name
|
||||
assert any(dropped_table in str(call) for call in delete_calls), (
|
||||
f"Expected to drop empty legacy table '{dropped_table}'"
|
||||
)
|
||||
|
||||
print(
|
||||
f"✅ Case 1a: Empty legacy table '{dropped_table}' auto-deleted successfully"
|
||||
)
|
||||
|
||||
|
||||
async def test_case1_nonempty_legacy_warning(
|
||||
mock_client_manager, mock_pg_db, mock_embedding_func
|
||||
):
|
||||
"""
|
||||
Case 1b: Both new and legacy tables exist, and legacy HAS DATA
|
||||
Expected: Log warning, do not delete legacy (preserve data)
|
||||
"""
|
||||
config = {
|
||||
"embedding_batch_num": 10,
|
||||
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
||||
}
|
||||
|
||||
embedding_func = EmbeddingFunc(
|
||||
embedding_dim=1536,
|
||||
func=mock_embedding_func.func,
|
||||
model_name="test-model",
|
||||
)
|
||||
|
||||
storage = PGVectorStorage(
|
||||
namespace=NameSpace.VECTOR_STORE_CHUNKS,
|
||||
global_config=config,
|
||||
embedding_func=embedding_func,
|
||||
workspace="test_ws",
|
||||
)
|
||||
|
||||
# Mock: Both tables exist
|
||||
async def mock_check_table_exists(table_name):
|
||||
return True # Both new and legacy exist
|
||||
|
||||
mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists)
|
||||
|
||||
# Mock: Legacy table has data (50 records)
|
||||
async def mock_query(sql, params=None, multirows=False, **kwargs):
|
||||
if "COUNT(*)" in sql:
|
||||
if storage.legacy_table_name in sql:
|
||||
return {"count": 50} # Legacy has data
|
||||
else:
|
||||
return {"count": 100} # New table has data
|
||||
return {}
|
||||
|
||||
mock_pg_db.query = AsyncMock(side_effect=mock_query)
|
||||
|
||||
with patch("lightrag.kg.postgres_impl.logger"):
|
||||
await storage.initialize()
|
||||
|
||||
# Verify: Legacy table with data should be preserved
|
||||
# We never auto-delete tables that contain data to prevent accidental data loss
|
||||
delete_calls = [
|
||||
call
|
||||
for call in mock_pg_db.execute.call_args_list
|
||||
if call[0][0] and "DROP TABLE" in call[0][0]
|
||||
]
|
||||
# Check if legacy table was deleted (it should not be)
|
||||
dropped_table = storage.legacy_table_name
|
||||
legacy_deleted = any(dropped_table in str(call) for call in delete_calls)
|
||||
assert not legacy_deleted, "Legacy table with data should NOT be auto-deleted"
|
||||
|
||||
print(
|
||||
f"✅ Case 1b: Legacy table '{dropped_table}' with data preserved (warning only)"
|
||||
)
|
||||
|
||||
|
||||
async def test_case1_sequential_workspace_migration(
|
||||
mock_client_manager, mock_pg_db, mock_embedding_func
|
||||
):
|
||||
"""
|
||||
Case 1c: Sequential workspace migration (Multi-tenant scenario)
|
||||
|
||||
Critical bug fix verification:
|
||||
Timeline:
|
||||
1. Legacy table has workspace_a (3 records) + workspace_b (3 records)
|
||||
2. Workspace A initializes first → Case 3 (only legacy exists) → migrates A's data
|
||||
3. Workspace B initializes later → Case 3 (both tables exist, legacy has B's data) → should migrate B's data
|
||||
4. Verify workspace B's data is correctly migrated to new table
|
||||
|
||||
This test verifies the migration logic correctly handles multi-tenant scenarios
|
||||
where different workspaces migrate sequentially.
|
||||
"""
|
||||
config = {
|
||||
"embedding_batch_num": 10,
|
||||
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
|
||||
}
|
||||
|
||||
embedding_func = EmbeddingFunc(
|
||||
embedding_dim=1536,
|
||||
func=mock_embedding_func.func,
|
||||
model_name="test-model",
|
||||
)
|
||||
|
||||
# Mock data: Legacy table has 6 records total (3 from workspace_a, 3 from workspace_b)
|
||||
mock_rows_a = [
|
||||
{"id": f"a_{i}", "content": f"A content {i}", "workspace": "workspace_a"}
|
||||
for i in range(3)
|
||||
]
|
||||
mock_rows_b = [
|
||||
{"id": f"b_{i}", "content": f"B content {i}", "workspace": "workspace_b"}
|
||||
for i in range(3)
|
||||
]
|
||||
|
||||
# Track migration state
|
||||
migration_state = {
|
||||
"new_table_exists": False,
|
||||
"workspace_a_migrated": False,
|
||||
"workspace_a_migration_count": 0,
|
||||
"workspace_b_migration_count": 0,
|
||||
}
|
||||
|
||||
# Step 1: Simulate workspace_a initialization (Case 3 - only legacy exists)
|
||||
# CRITICAL: Set db.workspace to workspace_a
|
||||
mock_pg_db.workspace = "workspace_a"
|
||||
|
||||
storage_a = PGVectorStorage(
|
||||
namespace=NameSpace.VECTOR_STORE_CHUNKS,
|
||||
global_config=config,
|
||||
embedding_func=embedding_func,
|
||||
workspace="workspace_a",
|
||||
)
|
||||
|
||||
# Mock table_exists for workspace_a
|
||||
async def mock_check_table_exists_a(table_name):
|
||||
if table_name == storage_a.legacy_table_name:
|
||||
return True
|
||||
if table_name == storage_a.table_name:
|
||||
return migration_state["new_table_exists"]
|
||||
return False
|
||||
|
||||
mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists_a)
|
||||
|
||||
# Mock query for workspace_a (Case 3)
|
||||
async def mock_query_a(sql, params=None, multirows=False, **kwargs):
|
||||
sql_upper = sql.upper()
|
||||
base_name = storage_a.legacy_table_name.upper()
|
||||
|
||||
if "COUNT(*)" in sql:
|
||||
has_model_suffix = "TEST_MODEL_1536D" in sql_upper
|
||||
is_legacy = base_name in sql_upper and not has_model_suffix
|
||||
has_workspace_filter = "WHERE workspace" in sql
|
||||
|
||||
if is_legacy and has_workspace_filter:
|
||||
workspace = params[0] if params and len(params) > 0 else None
|
||||
if workspace == "workspace_a":
|
||||
return {"count": 3}
|
||||
elif workspace == "workspace_b":
|
||||
return {"count": 3}
|
||||
elif is_legacy and not has_workspace_filter:
|
||||
# Global count in legacy table
|
||||
return {"count": 6}
|
||||
elif has_model_suffix:
|
||||
if has_workspace_filter:
|
||||
workspace = params[0] if params and len(params) > 0 else None
|
||||
if workspace == "workspace_a":
|
||||
return {"count": migration_state["workspace_a_migration_count"]}
|
||||
if workspace == "workspace_b":
|
||||
return {"count": migration_state["workspace_b_migration_count"]}
|
||||
return {
|
||||
"count": migration_state["workspace_a_migration_count"]
|
||||
+ migration_state["workspace_b_migration_count"]
|
||||
}
|
||||
elif multirows and "SELECT *" in sql:
|
||||
if "WHERE workspace" in sql:
|
||||
workspace = params[0] if params and len(params) > 0 else None
|
||||
if workspace == "workspace_a":
|
||||
# Handle keyset pagination
|
||||
if "id >" in sql:
|
||||
# params = [workspace, last_id, limit]
|
||||
last_id = params[1] if len(params) > 1 else None
|
||||
start_idx = 0
|
||||
for i, row in enumerate(mock_rows_a):
|
||||
if row["id"] == last_id:
|
||||
start_idx = i + 1
|
||||
break
|
||||
limit = params[2] if len(params) > 2 else 500
|
||||
else:
|
||||
# First batch: params = [workspace, limit]
|
||||
start_idx = 0
|
||||
limit = params[1] if len(params) > 1 else 500
|
||||
end = min(start_idx + limit, len(mock_rows_a))
|
||||
return mock_rows_a[start_idx:end]
|
||||
return {}
|
||||
|
||||
mock_pg_db.query = AsyncMock(side_effect=mock_query_a)
|
||||
|
||||
# Track migration via _run_with_retry (batch migration uses this)
|
||||
migration_a_executed = []
|
||||
|
||||
async def mock_run_with_retry_a(operation, **kwargs):
|
||||
migration_a_executed.append(True)
|
||||
migration_state["workspace_a_migration_count"] = len(mock_rows_a)
|
||||
return None
|
||||
|
||||
mock_pg_db._run_with_retry = AsyncMock(side_effect=mock_run_with_retry_a)
|
||||
|
||||
# Initialize workspace_a (Case 3)
|
||||
with patch("lightrag.kg.postgres_impl.logger"):
|
||||
await storage_a.initialize()
|
||||
migration_state["new_table_exists"] = True
|
||||
migration_state["workspace_a_migrated"] = True
|
||||
|
||||
print("✅ Step 1: Workspace A initialized")
|
||||
# Verify migration was executed via _run_with_retry (batch migration uses executemany)
|
||||
assert len(migration_a_executed) > 0, (
|
||||
"Migration should have been executed for workspace_a"
|
||||
)
|
||||
print(f"✅ Step 1: Migration executed {len(migration_a_executed)} batch(es)")
|
||||
|
||||
# Step 2: Simulate workspace_b initialization (Case 3 - both exist, but legacy has B's data)
|
||||
# CRITICAL: Set db.workspace to workspace_b
|
||||
mock_pg_db.workspace = "workspace_b"
|
||||
|
||||
storage_b = PGVectorStorage(
|
||||
namespace=NameSpace.VECTOR_STORE_CHUNKS,
|
||||
global_config=config,
|
||||
embedding_func=embedding_func,
|
||||
workspace="workspace_b",
|
||||
)
|
||||
|
||||
mock_pg_db.reset_mock()
|
||||
|
||||
# Mock table_exists for workspace_b (both exist)
|
||||
async def mock_check_table_exists_b(table_name):
|
||||
return True # Both tables exist
|
||||
|
||||
mock_pg_db.check_table_exists = AsyncMock(side_effect=mock_check_table_exists_b)
|
||||
|
||||
# Mock query for workspace_b (Case 3)
|
||||
async def mock_query_b(sql, params=None, multirows=False, **kwargs):
|
||||
sql_upper = sql.upper()
|
||||
base_name = storage_b.legacy_table_name.upper()
|
||||
|
||||
if "COUNT(*)" in sql:
|
||||
has_model_suffix = "TEST_MODEL_1536D" in sql_upper
|
||||
is_legacy = base_name in sql_upper and not has_model_suffix
|
||||
has_workspace_filter = "WHERE workspace" in sql
|
||||
|
||||
if is_legacy and has_workspace_filter:
|
||||
workspace = params[0] if params and len(params) > 0 else None
|
||||
if workspace == "workspace_b":
|
||||
return {"count": 3} # workspace_b still has data in legacy
|
||||
elif workspace == "workspace_a":
|
||||
return {"count": 0} # workspace_a already migrated
|
||||
elif is_legacy and not has_workspace_filter:
|
||||
# Global count: only workspace_b data remains
|
||||
return {"count": 3}
|
||||
elif has_model_suffix:
|
||||
if has_workspace_filter:
|
||||
workspace = params[0] if params and len(params) > 0 else None
|
||||
if workspace == "workspace_b":
|
||||
return {"count": migration_state["workspace_b_migration_count"]}
|
||||
elif workspace == "workspace_a":
|
||||
return {"count": 3}
|
||||
else:
|
||||
return {"count": 3 + migration_state["workspace_b_migration_count"]}
|
||||
elif multirows and "SELECT *" in sql:
|
||||
if "WHERE workspace" in sql:
|
||||
workspace = params[0] if params and len(params) > 0 else None
|
||||
if workspace == "workspace_b":
|
||||
# Handle keyset pagination
|
||||
if "id >" in sql:
|
||||
# params = [workspace, last_id, limit]
|
||||
last_id = params[1] if len(params) > 1 else None
|
||||
start_idx = 0
|
||||
for i, row in enumerate(mock_rows_b):
|
||||
if row["id"] == last_id:
|
||||
start_idx = i + 1
|
||||
break
|
||||
limit = params[2] if len(params) > 2 else 500
|
||||
else:
|
||||
# First batch: params = [workspace, limit]
|
||||
start_idx = 0
|
||||
limit = params[1] if len(params) > 1 else 500
|
||||
end = min(start_idx + limit, len(mock_rows_b))
|
||||
return mock_rows_b[start_idx:end]
|
||||
return {}
|
||||
|
||||
mock_pg_db.query = AsyncMock(side_effect=mock_query_b)
|
||||
|
||||
# Track migration via _run_with_retry for workspace_b
|
||||
migration_b_executed = []
|
||||
|
||||
async def mock_run_with_retry_b(operation, **kwargs):
|
||||
migration_b_executed.append(True)
|
||||
migration_state["workspace_b_migration_count"] = len(mock_rows_b)
|
||||
return None
|
||||
|
||||
mock_pg_db._run_with_retry = AsyncMock(side_effect=mock_run_with_retry_b)
|
||||
|
||||
# Initialize workspace_b (Case 3 - both tables exist)
|
||||
with patch("lightrag.kg.postgres_impl.logger"):
|
||||
await storage_b.initialize()
|
||||
|
||||
print("✅ Step 2: Workspace B initialized")
|
||||
|
||||
# Verify workspace_b migration happens when new table has no workspace_b data
|
||||
# but legacy table still has workspace_b data.
|
||||
assert len(migration_b_executed) > 0, (
|
||||
"Migration should have been executed for workspace_b"
|
||||
)
|
||||
print("✅ Step 2: Migration executed for workspace_b")
|
||||
|
||||
print("\n🎉 Case 1c: Sequential workspace migration verification complete!")
|
||||
print(" - Workspace A: Migrated successfully (only legacy existed)")
|
||||
print(" - Workspace B: Migrated successfully (new table empty for workspace_b)")
|
||||
@@ -0,0 +1,138 @@
|
||||
import importlib
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import lightrag.utils as utils_module
|
||||
from lightrag.kg.postgres_impl import PGGraphStorage, PostgreSQLDB
|
||||
from lightrag.namespace import NameSpace
|
||||
|
||||
|
||||
def make_db() -> PostgreSQLDB:
|
||||
return PostgreSQLDB(
|
||||
{
|
||||
"host": "localhost",
|
||||
"port": 5432,
|
||||
"user": "postgres",
|
||||
"password": "postgres",
|
||||
"database": "postgres",
|
||||
"workspace": "test_ws",
|
||||
"max_connections": 10,
|
||||
"connection_retry_attempts": 3,
|
||||
"connection_retry_backoff": 0,
|
||||
"connection_retry_backoff_max": 0,
|
||||
"pool_close_timeout": 5.0,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_timing_logs_success():
|
||||
db = make_db()
|
||||
|
||||
async def fake_run_with_retry(operation, **kwargs):
|
||||
conn = AsyncMock()
|
||||
conn.execute = AsyncMock(return_value="INSERT 0 1")
|
||||
await operation(conn)
|
||||
|
||||
db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry)
|
||||
|
||||
with patch("lightrag.kg.postgres_impl.performance_timing_log") as timing_log:
|
||||
await db.execute("SELECT 1", timing_label="test label")
|
||||
|
||||
assert any(
|
||||
"connection.execute completed" in call.args[0]
|
||||
for call in timing_log.call_args_list
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_timing_logs_failure():
|
||||
db = make_db()
|
||||
|
||||
async def fake_run_with_retry(operation, **kwargs):
|
||||
conn = AsyncMock()
|
||||
conn.execute = AsyncMock(side_effect=RuntimeError("boom"))
|
||||
await operation(conn)
|
||||
|
||||
db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry)
|
||||
|
||||
with patch("lightrag.kg.postgres_impl.performance_timing_log") as timing_log:
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
await db.execute("SELECT 1", timing_label="test label")
|
||||
|
||||
assert any(
|
||||
"connection.execute failed" in call.args[0]
|
||||
for call in timing_log.call_args_list
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_graph_upsert_node_passes_timing_label():
|
||||
storage = PGGraphStorage(
|
||||
namespace=NameSpace.GRAPH_STORE_CHUNK_ENTITY_RELATION,
|
||||
workspace="test_ws",
|
||||
global_config={},
|
||||
embedding_func=AsyncMock(),
|
||||
)
|
||||
storage.graph_name = "test_graph"
|
||||
storage._query = AsyncMock(return_value=[])
|
||||
|
||||
await storage.upsert_node(
|
||||
"node-1",
|
||||
{
|
||||
"entity_id": "node-1",
|
||||
"description": "desc",
|
||||
},
|
||||
)
|
||||
|
||||
assert storage._query.await_args.kwargs["timing_label"] == (
|
||||
"test_ws PGGraphStorage.upsert_node"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_graph_upsert_edge_passes_timing_label():
|
||||
storage = PGGraphStorage(
|
||||
namespace=NameSpace.GRAPH_STORE_CHUNK_ENTITY_RELATION,
|
||||
workspace="test_ws",
|
||||
global_config={},
|
||||
embedding_func=AsyncMock(),
|
||||
)
|
||||
storage.graph_name = "test_graph"
|
||||
# upsert_edge drives the lock + cypher via db._run_with_retry, not _query.
|
||||
storage.db = MagicMock()
|
||||
storage.db._run_with_retry = AsyncMock(return_value=None)
|
||||
|
||||
await storage.upsert_edge(
|
||||
"node-1",
|
||||
"node-2",
|
||||
{
|
||||
"weight": 1.0,
|
||||
"description": "desc",
|
||||
},
|
||||
)
|
||||
|
||||
assert storage.db._run_with_retry.await_args.kwargs["timing_label"] == (
|
||||
"test_ws PGGraphStorage.upsert_edge"
|
||||
)
|
||||
|
||||
|
||||
def test_performance_timing_logs_reads_new_env_only(monkeypatch):
|
||||
with monkeypatch.context() as m:
|
||||
m.setenv("LIGHTRAG_DOC_QUERY_TIMING_LOGS", "false")
|
||||
m.setenv("LIGHTRAG_PERFORMANCE_TIMING_LOGS", "true")
|
||||
reloaded = importlib.reload(utils_module)
|
||||
assert reloaded.PERFORMANCE_TIMING_LOGS is True
|
||||
|
||||
importlib.reload(utils_module)
|
||||
|
||||
|
||||
def test_performance_timing_logs_ignores_old_env(monkeypatch):
|
||||
with monkeypatch.context() as m:
|
||||
m.setenv("LIGHTRAG_DOC_QUERY_TIMING_LOGS", "true")
|
||||
m.setenv("LIGHTRAG_PERFORMANCE_TIMING_LOGS", "false")
|
||||
reloaded = importlib.reload(utils_module)
|
||||
assert reloaded.PERFORMANCE_TIMING_LOGS is False
|
||||
|
||||
importlib.reload(utils_module)
|
||||
@@ -0,0 +1,350 @@
|
||||
"""
|
||||
Integration test suite for PostgreSQL retry mechanism using real database.
|
||||
|
||||
This test suite connects to a real PostgreSQL database using credentials from .env
|
||||
and tests the retry mechanism with actual network failures.
|
||||
|
||||
Prerequisites:
|
||||
1. PostgreSQL server running and accessible
|
||||
2. .env file with POSTGRES_* configuration
|
||||
3. asyncpg installed: pip install asyncpg
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
from dotenv import load_dotenv
|
||||
from unittest.mock import patch
|
||||
from lightrag.kg.postgres_impl import PostgreSQLDB
|
||||
|
||||
asyncpg = pytest.importorskip("asyncpg")
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv(dotenv_path=".env", override=False)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.requires_db
|
||||
class TestPostgresRetryIntegration:
|
||||
"""Integration tests for PostgreSQL retry mechanism with real database."""
|
||||
|
||||
@pytest.fixture
|
||||
def db_config(self):
|
||||
"""Load database configuration from environment variables.
|
||||
|
||||
Uses new HA-optimized defaults that match postgres_impl.py ClientManager.get_config():
|
||||
- 10 retry attempts (up from 3)
|
||||
- 3.0s initial backoff (up from 0.5s)
|
||||
- 30.0s max backoff (up from 5.0s)
|
||||
"""
|
||||
return {
|
||||
"host": os.getenv("POSTGRES_HOST", "localhost"),
|
||||
"port": int(os.getenv("POSTGRES_PORT", "5432")),
|
||||
"user": os.getenv("POSTGRES_USER", "postgres"),
|
||||
"password": os.getenv("POSTGRES_PASSWORD", ""),
|
||||
"database": os.getenv("POSTGRES_DATABASE", "postgres"),
|
||||
"workspace": os.getenv("POSTGRES_WORKSPACE", "test_retry"),
|
||||
"max_connections": int(os.getenv("POSTGRES_MAX_CONNECTIONS", "10")),
|
||||
# Connection retry configuration - mirrors postgres_impl.py ClientManager.get_config()
|
||||
# NEW DEFAULTS optimized for HA deployments
|
||||
"connection_retry_attempts": min(
|
||||
100,
|
||||
int(os.getenv("POSTGRES_CONNECTION_RETRIES", "10")), # 3 → 10
|
||||
),
|
||||
"connection_retry_backoff": min(
|
||||
300.0,
|
||||
float(
|
||||
os.getenv("POSTGRES_CONNECTION_RETRY_BACKOFF", "3.0")
|
||||
), # 0.5 → 3.0
|
||||
),
|
||||
"connection_retry_backoff_max": min(
|
||||
600.0,
|
||||
float(
|
||||
os.getenv("POSTGRES_CONNECTION_RETRY_BACKOFF_MAX", "30.0")
|
||||
), # 5.0 → 30.0
|
||||
),
|
||||
"pool_close_timeout": min(
|
||||
30.0, float(os.getenv("POSTGRES_POOL_CLOSE_TIMEOUT", "5.0"))
|
||||
),
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_connection_success(self, db_config):
|
||||
"""
|
||||
Test successful connection to real PostgreSQL database.
|
||||
|
||||
This validates that:
|
||||
1. Database credentials are correct
|
||||
2. Connection pool initializes properly
|
||||
3. Basic query works
|
||||
"""
|
||||
print("\n" + "=" * 80)
|
||||
print("INTEGRATION TEST 1: Real Database Connection")
|
||||
print("=" * 80)
|
||||
print(
|
||||
f" → Connecting to {db_config['host']}:{db_config['port']}/{db_config['database']}"
|
||||
)
|
||||
|
||||
db = PostgreSQLDB(db_config)
|
||||
|
||||
try:
|
||||
# Initialize database connection
|
||||
await db.initdb()
|
||||
print(" ✓ Connection successful")
|
||||
|
||||
# Test simple query
|
||||
result = await db.query("SELECT 1 as test", multirows=False)
|
||||
assert result is not None
|
||||
assert result.get("test") == 1
|
||||
print(" ✓ Query executed successfully")
|
||||
|
||||
print("\n✅ Test passed: Real database connection works")
|
||||
print("=" * 80)
|
||||
finally:
|
||||
if db.pool:
|
||||
await db.pool.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_simulated_transient_error_with_real_db(self, db_config):
|
||||
"""
|
||||
Test retry mechanism with simulated transient errors on real database.
|
||||
|
||||
Simulates connection failures on first 2 attempts, then succeeds.
|
||||
Uses new HA defaults (10 retries, 3s backoff).
|
||||
"""
|
||||
print("\n" + "=" * 80)
|
||||
print("INTEGRATION TEST 2: Simulated Transient Errors")
|
||||
print("=" * 80)
|
||||
|
||||
db = PostgreSQLDB(db_config)
|
||||
attempt_count = {"value": 0}
|
||||
|
||||
# Original create_pool function
|
||||
original_create_pool = asyncpg.create_pool
|
||||
|
||||
async def mock_create_pool_with_failures(*args, **kwargs):
|
||||
"""Mock that fails first 2 times, then calls real create_pool."""
|
||||
attempt_count["value"] += 1
|
||||
print(f" → Connection attempt {attempt_count['value']}")
|
||||
|
||||
if attempt_count["value"] <= 2:
|
||||
print(" ✗ Simulating connection failure")
|
||||
raise asyncpg.exceptions.ConnectionFailureError(
|
||||
f"Simulated failure on attempt {attempt_count['value']}"
|
||||
)
|
||||
|
||||
print(" ✓ Allowing real connection")
|
||||
return await original_create_pool(*args, **kwargs)
|
||||
|
||||
try:
|
||||
# Patch create_pool to simulate failures
|
||||
with patch(
|
||||
"asyncpg.create_pool", side_effect=mock_create_pool_with_failures
|
||||
):
|
||||
await db.initdb()
|
||||
|
||||
assert attempt_count["value"] == 3, (
|
||||
f"Expected 3 attempts, got {attempt_count['value']}"
|
||||
)
|
||||
assert db.pool is not None, "Pool should be initialized after retries"
|
||||
|
||||
# Verify database is actually working
|
||||
result = await db.query("SELECT 1 as test", multirows=False)
|
||||
assert result.get("test") == 1
|
||||
|
||||
print(
|
||||
f"\n✅ Test passed: Retry mechanism worked, connected after {attempt_count['value']} attempts"
|
||||
)
|
||||
print("=" * 80)
|
||||
finally:
|
||||
if db.pool:
|
||||
await db.pool.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_retry_with_real_db(self, db_config):
|
||||
"""
|
||||
Test query-level retry with simulated connection issues.
|
||||
|
||||
Tests that queries retry on transient failures by simulating
|
||||
a temporary database unavailability.
|
||||
Uses new HA defaults (10 retries, 3s backoff).
|
||||
"""
|
||||
print("\n" + "=" * 80)
|
||||
print("INTEGRATION TEST 3: Query-Level Retry")
|
||||
print("=" * 80)
|
||||
|
||||
db = PostgreSQLDB(db_config)
|
||||
|
||||
try:
|
||||
# First initialize normally
|
||||
await db.initdb()
|
||||
print(" ✓ Database initialized")
|
||||
|
||||
# Close the pool to simulate connection loss
|
||||
print(" → Simulating connection loss (closing pool)...")
|
||||
await db.pool.close()
|
||||
db.pool = None
|
||||
|
||||
# Now query should trigger pool recreation and retry
|
||||
print(" → Attempting query (should auto-reconnect)...")
|
||||
result = await db.query("SELECT 1 as test", multirows=False)
|
||||
|
||||
assert result.get("test") == 1, "Query should succeed after reconnection"
|
||||
assert db.pool is not None, "Pool should be recreated"
|
||||
|
||||
print(" ✓ Query succeeded after automatic reconnection")
|
||||
print("\n✅ Test passed: Auto-reconnection works correctly")
|
||||
print("=" * 80)
|
||||
finally:
|
||||
if db.pool:
|
||||
await db.pool.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_queries_with_real_db(self, db_config):
|
||||
"""
|
||||
Test concurrent queries to validate thread safety and connection pooling.
|
||||
|
||||
Runs multiple concurrent queries to ensure no deadlocks or race conditions.
|
||||
Uses new HA defaults (10 retries, 3s backoff).
|
||||
"""
|
||||
print("\n" + "=" * 80)
|
||||
print("INTEGRATION TEST 4: Concurrent Queries")
|
||||
print("=" * 80)
|
||||
|
||||
db = PostgreSQLDB(db_config)
|
||||
|
||||
try:
|
||||
await db.initdb()
|
||||
print(" ✓ Database initialized")
|
||||
|
||||
# Launch 10 concurrent queries
|
||||
num_queries = 10
|
||||
print(f" → Launching {num_queries} concurrent queries...")
|
||||
|
||||
async def run_query(query_id):
|
||||
result = await db.query(
|
||||
f"SELECT {query_id} as id, pg_sleep(0.1)", multirows=False
|
||||
)
|
||||
return result.get("id")
|
||||
|
||||
start_time = time.time()
|
||||
tasks = [run_query(i) for i in range(num_queries)]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
# Check results
|
||||
successful = sum(1 for r in results if not isinstance(r, Exception))
|
||||
failed = sum(1 for r in results if isinstance(r, Exception))
|
||||
|
||||
print(f" → Completed in {elapsed:.2f}s")
|
||||
print(f" → Results: {successful} successful, {failed} failed")
|
||||
|
||||
assert successful == num_queries, (
|
||||
f"All {num_queries} queries should succeed"
|
||||
)
|
||||
assert failed == 0, "No queries should fail"
|
||||
|
||||
print("\n✅ Test passed: All concurrent queries succeeded, no deadlocks")
|
||||
print("=" * 80)
|
||||
finally:
|
||||
if db.pool:
|
||||
await db.pool.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pool_close_timeout_real(self, db_config):
|
||||
"""
|
||||
Test pool close timeout protection with real database.
|
||||
Uses new HA defaults (10 retries, 3s backoff).
|
||||
"""
|
||||
print("\n" + "=" * 80)
|
||||
print("INTEGRATION TEST 5: Pool Close Timeout")
|
||||
print("=" * 80)
|
||||
|
||||
db = PostgreSQLDB(db_config)
|
||||
|
||||
try:
|
||||
await db.initdb()
|
||||
print(" ✓ Database initialized")
|
||||
|
||||
# Trigger pool reset (which includes close)
|
||||
print(" → Triggering pool reset...")
|
||||
start_time = time.time()
|
||||
await db._reset_pool()
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
print(f" ✓ Pool reset completed in {elapsed:.2f}s")
|
||||
assert db.pool is None, "Pool should be None after reset"
|
||||
assert elapsed < db.pool_close_timeout + 1, (
|
||||
"Reset should complete within timeout"
|
||||
)
|
||||
|
||||
print("\n✅ Test passed: Pool reset handled correctly")
|
||||
print("=" * 80)
|
||||
finally:
|
||||
# Already closed in test
|
||||
pass
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_configuration_from_env(self, db_config):
|
||||
"""
|
||||
Test that configuration is correctly loaded from environment variables.
|
||||
"""
|
||||
print("\n" + "=" * 80)
|
||||
print("INTEGRATION TEST 6: Environment Configuration")
|
||||
print("=" * 80)
|
||||
|
||||
db = PostgreSQLDB(db_config)
|
||||
|
||||
print(" → Configuration loaded:")
|
||||
print(f" • Host: {db.host}")
|
||||
print(f" • Port: {db.port}")
|
||||
print(f" • Database: {db.database}")
|
||||
print(f" • User: {db.user}")
|
||||
print(f" • Workspace: {db.workspace}")
|
||||
print(f" • Max Connections: {db.max}")
|
||||
print(f" • Retry Attempts: {db.connection_retry_attempts}")
|
||||
print(f" • Retry Backoff: {db.connection_retry_backoff}s")
|
||||
print(f" • Max Backoff: {db.connection_retry_backoff_max}s")
|
||||
print(f" • Pool Close Timeout: {db.pool_close_timeout}s")
|
||||
|
||||
# Verify required fields are present
|
||||
assert db.host, "Host should be configured"
|
||||
assert db.port, "Port should be configured"
|
||||
assert db.user, "User should be configured"
|
||||
assert db.database, "Database should be configured"
|
||||
|
||||
print("\n✅ Test passed: All configuration loaded correctly from .env")
|
||||
print("=" * 80)
|
||||
|
||||
|
||||
def run_integration_tests():
|
||||
"""Run all integration tests with detailed output."""
|
||||
print("\n" + "=" * 80)
|
||||
print("POSTGRESQL RETRY MECHANISM - INTEGRATION TESTS")
|
||||
print("Testing with REAL database from .env configuration")
|
||||
print("=" * 80)
|
||||
|
||||
# Check if database configuration exists
|
||||
if not os.getenv("POSTGRES_HOST"):
|
||||
print("\n⚠️ WARNING: No POSTGRES_HOST in .env file")
|
||||
print("Please ensure .env file exists with PostgreSQL configuration.")
|
||||
return
|
||||
|
||||
print("\nRunning integration tests...\n")
|
||||
|
||||
# Run pytest with verbose output
|
||||
pytest.main(
|
||||
[
|
||||
__file__,
|
||||
"-v",
|
||||
"-s", # Don't capture output
|
||||
"--tb=short", # Short traceback format
|
||||
"--color=yes",
|
||||
"-x", # Stop on first failure
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_integration_tests()
|
||||
@@ -0,0 +1,35 @@
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from lightrag.kg.postgres_impl import PostgreSQLDB
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_table_exists_uses_search_path_visible_regclass():
|
||||
db = PostgreSQLDB.__new__(PostgreSQLDB)
|
||||
db.query = AsyncMock(return_value={"exists": True})
|
||||
|
||||
assert await db.check_table_exists("LIGHTRAG_DOC_FULL") is True
|
||||
|
||||
db.query.assert_awaited_once_with(
|
||||
"SELECT to_regclass($1) IS NOT NULL AS exists",
|
||||
["lightrag_doc_full"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_entity_relation_migration_reuses_table_exists_helper():
|
||||
db = PostgreSQLDB.__new__(PostgreSQLDB)
|
||||
db.check_table_exists = AsyncMock(return_value=True)
|
||||
db.query = AsyncMock()
|
||||
db.execute = AsyncMock()
|
||||
|
||||
await db._migrate_create_full_entities_relations_tables()
|
||||
|
||||
assert [call.args[0] for call in db.check_table_exists.await_args_list] == [
|
||||
"LIGHTRAG_FULL_ENTITIES",
|
||||
"LIGHTRAG_FULL_RELATIONS",
|
||||
]
|
||||
db.query.assert_not_called()
|
||||
db.execute.assert_not_called()
|
||||
@@ -0,0 +1,844 @@
|
||||
"""
|
||||
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 asyncio
|
||||
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] = []
|
||||
captured_execute: 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()
|
||||
tx = AsyncMock()
|
||||
tx.__aenter__.return_value = tx
|
||||
tx.__aexit__.return_value = False
|
||||
mock_conn.transaction = MagicMock(return_value=tx)
|
||||
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]))
|
||||
# Store (sql, args) from each execute call (chunked delete path).
|
||||
for call in mock_conn.execute.call_args_list:
|
||||
captured_execute.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._captured_execute = captured_execute
|
||||
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] = []
|
||||
captured_execute: 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]))
|
||||
for call in mock_conn.execute.call_args_list:
|
||||
captured_execute.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.__post_init__() # resolves batch-limit attrs (payload/records caps)
|
||||
storage._captured = captured
|
||||
storage._captured_execute = captured_execute
|
||||
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()
|
||||
# connection.transaction() is a synchronous factory in asyncpg that
|
||||
# returns an async context manager. AsyncMock would return a coroutine,
|
||||
# so swap in a MagicMock that returns an AsyncMock-backed context.
|
||||
tx_cm = AsyncMock()
|
||||
tx_cm.__aenter__.return_value = None
|
||||
tx_cm.__aexit__.return_value = None
|
||||
mock_conn.transaction = MagicMock(return_value=tx_cm)
|
||||
await operation(mock_conn)
|
||||
for call in mock_conn.executemany.call_args_list:
|
||||
captured.append((call.args[0], call.args[1]))
|
||||
# Also capture connection.execute calls (delete SQL inside flush).
|
||||
for call in mock_conn.execute.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._flush_lock = asyncio.Lock()
|
||||
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.[native-Fi].txt",
|
||||
"llm_cache_list": ["cache-key"],
|
||||
"heading": {"level": 2, "text": "Section A"},
|
||||
"sidecar": {"type": "drawing", "id": "img-1", "refs": []},
|
||||
}
|
||||
}
|
||||
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, heading, sidecar,
|
||||
# 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.[native-Fi].txt" # file_path
|
||||
assert json.loads(row[7]) == ["cache-key"] # llm_cache_list
|
||||
assert json.loads(row[8]) == {"level": 2, "text": "Section A"} # heading
|
||||
assert json.loads(row[9]) == {
|
||||
"type": "drawing",
|
||||
"id": "img-1",
|
||||
"refs": [],
|
||||
} # sidecar
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_text_chunks_missing_heading_sidecar_defaults_to_empty_dict():
|
||||
"""Plain-text chunks without heading/sidecar should serialize to '{}'."""
|
||||
storage = make_storage(NameSpace.KV_STORE_TEXT_CHUNKS)
|
||||
data = {
|
||||
"chunk-1": {
|
||||
"tokens": 10,
|
||||
"chunk_order_index": 0,
|
||||
"full_doc_id": "doc-1",
|
||||
"content": "plain text",
|
||||
"file_path": "/a/b.txt",
|
||||
}
|
||||
}
|
||||
await storage.upsert(data)
|
||||
|
||||
_, rows = storage._captured[0]
|
||||
row = rows[0]
|
||||
assert json.loads(row[8]) == {} # heading
|
||||
assert json.loads(row[9]) == {} # sidecar
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_text_chunks_none_heading_sidecar_defaults_to_empty_dict():
|
||||
"""Explicit None values should be coerced to '{}' to avoid type errors."""
|
||||
storage = make_storage(NameSpace.KV_STORE_TEXT_CHUNKS)
|
||||
data = {
|
||||
"chunk-1": {
|
||||
"tokens": 10,
|
||||
"chunk_order_index": 0,
|
||||
"full_doc_id": "doc-1",
|
||||
"content": "plain text",
|
||||
"file_path": "/a/b.txt",
|
||||
"heading": None,
|
||||
"sidecar": None,
|
||||
}
|
||||
}
|
||||
await storage.upsert(data)
|
||||
|
||||
_, rows = storage._captured[0]
|
||||
row = rows[0]
|
||||
assert json.loads(row[8]) == {}
|
||||
assert json.loads(row[9]) == {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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.[mineru-Fi].pdf",
|
||||
"sidecar_location": "lightrag://sidecar/doc-1",
|
||||
"parse_format": "lightrag",
|
||||
"content_hash": "deadbeef",
|
||||
"process_options": "Fi",
|
||||
"chunk_options": {"chunk_token_size": 1200, "chunk_overlap": 100},
|
||||
"parse_engine": "mineru",
|
||||
}
|
||||
}
|
||||
await storage.upsert(data)
|
||||
|
||||
assert len(storage._captured) == 1
|
||||
_, rows = storage._captured[0]
|
||||
row = rows[0]
|
||||
# SQL: (id, content, doc_name, workspace, sidecar_location, parse_format,
|
||||
# content_hash, process_options, chunk_options, parse_engine)
|
||||
assert row[0] == "doc-1"
|
||||
assert row[1] == "full text"
|
||||
assert row[2] == "/path/doc.[mineru-Fi].pdf"
|
||||
assert row[3] == "test_ws"
|
||||
assert row[4] == "lightrag://sidecar/doc-1"
|
||||
assert row[5] == "lightrag"
|
||||
assert row[6] == "deadbeef"
|
||||
assert row[7] == "Fi"
|
||||
assert json.loads(row[8]) == {"chunk_token_size": 1200, "chunk_overlap": 100}
|
||||
assert row[9] == "mineru"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_full_docs_missing_pipeline_fields_pass_through_as_none():
|
||||
"""Missing pipeline-derived fields must serialize as None at the Python
|
||||
layer so the SQL-level COALESCE guard can distinguish "caller did not
|
||||
supply" from "caller supplied a real value".
|
||||
|
||||
The 'raw' default for parse_format is provided by the column DDL on
|
||||
initial insert; the Python layer must NOT inject it, otherwise the
|
||||
COALESCE guard never triggers on subsequent partial writes (a follow-up
|
||||
upsert with no parse_format would re-stamp the column with 'raw' and
|
||||
blow away a previously-set 'lightrag').
|
||||
"""
|
||||
storage = make_storage(NameSpace.KV_STORE_FULL_DOCS)
|
||||
data = {"doc-1": {"content": "full text", "file_path": "/path/doc.pdf"}}
|
||||
await storage.upsert(data)
|
||||
|
||||
_, rows = storage._captured[0]
|
||||
row = rows[0]
|
||||
assert row[4] is None # sidecar_location
|
||||
assert row[5] is None # parse_format — DDL supplies 'raw' default on insert
|
||||
assert row[6] is None # content_hash
|
||||
assert row[7] is None # process_options
|
||||
assert json.loads(row[8]) == {} # chunk_options default
|
||||
assert row[9] is None # parse_engine
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_full_docs_none_chunk_options_defaults_to_empty_dict():
|
||||
storage = make_storage(NameSpace.KV_STORE_FULL_DOCS)
|
||||
data = {
|
||||
"doc-1": {
|
||||
"content": "full text",
|
||||
"file_path": "/path/doc.pdf",
|
||||
"chunk_options": None,
|
||||
}
|
||||
}
|
||||
await storage.upsert(data)
|
||||
|
||||
_, rows = storage._captured[0]
|
||||
assert json.loads(rows[0][8]) == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_full_docs_sql_protects_partial_writes():
|
||||
"""The ON CONFLICT clause must COALESCE+NULLIF every pipeline-derived
|
||||
column so a follow-up upsert that only carries ``content`` + ``doc_name``
|
||||
does not silently overwrite previously-recorded metadata back to defaults.
|
||||
|
||||
We assert this at the SQL-template level since the actual COALESCE
|
||||
behavior is executed by Postgres. The presence of the protective
|
||||
expression in the SQL is the single source of truth for the guarantee.
|
||||
"""
|
||||
storage = make_storage(NameSpace.KV_STORE_FULL_DOCS)
|
||||
await storage.upsert(
|
||||
{"doc-1": {"content": "full text", "file_path": "/path/doc.pdf"}}
|
||||
)
|
||||
sql, _ = storage._captured[0]
|
||||
normalized = " ".join(sql.split()).lower()
|
||||
|
||||
# Each pipeline-derived string column must be COALESCE/NULLIF-guarded
|
||||
for col in (
|
||||
"sidecar_location",
|
||||
"parse_format",
|
||||
"content_hash",
|
||||
"process_options",
|
||||
"parse_engine",
|
||||
):
|
||||
assert f"coalesce( nullif(excluded.{col}, '')" in normalized, (
|
||||
f"upsert_doc_full must guard {col} via COALESCE+NULLIF"
|
||||
)
|
||||
assert f"lightrag_doc_full.{col}" in normalized, (
|
||||
f"upsert_doc_full must preserve existing {col} on partial write"
|
||||
)
|
||||
|
||||
# chunk_options (JSONB) is guarded via CASE on NULL/empty-object literal
|
||||
assert "excluded.chunk_options is null" in normalized
|
||||
assert "excluded.chunk_options = '{}'::jsonb" in normalized
|
||||
assert "lightrag_doc_full.chunk_options" in normalized
|
||||
|
||||
# content / doc_name remain straight overwrites — they ARE the payload
|
||||
assert "content = excluded.content" in normalized
|
||||
assert "doc_name = excluded.doc_name" in normalized
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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)
|
||||
# upsert splits by _max_upsert_records_per_batch (byte cap left at default).
|
||||
storage._max_upsert_records_per_batch = 3
|
||||
|
||||
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_upsert_records_per_batch = 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"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# doc_status: content_hash tuple + COALESCE SQL guard
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_doc_status_upsert_includes_content_hash():
|
||||
storage = make_doc_status_storage()
|
||||
await storage.upsert(
|
||||
{
|
||||
"doc-1": {
|
||||
"content_summary": "summary",
|
||||
"content_length": 12,
|
||||
"chunks_count": 1,
|
||||
"status": "processed",
|
||||
"file_path": "/uploads/a.[native-Fi].txt",
|
||||
"chunks_list": ["chunk-1"],
|
||||
"metadata": {"source": "test"},
|
||||
"content_hash": "abc123",
|
||||
"created_at": "2024-01-01T00:00:00+00:00",
|
||||
"updated_at": "2024-01-01T00:00:00+00:00",
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
sql, rows = storage._captured[0]
|
||||
# content_hash should be present in the INSERT column list and tuple
|
||||
assert "content_hash" in sql
|
||||
row = rows[0]
|
||||
# Tuple layout: workspace, id, content_summary, content_length, chunks_count,
|
||||
# status, file_path, chunks_list, track_id, metadata, error_msg,
|
||||
# content_hash, created_at, updated_at
|
||||
assert row[6] == "/uploads/a.[native-Fi].txt"
|
||||
assert row[11] == "abc123"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_doc_status_upsert_missing_content_hash_is_none():
|
||||
"""Existing callers that do not pass content_hash still produce valid tuples."""
|
||||
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",
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
_, rows = storage._captured[0]
|
||||
assert rows[0][11] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_doc_status_upsert_sql_protects_existing_content_hash():
|
||||
"""The ON CONFLICT clause must COALESCE+NULLIF to preserve a previously
|
||||
set content_hash when a subsequent state-transition upsert carries no
|
||||
hash (None) or an empty string.
|
||||
|
||||
We assert this at the SQL-template level since the actual COALESCE
|
||||
behavior is executed by Postgres. The presence of the protective
|
||||
expression in the SQL is the single source of truth for the guarantee.
|
||||
"""
|
||||
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": [],
|
||||
"metadata": {},
|
||||
"created_at": "2024-01-01T00:00:00+00:00",
|
||||
"updated_at": "2024-01-01T00:00:00+00:00",
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
sql, _ = storage._captured[0]
|
||||
normalized = " ".join(sql.split()).lower()
|
||||
assert "coalesce(" in normalized
|
||||
assert "nullif(excluded.content_hash, '')" in normalized
|
||||
assert "lightrag_doc_status.content_hash" in normalized
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vector_flush_passes_timing_label():
|
||||
"""upsert() now buffers; the timing_label is emitted by the flush path
|
||||
that runs from index_done_callback() / finalize().
|
||||
"""
|
||||
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",
|
||||
}
|
||||
}
|
||||
)
|
||||
# No retry call until flush.
|
||||
assert storage._retry_kwargs == []
|
||||
|
||||
await storage.index_done_callback()
|
||||
|
||||
assert storage._retry_kwargs[0]["timing_label"] == (
|
||||
f"test_ws PGVectorStorage.flush[{NameSpace.VECTOR_STORE_CHUNKS}]"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Batch limits: KV byte-budget split + delete chunking
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_kv_upsert_splits_by_payload_bytes():
|
||||
storage = make_storage(NameSpace.KV_STORE_FULL_DOCS)
|
||||
# Disable the record cap; force a byte cap that no two rows can share.
|
||||
storage._max_upsert_records_per_batch = 0
|
||||
storage._max_upsert_payload_bytes = 200
|
||||
|
||||
data = {f"doc-{i}": {"content": "X" * 150, "file_path": ""} for i in range(3)}
|
||||
await storage.upsert(data)
|
||||
|
||||
# Each ~150-byte content row exceeds half the budget -> 3 separate batches.
|
||||
assert len(storage._captured) == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_kv_delete_splits_by_id_cap():
|
||||
storage = make_storage(NameSpace.KV_STORE_FULL_DOCS)
|
||||
storage._max_delete_records_per_batch = 2
|
||||
|
||||
await storage.delete([f"doc-{i}" for i in range(5)])
|
||||
|
||||
# 5 ids / cap 2 => 3 bounded ANY($2) DELETE statements, all in ONE
|
||||
# transaction (a single _run_with_retry closure) for all-or-nothing semantics.
|
||||
assert len(storage._retry_kwargs) == 1
|
||||
assert len(storage._captured_execute) == 3
|
||||
slices = [args[1] for _, args in storage._captured_execute]
|
||||
assert [len(s) for s in slices] == [2, 2, 1]
|
||||
assert all("DELETE FROM" in sql for sql, _ in storage._captured_execute)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_kv_delete_chunks_share_one_transaction():
|
||||
"""All delete chunks run inside a single connection.transaction()."""
|
||||
storage = make_storage(NameSpace.KV_STORE_FULL_DOCS)
|
||||
storage._max_delete_records_per_batch = 2
|
||||
|
||||
tx_calls = {"n": 0}
|
||||
|
||||
async def run(operation, **kwargs):
|
||||
mock_conn = AsyncMock()
|
||||
tx = AsyncMock()
|
||||
tx.__aenter__.return_value = tx
|
||||
tx.__aexit__.return_value = False
|
||||
|
||||
def _make_tx():
|
||||
tx_calls["n"] += 1
|
||||
return tx
|
||||
|
||||
mock_conn.transaction = MagicMock(side_effect=_make_tx)
|
||||
await operation(mock_conn)
|
||||
|
||||
storage.db._run_with_retry = AsyncMock(side_effect=run)
|
||||
|
||||
await storage.delete([f"doc-{i}" for i in range(5)])
|
||||
|
||||
# One transaction wrapping all 3 chunks, not one per chunk.
|
||||
assert tx_calls["n"] == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Batch limits: DocStatus record-cap split + delete chunking
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _doc_status_payload(i: int) -> dict:
|
||||
return {
|
||||
"content_summary": f"summary {i}",
|
||||
"content_length": 10,
|
||||
"chunks_count": 1,
|
||||
"status": "processed",
|
||||
"file_path": f"/{i}.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",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_doc_status_upsert_splits_by_record_cap():
|
||||
storage = make_doc_status_storage()
|
||||
storage._max_upsert_records_per_batch = 2
|
||||
|
||||
data = {f"doc-{i}": _doc_status_payload(i) for i in range(5)}
|
||||
await storage.upsert(data)
|
||||
|
||||
# 5 records / cap 2 => 3 executemany calls.
|
||||
assert len(storage._captured) == 3
|
||||
assert [len(rows) for _, rows in storage._captured] == [2, 2, 1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_doc_status_delete_splits_by_id_cap():
|
||||
storage = make_doc_status_storage()
|
||||
storage._max_delete_records_per_batch = 2
|
||||
|
||||
await storage.delete([f"doc-{i}" for i in range(5)])
|
||||
|
||||
# One transaction (single _run_with_retry) wrapping 3 bounded DELETEs.
|
||||
assert len(storage._retry_kwargs) == 1
|
||||
assert len(storage._captured_execute) == 3
|
||||
slices = [args[1] for _, args in storage._captured_execute]
|
||||
assert [len(s) for s in slices] == [2, 2, 1]
|
||||
@@ -0,0 +1,393 @@
|
||||
"""
|
||||
Unit tests for PGGraphStorage.upsert_edge Cypher query generation.
|
||||
|
||||
Verifies the Cypher query sent to AGE uses the OPTIONAL MATCH + DELETE +
|
||||
CREATE pattern with inline edge properties — the only reliable way to write
|
||||
edge properties in Apache AGE (SET r += {...}, ON CREATE/ON MATCH SET, and
|
||||
SET r.key = value all silently fail for DIRECTED edges).
|
||||
"""
|
||||
|
||||
import json
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import asyncpg
|
||||
from tenacity import wait_none
|
||||
|
||||
from lightrag.kg.postgres_impl import (
|
||||
PGGraphQueryException,
|
||||
PGGraphStorage,
|
||||
_is_transient_graph_write_error,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def make_graph_storage() -> PGGraphStorage:
|
||||
"""Construct a PGGraphStorage instance with a mocked db."""
|
||||
storage = PGGraphStorage.__new__(PGGraphStorage)
|
||||
storage.workspace = "test_ws"
|
||||
storage.namespace = "test_graph"
|
||||
storage.graph_name = "test_graph"
|
||||
storage.__post_init__() # resolves chunk-level batch-limit attrs
|
||||
storage.db = MagicMock()
|
||||
return storage
|
||||
|
||||
|
||||
class _FakeConnection:
|
||||
"""Captures statements + args passed to a fake asyncpg connection."""
|
||||
|
||||
def __init__(self):
|
||||
self.calls: list[dict] = []
|
||||
|
||||
def transaction(self):
|
||||
return _FakeTransaction()
|
||||
|
||||
async def execute(self, sql, *args):
|
||||
self.calls.append({"sql": sql, "args": args})
|
||||
return ""
|
||||
|
||||
|
||||
class _FakeTransaction:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
|
||||
async def _capture_upsert_edge(storage: PGGraphStorage, src: str, tgt: str, edge_data):
|
||||
"""Invoke upsert_edge against a fake connection and return the captured calls."""
|
||||
conn = _FakeConnection()
|
||||
|
||||
async def fake_run_with_retry(operation, **_kwargs):
|
||||
return await operation(conn)
|
||||
|
||||
storage.db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry)
|
||||
await storage.upsert_edge(src, tgt, edge_data)
|
||||
return conn.calls
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# upsert_edge: Cypher query correctness
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edge_uses_delete_create_not_set():
|
||||
"""Cypher must use OPTIONAL MATCH + DELETE + CREATE — not SET-based update.
|
||||
|
||||
Apache AGE silently drops edge properties written via SET r += {...},
|
||||
SET r.key = value, and ON CREATE/ON MATCH SET. The only reliable pattern
|
||||
is to delete any existing edge and CREATE a new one with inline props.
|
||||
"""
|
||||
storage = make_graph_storage()
|
||||
calls = await _capture_upsert_edge(
|
||||
storage, "NodeA", "NodeB", {"weight": "1.0", "description": "test edge"}
|
||||
)
|
||||
|
||||
# The cypher statement is the third one (after the per-edge + shared locks).
|
||||
cypher_sql = calls[2]["sql"]
|
||||
|
||||
# The new query must not contain any SET-based edge update — those silently
|
||||
# fail against AGE. All edge props live inline in the CREATE clause.
|
||||
assert "SET r" not in cypher_sql, (
|
||||
f"Edge SET clauses are silently dropped by AGE: {cypher_sql}"
|
||||
)
|
||||
assert "ON CREATE SET" not in cypher_sql
|
||||
assert "ON MATCH SET" not in cypher_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edge_contains_optional_match_delete_create():
|
||||
"""The Cypher query must use OPTIONAL MATCH + DELETE + CREATE with inline props."""
|
||||
storage = make_graph_storage()
|
||||
calls = await _capture_upsert_edge(storage, "Alice", "Bob", {"weight": "0.5"})
|
||||
|
||||
cypher_sql = calls[2]["sql"]
|
||||
assert "OPTIONAL MATCH (source)-[old:DIRECTED]-(target)" in cypher_sql
|
||||
assert "DELETE old" in cypher_sql
|
||||
assert "CREATE (source)-[r:DIRECTED" in cypher_sql
|
||||
assert "]->(target)" in cypher_sql
|
||||
# Edge properties must be inlined into the CREATE clause as a literal map.
|
||||
assert '`weight`: "0.5"' in cypher_sql
|
||||
assert "RETURN r" in cypher_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edge_handles_empty_props():
|
||||
"""Empty edge_data must inline an empty literal map, not crash."""
|
||||
storage = make_graph_storage()
|
||||
calls = await _capture_upsert_edge(storage, "Alice", "Bob", {})
|
||||
|
||||
cypher_sql = calls[2]["sql"]
|
||||
assert "CREATE (source)-[r:DIRECTED {}]->(target)" in cypher_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edge_uses_parameterized_match_ids():
|
||||
"""Source and target IDs must flow through Cypher parameters as agtype JSON."""
|
||||
storage = make_graph_storage()
|
||||
calls = await _capture_upsert_edge(storage, "Node A", "Node B", {"weight": "1.0"})
|
||||
|
||||
cypher_call = calls[2]
|
||||
cypher_sql = cypher_call["sql"]
|
||||
assert "entity_id: $src_id" in cypher_sql
|
||||
assert "entity_id: $tgt_id" in cypher_sql
|
||||
# Cypher params arrive as a single positional agtype JSON arg.
|
||||
params_json = cypher_call["args"][0]
|
||||
params = json.loads(params_json)
|
||||
assert params["src_id"] == "Node A"
|
||||
assert params["tgt_id"] == "Node B"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edge_serialises_with_advisory_lock():
|
||||
"""Concurrent upserts on the same edge must be serialised via pg_advisory_xact_lock.
|
||||
|
||||
OPTIONAL MATCH + DELETE + CREATE is not atomic on its own: two transactions
|
||||
could both pass the OPTIONAL MATCH and both run CREATE, leaving duplicate
|
||||
DIRECTED rows. We open a transaction and acquire a transaction-scoped
|
||||
advisory lock keyed on (graph_name, ordered (src_id, tgt_id)) before running
|
||||
the cypher upsert, so concurrent upserts of the same logical edge run
|
||||
serially without serialising independent graphs.
|
||||
|
||||
AGE refuses to plan a join against a cypher() call that contains a CREATE
|
||||
clause, so the lock cannot live in a CTE — it must be a separate statement
|
||||
on the same connection inside an explicit transaction.
|
||||
"""
|
||||
storage = make_graph_storage()
|
||||
calls = await _capture_upsert_edge(storage, "Alice", "Bob", {"weight": "0.5"})
|
||||
|
||||
# Three statements: per-edge exclusive lock, graph-wide shared lock, cypher.
|
||||
assert len(calls) == 3
|
||||
|
||||
lock_sql = calls[0]["sql"]
|
||||
assert "pg_advisory_xact_lock(" in lock_sql # exclusive (not _shared)
|
||||
# graph_name flows as $1 so independent AGE graphs in the same DB do not
|
||||
# serialise each other's edges.
|
||||
assert "$1::text || E'\\x01' ||" in lock_sql
|
||||
# Key must be order-independent for (src, tgt) so {A, B} and {B, A} collide
|
||||
# on the same lock (the OPTIONAL MATCH is undirected).
|
||||
assert "LEAST($2::text, $3::text)" in lock_sql
|
||||
assert "GREATEST($2::text, $3::text)" in lock_sql
|
||||
# Raw graph_name + node IDs flow as positional params — never interpolated.
|
||||
assert "test_graph" not in lock_sql
|
||||
assert "Alice" not in lock_sql and "Bob" not in lock_sql
|
||||
assert calls[0]["args"] == ("test_graph", "Alice", "Bob")
|
||||
|
||||
# Second statement: graph-wide SHARED lock keyed on graph_name only, so it
|
||||
# conflicts with the batch path's exclusive graph lock but not with other
|
||||
# single writers.
|
||||
shared_sql = calls[1]["sql"]
|
||||
assert "pg_advisory_xact_lock_shared" in shared_sql
|
||||
assert calls[1]["args"] == ("test_graph",)
|
||||
|
||||
# The cypher statement must not contain a lock — that would cause AGE to
|
||||
# reject the plan with "cypher create clause cannot be rescanned".
|
||||
cypher_sql = calls[2]["sql"]
|
||||
assert "pg_advisory_xact_lock" not in cypher_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edge_lock_key_includes_graph_name():
|
||||
"""Advisory lock key must include graph_name so independent AGE graphs in
|
||||
the same PostgreSQL database don't serialise each other's edges.
|
||||
|
||||
Regression for the codex review on PR #3056: pg_advisory_xact_lock is
|
||||
database-wide, so hashing only (src, tgt) would make {Alice, Bob} in
|
||||
`graph_a` block {Alice, Bob} in `graph_b` even though they touch different
|
||||
AGE graph tables.
|
||||
"""
|
||||
storage_a = make_graph_storage()
|
||||
storage_a.graph_name = "graph_a"
|
||||
storage_b = make_graph_storage()
|
||||
storage_b.graph_name = "graph_b"
|
||||
|
||||
calls_a = await _capture_upsert_edge(storage_a, "Alice", "Bob", {})
|
||||
calls_b = await _capture_upsert_edge(storage_b, "Alice", "Bob", {})
|
||||
|
||||
# graph_name flows as the first positional arg into the lock SQL so the
|
||||
# hashed lock key differs between graphs even when (src, tgt) match.
|
||||
assert calls_a[0]["args"] == ("graph_a", "Alice", "Bob")
|
||||
assert calls_b[0]["args"] == ("graph_b", "Alice", "Bob")
|
||||
|
||||
# And the lock template references graph_name as $1, with the ID pair as
|
||||
# $2/$3 — keep the param order pinned so future refactors don't silently
|
||||
# swap them.
|
||||
lock_sql = calls_a[0]["sql"]
|
||||
assert "$1::text" in lock_sql
|
||||
assert "LEAST($2::text, $3::text)" in lock_sql
|
||||
assert "GREATEST($2::text, $3::text)" in lock_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edge_lock_key_is_endpoint_order_independent():
|
||||
"""{A, B} and {B, A} must produce the same advisory lock key.
|
||||
|
||||
The lock SQL itself is identical across both call directions; only the
|
||||
positional args differ. LEAST/GREATEST inside the template then normalises
|
||||
them to the same hash input, so concurrent {A,B} and {B,A} writes collide
|
||||
on a single lock (matching the undirected OPTIONAL MATCH).
|
||||
"""
|
||||
storage = make_graph_storage()
|
||||
forward = await _capture_upsert_edge(storage, "Alice", "Bob", {})
|
||||
reverse = await _capture_upsert_edge(storage, "Bob", "Alice", {})
|
||||
|
||||
# Same lock SQL template for both directions.
|
||||
assert forward[0]["sql"] == reverse[0]["sql"]
|
||||
# graph_name first, then the endpoint pair in whatever order the caller
|
||||
# passed — LEAST/GREATEST canonicalises inside the SQL.
|
||||
assert forward[0]["args"][0] == reverse[0]["args"][0] == "test_graph"
|
||||
assert (
|
||||
set(forward[0]["args"][1:])
|
||||
== set(reverse[0]["args"][1:])
|
||||
== {
|
||||
"Alice",
|
||||
"Bob",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edge_wraps_transient_errors_for_retry(monkeypatch):
|
||||
"""Query-level transient errors must be wrapped in PGGraphQueryException so
|
||||
the outer @retry predicate can identify them and retry.
|
||||
|
||||
Regression: when upsert_edge was moved off self._query onto
|
||||
self.db._run_with_retry (to run the advisory lock + cypher in one
|
||||
transaction), the _query exception-wrapping path was bypassed. A raw
|
||||
asyncpg.DeadlockDetectedError surfacing from connection.execute would
|
||||
therefore fail _is_transient_graph_write_error's first guard
|
||||
(isinstance(exc, PGGraphQueryException)) and skip the retry loop, silently
|
||||
degrading concurrent ingestion under contention. This test pins the
|
||||
wrapping back in place and asserts the retry loop actually fires.
|
||||
"""
|
||||
# Make the retry loop fire with zero backoff so the test stays fast.
|
||||
monkeypatch.setattr(PGGraphStorage.upsert_edge.retry, "wait", wait_none())
|
||||
|
||||
storage = make_graph_storage()
|
||||
deadlock = asyncpg.exceptions.DeadlockDetectedError("simulated deadlock")
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def fake_run_with_retry(_operation, **_kwargs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
raise deadlock
|
||||
|
||||
storage.db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry)
|
||||
|
||||
with pytest.raises(PGGraphQueryException) as excinfo:
|
||||
await storage.upsert_edge("Alice", "Bob", {"weight": "1.0"})
|
||||
|
||||
# The original asyncpg exception is preserved as __cause__ so the predicate
|
||||
# can introspect it via exc.__cause__.
|
||||
assert excinfo.value.__cause__ is deadlock
|
||||
# And the predicate now recognises this exception as retryable.
|
||||
assert _is_transient_graph_write_error(excinfo.value) is True
|
||||
# Retried up to stop_after_attempt(3) — proves the wrapping actually
|
||||
# engages the @retry loop rather than failing fast on the first attempt.
|
||||
assert call_count == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edge_does_not_retry_non_transient_errors(monkeypatch):
|
||||
"""Non-transient errors must not be retried by the @retry loop.
|
||||
|
||||
The wrapping in upsert_edge unconditionally re-raises as
|
||||
PGGraphQueryException, but _is_transient_graph_write_error only returns
|
||||
True for a small set of asyncpg transient causes. A plain ValueError
|
||||
bubbling out of _run_with_retry should fail fast, not loop 3 times.
|
||||
"""
|
||||
monkeypatch.setattr(PGGraphStorage.upsert_edge.retry, "wait", wait_none())
|
||||
|
||||
storage = make_graph_storage()
|
||||
boom = ValueError("not a transient db error")
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def fake_run_with_retry(_operation, **_kwargs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
raise boom
|
||||
|
||||
storage.db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry)
|
||||
|
||||
with pytest.raises(PGGraphQueryException) as excinfo:
|
||||
await storage.upsert_edge("Alice", "Bob", {"weight": "1.0"})
|
||||
|
||||
assert excinfo.value.__cause__ is boom
|
||||
assert _is_transient_graph_write_error(excinfo.value) is False
|
||||
assert call_count == 1
|
||||
|
||||
|
||||
async def _capture_upsert_edges_batch(storage: PGGraphStorage, edges):
|
||||
"""Run upsert_edges_batch against a fake connection; return captured calls."""
|
||||
conn = _FakeConnection()
|
||||
|
||||
async def fake_run_with_retry(operation, **_kwargs):
|
||||
return await operation(conn)
|
||||
|
||||
storage.db._run_with_retry = AsyncMock(side_effect=fake_run_with_retry)
|
||||
await storage.upsert_edges_batch(edges)
|
||||
return conn.calls
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edges_batch_iterates_in_sorted_order():
|
||||
"""A chunk emits edge cypher in canonical (LEAST, GREATEST) order regardless
|
||||
of insertion order (deterministic dedup / reproducible replay)."""
|
||||
storage = make_graph_storage()
|
||||
|
||||
# Insertion order intentionally non-canonical: B-A, C-A, D-A.
|
||||
calls = await _capture_upsert_edges_batch(
|
||||
storage,
|
||||
[
|
||||
("B", "A", {"weight": "1"}),
|
||||
("C", "A", {"weight": "2"}),
|
||||
("D", "A", {"weight": "3"}),
|
||||
],
|
||||
)
|
||||
|
||||
# The batch chunk takes ONE graph-wide advisory lock (keyed on graph_name
|
||||
# only, no per-edge endpoints), not a per-edge lock. Order is observed from
|
||||
# the cypher calls' bound (src_id, tgt_id) params, canonicalised.
|
||||
lock_calls = [c for c in calls if "pg_advisory_xact_lock" in c["sql"]]
|
||||
assert len(lock_calls) == 1
|
||||
assert lock_calls[0]["args"] == ("test_graph",)
|
||||
cypher_calls = [c for c in calls if "CREATE (source)-[r:DIRECTED" in c["sql"]]
|
||||
edge_pairs = [
|
||||
(json.loads(c["args"][0])["src_id"], json.loads(c["args"][0])["tgt_id"])
|
||||
for c in cypher_calls
|
||||
]
|
||||
canonical_keys = [tuple(sorted(pair)) for pair in edge_pairs]
|
||||
assert canonical_keys == sorted(canonical_keys)
|
||||
assert canonical_keys == [("A", "B"), ("A", "C"), ("A", "D")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_edges_batch_dedupes_last_write_wins():
|
||||
"""Reciprocal duplicates collapse to a single edge upsert carrying the
|
||||
LATEST edge_data, regardless of which orientation arrives last."""
|
||||
storage = make_graph_storage()
|
||||
|
||||
calls = await _capture_upsert_edges_batch(
|
||||
storage,
|
||||
[
|
||||
("A", "B", {"weight": "first"}),
|
||||
("B", "A", {"weight": "second"}), # reciprocal, wins
|
||||
],
|
||||
)
|
||||
|
||||
# Exactly one edge cypher (CREATE) ran, with the latest payload inlined and
|
||||
# the last write's orientation in the bound params.
|
||||
cypher_calls = [c for c in calls if "CREATE (source)-[r:DIRECTED" in c["sql"]]
|
||||
assert len(cypher_calls) == 1
|
||||
assert '"second"' in cypher_calls[0]["sql"]
|
||||
assert '"first"' not in cypher_calls[0]["sql"]
|
||||
params_json = cypher_calls[0]["args"][0]
|
||||
assert json.loads(params_json) == {"src_id": "B", "tgt_id": "A"}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,85 @@
|
||||
"""PGVectorStorage.get_by_id/get_by_ids must not leak the embedding column.
|
||||
|
||||
The SQL-fallback path used ``SELECT *`` and returned ``content_vector`` — a
|
||||
pgvector value the asyncpg codec materializes as a numpy array, which is not
|
||||
JSON-serializable. ``get_entity_info(include_vector_data=True)`` puts that row
|
||||
into the API response, so ``POST /graph/entity/edit`` raised a 500 during
|
||||
response serialization when every storage was PostgreSQL. Both read paths must
|
||||
strip ``content_vector`` to match the buffered shape and the other vector
|
||||
backends.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from lightrag.kg.postgres_impl import PGVectorStorage
|
||||
from lightrag.namespace import NameSpace
|
||||
from lightrag.utils import EmbeddingFunc
|
||||
|
||||
pytestmark = pytest.mark.offline
|
||||
|
||||
|
||||
async def _embed(texts, **kwargs):
|
||||
return np.array([[0.0, 0.0, 0.0] for _ in texts], dtype=np.float32)
|
||||
|
||||
|
||||
def _make_storage(row):
|
||||
storage = PGVectorStorage(
|
||||
namespace=NameSpace.VECTOR_STORE_ENTITIES,
|
||||
workspace="ws",
|
||||
global_config={
|
||||
"embedding_batch_num": 10,
|
||||
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.5},
|
||||
},
|
||||
embedding_func=EmbeddingFunc(
|
||||
embedding_dim=3, func=_embed, model_name="test_model"
|
||||
),
|
||||
)
|
||||
|
||||
async def _query(sql, params, multirows=False):
|
||||
return [row] if multirows else row
|
||||
|
||||
storage.db = AsyncMock()
|
||||
storage.db.query = AsyncMock(side_effect=_query)
|
||||
storage._flush_lock = asyncio.Lock()
|
||||
return storage
|
||||
|
||||
|
||||
def _row():
|
||||
# content_vector is a numpy array in production (pgvector codec).
|
||||
return {
|
||||
"id": "ent-1",
|
||||
"content": "Alice\ndescription",
|
||||
"entity_name": "Alice",
|
||||
"source_id": "chunk-1",
|
||||
"file_path": "doc.txt",
|
||||
"content_vector": np.array([0.1, 0.2, 0.3], dtype=np.float32),
|
||||
"created_at": 1,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_by_id_strips_content_vector():
|
||||
storage = _make_storage(_row())
|
||||
|
||||
result = await storage.get_by_id("ent-1")
|
||||
|
||||
assert result is not None
|
||||
assert "content_vector" not in result
|
||||
assert result["content"] == "Alice\ndescription"
|
||||
assert result["id"] == "ent-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_by_ids_strips_content_vector():
|
||||
storage = _make_storage(_row())
|
||||
|
||||
results = await storage.get_by_ids(["ent-1"])
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0] is not None
|
||||
assert "content_vector" not in results[0]
|
||||
assert results[0]["content"] == "Alice\ndescription"
|
||||
Reference in New Issue
Block a user