Coverage for src/local_deep_research/journal_quality/data_sources/_openalex_common.py: 97%

1"""Shared helpers for the two OpenAlex snapshot fetchers. 

2 

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""" 

12 

13from __future__ import annotations 

14 

15import gzip 

16import json 

17from pathlib import Path 

18from typing import Callable, Iterator, Tuple 

19 

20from loguru import logger 

21 

22# Public OpenAlex snapshot — CC0, no auth, no rate limits. 

23# Manifest format documented at: 

24# https://docs.openalex.org/download-all-data/snapshot-data-format 

25# Each entry in ``manifest["files"]`` (``manifest["entries"]`` before the 

26# 2026-06 standard-format snapshot) has ``url`` (s3://...) and 

27# ``meta.content_length`` / ``meta.record_count``. We translate s3:// to 

28# the public HTTPS gateway so we don't need boto3. 

29OPENALEX_S3_BASE = "https://openalex.s3.amazonaws.com" 

30 

31# Only fetch parts hosted under the OpenAlex public S3 bucket — defense 

32# in depth on top of safe_get's private-IP block. A compromised or 

33# malformed manifest could otherwise list arbitrary attacker-controlled 

34# URLs. 

35OPENALEX_MANIFEST_ALLOWED_PREFIX = "s3://openalex/" 

36 

37 

38def s3_to_https(s3_url: str) -> str: 

39 """Translate ``s3://openalex/...`` to the public HTTPS gateway.""" 

40 return s3_url.replace( 

41 OPENALEX_MANIFEST_ALLOWED_PREFIX, OPENALEX_S3_BASE + "/", 1 

42 ) 

43 

44 

45def validate_manifest_entries(entries: list[dict], label: str) -> None: 

46 """Refuse to fetch if any manifest entry escapes the S3 allowlist. 

47 

48 Defense-in-depth: a compromised or tampered manifest could list 

49 URLs outside the OpenAlex bucket. Refusing the whole fetch rather 

50 than fetching some-and-not-others keeps failure modes simple. 

51 """ 

52 for entry in entries: 

53 raw = entry.get("url", "") 

54 if not raw.startswith(OPENALEX_MANIFEST_ALLOWED_PREFIX): 

55 raise ValueError( 

56 f"{label} manifest contains disallowed URL " 

57 f"(must start with {OPENALEX_MANIFEST_ALLOWED_PREFIX!r}): " 

58 f"{raw!r}" 

59 ) 

60 

61 

62# Per-partition retry budget. The default ``safe_get_with_retries`` 

63# budget (3 retries, 1-2-4 s backoff = ~7 s total) is sized for small 

64# request bodies and trips on a sustained mid-stream S3 hiccup: every 

65# retry of a ~5–10 MB partition that lands inside the same bad window 

66# fails the same way, exhausts the budget in seconds, and aborts the 

67# whole 30-partition pull. The release-gate workflow saw this twice in 

68# a row on 2026-04-26. 

69# 

70# 5 retries with 2-5-10-20-40 s backoff rides out a ~75 s S3 blip 

71# instead, while still bounding total wall-clock per partition at 

72# roughly ``timeout * 6 + 77 s`` — well inside the 45 min job timeout 

73# even if every partition needed all retries. 

74_PARTITION_MAX_RETRIES = 5 

75_PARTITION_BACKOFF_SECONDS = (2, 5, 10, 20, 40) 

76 

77 

78def iter_partitions( 

79 entries: list[dict], 

80 data_dir: Path, 

81 *, 

82 file_prefix: str, 

83 label: str, 

84 safe_get: Callable, 

85 timeout: int = 120, 

86 max_retries: int = _PARTITION_MAX_RETRIES, 

87 backoff_times: tuple = _PARTITION_BACKOFF_SECONDS, 

88) -> Iterator[Tuple[int, int, Iterator[dict]]]: 

89 """Download each partition, yielding ``(idx, total_parts, records)``. 

90 

91 Shared between ``openalex.py`` and ``institutions.py`` so the 

92 tmp-file lifecycle and malformed-JSON suppression (first-10 

93 warnings + one "further suppressed" notice) are defined once. 

94 

95 The caller iterates ``records`` for per-record work and is 

96 responsible for per-partition progress logging and ``progress_cb`` 

97 invocations — those need caller-specific state (running record 

98 count, schema-drift counters) that doesn't belong in the helper. 

99 

100 ``records`` decodes lazily and is valid only until this generator 

101 is advanced — the gzip handle closes and the tmp file is removed 

102 when the next partition is requested, so consume it in place. 

103 Deferring it raises ``ValueError: I/O operation on closed file``. 

104 

105 Args: 

106 entries: ``manifest["files"]`` — each dict has ``url`` 

107 starting with ``s3://openalex/``. 

108 data_dir: Directory used for the transient ``.<prefix>_part_<n>.gz`` 

109 files. Cleaned up even on exception. 

110 file_prefix: Leaf prefix for tmp files 

111 (e.g. ``openalex_sources`` / ``openalex_institutions``). 

112 label: Human-readable label used in log messages 

113 (e.g. ``"OpenAlex sources"`` / ``"Institutions"``). 

114 safe_get: Dependency-injected HTTP getter (lets the caller 

115 pick ``safe_get_with_retries`` without forcing a global 

116 import at module load). Must accept ``consume_body=True`` 

117 so body-stream transients (``ChunkedEncodingError``, 

118 ``ReadTimeout``) raised during ``resp.content`` are 

119 retried inside the wrapper, not propagated to abort the 

120 whole multi-partition pull. 

121 timeout: Per-partition HTTP timeout (seconds). 

122 max_retries: Per-partition retry budget. Defaults higher than 

123 ``safe_get_with_retries``' generic 3 because partition 

124 bodies are MB-sized and a mid-stream IncompleteRead aborts 

125 the whole multi-partition pull on exhaustion. 

126 backoff_times: Per-attempt sleep schedule. Defaults to a 

127 longer schedule than the generic ``safe_get_with_retries`` 

128 (1, 2, 4) so we ride out a sustained S3 blip instead of 

129 burning all retries inside the same bad window. 

130 """ 

131 malformed_total = 0 

132 total_parts = len(entries) 

133 

134 def decode(fh, idx: int) -> Iterator[dict]: 

135 nonlocal malformed_total 

136 for line in fh: 

137 line = line.strip() 

138 if not line: 138 ↛ 139line 138 didn't jump to line 139 because the condition on line 138 was never true

139 continue 

140 try: 

141 rec = json.loads(line) 

142 except (json.JSONDecodeError, ValueError): 

143 malformed_total += 1 

144 if malformed_total <= 10: 

145 logger.warning( 

146 f"{label} partition {idx}: skipping malformed JSON line" 

147 ) 

148 elif malformed_total == 11: 

149 logger.warning( 

150 f"{label} partition {idx}: further " 

151 "malformed lines suppressed" 

152 ) 

153 continue 

154 yield rec 

155 

156 for idx, entry in enumerate(entries): 

157 part_url = s3_to_https(entry["url"]) 

158 tmp_part = data_dir / f".{file_prefix}_part_{idx}.gz" 

159 

160 try: 

161 # consume_body=True: an OpenAlex S3 partition is ~10 MB 

162 # gzipped. A mid-stream ChunkedEncodingError / 

163 # IncompleteRead would otherwise abort the whole 30+ 

164 # partition pull. With consume_body, safe_get_with_retries 

165 # reads resp.content inside its retry loop and retries 

166 # body-stream transients the same way it retries 

167 # header-stage failures. 

168 resp = safe_get( 

169 part_url, 

170 timeout=timeout, 

171 consume_body=True, 

172 max_retries=max_retries, 

173 backoff_times=backoff_times, 

174 ) 

175 resp.raise_for_status() 

176 tmp_part.write_bytes(resp.content) 

177 # Unpin the compressed body; the suspended frame would 

178 # otherwise hold it until the caller drains the partition. 

179 del resp 

180 

181 with gzip.open(tmp_part, "rt", encoding="utf-8") as fh: 

182 yield idx, total_parts, decode(fh, idx) 

183 finally: 

184 tmp_part.unlink(missing_ok=True) 

185 

186 if malformed_total: 

187 logger.warning( 

188 f"{label}: {malformed_total:,} malformed lines skipped across " 

189 f"{total_parts} partitions" 

190 )