Coverage for src/local_deep_research/journal_quality/data_sources/_openalex_common.py: 98%
114 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""Shared helpers for the two OpenAlex snapshot fetchers.
3Sources and institutions both pull from the OpenAlex S3 bucket, translate
4``s3://`` URLs to the public HTTPS gateway, and defend-in-depth against
5a compromised manifest by allowlisting the ``s3://openalex/`` prefix.
6This file owns those three shared symbols so ``openalex.py`` and
7``institutions.py`` don't duplicate them (and can't drift). It also
8owns the per-partition streaming helper they both use to iterate
9records with consistent malformed-line suppression and tmp-file
10lifecycle.
11"""
13from __future__ import annotations
15import gzip
16import hashlib
17import json
18from pathlib import Path
19from typing import Callable, Iterator, Tuple
21from loguru import logger
23# Public OpenAlex snapshot — CC0, no auth, no rate limits.
24# Manifest format documented at:
25# https://docs.openalex.org/download-all-data/snapshot-data-format
26# Each entry in ``manifest["files"]`` (``manifest["entries"]`` before the
27# 2026-06 standard-format snapshot) has ``url`` (s3://...) and
28# ``meta.content_length`` / ``meta.record_count``. We translate s3:// to
29# the public HTTPS gateway so we don't need boto3.
30OPENALEX_S3_BASE = "https://openalex.s3.amazonaws.com"
32# Only fetch parts hosted under the OpenAlex public S3 bucket — defense
33# in depth on top of safe_get's private-IP block. A compromised or
34# malformed manifest could otherwise list arbitrary attacker-controlled
35# URLs.
36OPENALEX_MANIFEST_ALLOWED_PREFIX = "s3://openalex/"
39def s3_to_https(s3_url: str) -> str:
40 """Translate ``s3://openalex/...`` to the public HTTPS gateway."""
41 return s3_url.replace(
42 OPENALEX_MANIFEST_ALLOWED_PREFIX, OPENALEX_S3_BASE + "/", 1
43 )
46def validate_manifest_entries(entries: list[dict], label: str) -> None:
47 """Refuse to fetch if any manifest entry escapes the S3 allowlist.
49 Defense-in-depth: a compromised or tampered manifest could list
50 URLs outside the OpenAlex bucket. Refusing the whole fetch rather
51 than fetching some-and-not-others keeps failure modes simple.
52 """
53 for entry in entries:
54 raw = entry.get("url", "")
55 if not raw.startswith(OPENALEX_MANIFEST_ALLOWED_PREFIX):
56 raise ValueError(
57 f"{label} manifest contains disallowed URL "
58 f"(must start with {OPENALEX_MANIFEST_ALLOWED_PREFIX!r}): "
59 f"{raw!r}"
60 )
63# Per-partition retry budget. The default ``safe_get_with_retries``
64# budget (3 retries, 1-2-4 s backoff = ~7 s total) is sized for small
65# request bodies and trips on a sustained mid-stream S3 hiccup: every
66# retry of a ~5–10 MB partition that lands inside the same bad window
67# fails the same way, exhausts the budget in seconds, and aborts the
68# whole 30-partition pull. The release-gate workflow saw this twice in
69# a row on 2026-04-26.
70#
71# 5 retries with 2-5-10-20-40 s backoff rides out a ~75 s S3 blip
72# instead, while still bounding total wall-clock per partition at
73# roughly ``timeout * 6 + 77 s`` — well inside the 45 min job timeout
74# even if every partition needed all retries.
75_PARTITION_MAX_RETRIES = 5
76_PARTITION_BACKOFF_SECONDS = (2, 5, 10, 20, 40)
78# Absolute line cap: a real OpenAlex JSONL record is a few KB; even the
79# largest institution records stay far below 10 MB. A line longer than
80# this is hostile padding, not data — skipped as malformed without
81# being materialized.
82_MAX_JSONL_LINE_CHARS = 10_000_000
84# Absolute ceiling on a partition's decompressed stream, in characters.
85# Doubles as the fallback for entries whose manifest entry carries no
86# meta.content_length: the largest partition in the live fleet (an
87# institutions one, 176,715,737 compressed bytes) decompresses to
88# 1,979,631,761 bytes — 1.84 GiB — so 4 GiB keeps ~2.2x over today's
89# snapshot while still bounding a gzip ratio bomb. The manifest itself
90# is untrusted, so an anchored cap can never exceed it.
91_DEFAULT_MAX_PARTITION_DECOMPRESSED_BYTES = 4 * 1024**3
93# Manifest ``meta.content_length`` counts COMPRESSED gz bytes; the
94# enforced bound counts DECOMPRESSED characters — so the declaration
95# must be scaled by an expansion allowance, and the allowance has to
96# clear the whole live fleet or a real snapshot aborts as a "bomb".
97# Census of every live partition on 2026-09-24, read from each file's
98# gzip ISIZE trailer (decompressed bytes, an upper bound on the
99# characters this bound counts): the 196 sources partitions expand at
100# a median 14.23x and a maximum 18.69x — files[0], 560,290 compressed
101# bytes, decompresses to 10,470,101 bytes / 10,461,631 characters —
102# and the 80 institutions partitions top out at 12.17x. 32 leaves
103# ~1.7x over that worst live ratio while still catching gzip ratio
104# bombs (which expand ~1032:1).
105_COMPRESSED_BYTES_TO_DECOMPRESSED_CHARS = 32
107# Slack added to an anchored cap, in characters, so format drift on
108# small partitions is not false-flagged.
109_PARTITION_CAP_SLACK_CHARS = 1024**2
112def _partition_decompressed_cap(entry: dict) -> int:
113 """Upper bound for one partition's decompressed stream, in characters.
115 Anchored to the manifest's declared ``meta.content_length`` (the
116 compressed gz size) scaled by the expansion allowance plus slack:
117 a stream expanding far past what the manifest declared is a
118 decompression bomb, not a snapshot. The anchor is clamped to the
119 absolute ceiling — a hostile manifest can declare any size, so a
120 huge declaration must not buy a bigger bomb budget than a
121 meta-less entry gets.
122 """
123 meta = entry.get("meta") or {}
124 declared = meta.get("content_length")
125 # ``bool`` is an ``int``: a declaration of ``True`` would otherwise
126 # anchor the cap at 1 * 32 + 1 MiB and abort every real partition.
127 # The transport-side content_length check excludes bool the same
128 # way, so a non-integer declaration is no declaration at all.
129 if (
130 isinstance(declared, int)
131 and not isinstance(declared, bool)
132 and declared > 0
133 ):
134 return min(
135 _DEFAULT_MAX_PARTITION_DECOMPRESSED_BYTES,
136 declared * _COMPRESSED_BYTES_TO_DECOMPRESSED_CHARS
137 + _PARTITION_CAP_SLACK_CHARS,
138 )
139 return _DEFAULT_MAX_PARTITION_DECOMPRESSED_BYTES
142_DIGEST_HEX_LENGTH = 32
143_HEX_DIGIT_CHARS = frozenset("0123456789abcdefABCDEF")
146def _normalize_digest(declared: object) -> str | None:
147 """Return a lowercase 32-hex digest from a manifest ``meta.md5`` value. # DevSkim: ignore DS126858
149 Accepts a bare hex digest or the same digest wrapped in one matching
150 pair of double quotes (the S3 ETag wire form for a non-multipart
151 object). Returns ``None`` for anything else — a multipart ETag
152 (``<hex>-<n>``), a base64 digest, the wrong length, non-hex
153 characters, or a non-string value — so the caller can skip
154 verification for this partition instead of comparing bytes against
155 a value that was never a valid digest.
156 """
157 if not isinstance(declared, str):
158 return None
159 candidate = declared.strip()
160 if (
161 len(candidate) == _DIGEST_HEX_LENGTH + 2
162 and candidate[0] == '"'
163 and candidate[-1] == '"'
164 ):
165 candidate = candidate[1:-1]
166 if len(candidate) == _DIGEST_HEX_LENGTH and all(
167 c in _HEX_DIGIT_CHARS for c in candidate
168 ):
169 return candidate.lower()
170 return None
173def iter_partitions(
174 entries: list[dict],
175 data_dir: Path,
176 *,
177 file_prefix: str,
178 label: str,
179 safe_get: Callable,
180 timeout: int = 120,
181 max_retries: int = _PARTITION_MAX_RETRIES,
182 backoff_times: tuple = _PARTITION_BACKOFF_SECONDS,
183 max_line_chars: int = _MAX_JSONL_LINE_CHARS,
184) -> Iterator[Tuple[int, int, Iterator[dict]]]:
185 """Download each partition, yielding ``(idx, total_parts, records)``.
187 Shared between ``openalex.py`` and ``institutions.py`` so the
188 tmp-file lifecycle and malformed-JSON suppression (first-10
189 warnings + one "further suppressed" notice) are defined once.
191 The caller iterates ``records`` for per-record work and is
192 responsible for per-partition progress logging and ``progress_cb``
193 invocations — those need caller-specific state (running record
194 count, schema-drift counters) that doesn't belong in the helper.
196 ``records`` decodes lazily and is valid only until this generator
197 is advanced — the gzip handle closes and the tmp file is removed
198 when the next partition is requested, so consume it in place.
199 Deferring it raises ``ValueError: I/O operation on closed file``.
201 Args:
202 entries: ``manifest["files"]`` — each dict has ``url``
203 starting with ``s3://openalex/``.
204 data_dir: Directory used for the transient ``.<prefix>_part_<n>.gz``
205 files. Cleaned up even on exception.
206 file_prefix: Leaf prefix for tmp files
207 (e.g. ``openalex_sources`` / ``openalex_institutions``).
208 label: Human-readable label used in log messages
209 (e.g. ``"OpenAlex sources"`` / ``"Institutions"``).
210 safe_get: Dependency-injected HTTP getter (lets the caller
211 pick ``safe_get_with_retries`` without forcing a global
212 import at module load). Must accept ``consume_body=True``
213 so body-stream transients (``ChunkedEncodingError``,
214 ``ReadTimeout``) raised during ``resp.content`` are
215 retried inside the wrapper, not propagated to abort the
216 whole multi-partition pull. Must also accept
217 ``require_https=True``: partitions are fetched with the
218 scheme pinned so no redirect hop can move them to cleartext.
219 timeout: Per-partition HTTP timeout (seconds).
220 max_retries: Per-partition retry budget. Defaults higher than
221 ``safe_get_with_retries``' generic 3 because partition
222 bodies are MB-sized and a mid-stream IncompleteRead aborts
223 the whole multi-partition pull on exhaustion.
224 backoff_times: Per-attempt sleep schedule. Defaults to a
225 longer schedule than the generic ``safe_get_with_retries``
226 (1, 2, 4) so we ride out a sustained S3 blip instead of
227 burning all retries inside the same bad window.
229 Raises:
230 ValueError: If a partition's declared ``meta.content_length``
231 doesn't match the number of bytes actually received, or
232 its declared ``meta.md5`` (once normalized to a bare # DevSkim: ignore DS126858
233 digest) doesn't match the digest of the bytes received —
234 either aborts before anything is written to disk, so the
235 previous snapshot is left in place. Also if a partition's
236 decompressed stream runs past the bound
237 ``_partition_decompressed_cap`` derives from its manifest
238 entry: that one fires later, while the caller consumes the
239 yielded records, so the compressed body has already been
240 written to the temp file — which the ``finally`` still
241 removes.
242 """
243 malformed_total = 0
244 total_parts = len(entries)
246 def decode(fh, idx: int, max_chars: int, max_bytes: int) -> Iterator[dict]:
247 nonlocal malformed_total
249 def bomb_error() -> ValueError:
250 return ValueError(
251 f"{label} partition {idx}: decompressed stream exceeded "
252 f"its bound ({max_bytes} characters) — refusing a "
253 "possible decompression bomb"
254 )
256 seen_chars = 0
257 while True:
258 # readline(size) never buffers more than size characters, so
259 # an over-long line is detected — and its remainder discarded
260 # in bounded chunks — without ever being materialized whole.
261 # Every chunk read, kept or discarded, counts toward the
262 # bound so an all-overlong-line bomb cannot expand past it
263 # uncounted.
264 line = fh.readline(max_chars + 1)
265 if not line:
266 return
267 seen_chars += len(line)
268 if seen_chars > max_bytes:
269 raise bomb_error()
270 if len(line) > max_chars and not line.endswith("\n"):
271 while True:
272 rest = fh.readline(max_chars + 1)
273 if not rest: 273 ↛ 274line 273 didn't jump to line 274 because the condition on line 273 was never true
274 break
275 seen_chars += len(rest)
276 if seen_chars > max_bytes:
277 raise bomb_error()
278 if rest.endswith("\n"):
279 break
280 malformed_total += 1
281 if malformed_total <= 10:
282 logger.warning(
283 f"{label} partition {idx}: skipping over-long line"
284 )
285 elif malformed_total == 11:
286 logger.warning(
287 f"{label} partition {idx}: further "
288 "malformed lines suppressed"
289 )
290 continue
291 line = line.strip()
292 if not line: 292 ↛ 293line 292 didn't jump to line 293 because the condition on line 292 was never true
293 continue
294 try:
295 rec = json.loads(line)
296 except (json.JSONDecodeError, ValueError):
297 malformed_total += 1
298 if malformed_total <= 10:
299 logger.warning(
300 f"{label} partition {idx}: skipping malformed JSON line"
301 )
302 elif malformed_total == 11:
303 logger.warning(
304 f"{label} partition {idx}: further "
305 "malformed lines suppressed"
306 )
307 continue
308 yield rec
310 for idx, entry in enumerate(entries):
311 part_url = s3_to_https(entry["url"])
312 tmp_part = data_dir / f".{file_prefix}_part_{idx}.gz"
314 try:
315 # consume_body=True: an OpenAlex S3 partition is ~10 MB
316 # gzipped. A mid-stream ChunkedEncodingError /
317 # IncompleteRead would otherwise abort the whole 30+
318 # partition pull. With consume_body, safe_get_with_retries
319 # reads resp.content inside its retry loop and retries
320 # body-stream transients the same way it retries
321 # header-stage failures.
322 resp = safe_get(
323 part_url,
324 timeout=timeout,
325 consume_body=True,
326 max_retries=max_retries,
327 backoff_times=backoff_times,
328 require_https=True,
329 )
330 resp.raise_for_status()
331 meta = entry.get("meta") or {}
333 # Transport-corruption check when the manifest declares an
334 # expected size: compressed bytes whose length doesn't match
335 # meta.content_length were truncated or corrupted in
336 # transit — refuse before writing them to disk. This is not
337 # an authentication check: the declared length comes from
338 # the same manifest fetch as the bytes it describes, so it
339 # can't detect a hostile origin, only a garbled download. A
340 # declared value that isn't a plain integer (a string,
341 # float, or bool) is logged and skipped rather than treated
342 # as a mismatch.
343 declared_length = meta.get("content_length")
344 if declared_length is not None:
345 if isinstance(declared_length, int) and not isinstance(
346 declared_length, bool
347 ):
348 actual_length = len(resp.content)
349 if actual_length != declared_length:
350 raise ValueError(
351 f"{label} partition {idx}: content_length "
352 f"mismatch — manifest declares "
353 f"{declared_length} bytes but received "
354 f"{actual_length} bytes; "
355 "refusing possibly corrupted or truncated "
356 "partition"
357 )
358 else:
359 shape = (
360 f"{len(declared_length)}-character string"
361 if isinstance(declared_length, str)
362 else type(declared_length).__name__
363 )
364 logger.warning(
365 f"{label} partition {idx}: manifest "
366 f"content_length is not an integer ({shape}); "
367 "skipping the size check for this partition"
368 )
370 # Transport-corruption check when the manifest declares a
371 # digest: bytes that hash differently than meta.md5 were
372 # altered or corrupted in transit — refuse before writing
373 # them to disk. Like the content_length check above, this
374 # is not an authentication check: the digest comes from the
375 # same manifest fetch as the bytes it certifies, so a party
376 # able to substitute the partition can substitute the
377 # declared md5 too. A declared value that isn't a 32-hex
378 # digest (plain or ETag-quoted) is logged and skipped
379 # rather than treated as a mismatch.
380 declared_digest = meta.get("md5") # DevSkim: ignore DS126858
381 if declared_digest is not None:
382 normalized_digest = _normalize_digest(declared_digest)
383 if normalized_digest is None:
384 shape = (
385 f"{len(declared_digest)}-character string"
386 if isinstance(declared_digest, str)
387 else type(declared_digest).__name__
388 )
389 logger.warning(
390 f"{label} partition {idx}: manifest md5 is not " # DevSkim: ignore DS126858
391 f"a 32-character hex digest ({shape}); skipping "
392 "the integrity check for this partition"
393 )
394 else:
395 actual_digest = hashlib.md5( # DevSkim: ignore DS126858
396 resp.content, usedforsecurity=False
397 ).hexdigest()
398 if actual_digest != normalized_digest:
399 raise ValueError(
400 f"{label} partition {idx}: md5 mismatch — " # DevSkim: ignore DS126858
401 f"manifest declares {normalized_digest!r} but "
402 f"received {actual_digest!r}; refusing possibly "
403 "corrupted partition"
404 )
406 tmp_part.write_bytes(resp.content)
407 # Unpin the compressed body; the suspended frame would
408 # otherwise hold it until the caller drains the partition.
409 del resp
411 with gzip.open(tmp_part, "rt", encoding="utf-8") as fh:
412 cap = _partition_decompressed_cap(entry)
413 yield idx, total_parts, decode(fh, idx, max_line_chars, cap)
414 finally:
415 tmp_part.unlink(missing_ok=True)
417 if malformed_total:
418 logger.warning(
419 f"{label}: {malformed_total:,} malformed lines skipped across "
420 f"{total_parts} partitions"
421 )