37 lines
1.2 KiB
Python
37 lines
1.2 KiB
Python
|
"""kb-postgres access for kb-query beyond what `kb_retrieval.retrieval` covers -- the envelope
|
||
|
|
join (`source` + `entities` for link/metadata building, app/links.py) is out of scope for the
|
||
|
|
shared retrieval package (module 5 phase 4 plan §4: "to jest nowy kod w warstwie
|
||
|
|
HTTP-handlera kb-query, nie zmiana w kb_retrieval")."""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
|
||
|
|
import asyncpg
|
||
|
|
|
||
|
|
|
||
|
|
async def create_pool(dsn: str) -> asyncpg.Pool:
|
||
|
|
return await asyncpg.create_pool(dsn, min_size=1, max_size=5)
|
||
|
|
|
||
|
|
|
||
|
|
def _decode_jsonb(value: object) -> object:
|
||
|
|
"""asyncpg may return jsonb as a str or an already-decoded object."""
|
||
|
|
if value is None:
|
||
|
|
return None
|
||
|
|
if isinstance(value, str):
|
||
|
|
return json.loads(value)
|
||
|
|
return value
|
||
|
|
|
||
|
|
|
||
|
|
async def fetch_envelopes(conn: asyncpg.Connection, envelope_ids: list[str]) -> dict[str, dict]:
|
||
|
|
"""Batch-fetch `source`/`entities` for a set of chunk hits, keyed by envelope id."""
|
||
|
|
if not envelope_ids:
|
||
|
|
return {}
|
||
|
|
rows = await conn.fetch(
|
||
|
|
"SELECT id, source, entities FROM envelope WHERE id = ANY($1::text[])",
|
||
|
|
envelope_ids,
|
||
|
|
)
|
||
|
|
return {
|
||
|
|
r["id"]: {"source": r["source"], "entities": _decode_jsonb(r["entities"]) or []}
|
||
|
|
for r in rows
|
||
|
|
}
|