perf(kg-extract): migrate KG triple worker to psycopg3 AsyncConnectionPool

[jphein]

Summary

Migrates the KG triple-extraction worker from psycopg2 + asyncio.Lock to psycopg_pool.AsyncConnectionPool. After PR #206 raised LLM throughput ~10x by killing the JSON-blob ReDoS, the worker's single-shared-connection serialisation lock became the new throughput ceiling. This PR removes it.

What changed

Worker hot path (mempalace/kg_triple_worker.py)

• _SyncConnPool keeps its name (so test factories don't need to be renamed) but its internals are now a psycopg_pool.AsyncConnectionPool. The pool's configure= callback runs LOAD 'age' and SET search_path = ag_catalog, "$user", public once per fresh connection, paid once instead of per-cursor.
• _KGHandle.add_triple is async. No asyncio.Lock anywhere. Each coroutine grabs its own connection from the pool and writes concurrently.
• Pool sizing: min=max(4, max_concurrency // 2), max=max_concurrency + 2. Modest steady-state, enough headroom for the LLM-writer fleet plus the claim loop.
• Added async variants of _claim_batch, _fetch_drawer_text, _mark_completed, _mark_error, _seed_backfill. The sync versions are kept for the CLI --status path (a single short-lived connection — async pool would be overkill).

Driver swap elsewhere is mechanical

import psycopg as psycopg2 aliasing so existing monkeypatch.setattr(..., "_load_psycopg2", ...) and psycopg2.connect test seams keep working:

• mempalace/backends/postgres.py
• mempalace/knowledge_graph_age.py
• mempalace/searcher.py
• mempalace/migrate_to_postgres.py
• mempalace/cli.py
• mempalace/kg_writethrough.py

pyproject.toml — only psycopg lines touched:

postgres = ["psycopg[binary]>=3.2,<4", "psycopg-pool>=3.2,<4"]
kg-extract = ["httpx>=0.27", "psycopg[binary]>=3.2,<4", "psycopg-pool>=3.2,<4"]

Pool sizing rationale

Knob	Value	Why
min_size	max(4, max_concurrency // 2)	Steady idle capacity; avoids reconnect churn during quiet stretches
max_size	max_concurrency + 2	One conn per LLM-writer + headroom for the queue-claim/seed loops

max_concurrency is the LLM concurrency cap. Production runs with 24 → pool min 12, max 26. AGE per-connection setup pays once at first checkout, never again on that physical connection's lifetime.

Test plan

• tests/test_kg_triple_worker.py — fakes upgraded to dual-mode sync/async via small _NoopAwaitable and _SyncResult helpers. All 14 tests pass.
• tests/test_palace_graph.py, tests/test_age_kg_units.py, tests/test_kg_extraction_queue.py, tests/test_backfill_kg_triples.py, tests/test_backends_postgres.py, tests/test_migrate_to_postgres.py — all 120 pass (22 skipped, postgres-required).
• Full suite: 3157 passed (up from 3151 — gained the 6 worker tests), 38 pre-existing failures unrelated to psycopg (CLI search-output formatting, daemon routing, readme parity, source-adapter entry-point ordering).
• Production cutover smoke test (out of PR scope; tracked in deploy script under kg-extract-aurora).

References

• fix(kg-extract): drop ReDoS-prone regex in _parse_json_blob #206 — JSON-blob ReDoS fix that lifted LLM throughput and made the worker's lock the new bottleneck.

[Copilot]

Pull request overview

This PR removes the KG triple-extraction worker’s psycopg2 + global asyncio.Lock bottleneck by migrating the hot path to psycopg_pool.AsyncConnectionPool, enabling concurrent DB writes per coroutine and avoiding asyncio.to_thread overhead.

Changes:

• Refactors mempalace/kg_triple_worker.py to use an async connection pool for queue ops + AGE writes, and introduces async variants of queue helper functions.
• Mechanically swaps imports across the codebase to use psycopg (psycopg3) while preserving legacy _load_psycopg2 naming for test seams.
• Updates extras in pyproject.toml to depend on psycopg[binary] and psycopg-pool, and updates worker tests/fakes to support dual sync/async call patterns.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
mempalace/kg_triple_worker.py	Migrates worker DB hot path to AsyncConnectionPool, adds async queue helpers, and implements direct async AGE triple writes.
tests/test_kg_triple_worker.py	Updates fakes to be usable from both sync and async code paths; adapts tests to new async worker helpers.
pyproject.toml	Switches postgres/kg-extract extras from psycopg2-binary to psycopg[binary] + psycopg-pool.
mempalace/backends/postgres.py	Updates lazy driver import helper to return psycopg3 + psycopg3 sql.
mempalace/knowledge_graph_age.py	Updates lazy driver import to psycopg3 and documents legacy naming compatibility.
mempalace/searcher.py	Replaces psycopg2 imports with psycopg as psycopg2 aliasing and updates error text.
mempalace/migrate_to_postgres.py	Replaces psycopg2 imports with psycopg as psycopg2 aliasing and updates fatal install guidance.
mempalace/cli.py	Updates rooms command dependency check/import from psycopg2 to psycopg3.
mempalace/kg_writethrough.py	Updates documentation comment to reflect psycopg3 under legacy helper name.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[gemini-code-assist]

Warning

Gemini encountered an error creating the summary. You can try again by commenting /gemini summary.