Coverage for src/lumen/search/bundle.py: 88%
133 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-31 18:46 +0000
1"""
2This module builds FHIR searchset Bundles from database query results.
3It handles pagination, field filtering, summary options, and resource inclusion.
4"""
6from __future__ import annotations
8from datetime import UTC, datetime
9from typing import Any, cast
10from uuid import UUID
12import asyncpg
14from .parser import ParsedQuery
15from .query import SearchResult
17_SUMMARY_TRUE_FIELDS: frozenset[str] = frozenset(
18 {
19 "id",
20 "meta",
21 "resourceType",
22 "text",
23 "identifier",
24 "status",
25 "subject",
26 "patient",
27 "code",
28 "category",
29 "gender",
30 "birthDate",
31 "name",
32 }
33)
36def build_searchset_bundle(
37 base_url: str,
38 resource_type: str,
39 result: SearchResult,
40 query: ParsedQuery,
41 request_url: str,
42) -> dict[str, Any]:
43 """Construct a FHIR searchset Bundle dict."""
44 entries: list[dict[str, Any]] = []
46 for resource, vid, ts in zip(
47 result.resources, result.version_ids, result.last_updated, strict=True
48 ):
49 res = _apply_summary(resource, query)
50 res = _inject_meta(res, vid, ts)
51 rid = resource.get("id", "")
52 entries.append(
53 {
54 "fullUrl": f"{base_url}/{resource_type}/{rid}",
55 "resource": res,
56 "search": {"mode": "match"},
57 }
58 )
60 if result.warnings:
61 entries.append(
62 {
63 "fullUrl": f"{base_url}/OperationOutcome/search-warnings",
64 "resource": {
65 "resourceType": "OperationOutcome",
66 "issue": [
67 {
68 "severity": "warning",
69 "code": "not-supported",
70 "diagnostics": w,
71 }
72 for w in result.warnings
73 ],
74 },
75 "search": {"mode": "outcome"},
76 }
77 )
79 bundle: dict[str, Any] = {
80 "resourceType": "Bundle",
81 "type": "searchset",
82 "entry": entries,
83 }
85 if result.total is not None:
86 bundle["total"] = result.total
87 elif query.summary != "count":
88 # Exact total is expensive; omit unless already computed.
89 pass
91 # Pagination links
92 links: list[dict[str, str]] = [{"relation": "self", "url": request_url}]
93 if result.next_cursor:
94 next_url = _replace_page_token(request_url, result.next_cursor)
95 links.append({"relation": "next", "url": next_url})
96 bundle["link"] = links
98 return bundle
101def build_count_bundle(
102 total: int,
103 request_url: str,
104) -> dict[str, Any]:
105 return {
106 "resourceType": "Bundle",
107 "type": "searchset",
108 "total": total,
109 "entry": [],
110 "link": [{"relation": "self", "url": request_url}],
111 }
114async def _process_single_resource_includes(
115 entries: list[dict[str, Any]],
116 seen_ids: set[str],
117 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
118 tenant_id: UUID,
119 base_url: str,
120 resource: dict[str, Any],
121 param_name: str,
122 registry: Any,
123) -> None:
124 import json as _json
126 resource_type = resource.get("resourceType", "")
127 params = await registry.get_params(conn, tenant_id, resource_type)
128 refs = _extract_references_for_param(resource, param_name, params)
129 for ref_type, ref_id in refs:
130 key = f"{ref_type}/{ref_id}"
131 if key in seen_ids:
132 continue
133 # is_current (migration 009): a single current row, filtered directly. The old
134 # MAX(version_id) subquery aggregated the whole type even though only one id is wanted.
135 row = await conn.fetchrow(
136 """
137 SELECT r.resource, r.version_id, r.last_updated
138 FROM fhir_resources r
139 WHERE r.tenant_id=$1 AND r.resource_type=$2 AND r.resource_id=$3
140 AND r.is_current AND NOT r.deleted
141 """,
142 tenant_id,
143 ref_type,
144 ref_id,
145 )
146 if row:
147 raw = row["resource"]
148 res = _json.loads(raw) if isinstance(raw, str) else dict(raw)
149 res = _inject_meta(res, int(row["version_id"]), row["last_updated"])
150 seen_ids.add(key)
151 entries.append(
152 {
153 "fullUrl": f"{base_url}/{ref_type}/{ref_id}",
154 "resource": res,
155 "search": {"mode": "include"},
156 }
157 )
160async def _process_includes(
161 entries: list[dict[str, Any]],
162 seen_ids: set[str],
163 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
164 tenant_id: UUID,
165 base_url: str,
166 includes: list[str],
167 matched_resources: list[dict[str, Any]],
168) -> None:
169 from .registry import get_registry
171 registry = get_registry()
172 for include_spec in includes:
173 parts = include_spec.split(":")
174 if len(parts) < 2:
175 continue
176 param_name = parts[1]
178 for resource in matched_resources:
179 await _process_single_resource_includes(
180 entries, seen_ids, conn, tenant_id, base_url, resource, param_name, registry
181 )
184async def _process_revincludes(
185 entries: list[dict[str, Any]],
186 seen_ids: set[str],
187 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
188 tenant_id: UUID,
189 base_url: str,
190 revincludes: list[str],
191 matched_resources: list[dict[str, Any]],
192) -> None:
193 import json as _json
195 for revinclude_spec in revincludes:
196 parts = revinclude_spec.split(":")
197 if len(parts) < 2:
198 continue
199 source_type = parts[0]
200 param_name = parts[1]
202 match_ids = [r.get("id") for r in matched_resources if r.get("id")]
203 if not match_ids:
204 continue
206 # is_current (migration 009) replaces the O(table) MAX(version_id) aggregation; the
207 # spidx_reference join stays. DISTINCT collapses rows a resource matches more than once.
208 rows = await conn.fetch(
209 """
210 SELECT DISTINCT r.resource, r.version_id, r.last_updated, r.resource_id
211 FROM fhir_resources r
212 JOIN spidx_reference sr
213 ON sr.tenant_id=$1 AND sr.resource_type=$2
214 AND sr.resource_id=r.resource_id
215 AND sr.param_name=$3
216 AND sr.target_resource_id = ANY($4::text[])
217 WHERE r.tenant_id=$1 AND r.resource_type=$2
218 AND r.is_current AND NOT r.deleted
219 """,
220 tenant_id,
221 source_type,
222 param_name,
223 match_ids,
224 )
225 for row in rows:
226 raw = row["resource"]
227 res = _json.loads(raw) if isinstance(raw, str) else dict(raw)
228 key = f"{source_type}/{row['resource_id']}"
229 if key in seen_ids:
230 continue
231 res = _inject_meta(res, int(row["version_id"]), row["last_updated"])
232 seen_ids.add(key)
233 entries.append(
234 {
235 "fullUrl": f"{base_url}/{source_type}/{row['resource_id']}",
236 "resource": res,
237 "search": {"mode": "include"},
238 }
239 )
242async def attach_includes(
243 bundle: dict[str, Any],
244 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
245 tenant_id: UUID,
246 base_url: str,
247 includes: list[str],
248 revincludes: list[str],
249 matched_resources: list[dict[str, Any]],
250) -> None:
251 """Fetch and attach _include / _revinclude resources to the bundle."""
252 if not includes and not revincludes:
253 return
255 entries: list[dict[str, Any]] = list(bundle.get("entry", []))
256 seen_ids: set[str] = {
257 f"{e['resource']['resourceType']}/{e['resource'].get('id', '')}"
258 for e in entries
259 if "resource" in e
260 }
262 if includes:
263 await _process_includes(
264 entries, seen_ids, conn, tenant_id, base_url, includes, matched_resources
265 )
266 if revincludes:
267 await _process_revincludes(
268 entries, seen_ids, conn, tenant_id, base_url, revincludes, matched_resources
269 )
271 bundle["entry"] = entries
274# ── private helpers ────────────────────────────────────────────────────────────
277def _apply_summary(resource: dict[str, Any], query: ParsedQuery) -> dict[str, Any]:
278 if query.summary == "true":
279 return {k: v for k, v in resource.items() if k in _SUMMARY_TRUE_FIELDS}
280 if query.elements:
281 keep = set(query.elements) | {"resourceType", "id", "meta"}
282 return {k: v for k, v in resource.items() if k in keep}
283 return resource
286def _inject_meta(
287 resource: dict[str, Any], version_id: int, last_updated: datetime
288) -> dict[str, Any]:
289 meta: dict[str, Any] = dict(resource.get("meta") or {})
290 meta["versionId"] = str(version_id)
291 ts = last_updated
292 if ts.tzinfo is None:
293 ts = ts.replace(tzinfo=UTC)
294 meta["lastUpdated"] = ts.isoformat()
295 return {**resource, "meta": meta}
298def _replace_page_token(url: str, token: str) -> str:
299 import urllib.parse
301 parsed = urllib.parse.urlparse(url)
302 params = urllib.parse.parse_qs(parsed.query, keep_blank_values=True)
303 params["_page_token"] = [token]
304 new_query = urllib.parse.urlencode(params, doseq=True)
305 return urllib.parse.urlunparse(parsed._replace(query=new_query))
308def _parse_reference_string(val: Any) -> tuple[str, str] | None:
309 if isinstance(val, dict):
310 ref_str: str = cast("str", cast("dict[str, Any]", val).get("reference", ""))
311 elif isinstance(val, str):
312 ref_str = val
313 else:
314 return None
315 if not ref_str:
316 return None
317 parts: list[str] = ref_str.split("/")
318 if len(parts) >= 2:
319 return parts[-2], parts[-1]
320 return None
323def _extract_references_for_param(
324 resource: dict[str, Any], param_name: str, params: list[Any]
325) -> list[tuple[str, str]]:
326 """Naively extract reference values for a given parameter name from a resource dict."""
327 sp = next((p for p in params if p.name == param_name), None)
328 if sp is None or sp.type != "reference":
329 return []
331 from ..fhir.path import evaluate_path
333 refs: list[tuple[str, str]] = []
334 for path in sp.paths:
335 for val in evaluate_path(path, resource):
336 parsed = _parse_reference_string(val)
337 if parsed:
338 refs.append(parsed)
339 return refs