fix(effects): paginate FHIR and REST read-backs and verify total count (#1476)

packages/flow/openadapt_flow/runtime/effects/fhir.py

@@ -46,6 +46,7 @@
 
 import time
 from typing import Any, Optional, Protocol
+from urllib.parse import urljoin
 
 from openadapt_flow.runtime.effects._common import drop_record_fields, judge_records
 from openadapt_flow.runtime.effects.adapter import VerifierAdapterBase
@@ -135,9 +136,11 @@ def __init__(
         verify_tls: bool = True,
         poll_interval_s: float = 0.3,
         drop_fields: tuple[str, ...] = (),
+        max_pages: int = 50,
     ) -> None:
         self.base_url = base_url.rstrip("/")
         self.drop_fields = tuple(drop_fields)
+        self.max_pages = max_pages
         self.resource_type = resource_type
         self.search_params = dict(search_params or {})
         self.field_paths = dict(field_paths or DEFAULT_OBSERVATION_PATHS)
@@ -209,50 +212,85 @@ def _flatten(self, resource: dict[str, Any]) -> dict[str, Any]:
     def _search(self) -> Optional[list[dict[str, Any]]]:
         """Run the configured FHIR search; return flattened resources.
 
-        Returns ``None`` (read as unreadable -> INDETERMINATE) on any
-        transport error, non-2xx (incl. auth 401/403), non-Bundle body, or
-        malformed entries. Never raises.
+        Follows pagination (Bundle.link[rel=next]) up to ``max_pages`` and
+        verifies that the number of resources read matches Bundle.total when
+        provided. Returns ``None`` (read as unreadable -> INDETERMINATE) on
+        any transport error, non-2xx (incl. auth 401/403), non-Bundle body,
+        malformed entries, pagination cap exhaustion, or total count mismatch.
+        Never raises.
         """
         url = f"{self.base_url}/{self.resource_type}"
         headers = self._headers()
         if headers is None:
             return None
         session = self._get_session()
-        # ``verify`` is passed per request only for a ``requests``-style client;
-        # the httpx fallback already bound it at construction (see
-        # ``_get_session``), and passing it again would raise.
         tls_kwargs = {"verify": self.verify_tls} if self._verify_per_request else {}
-        try:
-            resp = session.get(
-                url,
-                params=self.search_params,
-                headers=headers,
-                timeout=self.timeout_s,
-                **tls_kwargs,
-            )
-        except Exception:  # noqa: BLE001 - transport failure is unreadable
-            return None
-        if resp.status_code == 401 and self.token_source is not None:
-            self.token_source.invalidate()
-        if resp.status_code // 100 != 2:
-            return None
-        try:
-            bundle = resp.json()
-        except Exception:  # noqa: BLE001 - unparseable body is unreadable
-            return None
-        if not isinstance(bundle, dict) or bundle.get("resourceType") != "Bundle":
-            return None
-        entries = bundle.get("entry", [])
-        if not isinstance(entries, list):
-            return None
+
         out: list[dict[str, Any]] = []
-        for entry in entries:
-            if not isinstance(entry, dict):
+        expected_total: Optional[int] = None
+        next_url: Optional[str] = url
+        params: Optional[dict[str, Any]] = self.search_params
+        pages_read = 0
+
+        while next_url is not None:
+            if pages_read >= self.max_pages:
+                return None
+            try:
+                resp = session.get(
+                    next_url,
+                    params=params,
+                    headers=headers,
+                    timeout=self.timeout_s,
+                    **tls_kwargs,
+                )
+            except Exception:  # noqa: BLE001 - transport failure is unreadable
                 return None
-            resource = entry.get("resource")
-            if not isinstance(resource, dict):
+            if resp.status_code == 401 and self.token_source is not None:
+                self.token_source.invalidate()
+            if resp.status_code // 100 != 2:
                 return None
-            out.append(self._flatten(resource))
+            try:
+                bundle = resp.json()
+            except Exception:  # noqa: BLE001 - unparseable body is unreadable
+                return None
+            if not isinstance(bundle, dict) or bundle.get("resourceType") != "Bundle":
+                return None
+
+            if expected_total is None and isinstance(bundle.get("total"), int):
+                expected_total = bundle["total"]
+
+            entries = bundle.get("entry", [])
+            if not isinstance(entries, list):
+                return None
+            for entry in entries:
+                if not isinstance(entry, dict):
+                    return None
+                resource = entry.get("resource")
+                if not isinstance(resource, dict):
+                    return None
+                out.append(self._flatten(resource))
+
+            pages_read += 1
+
+            # Check for next page in bundle.link
+            links = bundle.get("link", [])
+            next_link_url: Optional[str] = None
+            if isinstance(links, list):
+                for link in links:
+                    if isinstance(link, dict) and link.get("relation") == "next":
+                        raw_link = link.get("url")
+                        if isinstance(raw_link, str) and raw_link.strip():
+                            next_link_url = urljoin(
+                                str(getattr(resp, "url", next_url)), raw_link.strip()
+                            )
+                        break
+
+            next_url = next_link_url
+            params = None  # query parameters are already encoded in next_url
+
+        if expected_total is not None and len(out) != expected_total:
+            return None
+
         return drop_record_fields(out, self.drop_fields)
 
     # -- EffectVerifier protocol --------------------------------------------

packages/flow/openadapt_flow/runtime/effects/rest.py

@@ -22,6 +22,7 @@
 
 import time
 from typing import Any, Optional
+from urllib.parse import urljoin
 
 from openadapt_flow.runtime.effects._common import drop_record_fields, judge_records
 from openadapt_flow.runtime.effects.adapter import VerifierAdapterBase
@@ -71,6 +72,7 @@ def __init__(
         timeout_s: float = 5.0,
         poll_interval_s: float = 0.2,
         drop_fields: tuple[str, ...] = (),
+        max_pages: int = 50,
     ) -> None:
         self.base_url = base_url.rstrip("/")
         self.drop_fields = tuple(drop_fields)
@@ -80,6 +82,7 @@ def __init__(
         self.timeout_s = timeout_s
         self.poll_interval_s = poll_interval_s
         self._session = session
+        self.max_pages = max_pages
 
     # -- transport ----------------------------------------------------------
 
@@ -110,39 +113,86 @@ def _get_session(self) -> Any:
     def _fetch_records(self) -> Optional[list[dict[str, Any]]]:
         """GET the system-of-record document and extract the records list.
 
+        Follows pagination (RFC 5988 Link header rel=next or body next/links)
+        up to ``max_pages`` and validates that the number of records read
+        matches ``total`` when provided.
+
         Returns ``None`` -- read as unreadable, forcing INDETERMINATE -- on
-        any transport error, non-2xx status, or shape mismatch. Never raises.
+        any transport error, non-2xx status, shape mismatch, pagination cap
+        exhaustion, or total count mismatch. Never raises.
         """
         url = f"{self.base_url}{self.records_path}"
-        try:
-            # ``headers`` is passed ONLY when configured, so injected
-            # ``requests``-style sessions with a narrower ``get`` signature
-            # (tests, custom oracles) keep working exactly as before.
-            if self.headers:
-                resp = self._get_session().get(
-                    url, timeout=self.timeout_s, headers=self.headers
-                )
+        all_records: list[dict[str, Any]] = []
+        expected_total: Optional[int] = None
+        next_url: Optional[str] = url
+        pages_read = 0
+
+        while next_url is not None:
+            if pages_read >= self.max_pages:
+                return None
+            try:
+                # ``headers`` is passed ONLY when configured, so injected
+                # ``requests``-style sessions with a narrower ``get`` signature
+                # (tests, custom oracles) keep working exactly as before.
+                if self.headers:
+                    resp = self._get_session().get(
+                        next_url, timeout=self.timeout_s, headers=self.headers
+                    )
+                else:
+                    resp = self._get_session().get(next_url, timeout=self.timeout_s)
+            except Exception:  # noqa: BLE001 - any transport failure is unreadable
+                return None
+            if resp.status_code // 100 != 2:
+                return None
+            try:
+                body = resp.json()
+            except Exception:  # noqa: BLE001 - unparseable body is unreadable
+                return None
+
+            if self.records_key is None:
+                records = body
+            elif isinstance(body, dict):
+                records = body.get(self.records_key)
+                if expected_total is None and isinstance(body.get("total"), int):
+                    expected_total = body["total"]
+                elif expected_total is None and isinstance(
+                    body.get("total_count"), int
+                ):
+                    expected_total = body["total_count"]
             else:
-                resp = self._get_session().get(url, timeout=self.timeout_s)
-        except Exception:  # noqa: BLE001 - any transport failure is unreadable
-            return None
-        if resp.status_code // 100 != 2:
-            return None
-        try:
-            body = resp.json()
-        except Exception:  # noqa: BLE001 - unparseable body is unreadable
-            return None
-        if self.records_key is None:
-            records = body
-        elif isinstance(body, dict):
-            records = body.get(self.records_key)
-        else:
-            return None
-        if not isinstance(records, list):
+                return None
+
+            if not isinstance(records, list):
+                return None
+            all_records.extend([r for r in records if isinstance(r, dict)])
+            pages_read += 1
+
+            # Check next link: from resp.links (RFC 5988 / 8288 Link header) or body dictionary
+            next_link: Optional[str] = None
+            if hasattr(resp, "links") and isinstance(resp.links, dict):
+                next_entry = resp.links.get("next")
+                if isinstance(next_entry, dict) and next_entry.get("url"):
+                    next_link = next_entry["url"]
+                elif isinstance(next_entry, str):
+                    next_link = next_entry
+            if next_link is None and isinstance(body, dict):
+                raw_next = body.get("next") or body.get("next_url")
+                if isinstance(raw_next, str) and raw_next.strip():
+                    next_link = raw_next.strip()
+                elif isinstance(body.get("links"), dict):
+                    links_obj = body["links"]
+                    if isinstance(links_obj.get("next"), str):
+                        next_link = links_obj["next"].strip()
+
+            if next_link is not None:
+                next_url = urljoin(str(getattr(resp, "url", next_url)), next_link)
+            else:
+                next_url = None
+
+        if expected_total is not None and len(all_records) != expected_total:
             return None
-        return drop_record_fields(
-            [r for r in records if isinstance(r, dict)], self.drop_fields
-        )
+
+        return drop_record_fields(all_records, self.drop_fields)
 
     # -- EffectVerifier protocol --------------------------------------------
 

packages/flow/tests/_fhir_fake.py

@@ -97,13 +97,36 @@ def search(self, patient: Optional[str]) -> list[dict]:
             return out
 
 
-def _bundle(resources: list[dict]) -> dict:
+def _bundle(
+    resources: list[dict],
+    *,
+    count: Optional[int] = None,
+    offset: int = 0,
+    base_url: str = "http://fake/Observation",
+    extra_params: Optional[dict[str, str]] = None,
+    total_override: Optional[int] = None,
+) -> dict:
+    total = len(resources) if total_override is None else total_override
+    if count is not None and count > 0:
+        paged = resources[offset : offset + count]
+    else:
+        paged = resources[offset:]
+
+    links = [{"relation": "self", "url": base_url}]
+    if count is not None and offset + count < len(resources):
+        params = dict(extra_params or {})
+        params["_count"] = str(count)
+        params["_offset"] = str(offset + count)
+        query_str = "&".join(f"{k}={v}" for k, v in params.items())
+        next_url = f"{base_url}?{query_str}" if query_str else base_url
+        links.append({"relation": "next", "url": next_url})
+
     return {
         "resourceType": "Bundle",
         "type": "searchset",
-        "total": len(resources),
-        "link": [{"relation": "self", "url": "http://fake/Observation"}],
-        "entry": [{"fullUrl": f"urn:uuid:{r['id']}", "resource": r} for r in resources],
+        "total": total,
+        "link": links,
+        "entry": [{"fullUrl": f"urn:uuid:{r['id']}", "resource": r} for r in paged],
     }
 
 
@@ -139,7 +162,26 @@ def do_GET(self) -> None:  # noqa: N802
             if parsed.path.rstrip("/").endswith("/Observation"):
                 qs = parse_qs(parsed.query)
                 patient = (qs.get("patient", [None]) or [None])[0]
-                self._send(200, _bundle(store.search(patient)))
+                count_raw = (qs.get("_count", [None]) or [None])[0]
+                offset_raw = (qs.get("_offset", [None]) or [None])[0]
+                bad_total = (qs.get("_bad_total", [None]) or [None])[0]
+                count = int(count_raw) if count_raw is not None else None
+                offset = int(offset_raw) if offset_raw is not None else 0
+                total_override = int(bad_total) if bad_total is not None else None
+                all_res = store.search(patient)
+                host = self.headers.get("Host", "127.0.0.1")
+                extra = {"patient": patient} if patient else {}
+                self._send(
+                    200,
+                    _bundle(
+                        all_res,
+                        count=count,
+                        offset=offset,
+                        base_url=f"http://{host}{parsed.path}",
+                        extra_params=extra,
+                        total_override=total_override,
+                    ),
+                )
                 return
             self._send(
                 404,