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, list[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 Args: 

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

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

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

104 files. Cleaned up even on exception. 

105 file_prefix: Leaf prefix for tmp files 

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

107 label: Human-readable label used in log messages 

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

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

110 pick ``safe_get_with_retries`` without forcing a global 

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

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

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

114 retried inside the wrapper, not propagated to abort the 

115 whole multi-partition pull. 

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

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

118 ``safe_get_with_retries``' generic 3 because partition 

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

120 the whole multi-partition pull on exhaustion. 

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

122 longer schedule than the generic ``safe_get_with_retries`` 

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

124 burning all retries inside the same bad window. 

125 """ 

126 malformed_total = 0 

127 total_parts = len(entries) 

128 

129 for idx, entry in enumerate(entries): 

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

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

132 records: list[dict] = [] 

133 

134 try: 

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

136 # gzipped. A mid-stream ChunkedEncodingError / 

137 # IncompleteRead would otherwise abort the whole 30+ 

138 # partition pull. With consume_body, safe_get_with_retries 

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

140 # body-stream transients the same way it retries 

141 # header-stage failures. 

142 resp = safe_get( 

143 part_url, 

144 timeout=timeout, 

145 consume_body=True, 

146 max_retries=max_retries, 

147 backoff_times=backoff_times, 

148 ) 

149 resp.raise_for_status() 

150 tmp_part.write_bytes(resp.content) 

151 

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

153 for line in fh: 

154 line = line.strip() 

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

156 continue 

157 try: 

158 rec = json.loads(line) 

159 except (json.JSONDecodeError, ValueError): 

160 malformed_total += 1 

161 if malformed_total <= 10: 

162 logger.warning( 

163 f"{label} partition {idx}: skipping " 

164 f"malformed JSON line" 

165 ) 

166 elif malformed_total == 11: 

167 logger.warning( 

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

169 "malformed lines suppressed" 

170 ) 

171 continue 

172 records.append(rec) 

173 finally: 

174 tmp_part.unlink(missing_ok=True) 

175 

176 yield idx, total_parts, records 

177 

178 if malformed_total: 

179 logger.warning( 

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

181 f"{total_parts} partitions" 

182 )