Coverage for src/lumen/search/query.py: 74%
618 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-08-31 18:46 +0000
1"""
2This file compiles parsed FHIR search parameters into executable SQL queries.
3It handles table joins, filtering predicates, sorting, and cursor-based pagination.
4"""
6from __future__ import annotations
8import base64
9import contextlib
10import json
11import time
12from dataclasses import dataclass
13from datetime import UTC, datetime
14from decimal import Decimal
15from typing import Any
16from uuid import UUID
18import asyncpg
20from ..fhir.params import SearchParamDef, SearchParamType
21from ..observability import record_interaction, record_search_duration
22from .parser import (
23 ChainedParam,
24 HasParam,
25 ParsedParam,
26 ParsedQuery,
27 Prefix,
28 SortSpec,
29)
31# ── public API ─────────────────────────────────────────────────────────────────
34class SearchResult:
35 def __init__(
36 self,
37 resources: list[dict[str, Any]],
38 version_ids: list[int],
39 last_updated: list[datetime],
40 total: int | None,
41 next_cursor: str | None,
42 warnings: list[str] | None = None,
43 ) -> None:
44 self.resources = resources
45 self.version_ids = version_ids
46 self.last_updated = last_updated
47 self.total = total
48 self.next_cursor = next_cursor
49 self.warnings = warnings or []
52def _filter_known_params(
53 query: ParsedQuery,
54 registry: dict[str, SearchParamDef],
55 ignore_unknown_parameters: bool,
56) -> tuple[list[ParsedParam], list[str]]:
57 known_params: list[ParsedParam] = []
58 warnings: list[str] = []
59 for p in query.params:
60 if p.name in registry:
61 known_params.append(p)
62 else:
63 if not ignore_unknown_parameters:
64 from ..core.exceptions import OperationOutcomeError
66 raise OperationOutcomeError(
67 status_code=400,
68 code="not-supported",
69 diagnostics=f"Unknown search parameter '{p.name}' is not supported.",
70 )
71 warnings.append(f"Unknown search parameter '{p.name}' was ignored.")
72 return known_params, warnings
75async def _resolve_valuesets(
76 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
77 tenant_id: UUID,
78 known_params: list[ParsedParam],
79) -> dict[str, list[dict[str, Any]]]:
80 expanded_valuesets: dict[str, list[dict[str, Any]]] = {}
81 for p in known_params:
82 if p.modifier in ("in", "not-in"):
83 for sv in p.values:
84 if sv.raw not in expanded_valuesets:
85 expanded_valuesets[sv.raw] = await _expand_valueset_by_url(
86 conn, tenant_id, sv.raw
87 )
88 return expanded_valuesets
91def _parse_search_rows(
92 rows: list[asyncpg.Record],
93) -> tuple[list[dict[str, Any]], list[int], list[datetime]]:
94 import json as _json
96 resources: list[dict[str, Any]] = []
97 version_ids: list[int] = []
98 timestamps: list[datetime] = []
100 for row in rows:
101 raw = row["resource"]
102 if isinstance(raw, str):
103 resources.append(_json.loads(raw))
104 else:
105 resources.append(dict(raw))
106 version_ids.append(int(row["version_id"]))
107 ts = row["last_updated"]
108 if ts.tzinfo is None:
109 ts = ts.replace(tzinfo=UTC)
110 timestamps.append(ts)
111 return resources, version_ids, timestamps
114async def execute_search(
115 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
116 tenant_id: UUID,
117 resource_type: str,
118 query: ParsedQuery,
119 params: list[SearchParamDef],
120 ignore_unknown_parameters: bool = True,
121) -> SearchResult:
122 """Time the search and emit domain metrics; delegates to the implementation."""
123 start = time.perf_counter()
124 try:
125 return await _execute_search_impl(
126 conn, tenant_id, resource_type, query, params, ignore_unknown_parameters
127 )
128 finally:
129 record_search_duration(resource_type, (time.perf_counter() - start) * 1000)
130 record_interaction(resource_type, "search")
133async def _execute_search_impl(
134 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
135 tenant_id: UUID,
136 resource_type: str,
137 query: ParsedQuery,
138 params: list[SearchParamDef],
139 ignore_unknown_parameters: bool = True,
140) -> SearchResult:
141 registry = {sp.name: sp for sp in params}
143 # Resolve chained params (target type + final field type) so the SQL builder
144 # can join the correct spidx_* table instead of assuming token.
145 _resolve_chained_params(query.chained_params)
147 # Split out known vs unknown params; unknown are ignored with a warning.
148 known_params, warnings = _filter_known_params(query, registry, ignore_unknown_parameters)
150 # Expand any ValueSets for :in / :not-in modifiers
151 expanded_valuesets = await _resolve_valuesets(conn, tenant_id, known_params)
153 # For _summary=count: just return the count, no rows.
154 if query.summary == "count":
155 count_sql, count_args = _build_count_sql(
156 tenant_id,
157 resource_type,
158 known_params,
159 registry,
160 expanded_valuesets,
161 query.has_params or None,
162 query.chained_params or None,
163 )
164 row = await conn.fetchrow(count_sql, *count_args)
165 total = int(row["total"]) if row else 0 # type: ignore[index]
166 return SearchResult([], [], [], total, None, warnings)
168 keys = _order_keys(query.sort, registry)
169 sql, args = _build_search_sql(
170 tenant_id,
171 resource_type,
172 known_params,
173 registry,
174 query.count,
175 query.cursor,
176 keys,
177 expanded_valuesets,
178 query.has_params or None,
179 query.chained_params or None,
180 )
181 rows: list[asyncpg.Record] = await conn.fetch(sql, *args)
183 resources, version_ids, timestamps = _parse_search_rows(rows)
185 next_cursor: str | None = None
186 if len(rows) == query.count and rows:
187 last = rows[-1]
188 sort_values = [_serialize_sort_value(last[f"sortk{i}"]) for i in range(len(keys))]
189 next_cursor = _encode_cursor(sort_values, str(last["resource_id"]))
191 # Bundle.total is normally omitted (an exact count is a separate,
192 # expensive query) — but on an unpaginated first page that doesn't fill
193 # the requested count, len(rows) already *is* the exact total, for free.
194 # Later pages can't infer it this way (we don't track how many resources
195 # came before the cursor), so total stays unset there as before.
196 total: int | None = None
197 if query.cursor is None and len(rows) < query.count:
198 total = len(rows)
200 return SearchResult(resources, version_ids, timestamps, total, next_cursor, warnings)
203async def _expand_valueset_by_url(
204 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record],
205 tenant_id: UUID,
206 url: str,
207) -> list[dict[str, Any]]:
208 """Resolve and expand a ValueSet by its canonical URL."""
209 from ..db.repository import ResourceRepository
210 from ..fhir import terminology as tx
211 from ..search.registry import get_registry
213 repo = ResourceRepository(conn, get_registry())
214 vs_record = await repo.find_current_by_url(tenant_id, "ValueSet", url)
215 if vs_record is None:
216 return []
218 value_set = vs_record["resource"]
219 systems: dict[str, dict[str, Any]] = {}
220 for system_url in tx.referenced_systems(value_set):
221 cs = await repo.find_current_by_url(tenant_id, "CodeSystem", system_url)
222 if cs is not None:
223 systems[system_url] = cs["resource"]
225 try:
226 return tx.expand_valueset(value_set, systems)
227 except tx.TerminologyError:
228 return []
231# ── SQL builders ───────────────────────────────────────────────────────────────
234def _build_count_sql(
235 tenant_id: UUID,
236 resource_type: str,
237 params: list[ParsedParam],
238 registry: dict[str, SearchParamDef],
239 expanded_valuesets: dict[str, list[dict[str, Any]]],
240 has_params: list[HasParam] | None = None,
241 chained_params: list[ChainedParam] | None = None,
242) -> tuple[str, list[Any]]:
243 args: list[Any] = [tenant_id, resource_type]
244 ctes, joins, where_clauses, args = _build_joins_and_where(
245 params,
246 registry,
247 args,
248 tenant_id,
249 resource_type,
250 expanded_valuesets,
251 has_params,
252 chained_params,
253 )
254 with_sql = ("WITH " + ",\n".join(ctes)) if ctes else ""
255 join_sql = "\n".join(joins)
256 where_sql = (" AND " + " AND ".join(where_clauses)) if where_clauses else ""
257 # Current rows are marked by is_current (migration 009); filtering them directly replaces
258 # the O(table) MAX(version_id) aggregation that this query used to run before every count.
259 sql = f"""
260 {with_sql}
261 SELECT COUNT(DISTINCT (r.resource_id)) AS total
262 FROM fhir_resources r
263 {join_sql}
264 WHERE r.tenant_id = $1 AND r.resource_type = $2
265 AND r.is_current AND NOT r.deleted{where_sql}
266 """
267 return sql, args
270def _build_search_sql(
271 tenant_id: UUID,
272 resource_type: str,
273 params: list[ParsedParam],
274 registry: dict[str, SearchParamDef],
275 count: int,
276 cursor: str | None,
277 keys: list[_OrderKey],
278 expanded_valuesets: dict[str, list[dict[str, Any]]],
279 has_params: list[HasParam] | None = None,
280 chained_params: list[ChainedParam] | None = None,
281) -> tuple[str, list[Any]]:
282 args: list[Any] = [tenant_id, resource_type]
284 ctes, joins, where_clauses, args = _build_joins_and_where(
285 params,
286 registry,
287 args,
288 tenant_id,
289 resource_type,
290 expanded_valuesets,
291 has_params,
292 chained_params,
293 )
295 # Cursor condition: keyset over the full sort tuple + resource_id tiebreaker.
296 if cursor is not None:
297 cursor_values, cursor_id = _decode_cursor(cursor)
298 if cursor_id is not None and len(cursor_values) == len(keys):
299 where_clauses.append(_keyset_predicate(keys, cursor_values, cursor_id, args))
301 with_sql = ("WITH " + ",\n".join(ctes)) if ctes else ""
302 join_sql = "\n".join(joins)
303 where_sql = (" AND " + " AND ".join(where_clauses)) if where_clauses else ""
305 # DISTINCT ON must list the leading ORDER BY expressions verbatim. The sort
306 # expressions are deterministic per resource_id (they don't depend on the
307 # filter joins), so DISTINCT ON (sort_exprs, resource_id) collapses the
308 # duplicate rows that one-to-many filter joins produce.
309 distinct_exprs = ", ".join([*(k.expr for k in keys), "r.resource_id"])
310 order_sql = _build_order(keys)
311 select_sortk = "".join(f", {k.expr} AS sortk{i}" for i, k in enumerate(keys))
313 n_limit = len(args) + 1
314 args.append(count)
316 # is_current (migration 009) replaces the O(table) MAX(version_id) aggregation. DISTINCT ON
317 # is still required: one-to-many filter joins (spidx_*) can emit a current row more than once.
318 sql = f"""
319 {with_sql}
320 SELECT DISTINCT ON ({distinct_exprs})
321 r.resource, r.version_id, r.last_updated, r.resource_id{select_sortk}
322 FROM fhir_resources r
323 {join_sql}
324 WHERE r.tenant_id = $1 AND r.resource_type = $2
325 AND r.is_current AND NOT r.deleted{where_sql}
326 {order_sql}
327 LIMIT ${n_limit}
328 """
329 return sql, args
332def _build_has_clause(
333 hp: HasParam,
334 args: list[Any],
335 cte_name: str,
336) -> tuple[str, str]:
337 """_has reverse chaining as a materialized CTE + join, returned as (cte_body, join).
339 Example: Patient?_has:Coverage:beneficiary:type=GKV
340 Finds patients where a Coverage with beneficiary→patient and type=GKV exists.
342 The set of matching target ids is built once (``AS MATERIALIZED``) and joined on
343 ``r.resource_id`` — deliberately NOT a correlated ``EXISTS (… = r.resource_id)``. Postgres
344 estimated that correlated form at rows=1 while it returned thousands, driving a nested loop
345 that re-probed the join per candidate (the _has degradation in
346 development/performance-befunde-20260719-production-week.md, section 9a). MATERIALIZED is
347 an optimization fence: it forces single evaluation and a hash join, which neither an
348 ``IN`` rewrite nor extended statistics achieved. Uses $1 (tenant_id) and $2 (resource_type).
349 """
350 or_parts: list[str] = []
351 for sv in hp.values:
352 raw = sv.raw
353 if "|" in raw:
354 sys_part, code_part = raw.split("|", 1)
355 n = len(args) + 1
356 if sys_part == "":
357 args.append(code_part)
358 or_parts.append(f"ft.code = ${n}")
359 else:
360 args.append(sys_part)
361 args.append(code_part)
362 or_parts.append(f"(ft.system = ${n} AND ft.code = ${n + 1})")
363 else:
364 n = len(args) + 1
365 args.append(raw)
366 or_parts.append(f"ft.code = ${n}")
368 filter_where = f"({' OR '.join(or_parts)})" if or_parts else "TRUE"
370 jtype_n = len(args) + 1
371 args.append(hp.joining_type)
372 jparam_n = len(args) + 1
373 args.append(hp.join_param)
374 fparam_n = len(args) + 1
375 args.append(hp.filter_param)
377 cte_body = (
378 f"{cte_name} AS MATERIALIZED ("
379 f" SELECT DISTINCT sr.target_resource_id AS rid"
380 f" FROM spidx_reference sr"
381 f" JOIN spidx_token ft"
382 f" ON ft.tenant_id = $1"
383 f" AND ft.resource_type = ${jtype_n}"
384 f" AND ft.resource_id = sr.resource_id"
385 f" AND ft.param_name = ${fparam_n}"
386 f" AND {filter_where}"
387 f" WHERE sr.tenant_id = $1"
388 f" AND sr.resource_type = ${jtype_n}"
389 f" AND sr.param_name = ${jparam_n}"
390 f" AND (sr.target_resource_type = $2 OR sr.target_resource_type IS NULL)"
391 f")"
392 )
393 join = f"JOIN {cte_name} ON {cte_name}.rid = r.resource_id"
394 return cte_body, join
397def _resolve_chained_params(
398 chained_params: list[ChainedParam],
399) -> None:
400 """Populate each ChainedParam's target_type / chained_field / chained_type.
402 Parses an optional resource-type prefix out of ``chained_param``
403 (``Patient.name`` → target "Patient", field "name") and looks up the
404 field's index type — on the named target when given, otherwise across all
405 base types (a standard param's type is consistent wherever defined). The
406 resolved type selects the spidx_* table in ``_build_chained_clause``;
407 ``None`` (unknown field, or an unsupported multi-hop ``a.b.c`` chain) leaves
408 the clause to fall back to token matching, preserving prior behavior.
409 """
410 from .registry import get_registry
412 reg = get_registry()
413 for cp in chained_params:
414 chained = cp.chained_param
415 target_type: str | None = None
416 field = chained
417 if "." in chained:
418 head, rest = chained.split(".", 1)
419 if head[:1].isupper():
420 target_type = head
421 field = rest
422 if "." in field:
423 # Multi-hop chain (a.b.c) — not modeled by the single-hop EXISTS;
424 # leave unresolved so it degrades to the prior token behavior.
425 cp.chained_field = field
426 continue
427 cp.target_type = target_type
428 cp.chained_field = field
429 if target_type is not None:
430 cp.chained_type = reg.base_param_type(target_type, field)
431 else:
432 cp.chained_type = reg.any_base_param_type(field)
435def _build_chained_clause(
436 cp: ChainedParam,
437 args: list[Any],
438 cte_name: str,
439) -> tuple[str, str]:
440 """Chained parameter as a materialized CTE + join: base_param.chained=value.
442 Resolves the base reference (spidx_reference) to its target, then filters
443 the target on ``chained_field`` using the spidx_* table matching the
444 resolved ``chained_type`` (reference / string / date / quantity / number /
445 token). The per-type matching predicate is produced by the same clause
446 builders used for direct params, aliased to ``ct``. Falls back to token
447 matching when the type is unresolved. Uses $1 (tenant_id) and $2
448 (resource_type) from the outer query.
449 """
450 chained_field = cp.chained_field or cp.chained_param
451 synthetic = ParsedParam(name=chained_field, modifier="", values=cp.values)
452 ctype = cp.chained_type or "token"
454 if ctype == "string":
455 table, _, wheres = _string_clause(synthetic, "ct", args)
456 elif ctype == "reference":
457 table, _, wheres = _reference_clause(synthetic, "ct", args)
458 elif ctype == "date":
459 table, _, wheres = _date_clause(synthetic, "ct", args)
460 elif ctype == "quantity":
461 table, _, wheres = _quantity_clause(synthetic, "ct", args)
462 elif ctype == "number":
463 table, _, wheres = _number_clause(synthetic, "ct", args)
464 else: # token (and unresolved fallback)
465 table, _, wheres = _token_clause(synthetic, "ct", args, {})
466 filter_where = " AND ".join(wheres) if wheres else "TRUE"
468 bparam_n = len(args) + 1
469 args.append(cp.base_param)
470 cparam_n = len(args) + 1
471 args.append(chained_field)
473 target_type_filter = ""
474 if cp.target_type is not None:
475 tt_n = len(args) + 1
476 args.append(cp.target_type)
477 target_type_filter = f" AND cr.target_resource_type = ${tt_n}"
479 # Materialized CTE + join for the same planner reason as _build_has_clause: the source id
480 # set is built once and hash-joined, instead of a correlated EXISTS re-probed per candidate.
481 cte_body = (
482 f"{cte_name} AS MATERIALIZED ("
483 f" SELECT DISTINCT cr.resource_id AS rid"
484 f" FROM spidx_reference cr"
485 f" JOIN {table} ct"
486 f" ON ct.tenant_id = $1"
487 f" AND ct.resource_type = cr.target_resource_type"
488 f" AND ct.resource_id = cr.target_resource_id"
489 f" AND ct.param_name = ${cparam_n}"
490 f" AND {filter_where}"
491 f" WHERE cr.tenant_id = $1"
492 f" AND cr.resource_type = $2"
493 f" AND cr.param_name = ${bparam_n}"
494 f" AND cr.target_resource_type IS NOT NULL"
495 f"{target_type_filter}"
496 f")"
497 )
498 join = f"JOIN {cte_name} ON {cte_name}.rid = r.resource_id"
499 return cte_body, join
502def _build_joins_and_where(
503 params: list[ParsedParam],
504 registry: dict[str, SearchParamDef],
505 args: list[Any],
506 tenant_id: UUID,
507 resource_type: str,
508 expanded_valuesets: dict[str, list[dict[str, Any]]],
509 has_params: list[HasParam] | None = None,
510 chained_params: list[ChainedParam] | None = None,
511) -> tuple[list[str], list[str], list[str], list[Any]]:
512 ctes: list[str] = []
513 joins: list[str] = []
514 where_clauses: list[str] = []
516 for i, p in enumerate(params):
517 sp = registry.get(p.name)
518 if sp is None:
519 continue
520 alias = f"sp{i}"
521 table, join_sql, param_wheres = _build_param_clause(
522 p, sp, alias, args, tenant_id, resource_type, expanded_valuesets
523 )
524 if table:
525 joins.append(
526 f"JOIN {table} {alias} ON {alias}.tenant_id = $1"
527 f" AND {alias}.resource_type = $2"
528 f" AND {alias}.resource_id = r.resource_id"
529 f" AND {alias}.param_name = '{sp.name}'"
530 )
531 if join_sql:
532 joins.append(join_sql)
533 where_clauses.extend(param_wheres)
535 # _has and chained each become a materialized CTE joined on r.resource_id (see
536 # _build_has_clause). CTE names must be unique within the statement, hence the shared index.
537 cte_idx = 0
538 for hp in has_params or []:
539 cte_body, join = _build_has_clause(hp, args, f"_hc{cte_idx}")
540 ctes.append(cte_body)
541 joins.append(join)
542 cte_idx += 1
544 for cp in chained_params or []:
545 cte_body, join = _build_chained_clause(cp, args, f"_hc{cte_idx}")
546 ctes.append(cte_body)
547 joins.append(join)
548 cte_idx += 1
550 return ctes, joins, where_clauses, args
553def _build_param_clause(
554 p: ParsedParam,
555 sp: SearchParamDef,
556 alias: str,
557 args: list[Any],
558 tenant_id: UUID,
559 resource_type: str,
560 expanded_valuesets: dict[str, list[dict[str, Any]]],
561) -> tuple[str, str, list[str]]:
562 """Return (table_name, extra_join_sql, where_clauses)."""
564 # :missing modifier — check absence in the spidx table using NOT EXISTS
565 if p.modifier == "missing":
566 is_missing = p.values and p.values[0].raw.lower() == "true"
567 table = _table_for_type(sp.type)
568 exists_clause = (
569 f"EXISTS (SELECT 1 FROM {table} x"
570 f" WHERE x.tenant_id=$1 AND x.resource_type=$2"
571 f" AND x.resource_id=r.resource_id AND x.param_name='{sp.name}')"
572 )
573 if is_missing:
574 return "", "", [f"NOT {exists_clause}"]
575 else:
576 return "", "", [exists_clause]
578 if sp.type == "string":
579 return _string_clause(p, alias, args)
580 elif sp.type == "token":
581 return _token_clause(p, alias, args, expanded_valuesets)
582 elif sp.type == "date":
583 return _date_clause(p, alias, args)
584 elif sp.type == "reference":
585 return _reference_clause(p, alias, args)
586 elif sp.type == "quantity":
587 return _quantity_clause(p, alias, args)
588 elif sp.type == "number":
589 return _number_clause(p, alias, args)
590 elif sp.type == "composite":
591 return _composite_clause(p, sp, alias, args)
592 elif sp.type == "geo":
593 return _geo_clause(p, alias, args)
595 return "", "", []
598# ── per-type clause builders ───────────────────────────────────────────────────
601def _string_clause(p: ParsedParam, alias: str, args: list[Any]) -> tuple[str, str, list[str]]:
602 or_parts: list[str] = []
603 for sv in p.values:
604 n = len(args) + 1
605 if p.modifier == "exact":
606 args.append(sv.raw)
607 or_parts.append(f"{alias}.value = ${n}")
608 elif p.modifier == "contains":
609 args.append(f"%{sv.raw.lower()}%")
610 or_parts.append(f"{alias}.value LIKE ${n}")
611 else: # default: starts-with (FHIR default for string)
612 args.append(f"{sv.raw.lower()}%")
613 or_parts.append(f"{alias}.value LIKE ${n}")
614 where = f"({' OR '.join(or_parts)})" if or_parts else "TRUE"
615 return "spidx_string", "", [where]
618def _build_concept_parts(
619 concepts: list[dict[str, Any]],
620 alias: str,
621 args: list[Any],
622) -> list[str]:
623 concept_parts: list[str] = []
624 for concept in concepts:
625 system = concept.get("system")
626 code = concept.get("code")
627 if not code:
628 continue
629 if system:
630 n = len(args) + 1
631 args.append(system)
632 args.append(code)
633 concept_parts.append(
634 f"(({alias}.system IS NULL OR {alias}.system = ${n}) AND {alias}.code = ${n + 1})"
635 )
636 else:
637 n = len(args) + 1
638 args.append(code)
639 concept_parts.append(f"{alias}.code = ${n}")
640 return concept_parts
643def _token_valueset_clause(
644 p: ParsedParam,
645 alias: str,
646 args: list[Any],
647 expanded_valuesets: dict[str, list[dict[str, Any]]],
648) -> tuple[str, str, list[str]]:
649 or_parts: list[str] = []
650 for sv in p.values:
651 concepts = expanded_valuesets.get(sv.raw, [])
652 concept_parts = _build_concept_parts(concepts, alias, args)
653 if concept_parts:
654 or_parts.append(f"({' OR '.join(concept_parts)})")
655 else:
656 or_parts.append("FALSE")
658 if p.modifier == "not-in":
659 where = f"NOT ({' OR '.join(or_parts)})" if or_parts else "TRUE"
660 else:
661 where = f"({' OR '.join(or_parts)})" if or_parts else "FALSE"
663 return "spidx_token", "", [where]
666def _token_direct_clause(
667 p: ParsedParam,
668 alias: str,
669 args: list[Any],
670) -> tuple[str, str, list[str]]:
671 or_parts: list[str] = []
672 for sv in p.values:
673 raw = sv.raw
674 # Format: [system]|code or |code or system| or code
675 if "|" in raw:
676 system_part, code_part = raw.split("|", 1)
677 n = len(args) + 1
678 if system_part == "" and code_part == "":
679 # | — degenerate, matches nothing meaningful; skip.
680 continue
681 if system_part == "":
682 # |code — match code, any system
683 args.append(code_part)
684 or_parts.append(f"{alias}.code = ${n}")
685 elif code_part == "":
686 # system| — match system, any code
687 args.append(system_part)
688 or_parts.append(f"{alias}.system = ${n}")
689 else:
690 args.append(system_part)
691 args.append(code_part)
692 or_parts.append(
693 f"({alias}.system = ${n} OR {alias}.system IS NULL) AND {alias}.code = ${n + 1}"
694 )
695 else:
696 n = len(args) + 1
697 args.append(raw)
698 or_parts.append(f"{alias}.code = ${n}")
700 if p.modifier == "not":
701 where = f"NOT ({' OR '.join(or_parts)})" if or_parts else "TRUE"
702 else:
703 where = f"({' OR '.join(or_parts)})" if or_parts else "TRUE"
705 return "spidx_token", "", [where]
708def _token_text_clause(
709 p: ParsedParam,
710 alias: str,
711 args: list[Any],
712) -> tuple[str, str, list[str]]:
713 """:text modifier — case-insensitive starts-with match against the token's
714 textual part (``CodeableConcept.text`` / ``Coding.display``)."""
715 or_parts: list[str] = []
716 for sv in p.values:
717 n = len(args) + 1
718 args.append(f"{sv.raw.lower()}%")
719 or_parts.append(f"{alias}.display IS NOT NULL AND lower({alias}.display) LIKE ${n}")
720 where = f"({' OR '.join(f'({part})' for part in or_parts)})" if or_parts else "FALSE"
721 return "spidx_token", "", [where]
724def _token_clause(
725 p: ParsedParam,
726 alias: str,
727 args: list[Any],
728 expanded_valuesets: dict[str, list[dict[str, Any]]],
729) -> tuple[str, str, list[str]]:
730 if p.modifier in ("in", "not-in"):
731 return _token_valueset_clause(p, alias, args, expanded_valuesets)
732 if p.modifier == "text":
733 return _token_text_clause(p, alias, args)
734 return _token_direct_clause(p, alias, args)
737def _date_clause(p: ParsedParam, alias: str, args: list[Any]) -> tuple[str, str, list[str]]:
738 or_parts: list[str] = []
739 for sv in p.values:
740 try:
741 lo, hi = _parse_search_date(sv.raw)
742 except ValueError:
743 continue
744 prefix = sv.prefix or "eq"
745 cond = _date_comparison(alias, prefix, lo, hi, args)
746 if cond:
747 or_parts.append(cond)
748 where = f"({' OR '.join(or_parts)})" if or_parts else "TRUE"
749 return "spidx_date", "", [where]
752def _date_comparison(
753 alias: str, prefix: Prefix | str, lo: datetime, hi: datetime, args: list[Any]
754) -> str:
755 # The indexed range: [value_low, value_high]
756 # The search range: [lo, hi]
757 if prefix in ("eq", ""):
758 # Ranges overlap
759 n = len(args) + 1
760 args.extend([hi, lo])
761 return (
762 f"({alias}.value_low <= ${n}::timestamptz"
763 f" AND COALESCE({alias}.value_high, {alias}.value_low) >= ${n + 1}::timestamptz)"
764 )
765 elif prefix == "ne":
766 n = len(args) + 1
767 args.extend([hi, lo])
768 return (
769 f"NOT ({alias}.value_low <= ${n}::timestamptz"
770 f" AND COALESCE({alias}.value_high, {alias}.value_low) >= ${n + 1}::timestamptz)"
771 )
772 elif prefix == "gt":
773 n = len(args) + 1
774 args.append(hi)
775 return f"COALESCE({alias}.value_high, {alias}.value_low) > ${n}::timestamptz"
776 elif prefix == "ge":
777 n = len(args) + 1
778 args.append(lo)
779 return f"COALESCE({alias}.value_high, {alias}.value_low) >= ${n}::timestamptz"
780 elif prefix == "lt":
781 n = len(args) + 1
782 args.append(lo)
783 return f"{alias}.value_low < ${n}::timestamptz"
784 elif prefix == "le":
785 n = len(args) + 1
786 args.append(hi)
787 return f"{alias}.value_low <= ${n}::timestamptz"
788 elif prefix == "sa": # starts after
789 n = len(args) + 1
790 args.append(hi)
791 return f"{alias}.value_low > ${n}::timestamptz"
792 elif prefix == "eb": # ends before
793 n = len(args) + 1
794 args.append(lo)
795 return f"COALESCE({alias}.value_high, {alias}.value_low) < ${n}::timestamptz"
796 return "TRUE"
799def _reference_clause(p: ParsedParam, alias: str, args: list[Any]) -> tuple[str, str, list[str]]:
800 or_parts: list[str] = []
801 for sv in p.values:
802 raw = sv.raw
803 # Format: ResourceType/id or id or ResourceType/id (absolute URL last segment)
804 if "/" in raw:
805 parts = raw.rsplit("/", 1)
806 target_type, target_id = parts[0].split("/")[-1], parts[-1]
807 n = len(args) + 1
808 args.extend([target_type, target_id])
809 or_parts.append(
810 f"({alias}.target_resource_type = ${n} AND {alias}.target_resource_id = ${n + 1})"
811 )
812 else:
813 n = len(args) + 1
814 args.append(raw)
815 or_parts.append(f"{alias}.target_resource_id = ${n}")
817 where = f"({' OR '.join(or_parts)})" if or_parts else "TRUE"
818 return "spidx_reference", "", [where]
821def _quantity_clause(p: ParsedParam, alias: str, args: list[Any]) -> tuple[str, str, list[str]]:
822 or_parts: list[str] = []
823 for sv in p.values:
824 # Format: value|system|code or value|system or value
825 parts = sv.raw.split("|")
826 try:
827 num = float(parts[0])
828 except (ValueError, IndexError):
829 continue
830 prefix = sv.prefix or "eq"
831 n = len(args) + 1
832 args.append(num)
833 cond = _numeric_comparison(alias, prefix, n)
834 if len(parts) >= 2 and parts[1]:
835 n2 = len(args) + 1
836 args.append(parts[1])
837 cond = f"({cond} AND {alias}.system = ${n2})"
838 if len(parts) >= 3 and parts[2]:
839 n3 = len(args) + 1
840 args.append(parts[2])
841 cond = f"({cond} AND {alias}.code = ${n3})"
842 or_parts.append(cond)
843 where = f"({' OR '.join(or_parts)})" if or_parts else "TRUE"
844 return "spidx_quantity", "", [where]
847def _number_clause(p: ParsedParam, alias: str, args: list[Any]) -> tuple[str, str, list[str]]:
848 or_parts: list[str] = []
849 for sv in p.values:
850 try:
851 num = float(sv.raw)
852 except ValueError:
853 continue
854 prefix = sv.prefix or "eq"
855 n = len(args) + 1
856 args.append(num)
857 or_parts.append(_numeric_comparison(alias, prefix, n))
858 where = f"({' OR '.join(or_parts)})" if or_parts else "TRUE"
859 return "spidx_number", "", [where]
862def _composite_token_component(
863 raw: str, alias: str, system_col: str, code_col: str, args: list[Any]
864) -> str:
865 """Same [system]|code / |code / system| / code matching as a plain token
866 param (see _token_direct_clause), targeting the given composite columns.
867 """
868 if "|" in raw:
869 system_part, code_part = raw.split("|", 1)
870 n = len(args) + 1
871 if system_part == "" and code_part == "":
872 return "FALSE"
873 if system_part == "":
874 args.append(code_part)
875 return f"{alias}.{code_col} = ${n}"
876 if code_part == "":
877 args.append(system_part)
878 return f"{alias}.{system_col} = ${n}"
879 args.append(system_part)
880 args.append(code_part)
881 return (
882 f"({alias}.{system_col} = ${n} OR {alias}.{system_col} IS NULL)"
883 f" AND {alias}.{code_col} = ${n + 1}"
884 )
885 n = len(args) + 1
886 args.append(raw)
887 return f"{alias}.{code_col} = ${n}"
890def _composite_quantity_component(raw: str, alias: str, args: list[Any]) -> str:
891 """Exact-match only (no gt/lt/... prefixes) — value[|system[|code]]."""
892 parts = raw.split("|")
893 try:
894 num = float(parts[0])
895 except (ValueError, IndexError):
896 return "FALSE"
897 n = len(args) + 1
898 args.append(num)
899 cond = f"{alias}.c2_qty_value = ${n}"
900 if len(parts) >= 2 and parts[1]:
901 n2 = len(args) + 1
902 args.append(parts[1])
903 cond = f"({cond} AND {alias}.c2_qty_system = ${n2})"
904 if len(parts) >= 3 and parts[2]:
905 n3 = len(args) + 1
906 args.append(parts[2])
907 cond = f"({cond} AND {alias}.c2_qty_code = ${n3})"
908 return cond
911def _composite_clause(
912 p: ParsedParam, sp: SearchParamDef, alias: str, args: list[Any]
913) -> tuple[str, str, list[str]]:
914 """Composite param: value is ``component1$component2`` (FHIR composite
915 search syntax). Both components must match on the *same* indexed row —
916 that row-level correlation is what spidx_composite exists for.
917 """
918 if sp.components is None or len(sp.components) != 2:
919 return "", "", ["FALSE"]
920 _, c2 = sp.components
922 or_parts: list[str] = []
923 for sv in p.values:
924 comp_parts = sv.raw.split("$", 1)
925 if len(comp_parts) != 2:
926 continue
927 c1_raw, c2_raw = comp_parts
928 c1_cond = _composite_token_component(
929 c1_raw, alias, "c1_token_system", "c1_token_code", args
930 )
931 c2_cond = (
932 _composite_token_component(c2_raw, alias, "c2_token_system", "c2_token_code", args)
933 if c2.type == "token"
934 else _composite_quantity_component(c2_raw, alias, args)
935 )
936 or_parts.append(f"({c1_cond} AND {c2_cond})")
938 where = f"({' OR '.join(or_parts)})" if or_parts else "TRUE"
939 return "spidx_composite", "", [where]
942# near=lat|lon[|distance|units] — distance defaults to this radius (km) when
943# omitted, matching common FHIR server defaults for an unqualified proximity
944# search.
945_DEFAULT_NEAR_DISTANCE_KM = 10.0
946_EARTH_RADIUS_KM = 6371.0
947_KM_PER_UNIT: dict[str, float] = {"km": 1.0, "mi": 1.609344, "[mi_us]": 1.609344}
950def _geo_clause(p: ParsedParam, alias: str, args: list[Any]) -> tuple[str, str, list[str]]:
951 """``near=latitude|longitude[|distance|units]`` — Haversine great-circle
952 distance computed in SQL (no PostGIS/earthdistance extension required).
953 """
954 or_parts: list[str] = []
955 for sv in p.values:
956 parts = sv.raw.split("|")
957 try:
958 lat = float(parts[0])
959 lon = float(parts[1])
960 except (ValueError, IndexError):
961 continue
962 distance_km = _DEFAULT_NEAR_DISTANCE_KM
963 if len(parts) >= 3 and parts[2]:
964 with contextlib.suppress(ValueError):
965 distance_km = float(parts[2]) * _KM_PER_UNIT.get(
966 parts[3] if len(parts) >= 4 and parts[3] else "km", 1.0
967 )
968 n = len(args) + 1
969 args.extend([lat, lon, distance_km])
970 # Clamp the acos() argument to [-1, 1] — floating-point rounding can
971 # push it just outside that domain for near-identical/antipodal points.
972 or_parts.append(
973 f"({_EARTH_RADIUS_KM} * acos(LEAST(1.0, GREATEST(-1.0,"
974 f" cos(radians(${n})) * cos(radians({alias}.latitude))"
975 f" * cos(radians({alias}.longitude) - radians(${n + 1}))"
976 f" + sin(radians(${n})) * sin(radians({alias}.latitude))"
977 f")))) <= ${n + 2}"
978 )
979 where = f"({' OR '.join(or_parts)})" if or_parts else "TRUE"
980 return "spidx_geo", "", [where]
983def _numeric_comparison(alias: str, prefix: str, n: int) -> str:
984 if prefix in ("eq", ""):
985 return f"{alias}.value = ${n}"
986 elif prefix == "ne":
987 return f"{alias}.value != ${n}"
988 elif prefix == "gt":
989 return f"{alias}.value > ${n}"
990 elif prefix == "ge":
991 return f"{alias}.value >= ${n}"
992 elif prefix == "lt":
993 return f"{alias}.value < ${n}"
994 elif prefix == "le":
995 return f"{alias}.value <= ${n}"
996 return f"{alias}.value = ${n}"
999# ── sort ───────────────────────────────────────────────────────────────────────
1002@dataclass(frozen=True)
1003class _OrderKey:
1004 """One resolved sort key: a deterministic per-resource SQL expression."""
1006 expr: str # scalar SQL expression (column or correlated subquery)
1007 direction: str # "asc" | "desc"
1008 pgtype: str # "timestamptz" | "numeric" | "text" — cast for cursor params
1011# Column carrying the sortable value in each spidx table, and the cursor cast.
1012_SORT_COLUMN: dict[SearchParamType, tuple[str, str]] = {
1013 "string": ("value", "text"),
1014 "token": ("code", "text"),
1015 "date": ("value_low", "timestamptz"),
1016 "reference": ("target_resource_id", "text"),
1017 "quantity": ("value", "numeric"),
1018 "number": ("value", "numeric"),
1019}
1022def _order_keys(sort: list[SortSpec], registry: dict[str, SearchParamDef]) -> list[_OrderKey]:
1023 """Resolve ``_sort`` specs into SQL order keys.
1025 ``_lastUpdated`` and ``_id`` map to columns on ``fhir_resources``; any other
1026 spec must name a registered search parameter, whose value is pulled from its
1027 ``spidx_*`` table via a correlated aggregate (MIN for asc, MAX for desc, so
1028 multi-valued parameters sort by the value that comes first in that
1029 direction). Unknown sort parameters are ignored (FHIR: servers MAY ignore).
1030 Defaults to ``_lastUpdated`` ascending when no usable sort is given.
1031 """
1032 keys: list[_OrderKey] = []
1033 for spec in sort:
1034 if spec.param == "_lastUpdated":
1035 keys.append(_OrderKey("r.last_updated", spec.direction, "timestamptz"))
1036 elif spec.param == "_id":
1037 keys.append(_OrderKey("r.resource_id", spec.direction, "text"))
1038 else:
1039 sp = registry.get(spec.param)
1040 if sp is None or sp.type not in _SORT_COLUMN:
1041 continue # unknown / unsortable parameter — ignore
1042 column, pgtype = _SORT_COLUMN[sp.type]
1043 table = _table_for_type(sp.type)
1044 agg = "MIN" if spec.direction == "asc" else "MAX"
1045 expr = (
1046 f"(SELECT {agg}(s.{column}) FROM {table} s"
1047 f" WHERE s.tenant_id = $1 AND s.resource_type = $2"
1048 f" AND s.resource_id = r.resource_id AND s.param_name = '{sp.name}')"
1049 )
1050 keys.append(_OrderKey(expr, spec.direction, pgtype))
1052 if not keys:
1053 keys.append(_OrderKey("r.last_updated", "asc", "timestamptz"))
1054 return keys
1057def _build_order(keys: list[_OrderKey]) -> str:
1058 parts: list[str] = []
1059 for k in keys:
1060 if k.direction == "asc":
1061 parts.append(f"{k.expr} ASC NULLS LAST")
1062 else:
1063 parts.append(f"{k.expr} DESC NULLS FIRST")
1064 parts.append("r.resource_id ASC")
1065 return "ORDER BY " + ", ".join(parts)
1068def _keyset_predicate(
1069 keys: list[_OrderKey], values: list[str | None], cursor_id: str, args: list[Any]
1070) -> str:
1071 """Build the lexicographic ``row > cursor`` predicate for keyset pagination.
1073 Handles per-key direction and NULL ordering (asc → NULLS LAST, desc → NULLS
1074 FIRST) so it matches the ORDER BY exactly. ``r.resource_id`` is the final,
1075 unique, non-null tiebreaker.
1076 """
1077 placeholders: list[str] = []
1078 for k, v in zip(keys, values, strict=True):
1079 args.append(v)
1080 # Cursor values travel as text; cast via text so asyncpg sends a string
1081 # (a bare ``$n::timestamptz`` would make asyncpg expect a datetime).
1082 placeholders.append(f"(${len(args)}::text)::{k.pgtype}")
1083 args.append(cursor_id)
1084 id_ph = f"${len(args)}"
1086 def after(k: _OrderKey, ph: str) -> str:
1087 if k.direction == "asc":
1088 # NULLS LAST: a row sorts after the cursor when the cursor value is
1089 # non-null and the row value is larger or null.
1090 larger = f"{k.expr} IS NULL OR {k.expr} > {ph}"
1091 return f"CASE WHEN {ph} IS NULL THEN FALSE ELSE ({larger}) END"
1092 # desc, NULLS FIRST: rows smaller than the cursor, or any non-null row
1093 # when the cursor value itself is null.
1094 smaller = f"{k.expr} IS NOT NULL AND {k.expr} < {ph}"
1095 return f"CASE WHEN {ph} IS NULL THEN {k.expr} IS NOT NULL ELSE ({smaller}) END"
1097 terms: list[str] = []
1098 eq_prefix: list[str] = []
1099 for k, ph in zip(keys, placeholders, strict=True):
1100 terms.append(" AND ".join([*eq_prefix, f"({after(k, ph)})"]))
1101 eq_prefix.append(f"({k.expr} IS NOT DISTINCT FROM {ph})")
1102 terms.append(" AND ".join([*eq_prefix, f"r.resource_id > {id_ph}"]))
1104 return "(" + " OR ".join(f"({t})" for t in terms) + ")"
1107def _table_for_type(sp_type: SearchParamType) -> str:
1108 return {
1109 "string": "spidx_string",
1110 "token": "spidx_token",
1111 "date": "spidx_date",
1112 "reference": "spidx_reference",
1113 "quantity": "spidx_quantity",
1114 "number": "spidx_number",
1115 "composite": "spidx_composite",
1116 "geo": "spidx_geo",
1117 }[sp_type]
1120# ── cursor ─────────────────────────────────────────────────────────────────────
1123def _serialize_sort_value(value: Any) -> str | None:
1124 """Render a sort-key value for the cursor; SQL re-casts it via ``pgtype``."""
1125 if value is None:
1126 return None
1127 if isinstance(value, datetime):
1128 ts = value if value.tzinfo is not None else value.replace(tzinfo=UTC)
1129 return ts.isoformat()
1130 if isinstance(value, Decimal):
1131 return str(value)
1132 return str(value)
1135def _encode_cursor(sort_values: list[str | None], resource_id: str) -> str:
1136 payload = json.dumps({"k": sort_values, "id": resource_id})
1137 return base64.urlsafe_b64encode(payload.encode()).decode()
1140def _decode_cursor(cursor: str) -> tuple[list[str | None], str | None]:
1141 try:
1142 payload = json.loads(base64.urlsafe_b64decode(cursor.encode()))
1143 values = [None if v is None else str(v) for v in payload["k"]]
1144 return values, str(payload["id"])
1145 except Exception:
1146 return [], None
1149# ── date parsing (search values) ──────────────────────────────────────────────
1152def _parse_search_date(value: str) -> tuple[datetime, datetime]:
1153 """Parse a FHIR search date string to (low, high) inclusive bounds (UTC)."""
1154 import re
1156 _DATE_ONLY_RE = re.compile(r"^\d{4}(-\d{2}(-\d{2})?)?$")
1157 if _DATE_ONLY_RE.match(value):
1158 parts = value.split("-")
1159 if len(parts) == 1:
1160 import calendar
1162 lo = datetime(int(parts[0]), 1, 1, tzinfo=UTC)
1163 hi = datetime(int(parts[0]), 12, 31, 23, 59, 59, tzinfo=UTC)
1164 elif len(parts) == 2:
1165 import calendar
1167 y, m = int(parts[0]), int(parts[1])
1168 last_day = calendar.monthrange(y, m)[1]
1169 lo = datetime(y, m, 1, tzinfo=UTC)
1170 hi = datetime(y, m, last_day, 23, 59, 59, tzinfo=UTC)
1171 else:
1172 y, m, d = int(parts[0]), int(parts[1]), int(parts[2])
1173 lo = datetime(y, m, d, tzinfo=UTC)
1174 hi = datetime(y, m, d, 23, 59, 59, tzinfo=UTC)
1175 return lo, hi
1176 clean = value.replace("Z", "+00:00")
1177 clean = re.sub(r"(\.\d{6})\d+", r"\1", clean)
1178 dt = datetime.fromisoformat(clean)
1179 if dt.tzinfo is None:
1180 dt = dt.replace(tzinfo=UTC)
1181 return dt, dt