feat(kg): add PGTableGraphStorage — PostgreSQL-native graph backend (no AGE dependency)

[ysys143]

Summary

PGTableGraphStorage — a GRAPH_STORAGE backend on plain PostgreSQL tables + B-tree indexes, with no Apache AGE, no Cypher, and no custom extensions. Closes #3095.

Why: AGE can't be installed on managed PostgreSQL (RDS, Cloud SQL, Supabase, Neon). This backend makes LightRAG's graph layer run on any stock PostgreSQL 14+.

Renamed from PgRcteGraphStorage per review — traversal is now a frontier-capped iterative BFS over indexed tables, not a recursive CTE, so "Table" describes it better.

Changes

• lightrag/kg/pgtable_impl.py — PGTableGraphStorage (all 18 BaseGraphStorage methods + 5 batch overrides)
• lightrag/kg/__init__.py — registers in STORAGE_IMPLEMENTATIONS / STORAGE_ENV_REQUIREMENTS / STORAGES
• lightrag/kg/postgres_impl.py — pgvector imported lazily (stock-PG boxes can load the module)
• .github/workflows/pg-smoke.yml — pg-smoke on PostgreSQL 18

Schema (auto-created on initialize())

lightrag_graph_nodes(workspace, namespace, id, properties JSONB)
  PRIMARY KEY (workspace, namespace, id)
lightrag_graph_edges(workspace, namespace, src_id, tgt_id, properties JSONB, updated_at)
  PRIMARY KEY (workspace, namespace, src_id, tgt_id)

Implementation notes

• Traversal: exact-label get_knowledge_graph is a frontier-capped iterative BFS bounded by max_nodes (degree-priority admission, seed pinned). Replaces a recursive CTE that blew up on hubs / deep max_depth — see the depth sweep.
• Undirected canonicalization via Python min/max (never SQL LEAST/GREATEST — avoids non-C collation divergence).
• Node upsert merges properties (jsonb ||) instead of replacing; entity_id required.
• Transient write retry (deadlock / serialization / lock / cancel).
• Same env vars as PGGraphStorage.

Benchmark (PostgreSQL 18, GCP n2-standard-8, BA 8k nodes / ~40k edges)

Backend	Seed	RPS	p50	p95
PGTable	3.0 s	1,431	3.5 ms	37.3 ms
Neo4j	4.0 s	285	18.6 ms	106 ms
AGE	434 s	73	9.6 ms	1,072 ms
NetworkX (in-mem)	0.1 s	3,679	2.1 ms	8.9 ms

get_knowledge_graph p50: PGTable 39 ms vs AGE 1,099 ms (~28×); seed ~145×; get_node_edges/upsert_edges_batch ~25×.

PG backends are VACUUM ANALYZEd after bulk-load before measuring — without it the first concurrent run races autovacuum and read latency spikes ~25× as a measurement artifact (not a backend property).

Cross-backend depth sweep — get_knowledge_graph latency (ms) as max_depth grows (BA n=3000, max_nodes=1000, hub seed degree 245):

depth	NetworkX	PGTable	AGE	Neo4j
1	12.5	35.7	370.8	25.9
2	38.9	129.4	547.6	355.9
3	175.9	110.2	525.4	463.5
5	43.2	234.1	564.1	446.1

PGTable's frontier BFS beats AGE at every depth and matches/beats Neo4j at deep hubs (depth 5: 234 vs 446 ms) — a plain-table BFS competing with a native graph DB. The /graphs API defaults to max_depth=3 with no upper bound on depth; the recursive-CTE approach this PR initially used (then dropped) timed out at hub max_depth ≥ 3, which is exactly why it was replaced by the frontier BFS. Raw results: gist.

Review follow-up

• Codex items 1–7 addressed (namespace isolation, lazy pgvector, test layout, smoke marker, seed-pin truncation, search guard + relevance [moved to Python _search_score, same ordering], edge canonical sort).
• Additional cross-review fixes: wildcard no longer seed-pins a node whose id is literally "*"; self-loop double-count in get_nodes_edges_batch fixed; SELECT ... FOR UPDATE in legacy edge normalization (serializes the one-time migration against concurrent upserts). Regression tests added.
• Sibling-backend alignment: constructor validate_workspace (matches PGGraphStorage / NetworkXStorage); wildcard get_knowledge_graph("*") is bounded in SQL (degree-ordered LIMIT, no full-table fetch); get_nodes_edges_batch OR predicate → index-friendly UNION; PGTableGraphStorage registered in the setup wizard (scripts/setup/lib/storage_requirements.sh).

Known constraints

• The legacy-edge FK cleanup in DDL operates table-wide (FK is a global composite constraint); one-time migration guarded by an advisory lock.
• When multiple LightRAG instances share one PostgreSQL, the connection-pool signature includes vector_storage, so mixing different non-vector vector backends against the same pool is rejected by design.

Tests

Unit (mock) + integration smoke (-m pg_smoke) on PG18 + NetworkX equivalence (19) + batch-ordering. Full offline suite green (2076 passed) on Python 3.12 and 3.14 after merging upstream/main.

[ysys143]

Applied ruff format (v0.6.4, matching the pre-commit config) to address the linting CI check. ruff check was clean; I ran the formatter locally against a newer version which didn't flag any differences — missed that the CI pins to v0.6.4. No logic changes. @danielaskdd

[ysys143]

Rebased onto the latest main and ran pre-commit run --all-files locally — all checks pass now.

The previous CI failure was caused by pre-existing lint issues in upstream files (env.example trailing whitespace, lightrag/pipeline.py ruff-format) that the CI catches via --all-files on the merge commit. Those weren't introduced by this PR.

[danielaskdd]

Thanks for this — a pure-SQL graph backend that drops the AGE requirement is genuinely valuable for managed Postgres (RDS / Cloud SQL / Supabase / Neon). The structure is clean and all 18 BaseGraphStorage methods are present. Below are the issues I think should be resolved before merge, plus notes on the benchmark framing.

Blocking

1. Independent connection pool bypasses ClientManager / PostgreSQLDB
_PgPool calls asyncpg.create_pool() directly instead of going through the shared, ref-counted pool that PGKVStorage / PGVectorStorage / PGGraphStorage all use (ClientManager in postgres_impl.py). Consequences when this backend runs alongside the other PG storages:

• Connection counts stack (each pool has its own max_size), making too many connections easy to hit against the default PG max_connections=100.
• You lose the shared infra that lives in PostgreSQLDB: SSL config (POSTGRES_SSL_*), connection retry/backoff (POSTGRES_CONNECTION_RETRIES), the POSTGRES_WORKSPACE precedence, and centralized table creation/migration (check_tables).
• Workspace handling diverges: this uses self.workspace directly and ignores the POSTGRES_WORKSPACE env var the rest of the codebase honors.

Recommend reusing ClientManager.get_client() / release_client() and PostgreSQLDB.query/execute rather than a private pool.

2. Undirected edges are not normalized on write
The contract is undirected, and PGGraphStorage canonicalizes edges (LEAST/GREATEST) so (A,B) and (B,A) are one edge. Here the PK is (workspace, src_id, tgt_id) with no normalization, so upsert_edge(A,B) followed by upsert_edge(B,A) produces two rows. That means:

• a duplicated undirected edge with possibly divergent properties (get_edge then returns an arbitrary one), and
• node_degree (src_id=$2 OR tgt_id=$2) double-counts that edge.

The "both directions" handling only covers the read path; the write path still needs normalization. Note the equivalence tests likely won't catch this since each backend is ingested independently.

3. Batch write methods are not overridden
Only the read batches are overridden. upsert_nodes_batch / upsert_edges_batch / has_nodes_batch fall back to the base-class N+1 serial loops, so ingestion does one acquire+execute per node/edge — much slower than PGGraphStorage's chunked batched MERGE. (This is also why the benchmark's single-upsert_node mix doesn't reveal it.)

Correctness / structure (non-blocking)

• get_knowledge_graph truncation is non-deterministic: SELECT DISTINCT ... LIMIT max_nodes has no ORDER BY, so which nodes survive truncation is arbitrary — whereas the AGE backend ranks by degree. The two are not equivalent here, and results aren't reproducible on large graphs. Consider ordering by degree.
• is_truncated = len(node_ids) >= max_nodes falsely reports truncation when the graph has exactly max_nodes nodes and nothing was cut.
• The DISTINCT comment ("deduplicates by id across depths") is inaccurate — DISTINCT applies to the whole (id, properties) row, not id alone.
• idx_lightrag_graph_nodes_props (GIN on properties) isn't used by any query in this file — pure write amplification unless a property-filter query is planned.
• Module docstring names the tables graph_nodes/graph_edges, but the DDL creates lightrag_graph_nodes/lightrag_graph_edges.

Benchmark credibility

The directional result is believable — plain indexed SQL > Neo4j > AGE, and AGE's per-call cypher() parse overhead blowing up p95 under concurrency is a known effect. Two caveats on the framing though:

• The workload is described as mirroring LightRAG's real call distribution, but it's point-queries only — it excludes get_knowledge_graph (the expensive recursive query) and the *_batch methods that LightRAG's retrieval and ingestion paths actually rely on (including the write batches not optimized here). So it tends to flatter this implementation specifically where it's weakest.
• It's not reproducible as presented: no script/raw data, request counts, variance, or the pool sizes/configs used for each backend. Equalizing pool size across backends matters a lot for an RPS comparison.
• The equivalence claim covers 5 read ops and explicitly not get_knowledge_graph, writes, or deletes — i.e. it doesn't exercise the normalization/truncation issues above.

Happy to help with any of these if useful.

---

Generated by Claude Code

[danielaskdd]

Thanks a lot for the contribution, @ysys143 — a pure-SQL graph backend that removes the AGE dependency is a genuinely useful direction for users on managed Postgres (RDS / Cloud SQL / Supabase / Neon), and the implementation is clean and well-organized.

Before we can merge, could you please refine a few things? I've left the details in the review above; the key ones are:

1. Use the shared connection pool. Please route through ClientManager / PostgreSQLDB instead of creating a private asyncpg pool, so this backend doesn't stack connections alongside the other PG storages and so it inherits SSL config, retry/backoff, POSTGRES_WORKSPACE handling, and centralized table creation. (This is the point raised earlier in the thread.)
2. Normalize undirected edges on write (e.g. LEAST/GREATEST, as PGGraphStorage does), so (A,B) and (B,A) don't become two rows and inflate node_degree.
3. Override the batch-write methods (upsert_nodes_batch / upsert_edges_batch / has_nodes_batch) — they currently fall back to N+1 serial calls, which will slow ingestion.
4. Smaller fixes: make get_knowledge_graph truncation deterministic (add an ORDER BY, ideally by degree to match the AGE backend), fix the is_truncated edge case, drop the unused GIN index, and align the docstring table names with the DDL.

On the benchmark: the direction is believable, but could you share the benchmark script and raw data, and extend the workload to include get_knowledge_graph and the *_batch methods (which the real retrieval/ingestion paths use)? Equalizing the pool size/config across the three backends would also make the RPS comparison more apples-to-apples. Similarly, it'd be great to extend the equivalence tests beyond the 5 read ops to cover writes/deletes and get_knowledge_graph.

Happy to help if any of these are unclear. Looking forward to the next revision!

---

Generated by Claude Code

[ysys143]

Revision — PgRcteGraphStorage

Thanks for the detailed review. All blocking and non-blocking items have been addressed.

---

What changed

1. Shared connection pool (ClientManager / PostgreSQLDB)

_PgPool / asyncpg.create_pool() removed. PgRcteGraphStorage now routes through ClientManager.get_client(), the same ref-counted pool used by PGKVStorage and PGVectorStorage. SSL config (POSTGRES_SSL_*), retry/backoff, POSTGRES_WORKSPACE precedence, and centralised table creation are all inherited automatically.

2. Undirected edge normalisation on write

All write paths now canonicalise src_id = min(a, b), tgt_id = max(a, b) before INSERT. upsert_edge(A, B) and upsert_edge(B, A) now write a single row — node_degree no longer double-counts.

3. Batch write overrides

Method	Implementation
upsert_nodes_batch	Single INSERT … SELECT FROM unnest($ids, $props) — last-write-wins dedup in Python before SQL
upsert_edges_batch	Same unnest INSERT after normalisation + dedup
has_nodes_batch	Delegates to get_nodes_batch (WHERE id = ANY($2))

4. Non-blocking fixes

• get_knowledge_graph: fetch max_nodes + 1 rows; is_truncated = len > max_nodes; sort by degree descending to match AGE behaviour
• is_truncated off-by-one fixed
• Unused GIN index removed from DDL
• Docstring table names corrected

---

Benchmark

Infrastructure: GCP n2-standard-8 (8 vCPU, 32 GB), Barabási-Albert graph seed=42 — 8,000 nodes / ~40,000 edges, 10 workers × 30 s, POSTGRES_MAX_CONNECTIONS=20 identical for all PG backends.

Reproduce with:

docker compose -f docker-compose.bench.yml up -d --wait
python bench/bench_graph_ops.py --nodes 8000 --workers 10 --duration 30 --output bench-results.json

Raw data: bench/bench-results.json

Backend	Seed	RPS	p50	p95	p99
PgRcteGraphStorage	1.3 s	591	1.1 ms	173.5 ms	224.7 ms
Neo4JStorage	6.6 s	282	17.7 ms	115.4 ms	263.7 ms
PGGraphStorage (AGE)	433 s	69	7.2 ms	1,201 ms	1,396 ms
NetworkXStorage (in-memory baseline)	0.1 s	3,740	2.0 ms	8.8 ms	15.0 ms

Per-operation breakdown — PgRcteGraphStorage vs PGGraphStorage (AGE):

Operation	RCTE p50	AGE p50	RCTE p95	AGE p95
get_node	0.6 ms	2.1 ms	2.2 ms	21 ms
has_node	0.6 ms	1.8 ms	2.6 ms	23 ms
node_degree	0.7 ms	2.8 ms	2.4 ms	24 ms
get_node_edges	1.3 ms	130 ms	4.8 ms	222 ms
get_nodes_batch	0.8 ms	2.1 ms	2.6 ms	18 ms
upsert_node	1.3 ms	5.0 ms	12.7 ms	34 ms
upsert_nodes_batch	1.4 ms	15.2 ms	13.5 ms	44 ms
upsert_edges_batch	1.5 ms	108 ms	13.2 ms	218 ms
get_knowledge_graph	181 ms	1,230 ms	245 ms	1,468 ms

RCTE p95 (173 ms) is dominated by get_knowledge_graph (8% of workload). Point queries are sub-1 ms p50. AGE seed time (433 s) vs RCTE (1.3 s) reflects per-row Cypher upsert vs chunked unnest INSERT.

---

Test coverage

Offline unit tests — TestPgRcteBatchOrdering (3 tests, no DB required):

• upsert_nodes_batch deduplicates, issues a single _execute call
• upsert_edges_batch collapses (A,B) + (B,A) to one row; last write wins
• Canonical order is lexicographic min/max regardless of call order

Offline pipeline integration test — TestPgRctePipelineIntegration (1 test, no DB required):

• Drives ainsert_custom_kg() against the real batch write logic with DB methods patched
• Verifies batch writes are issued (not N+1 serial upserts)
• Verifies a relation submitted as (EntityB, EntityA) reaches SQL as (EntityA, EntityB) — canonical normalisation flows end-to-end through the pipeline

Equivalence tests — 19 tests covering all BaseGraphStorage methods vs NetworkXStorage reference: read ops, write/update ops, delete ops, and all batch read/write methods including reversed-pair upsert_edges_batch + node_degree check.

E2E smoke test (run locally on GCP VM, not included in upstream test suite per project conventions):

• LightRAG.ainsert() → ALICE, BOB nodes + edge written to PostgreSQL → node_degree == 1 (normalisation correct) → LightRAG.aquery() completes — PASSED (1.52 s, Python 3.14, PostgreSQL 17)

[danielaskdd]

🔴 Blocking — get_node_edges return order violates the contract and drops data on entity merge

All three reference backends return get_node_edges(X) with the queried node always first — (X, neighbor):

• NetworkX: list(graph.edges(X)) always reports the nbunch node first;
• Neo4j (neo4j_impl.py L952): source_label = n.entity_id = X;
• AGE (postgres_impl.py L6352): RETURN n.entity_id AS source_id.

PgRcteGraphStorage.get_node_edges instead returns the canonical (min, max) pair, so the queried node can land in either position. A downstream consumer depends on the "queried node first" contract — _merge_entities_impl in lightrag/utils_graph.py#L1319-L1323:

for src, tgt in edges:
    if src == entity_name:          # no else branch
        edge_data = await ...get_edge(src, tgt)
        all_relations.append((src, tgt, edge_data))

Under this backend, every edge where the merged entity's id is lexicographically greater than its neighbor's gets src == neighbor → the condition is false → the relation is silently dropped. Net effect: merging an entity loses roughly half of its relations, with no error.

Note the rename path (utils_graph.py#L342-L351) handles both sides (if source == entity_name ... else: # target == entity_name), which confirms the ordering convention is real and relied upon.

Why the equivalence test didn't catch it: check_equiv in bench/compare_three_way.py compares get_node_edges via frozenset(e), which deliberately erases tuple order, so "ALL MATCH" does not validate this contract.

Suggested fix: normalize the queried node to position 0 in both get_node_edges and get_nodes_edges_batch, e.g.

return [(source_node_id, t if s == source_node_id else s) for s, t in rows]

and add a test that asserts ordering (not just frozenset equality).

[danielaskdd]

🟡 get_knowledge_graph truncation is non-deterministic and does not match AGE (and is untested)

In get_knowledge_graph, the SELECT DISTINCT ... LIMIT fetch_limit truncates over a recursive CTE whose row order is not deterministic. The subsequent node_rows.sort(key=lambda r: degrees...) runs after node_rows[:max_nodes], so it only re-orders an already-arbitrarily-truncated set.

The code comment claims this is done "to match AGE backend behaviour under truncation", but the set of nodes that survives truncation is arbitrary here — not degree-prioritized — so it does not actually match AGE when the graph exceeds max_nodes.

This path was never validated: the equivalence suite uses a 5-node graph with max_nodes=100 (no truncation ever triggers), and the perf path doesn't check correctness.

Suggested fix: push the ordering into SQL before the limit (ORDER BY degree DESC LIMIT ...) so truncation is deterministic and degree-prioritized — or, if matching AGE exactly isn't a goal, update the PR description / comment to state that the truncated node set is not guaranteed to match AGE.

[danielaskdd]

🟡 get_knowledge_graph traversal cost is not bounded by max_nodes

The PR describes get_knowledge_graph as "bounded BFS, max_depth + max_nodes respected", but max_nodes is only the outer LIMIT — it bounds the output, not the traversal cost. WITH RECURSIVE materializes the entire reachable connected component first, then applies LIMIT.

This is most visible for node_label = "*": the base case already selects every node in the workspace, so the recursive part is pure dead weight, and the whole node/edge set is walked regardless of max_nodes. On large graphs (tens of thousands of nodes) this can be slow and memory-heavy.

Suggested fix: for node_label = "*", use a flat SELECT ... LIMIT instead of the recursive CTE; for the seeded case, document that cost scales with the reachable component, not with max_nodes.

[danielaskdd]

🟡 The real-DB glue is not covered in CI

All unit tests mock _execute / _fetch / _fetchrow / _fetchval, and the pipeline integration test patches them too. The actual code paths — _execute → PostgreSQLDB.execute(data=dict) positional binding, and json.loads() on JSONB columns returned by asyncpg — are only exercised by the bench/ scripts, which require a multi-container setup and are not part of CI.

The param-passing glue is subtle (an {str(i): v} dict whose .values() order must line up with $1..$N). It is in fact correct today, but nothing in CI would catch a regression.

Suggested fix: add a real-PostgreSQL smoke test (e.g. testcontainers, CI-gated) covering at least upsert_node/get_node/upsert_edge/get_edge/get_node_edges against a live instance.

[danielaskdd]

🟢 Minor / style

• DDL comment vs implementation mismatch: the docstring/DDL comment says src_id = LEAST(a, b), but canonicalization is done entirely in Python via min/max. It is consistent today (all read/write paths use Python), but Python compares by Unicode code point while SQL LEAST/GREATEST uses the DB collation — if anyone later adds SQL-side canonicalization the two could disagree and produce duplicate edges. Suggest rewording the comment to "canonicalized in Python via min/max" and noting that SQL-side normalization must not be introduced.
• Missing validate_workspace: other PG storages call validate_workspace(self.workspace) in __post_init__ (postgres_impl.py L2512); PgRcteGraphStorage does not.
• import asyncio inside edge_degree should be a module-level import.
• bench/ (~1070 lines) committed to the repo: hardcoded credentials/ports, and AgeDirectStore builds Cypher via string interpolation. Acceptable for a throwaway benchmark, but it adds untested maintenance surface — could you confirm whether these belong in-repo or in the linked gist?

[danielaskdd]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: cb94c61651

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[danielaskdd]

Code review — PGTableGraphStorage

Reviewed the new backend end-to-end. The implementation is careful and conforms to BaseGraphStorage faithfully — most candidate issues turned out to be documented-intentional or consistent with the reference backends (NetworkX / PGGraphStorage), including the min/max canonical single-row undirected design. The findings below are the ones worth acting on, most-severe first.

1. [Blocking / correctness] Legacy-edge normalization can crash initialize()

lightrag/kg/pgtable_impl.py — _normalize_legacy_edges (~L531–554)

When a group contains both the canonical row (A,B) and its reversed duplicate (B,A), old_srcs/old_tgts include the canonical target (A,B) itself. The statement

WITH deleted AS (DELETE ... (A,B),(B,A)) INSERT (A,B) ON CONFLICT DO UPDATE

runs as a single command on one snapshot. Per the Postgres docs, the effects of a data-modifying statement in WITH are not visible to the rest of the query, so the INSERT still sees (A,B) as live, performs DO UPDATE on it, and that tuple is concurrently deleted by the same command →

ERROR: tuple to be updated was already modified by an operation triggered by the current command

The whole migration transaction rolls back, so the backend fails on every startup and cannot self-heal. This is exactly the migration's target scenario (a legacy table written before canonicalization that holds both directions). The offline unit tests mock _normalize_legacy_edges, so this real SQL path is uncovered.

Suggested fix: exclude the canonical target (src,tgt) from the delete set (delete only the non-canonical members), then plain INSERT ... ON CONFLICT. Please add a real-PG regression test with both (A,B) and (B,A) present.

2. [Medium / robustness] Read paths have no query-level transient retry

lightrag/kg/pgtable_impl.py — _fetch / _fetchrow / _fetchval (~L381–393)

Only _execute (writes) retries deadlock/serialization/lock/cancel via tenacity. Reads go straight through PostgreSQLDB.query, which only retries connection-level transients (_run_with_retry uses self._transient_exceptions), not serialization conflicts. get_knowledge_graph's BFS issues many query calls; under a SERIALIZABLE / REPEATABLE READ pool, a read racing concurrent ingestion can surface SerializationError straight to the caller and 500 the /graphs request. Rare under the default READ COMMITTED, but exposed at higher isolation. Suggest wrapping reads in the same transient-retry, or documenting that this backend assumes READ COMMITTED.

3. [Low / data fidelity] Nondeterministic survivor in legacy normalization

lightrag/kg/pgtable_impl.py (~L519–525)

Reversed duplicates written in the same legacy batch share updated_at, so max(group, key=...) returns whichever the DB ordered first — meaning which edge's properties (weight/description) survives the collapse is arbitrary and not reproducible across runs/replicas. Suggest a deterministic secondary tie-break (e.g. prefer the already-canonical row).

4. [Test coverage] Validate against the shared contract suite

tests/kg/test_graph_storage.py

Every graph backend must pass the shared end-to-end contract suite (basic / advanced / batch ops / undirected / special chars / string escaping / batch upsert / query helpers). The harness is registry-driven and this PR already registers PGTableGraphStorage, so no harness changes are needed.

Now (please attach before merge): run it locally and paste the result. It's a CLI program — run it and pick 9 (All Tests), especially 4 (Undirected Property Test), which is what validates that your min/max canonical single-row storage actually satisfies undirected semantics. PGGraphStorage passes it via mirrored bidirectional edges; you take a different route, so this test matters more here:

# interactive (the file's intended usage)
LIGHTRAG_GRAPH_STORAGE=PGTableGraphStorage POSTGRES_HOST=... POSTGRES_PORT=5432 \
POSTGRES_USER=... POSTGRES_PASSWORD=... POSTGRES_DATABASE=... \
python tests/kg/test_graph_storage.py        # choose 9 = All Tests (or 4 = Undirected)

# or non-interactive (pytest)
LIGHTRAG_GRAPH_STORAGE=PGTableGraphStorage POSTGRES_*=... \
pytest tests/kg/test_graph_storage.py -m integration --run-integration -v

Please also refresh the stale backend list in the file's docstring (tests/kg/test_graph_storage.py:8–13, currently missing PGTableGraphStorage).

Follow-up (suggested, non-blocking): wire the shared suite into CI as a per-backend gate. The current .github/workflows/pg-smoke.yml only runs the bespoke tests/kg/pgtable_impl/test_pgtable_smoke.py, so "every backend passes the same contract" isn't enforced in CI yet. Add a step running tests/kg/test_graph_storage.py with LIGHTRAG_GRAPH_STORAGE=PGTableGraphStorage.

When wiring it in, the false-green trap must be solved (both suites share it):

• test_pgtable_smoke.py:31–36 does pytest.skip(allow_module_level=True) when POSTGRES_PASSWORD is unset;
• test_graph_storage.py's storage() fixture skips when .env / the backend isn't configured.

Either way the job exits 0 green having run zero tests if the env ever fails to propagate. In CI (where the DB is guaranteed) a missing DB should fail loudly — e.g. assert skipped == 0 / passed > 0 via --junitxml, or add an explicit CI-only switch (PGTABLE_SMOKE_REQUIRED=1) that raises instead of skipping.

5. [Acknowledged — no action] Known constraints

The table-wide orphan-edge DELETE on every startup, and the per-hop BFS cost being unbounded by max_nodes, are both already documented in code comments / the PR description as known constraints / tracked follow-ups. Confirmed during review; noted only for completeness.

[ysys143]

Thanks for the careful review. Addressed in d625c921, dc022055, 00efdad3 (head 00efdad3).

1. Legacy-edge normalization

Fixed: the canonical (src, tgt) row is now excluded from the CTE delete set, so the delete set and the ON CONFLICT INSERT target are disjoint.

One note for transparency: while verifying I tried to reproduce the tuple to be updated was already modified by an operation triggered by the current command error directly, seeding both (A,B) and (B,A), and on PG 15 / 16 / 17 / 18 the original statement actually returned the correct single canonical row with no error. But the concern stands regardless — it relied on modifying the same row twice in one data-modifying CTE, which the docs leave unspecified, so the result isn't guaranteed across versions/plans. The fix makes it well-defined everywhere. Added a real-PostgreSQL regression test (tests/kg/pgtable_impl/test_pgtable_smoke.py) that seeds the reversed duplicate and asserts a single canonical survivor — this also exercises the migration path the offline unit tests mock away.

3. Non-deterministic survivor

Fixed: the survivor tie-break now prefers the already-canonical row when updated_at ties, so which properties survive is reproducible.

2. Read-path transient retry

Fixed: factored the transient-retry into a shared _with_retry helper used by _fetch / _fetchrow / _fetchval as well as _execute (predicate renamed _is_transient_error, since it now covers reads). No effect under the default READ COMMITTED, but a read racing concurrent ingestion under SERIALIZABLE / REPEATABLE READ no longer surfaces a SerializationError straight to the caller.

4. Shared contract suite

Ran it against PGTableGraphStorage on real PG18:

LIGHTRAG_GRAPH_STORAGE=PGTableGraphStorage pytest tests/kg/test_graph_storage.py -m integration --run-integration
=> 8 passed, 0 skipped

All 8 pass, including test_graph_undirected_property — confirming the min/max single-row design satisfies the undirected contract end-to-end. Wired it into pg-smoke.yml as a per-backend gate, with the false-green guard you flagged: it parses the JUnit XML and fails unless tests > 0 and skipped == 0. Also added PGTableGraphStorage to the suite's docstring backend list.

Additional fix found while reviewing

The namespaced-PK migration ran ADD PRIMARY KEY directly; on a legacy table that holds duplicate rows for the new key this fails (could not create unique index … is duplicate) and aborts startup — verified directly. Added a de-dupe (keep newest by updated_at, ctid tiebreak) before each PK rebuild, with a real-PG regression test.

Considered, no change needed

• get_knowledge_graph is_truncated: checked against NetworkXStorage directly — it returns is_truncated=False for a depth-limited cut (it only flags the max_nodes budget cut), so the current behaviour already matches the reference; I reverted an earlier attempt to flag depth cuts as it would have diverged.
• Edge-property update is REPLACE (matching PGGraphStorage): the abstract BaseGraphStorage.upsert_edge contract is silent on merge-vs-replace, and LightRAG's pipeline always passes the full property set and merges at the application layer (operate.py), so it's unobservable to callers.

Follow-ups (not in this PR)

The table-wide orphan cleanup and the per-hop BFS cost being unbounded by max_nodes remain documented as known constraints. The batch-size chunking PGGraphStorage does via _chunk_by_budget is a reasonable robustness follow-up for very large *_batch calls — I'll track it separately rather than grow this PR.

[danielaskdd]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 00efdad3a9

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[ysys143]

Follow-up on the contract-suite CI gate (commit 8c1bb071).

While double-checking the false-green guard I added for #4, I found a second false-green that's worth flagging for the shared suite itself: most tests in tests/kg/test_graph_storage.py wrap their body in try/except Exception and return False on failure (e.g. test_graph_basic, test_graph_advanced, test_graph_batch_operations, …). pytest records a non-None return as a passed test (only a PytestReturnNotNoneWarning), so a failed assertion inside one of those tests is swallowed and the run still reports N passed. This affects any backend that runs this suite in CI, not just PGTableGraphStorage.

Verified locally: injecting a failure into test_graph_basic still reports 8 passed, and a JUnit failures/errors check (my original guard) does not catch it.

For now I made the PGTable CI step robust to it: run the suite with -s so the except block's "An error occurred during the test" marker reaches the log even for a swallowed-then-passed test, and fail the step if that marker appears (in addition to the tests > 0 / skipped == 0 / JUnit check). A cleaner upstream fix would be to let those tests assert/raise instead of return False, so pytest reports real failures directly — happy to send that as a separate PR if you'd like.

[danielaskdd]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 3354c302f3

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[danielaskdd]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 1548d7a43c

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Use a one-character LIKE escape

When PGTableGraphStorage.search_labels() runs against PostgreSQL with the default standard_conforming_strings=on, the SQL literal in ESCAPE '\\' evaluates to two backslash characters, but PostgreSQL requires the ESCAPE value to be empty or exactly one character. This makes any label search on the new PGTable backend fail at execution time instead of returning matching labels; use a single-character escape literal such as ESCAPE E'\\' or otherwise avoid the doubled standard string literal.

Useful? React with 👍 / 👎.

[ysys143]

This looks like a false positive on the default configuration — verified on PG18.

The query is a regular (non-raw) Python triple-quoted string, so the source '\\' is already a single backslash by the time the SQL reaches Postgres (Python collapses \\ -> \), not two. Checked directly:

• the SQL text sent to the server contains exactly one backslash;
• standard_conforming_strings is on (the Postgres default since 9.1);
• ... LIKE $1 ESCAPE '\' executes with no "invalid escape string" error (returns a normal boolean).

The shared contract suite's test_graph_query_helpers also exercises search_labels against a live PG18 and passes, so the escape clause works end to end.

The only configuration where a single-backslash literal would be a problem is the non-default, deprecated standard_conforming_strings=off. If you'd like it to be robust there too I can switch to ESCAPE E'\\' (explicit escape-string syntax, unambiguously one character regardless of that setting), but no change is needed for default Postgres.

[danielaskdd]

Agreed — this is a false positive, your analysis is correct.

The literal here is a regular (non-raw) triple-quoted Python string, so the \ in the source collapses to a single backslash before the SQL ever reaches Postgres. The text on the wire is ... ESCAPE '', and under standard_conforming_strings=on (the default since 9.1) that's a one-character escape literal — exactly what ESCAPE expects. The bot conflated the source-level representation (two characters in the .py file) with the runtime string value (one backslash). It would only hold if the string were raw (r"""..."""), which it isn't. The contract suite exercising search_labels against live PG18 confirms it end to end.

So no change is needed for default Postgres. The one edge case is the non-default, deprecated standard_conforming_strings=off. Since E'\' is unambiguously one character regardless of that setting and the change is essentially free, I'd lean toward applying it purely as defensive hardening — but I'm happy either way; it's not blocking.

Resolving this one as non-actionable for the default configuration.

[danielaskdd]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Bravo.

Reviewed commit: 1548d7a43c

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[danielaskdd]

Review summary

Did a full read of pgtable_impl.py and cross-checked it against the NetworkX / PGGraphStorage references it mirrors. No correctness bugs found — the implementation is solid. Verified as correct (so reviewers needn't re-check): $N-positional param plumbing through both PostgreSQLDB.execute/.query; single-statement edge insert + FK validity in the WITH endpoints … CROSS JOIN endpoint_write upserts; the legacy-edge migration's disjoint DELETE/INSERT sets and updated_at <= $8 concurrency guard; self-loop degree consistency; the BFS budget off-by-one + degree-priority truncation; and _search_score mirroring NetworkX.

One thing I want to flag as a non-issue in case it comes up: is_truncated reflects only the max_nodes axis, never depth-limit truncation. That's correct — NetworkXStorage does exactly the same (networkx_impl.py:598; its depth-limit branch only logs), so the behavior is faithful to the reference.

The remaining items are all test-hardening suggestions — nothing blocking:

1. tests/kg/pgtable_impl/test_pgtable_graph_storage.py:436 — test_get_knowledge_graph_uses_global_config_when_max_nodes_none mocks _fetch → [], so it returns at the empty-result early-exit before the max_graph_nodes=42 cap is ever read. It only asserts kg.nodes == [] and would pass even if the None → global_config fallback were broken. Suggest returning >42 rows and asserting truncation at 42 so the fallback is actually exercised.

2. Edge-replace contract has no test. Node upsert merges (||) while edge upsert replaces properties. The merge side is asserted, but no test upserts the same edge twice with different keys to confirm old keys are dropped — a regression to || on edges wouldn't be caught. Worth adding (smoke or mock).

3. get_node_edges self-loop untested. Only the batch variant get_nodes_edges_batch is tested for "self-loop appears once"; the single-method get_node_edges (separate code path) is never exercised with an (A, A) edge.

4. Batch entity_id validation untested. upsert_nodes_batch's ValueError-when-entity_id-absent and forced canonicalization of a mismatched entity_id are only covered for the single upsert_node. The batch dedup test uses inputs whose entity_id already equals node_id.

5. test_pgtable_graph_storage.py:720 (minor) — test_self_loop_degree_consistency hardcodes _fetchval → 2, so the node_degree == 2 assertion is guaranteed by the mock rather than the SQL. The real CASE-clause logic is covered by the smoke test, so this is low priority.

Overall: nice work — clean, well-documented, and the integration smoke suite genuinely exercises the core paths against live PostgreSQL. The above are just suggestions to close coverage gaps.

---

Generated by Claude Code

[ysys143]

Pushed two follow-up commits addressing the test-coverage gaps from the review (head now 3417bd9e):

• 1dcee326 test(kg): close coverage gaps

  • test_get_knowledge_graph_uses_global_config_when_max_nodes_none was vacuous (it mocked _fetch -> [] and hit the empty-result early-exit before the max_graph_nodes cap was ever read). Rewritten to return 50 rows (> the 42 cap), asserting the node query requested cap+1 (43), the result truncates to 42, and is_truncated is True.
  • Added test_upsert_nodes_batch_raises_without_entity_id and test_upsert_nodes_batch_overwrites_mismatched_entity_id so the batch path enforces entity_id like the single-row upsert_node.
  • Added real-PG smoke tests: test_upsert_edge_replaces_properties (edge upsert REPLACES, unlike node merge) and test_get_node_edges_self_loop_appears_once.

• 3417bd9e fix(kg): use an E'' escape literal in search_labels LIKE

  • Changed ESCAPE '\' to ESCAPE E'\\' so it is well-defined regardless of standard_conforming_strings.

Verified: shared contract suite 8/8 (real PG18, 0 skipped), full offline suite green on 3.12 and 3.14, and pre-commit clean at the pinned ruff rev.

[dan1els]

I really looking forward for that. AGE messing up with planner predictions which causing huge loops and performance significantly degrading.
Btw are you going to provide any migration guides from AGE to this schema?