Spaces:
Sleeping
Sleeping
File size: 5,877 Bytes
b92b744 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | """ClinicalTrials.gov paging and shared atomic checkpoint helpers.
OpenAlex is read exclusively from the public S3 snapshot in snapshot.py.
"""
import asyncio
import hashlib
import json
import os
from collections.abc import Callable
from pathlib import Path
import httpx
CTGOV_URL = "https://clinicaltrials.gov/api/v2/studies"
CTGOV_FILTER = "AREA[StudyType]INTERVENTIONAL AND AREA[OverallStatus](COMPLETED OR TERMINATED OR WITHDRAWN OR SUSPENDED)"
def atomic_json(path: Path, data: dict) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix(path.suffix + ".tmp")
with temporary.open("w") as handle:
json.dump(data, handle)
handle.flush()
os.fsync(handle.fileno())
temporary.replace(path)
async def request_json(
client: httpx.AsyncClient,
url: str,
params: dict,
retries: int = 5,
sleep: Callable = asyncio.sleep,
) -> dict:
"""Bounded retries; exception messages never include URLs or credentials."""
for attempt in range(retries + 1):
try:
response = await client.get(url, params=params)
if response.status_code == 429 or response.status_code >= 500:
if attempt == retries:
raise RuntimeError(
f"Source unavailable after retries (HTTP {response.status_code})"
)
retry_after = response.headers.get("retry-after", "")
delay = (
min(60, float(retry_after)) if retry_after.isdigit() else min(60, 2**attempt)
)
await sleep(delay)
continue
if response.status_code >= 400:
raise RuntimeError(
f"Source rejected request (HTTP {response.status_code}); check source parameters"
)
data = response.json()
if not isinstance(data, dict):
raise TypeError("Source returned an unexpected response shape")
return data
except (httpx.TransportError, ValueError):
if attempt == retries:
raise RuntimeError("Source transport or JSON error after retries") from None
await sleep(min(60, 2**attempt))
raise RuntimeError("Unreachable retry state")
async def fetch_pages(
source: str,
output: Path,
*,
query: str = "",
filters: str | None = None,
limit: int | None = 1000,
resume: bool = False,
page_size: int = 1000,
client: httpx.AsyncClient | None = None,
) -> dict:
if source != "ctgov":
raise ValueError(
"Only ctgov uses API paging; use the snapshot command for OpenAlex public S3"
)
if limit is not None and limit <= 0:
raise ValueError("limit must be positive or None for source exhaustion")
if not 1 <= page_size <= 1000:
raise ValueError("page_size must be between 1 and 1000")
filters = CTGOV_FILTER if filters is None else filters
signature = hashlib.sha256(json.dumps([source, query, filters, page_size]).encode()).hexdigest()
checkpoint = output.with_suffix(output.suffix + ".checkpoint.json")
state = {
"signature": signature,
"source": source,
"cursor": None,
"count": 0,
"bytes": 0,
"done": False,
"total_count": None,
}
output.parent.mkdir(parents=True, exist_ok=True)
if resume and checkpoint.exists():
state = json.loads(checkpoint.read_text())
if state.get("signature") != signature:
raise ValueError("Resume parameters differ from the checkpoint")
if not output.exists() or output.stat().st_size < state["bytes"]:
raise ValueError("Checkpoint output is missing or truncated")
elif output.exists():
raise FileExistsError("Output already exists; use --resume or select a new output")
if not output.exists():
output.touch()
if output.stat().st_size != state["bytes"]:
with output.open("r+b") as handle:
handle.truncate(state["bytes"])
atomic_json(checkpoint, state)
if state["done"] or (limit is not None and state["count"] >= limit):
return state
owned = client is None
client = client or httpx.AsyncClient(timeout=60, follow_redirects=True)
try:
while not state["done"] and (limit is None or state["count"] < limit):
size = page_size if limit is None else min(page_size, limit - state["count"])
params = {"format": "json", "pageSize": size, "countTotal": "true"}
if filters:
params["filter.advanced"] = filters
if query:
params["query.term"] = query
if state["cursor"]:
params["pageToken"] = state["cursor"]
data = await request_json(client, CTGOV_URL, params)
rows, next_cursor = data.get("studies"), data.get("nextPageToken")
if not isinstance(rows, list):
raise TypeError("Source response has no study list")
if rows and next_cursor and next_cursor == state["cursor"]:
raise RuntimeError("Source repeated its paging cursor")
with output.open("ab") as handle:
for row in rows:
handle.write((json.dumps(row, ensure_ascii=False) + "\n").encode())
handle.flush()
os.fsync(handle.fileno())
state["bytes"] = handle.tell()
state.update(
count=state["count"] + len(rows),
cursor=next_cursor,
done=not rows or not next_cursor,
total_count=data.get("totalCount", state.get("total_count")),
)
atomic_json(checkpoint, state)
finally:
if owned:
await client.aclose()
return state
|