Coverage for src/lumen/search/indexer.py: 85%
227 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-31 18:46 +0000
1"""
2This module extracts indexable values from FHIR resources based on search parameter definitions.
3It populates the corresponding search index tables within the resource write transaction.
4"""
6from __future__ import annotations
8import re
9from datetime import UTC, datetime
10from typing import Any, cast
11from uuid import UUID
13import asyncpg
14import structlog
16from ..fhir.params import SearchParamDef
17from ..fhir.path import evaluate_path
19logger = structlog.get_logger(__name__)
21# ── public API ─────────────────────────────────────────────────────────────────
24async def index_resource(
25 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
26 tenant_id: UUID,
27 resource_type: str,
28 resource_id: str,
29 resource: dict[str, Any],
30 params: list[SearchParamDef],
31) -> None:
32 """Delete old spidx_* rows for this resource and insert fresh ones."""
33 await _delete_indexes(conn, tenant_id, resource_type, resource_id)
34 for sp in params:
35 await _index_param(conn, tenant_id, resource_type, resource_id, resource, sp)
38async def delete_indexes(
39 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
40 tenant_id: UUID,
41 resource_type: str,
42 resource_id: str,
43) -> None:
44 await _delete_indexes(conn, tenant_id, resource_type, resource_id)
47# ── private helpers ────────────────────────────────────────────────────────────
50async def _delete_indexes(
51 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
52 tenant_id: UUID,
53 resource_type: str,
54 resource_id: str,
55) -> None:
56 for table in (
57 "spidx_string",
58 "spidx_token",
59 "spidx_date",
60 "spidx_reference",
61 "spidx_quantity",
62 "spidx_number",
63 "spidx_composite",
64 "spidx_geo",
65 ):
66 await conn.execute(
67 f"DELETE FROM {table} WHERE tenant_id=$1 AND resource_type=$2 AND resource_id=$3",
68 tenant_id,
69 resource_type,
70 resource_id,
71 )
74async def _index_param(
75 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
76 tenant_id: UUID,
77 resource_type: str,
78 resource_id: str,
79 resource: dict[str, Any],
80 sp: SearchParamDef,
81) -> None:
82 values: list[Any] = []
83 for path in sp.paths:
84 values.extend(evaluate_path(path, resource))
86 if sp.type == "string":
87 await _insert_strings(conn, tenant_id, resource_type, resource_id, sp.name, values)
88 elif sp.type == "token":
89 await _insert_tokens(conn, tenant_id, resource_type, resource_id, sp.name, values)
90 elif sp.type == "date":
91 await _insert_dates(conn, tenant_id, resource_type, resource_id, sp.name, values)
92 elif sp.type == "reference":
93 await _insert_references(conn, tenant_id, resource_type, resource_id, sp.name, values)
94 elif sp.type == "quantity":
95 await _insert_quantities(conn, tenant_id, resource_type, resource_id, sp.name, values)
96 elif sp.type == "number":
97 await _insert_numbers(conn, tenant_id, resource_type, resource_id, sp.name, values)
98 elif sp.type == "composite":
99 await _insert_composite(conn, tenant_id, resource_type, resource_id, sp, values)
100 elif sp.type == "geo":
101 await _insert_geo(conn, tenant_id, resource_type, resource_id, sp.name, values)
104# ── string ─────────────────────────────────────────────────────────────────────
107def _string_values(node: Any) -> list[str]:
108 """Collect indexable strings from a node.
110 String search parameters often point at complex types (HumanName, Address);
111 collect every string leaf so e.g. ``Patient.name`` indexes family/given/text.
112 """
113 out: list[str] = []
114 if isinstance(node, str):
115 if node.strip():
116 out.append(node)
117 elif isinstance(node, dict):
118 for v in cast("dict[str, Any]", node).values():
119 out.extend(_string_values(v))
120 elif isinstance(node, list):
121 for v in cast("list[Any]", node):
122 out.extend(_string_values(v))
123 return out
126async def _insert_strings(
127 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
128 tenant_id: UUID,
129 resource_type: str,
130 resource_id: str,
131 param_name: str,
132 values: list[Any],
133) -> None:
134 for val in values:
135 for s in _string_values(val):
136 await conn.execute(
137 """
138 INSERT INTO spidx_string (tenant_id, resource_type, resource_id, param_name, value)
139 VALUES ($1,$2,$3,$4,$5)
140 """,
141 tenant_id,
142 resource_type,
143 resource_id,
144 param_name,
145 s.lower(),
146 )
149# ── token ──────────────────────────────────────────────────────────────────────
151# (system, code, display) — ``display`` is the textual part searched by ``:text``.
152TokenPair = tuple[str | None, str, str | None]
155def _extract_codeable_concept_pairs(d: dict[str, Any]) -> list[TokenPair]:
156 # ``display`` carries the textual part searched by the token ``:text`` modifier.
157 pairs: list[TokenPair] = []
158 text = d.get("text")
159 text_str = str(text) if text else None
160 if text_str:
161 pairs.append((None, text_str, text_str))
162 for coding in cast("list[Any]", d.get("coding", [])):
163 if isinstance(coding, dict):
164 cd: dict[str, Any] = cast("dict[str, Any]", coding)
165 code = cd.get("code", "")
166 if code:
167 display = cd.get("display")
168 pairs.append(
169 (
170 cast(str | None, cd.get("system")),
171 str(code),
172 str(display) if display else None,
173 )
174 )
175 return pairs
178def _extract_token_pairs_from_dict(d: dict[str, Any]) -> list[TokenPair]:
179 # CodeableConcept: has coding[] and/or text.
180 if "coding" in d or "text" in d:
181 return _extract_codeable_concept_pairs(d)
183 # Coding: system + code (+ display for :text).
184 if "code" in d:
185 code = d.get("code", "")
186 if code:
187 display = d.get("display")
188 return [
189 (
190 cast(str | None, d.get("system")),
191 str(code),
192 str(display) if display else None,
193 )
194 ]
196 # Identifier: system + value.
197 if "value" in d:
198 val = d.get("value", "")
199 if val:
200 return [(cast(str | None, d.get("system")), str(val), None)]
202 return []
205def _extract_token_pairs(value: Any) -> list[TokenPair]:
206 """Return (system, code, display) tuples for an extracted value, inferred from shape.
208 Handles the token-bearing datatypes: code (str), boolean, Coding,
209 CodeableConcept, and Identifier — distinguished by the JSON node's keys.
210 ``display`` is the textual part matched by the ``:text`` modifier and is
211 ``None`` when the datatype carries no text.
212 """
213 if isinstance(value, bool):
214 return [(None, "true" if value else "false", None)]
215 if isinstance(value, str):
216 return [(None, value, None)] if value else []
217 if not isinstance(value, dict):
218 return []
220 return _extract_token_pairs_from_dict(cast("dict[str, Any]", value))
223async def _insert_tokens(
224 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
225 tenant_id: UUID,
226 resource_type: str,
227 resource_id: str,
228 param_name: str,
229 values: list[Any],
230) -> None:
231 for val in values:
232 for system, code, display in _extract_token_pairs(val):
233 await conn.execute(
234 """
235 INSERT INTO spidx_token
236 (tenant_id, resource_type, resource_id, param_name, system, code, display)
237 VALUES ($1,$2,$3,$4,$5,$6,$7)
238 """,
239 tenant_id,
240 resource_type,
241 resource_id,
242 param_name,
243 system,
244 code,
245 display,
246 )
249# ── date ───────────────────────────────────────────────────────────────────────
251_DATE_ONLY_RE = re.compile(r"^\d{4}(-\d{2}(-\d{2})?)?$")
254def _parse_fhir_date(value: str) -> tuple[datetime, datetime]:
255 """Parse a FHIR date/dateTime string into (low, high) inclusive bounds."""
256 if _DATE_ONLY_RE.match(value):
257 parts = value.split("-")
258 if len(parts) == 1: # year only
259 lo = datetime(int(parts[0]), 1, 1, tzinfo=UTC)
260 hi = datetime(int(parts[0]), 12, 31, 23, 59, 59, tzinfo=UTC)
261 elif len(parts) == 2: # year-month
262 import calendar
264 y, m = int(parts[0]), int(parts[1])
265 last_day = calendar.monthrange(y, m)[1]
266 lo = datetime(y, m, 1, tzinfo=UTC)
267 hi = datetime(y, m, last_day, 23, 59, 59, tzinfo=UTC)
268 else: # full date
269 y, m, d = int(parts[0]), int(parts[1]), int(parts[2])
270 lo = datetime(y, m, d, tzinfo=UTC)
271 hi = datetime(y, m, d, 23, 59, 59, tzinfo=UTC)
272 return lo, hi
274 # dateTime / instant — strip timezone info and parse
275 clean = value.replace("Z", "+00:00")
276 # Remove sub-second precision beyond microseconds
277 clean = re.sub(r"(\.\d{6})\d+", r"\1", clean)
278 dt = datetime.fromisoformat(clean)
279 if dt.tzinfo is None:
280 dt = dt.replace(tzinfo=UTC)
281 return dt, dt
284def _parse_extracted_date(
285 val: Any,
286 resource_type: str,
287 resource_id: str,
288 param_name: str,
289) -> tuple[datetime | None, datetime | None]:
290 lo: datetime | None = None
291 hi: datetime | None = None
292 try:
293 # Period / Range: {start,end}; otherwise a date/dateTime/instant string.
294 if isinstance(val, dict) and ("start" in val or "end" in val):
295 period: dict[str, Any] = cast("dict[str, Any]", val)
296 start_str = period.get("start")
297 end_str = period.get("end")
298 if start_str:
299 lo, _ = _parse_fhir_date(str(start_str))
300 if end_str:
301 _, hi = _parse_fhir_date(str(end_str))
302 elif isinstance(val, str) and val:
303 lo, hi = _parse_fhir_date(val)
304 except (ValueError, KeyError) as e:
305 logger.warning(
306 "indexer.parse_error",
307 resource_type=resource_type,
308 resource_id=resource_id,
309 param_name=param_name,
310 value=val,
311 error=str(e),
312 )
313 return lo, hi
316async def _insert_dates(
317 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
318 tenant_id: UUID,
319 resource_type: str,
320 resource_id: str,
321 param_name: str,
322 values: list[Any],
323) -> None:
324 for val in values:
325 lo, hi = _parse_extracted_date(val, resource_type, resource_id, param_name)
326 if lo is not None or hi is not None:
327 await conn.execute(
328 """
329 INSERT INTO spidx_date
330 (tenant_id, resource_type, resource_id, param_name, value_low, value_high)
331 VALUES ($1,$2,$3,$4,$5,$6)
332 """,
333 tenant_id,
334 resource_type,
335 resource_id,
336 param_name,
337 lo,
338 hi,
339 )
342# ── reference ──────────────────────────────────────────────────────────────────
345def _parse_reference(value: Any) -> tuple[str | None, str] | None:
346 """Extract (resource_type, resource_id) from a FHIR Reference dict or string."""
347 if isinstance(value, dict):
348 ref: str = cast("str", cast("dict[str, Any]", value).get("reference", ""))
349 elif isinstance(value, str):
350 ref = value
351 else:
352 return None
354 if not ref:
355 return None
357 # Absolute URL: https://example.com/fhir/Patient/123
358 # Relative: Patient/123 or just 123
359 parts: list[str] = ref.split("/")
360 if len(parts) >= 2:
361 return parts[-2], parts[-1]
362 return None, parts[-1]
365async def _insert_references(
366 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
367 tenant_id: UUID,
368 resource_type: str,
369 resource_id: str,
370 param_name: str,
371 values: list[Any],
372) -> None:
373 for val in values:
374 parsed = _parse_reference(val)
375 if parsed is None:
376 continue
377 target_type, target_id = parsed
378 if target_id:
379 await conn.execute(
380 """
381 INSERT INTO spidx_reference
382 (tenant_id, resource_type, resource_id, param_name,
383 target_resource_type, target_resource_id)
384 VALUES ($1,$2,$3,$4,$5,$6)
385 """,
386 tenant_id,
387 resource_type,
388 resource_id,
389 param_name,
390 target_type,
391 target_id,
392 )
395# ── quantity ───────────────────────────────────────────────────────────────────
398async def _insert_quantities(
399 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
400 tenant_id: UUID,
401 resource_type: str,
402 resource_id: str,
403 param_name: str,
404 values: list[Any],
405) -> None:
406 for val in values:
407 if not isinstance(val, dict):
408 continue
409 qty: dict[str, Any] = cast("dict[str, Any]", val)
410 raw_value = qty.get("value")
411 if raw_value is None:
412 continue
413 try:
414 numeric = float(raw_value) # type: ignore[arg-type] # Any from FHIR JSON
415 except (TypeError, ValueError):
416 continue
417 await conn.execute(
418 """
419 INSERT INTO spidx_quantity
420 (tenant_id, resource_type, resource_id, param_name, system, code, value)
421 VALUES ($1,$2,$3,$4,$5,$6,$7)
422 """,
423 tenant_id,
424 resource_type,
425 resource_id,
426 param_name,
427 cast(str | None, qty.get("system")),
428 cast(str | None, qty.get("code") or qty.get("unit")),
429 numeric,
430 )
433# ── number ─────────────────────────────────────────────────────────────────────
436async def _insert_numbers(
437 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
438 tenant_id: UUID,
439 resource_type: str,
440 resource_id: str,
441 param_name: str,
442 values: list[Any],
443) -> None:
444 for val in values:
445 try:
446 numeric = float(val)
447 except (TypeError, ValueError) as e:
448 logger.warning(
449 "indexer.parse_error",
450 resource_type=resource_type,
451 resource_id=resource_id,
452 param_name=param_name,
453 value=val,
454 error=str(e),
455 )
456 continue
457 await conn.execute(
458 """
459 INSERT INTO spidx_number
460 (tenant_id, resource_type, resource_id, param_name, value)
461 VALUES ($1,$2,$3,$4,$5)
462 """,
463 tenant_id,
464 resource_type,
465 resource_id,
466 param_name,
467 numeric,
468 )
471# ── composite ──────────────────────────────────────────────────────────────────
474async def _insert_composite(
475 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
476 tenant_id: UUID,
477 resource_type: str,
478 resource_id: str,
479 sp: SearchParamDef,
480 occurrences: list[Any],
481) -> None:
482 """Index each occurrence of a composite param's base path (e.g. one
483 ``Observation.component`` entry, or the whole resource for the
484 ``Observation | Observation.component`` "either" branch) as a row pairing
485 both components — the correlation a plain per-type spidx table can't
486 express (see migration 004 for why).
487 """
488 if sp.components is None or len(sp.components) != 2:
489 return
490 c1, c2 = sp.components
491 for occurrence in occurrences:
492 if not isinstance(occurrence, dict):
493 continue
494 occ: dict[str, Any] = cast("dict[str, Any]", occurrence)
495 c1_pairs = [
496 (system, code)
497 for val in evaluate_path(c1.expression, occ)
498 for system, code, _ in _extract_token_pairs(val)
499 ]
500 if not c1_pairs:
501 continue
502 if c2.type == "token":
503 await _insert_composite_token_pairs(
504 conn, tenant_id, resource_type, resource_id, sp.name, c1_pairs, c2.expression, occ
505 )
506 else:
507 await _insert_composite_quantity_pairs(
508 conn, tenant_id, resource_type, resource_id, sp.name, c1_pairs, c2.expression, occ
509 )
512async def _insert_composite_token_pairs(
513 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
514 tenant_id: UUID,
515 resource_type: str,
516 resource_id: str,
517 param_name: str,
518 c1_pairs: list[tuple[str | None, str]],
519 c2_expression: str,
520 occurrence: dict[str, Any],
521) -> None:
522 c2_pairs = [
523 (system, code)
524 for val in evaluate_path(c2_expression, occurrence)
525 for system, code, _ in _extract_token_pairs(val)
526 ]
527 for c1_system, c1_code in c1_pairs:
528 for c2_system, c2_code in c2_pairs:
529 await conn.execute(
530 """
531 INSERT INTO spidx_composite
532 (tenant_id, resource_type, resource_id, param_name,
533 c1_token_system, c1_token_code, c2_token_system, c2_token_code)
534 VALUES ($1,$2,$3,$4,$5,$6,$7,$8)
535 """,
536 tenant_id,
537 resource_type,
538 resource_id,
539 param_name,
540 c1_system,
541 c1_code,
542 c2_system,
543 c2_code,
544 )
547async def _insert_composite_quantity_pairs(
548 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
549 tenant_id: UUID,
550 resource_type: str,
551 resource_id: str,
552 param_name: str,
553 c1_pairs: list[tuple[str | None, str]],
554 c2_expression: str,
555 occurrence: dict[str, Any],
556) -> None:
557 for c2_val in evaluate_path(c2_expression, occurrence):
558 if not isinstance(c2_val, dict):
559 continue
560 qty: dict[str, Any] = cast("dict[str, Any]", c2_val)
561 raw_value = qty.get("value")
562 if raw_value is None:
563 continue
564 try:
565 numeric = float(raw_value) # type: ignore[arg-type] # Any from FHIR JSON
566 except (TypeError, ValueError):
567 continue
568 for c1_system, c1_code in c1_pairs:
569 await conn.execute(
570 """
571 INSERT INTO spidx_composite
572 (tenant_id, resource_type, resource_id, param_name,
573 c1_token_system, c1_token_code, c2_qty_system, c2_qty_code, c2_qty_value)
574 VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9)
575 """,
576 tenant_id,
577 resource_type,
578 resource_id,
579 param_name,
580 c1_system,
581 c1_code,
582 cast(str | None, qty.get("system")),
583 cast(str | None, qty.get("code") or qty.get("unit")),
584 numeric,
585 )
588# ── geo ────────────────────────────────────────────────────────────────────────
591async def _insert_geo(
592 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
593 tenant_id: UUID,
594 resource_type: str,
595 resource_id: str,
596 param_name: str,
597 values: list[Any],
598) -> None:
599 for val in values:
600 if not isinstance(val, dict):
601 continue
602 position: dict[str, Any] = cast("dict[str, Any]", val)
603 lat, lon = position.get("latitude"), position.get("longitude")
604 if lat is None or lon is None:
605 continue
606 try:
607 lat_f, lon_f = float(lat), float(lon) # type: ignore[arg-type] # Any from FHIR JSON
608 except (TypeError, ValueError):
609 continue
610 await conn.execute(
611 """
612 INSERT INTO spidx_geo
613 (tenant_id, resource_type, resource_id, param_name, latitude, longitude)
614 VALUES ($1,$2,$3,$4,$5,$6)
615 """,
616 tenant_id,
617 resource_type,
618 resource_id,
619 param_name,
620 lat_f,
621 lon_f,
622 )