Coverage for oc_ds_converter / datasource / orcid_index.py: 81%

1# SPDX-FileCopyrightText: 2026 Arcangelo Massari <arcangelo.massari@unibo.it> 

2# 

3# SPDX-License-Identifier: ISC 

4 

5from __future__ import annotations 

6 

7import json 

8import os 

9from collections import defaultdict 

10from concurrent.futures import ProcessPoolExecutor, as_completed 

11from csv import DictReader 

12from multiprocessing import get_context 

13from os import cpu_count, sep, walk 

14from os.path import exists 

15from typing import Protocol, cast 

16 

17import fakeredis 

18 

19from oc_ds_converter.datasource.redis import RedisDataSource 

20from oc_ds_converter.lib.console import create_progress 

21from oc_ds_converter.oc_idmanager import DOIManager 

22 

23 

24class OrcidIndexInterface(Protocol): 

25 def get_value(self, id_string: str, /) -> set[str] | None: ... 

26 def get_values_batch(self, ids: list[str], /) -> dict[str, set[str]]: ... 

27 

28 

29class OrcidIndexRedis: 

30 def __init__(self, testing: bool = False) -> None: 

31 if testing: 

32 self._r = fakeredis.FakeStrictRedis(decode_responses=True) 

33 else: 

34 self._redis = RedisDataSource("DOI-ORCID-INDEX") 

35 self._r = self._redis._r 

36 

37 def get_value(self, doi: str) -> set[str] | None: 

38 result = cast(set[str], self._r.smembers(doi)) 

39 if result: 

40 return result 

41 return None 

42 

43 def get_values_batch(self, dois: list[str]) -> dict[str, set[str]]: 

44 if not dois: 

45 return {} 

46 pipe = self._r.pipeline() 

47 for doi in dois: 

48 pipe.smembers(doi) 

49 results = pipe.execute() 

50 return { 

51 doi: cast(set[str], members) 

52 for doi, members in zip(dois, results) 

53 if members 

54 } 

55 

56 def add_values_batch(self, data: dict[str, set[str]]) -> None: 

57 pipe = self._r.pipeline() 

58 for doi, values in data.items(): 

59 if values: 

60 pipe.sadd(doi, *values) 

61 pipe.execute() 

62 

63 def has_data(self) -> bool: 

64 return cast(int, self._r.dbsize()) > 0 

65 

66 def clear(self) -> None: 

67 self._r.flushdb() 

68 

69 

70def _process_csv_file(csv_path: str) -> dict[str, set[str]]: 

71 """Process a single CSV file and return DOI -> ORCID mappings. 

72 

73 This function runs in a separate process for parallelization. 

74 """ 

75 doi_manager = DOIManager() 

76 result: dict[str, set[str]] = defaultdict(set) 

77 

78 with open(csv_path, 'r', encoding='utf-8') as f: 

79 reader = DictReader(f) 

80 for row in reader: 

81 raw_doi = row['id'] 

82 doi = doi_manager.normalise(raw_doi, include_prefix=True) 

83 if doi: 

84 result[doi].add(row['value']) 

85 

86 return dict(result) 

87 

88 

89def load_orcid_index_to_redis( 

90 orcid_index_dir: str, 

91 orcid_index_redis: OrcidIndexRedis, 

92 batch_size: int = 1000000, 

93 max_workers: int | None = None, 

94) -> None: 

95 if not exists(orcid_index_dir): 

96 return 

97 

98 files_to_process: list[str] = [] 

99 for cur_dir, _, cur_files in walk(orcid_index_dir): 

100 for cur_file in cur_files: 

101 if cur_file.endswith('.csv'): 

102 files_to_process.append(cur_dir + sep + cur_file) 

103 

104 if not files_to_process: 

105 return 

106 

107 if max_workers is None: 

108 max_workers = min(cpu_count() or 4, len(files_to_process)) 

109 

110 batch: dict[str, set[str]] = {} 

111 count = 0 

112 

113 with create_progress() as progress: 

114 task = progress.add_task( 

115 "[green]Loading DOI-ORCID index files", total=len(files_to_process) 

116 ) 

117 

118 mp_method = "spawn" if os.name == "nt" else "forkserver" 

119 

120 with ProcessPoolExecutor( 

121 max_workers=max_workers, mp_context=get_context(mp_method) 

122 ) as executor: 

123 futures = { 

124 executor.submit(_process_csv_file, path): path 

125 for path in files_to_process 

126 } 

127 

128 for future in as_completed(futures): 

129 file_data = future.result() 

130 

131 # Merge results into batch 

132 for doi, values in file_data.items(): 

133 if doi in batch: 

134 batch[doi].update(values) 

135 else: 

136 batch[doi] = values 

137 count += len(values) 

138 

139 # Flush to Redis when batch is large enough 

140 if count >= batch_size: 

141 orcid_index_redis.add_values_batch(batch) 

142 batch = {} 

143 count = 0 

144 

145 progress.update(task, advance=1) 

146 

147 if batch: 

148 orcid_index_redis.add_values_batch(batch) 

149 

150 

151class PublishersRedis: 

152 MEMBER_PREFIX = "member:" 

153 DOI_PREFIX_KEY = "prefix:" 

154 

155 def __init__(self, testing: bool = False) -> None: 

156 if testing: 

157 self._r = fakeredis.FakeStrictRedis(decode_responses=True) 

158 else: 

159 self._redis = RedisDataSource("PUBLISHERS-INDEX") 

160 self._r = self._redis._r 

161 

162 def get_by_member(self, member_id: str) -> dict[str, str | set[str]] | None: 

163 key = f"{self.MEMBER_PREFIX}{member_id}" 

164 data = self._r.get(key) 

165 if data: 

166 result = json.loads(str(data)) 

167 result["prefixes"] = set(result["prefixes"]) 

168 return result 

169 return None 

170 

171 def get_by_prefix(self, prefix: str) -> dict[str, str | set[str]] | None: 

172 key = f"{self.DOI_PREFIX_KEY}{prefix}" 

173 member_id = self._r.get(key) 

174 if member_id: 

175 return self.get_by_member(str(member_id)) 

176 return None 

177 

178 def set_publisher(self, member_id: str, name: str, prefixes: set[str]) -> None: 

179 member_key = f"{self.MEMBER_PREFIX}{member_id}" 

180 data = {"name": name, "prefixes": list(prefixes)} 

181 self._r.set(member_key, json.dumps(data)) 

182 for prefix in prefixes: 

183 prefix_key = f"{self.DOI_PREFIX_KEY}{prefix}" 

184 self._r.set(prefix_key, member_id) 

185 

186 def set_publishers_batch(self, publishers: dict[str, dict[str, str | set[str]]]) -> None: 

187 pipe = self._r.pipeline() 

188 for member_id, data in publishers.items(): 

189 member_key = f"{self.MEMBER_PREFIX}{member_id}" 

190 prefixes_list = list(data["prefixes"]) 

191 pipe.set(member_key, json.dumps({"name": data["name"], "prefixes": prefixes_list})) 

192 for prefix in prefixes_list: 

193 prefix_key = f"{self.DOI_PREFIX_KEY}{prefix}" 

194 pipe.set(prefix_key, member_id) 

195 pipe.execute() 

196 

197 def has_data(self) -> bool: 

198 return cast(int, self._r.dbsize()) > 0 

199 

200 def clear(self) -> None: 

201 self._r.flushdb() 

202 

203 

204def load_publishers_to_redis( 

205 publishers_filepath: str, 

206 publishers_redis: PublishersRedis, 

207 batch_size: int = 5000, 

208) -> None: 

209 if not exists(publishers_filepath): 

210 return 

211 

212 batch: dict[str, dict[str, str | set[str]]] = {} 

213 count = 0 

214 

215 with open(publishers_filepath, 'r', encoding='utf-8') as f: 

216 reader = DictReader(f) 

217 for row in reader: 

218 pub_id = row['id'] 

219 if pub_id not in batch: 

220 batch[pub_id] = {'name': row['name'], 'prefixes': set()} 

221 prefixes = batch[pub_id]['prefixes'] 

222 prefixes.add(row['prefix']) # type: ignore[union-attr] 

223 count += 1 

224 

225 if count >= batch_size: 

226 publishers_redis.set_publishers_batch(batch) 

227 batch = {} 

228 count = 0 

229 

230 if batch: 

231 publishers_redis.set_publishers_batch(batch)