Coverage for src/crawler/cmds/augment/augment_cmd.py: 20%

1import logging 

2from concurrent.futures import Future, ThreadPoolExecutor 

3from typing import TYPE_CHECKING, Literal, NotRequired, TypedDict 

4 

5from history.model_data import HistoryChildDict, MatchingEventResult 

6from history.models import HistoryEventStatus 

7from history.utils import insert_history_event 

8from matching_back.backends import BACKENDS 

9from ptf import model_helpers 

10from ptf.cmds.base_cmds import baseCmd 

11from ptf.models import Article, Collection 

12 

13from crawler.cmds.augment import update_article_with_xarticle 

14from crawler.models.source import Source 

15 

16if TYPE_CHECKING: 

17 from collections.abc import Callable 

18 

19 from matching.backends import MatchingBackend 

20 from ptf.models import ArticleQuerySet 

21 

22 class AugmentArticlesParams(TypedDict): 

23 source_id: str 

24 collection_pid: NotRequired[str] 

25 issue_pids: NotRequired[list[str]] 

26 update_pdf_link: NotRequired[bool] 

27 

28 from typing import NotRequired, TypedDict 

29 

30 class MatchingOperationMessage(TypedDict): 

31 pid: str 

32 zbl_id: "NotRequired[str]" 

33 score: "NotRequired[int]" 

34 status: HistoryEventStatus 

35 message: "NotRequired[str]" 

36 status_message: "NotRequired[str]" 

37 

38 

39_logger = logging.getLogger(__name__) 

40 

41 

42def compile_history_events(event_messages: "list[MatchingOperationMessage]"): 

43 children = [] 

44 for msg in event_messages: 

45 resource = model_helpers.get_resource(msg["pid"]) 

46 if not resource: 

47 raise ValueError("Resource not found") 

48 data: "HistoryChildDict" = { 

49 "resource": resource, 

50 "type": "zbl-id", 

51 "status": msg.get("status"), 

52 } 

53 

54 if "zbl_id" in msg and msg["zbl_id"] != "": 

55 data["url"] = f"https://zbmath.org/{msg['zbl_id']}" 

56 

57 if "score" in msg: 

58 data["score"] = msg["score"] 

59 

60 if "message" in msg: 

61 data["message"] = msg["message"] 

62 

63 if "status_message" in msg: 

64 data["status_message"] = msg["status_message"] 

65 

66 children.append(data) 

67 return children 

68 

69 

70class AugmentArticlesCmd(baseCmd[None]): 

71 """ 

72 Metadata augment command. 

73 Calls the desired Matching Backend defined in ptf-back. 

74 """ 

75 

76 required_params = ["backend", "source_id"] 

77 

78 backend: "MatchingBackend" 

79 source_id: str 

80 collection_pid: str | None = None 

81 issue_pids: list[str] = [] 

82 update_pdf_link: bool = False 

83 collection: "Collection | None" = None 

84 queryset: "ArticleQuerySet" 

85 event_messages: list[dict] 

86 callback: "Callable" = lambda _: _ 

87 

88 def __init__(self, params: "AugmentArticlesParams", backend: Literal["zbmath", "crossref"]): 

89 super().__init__(params) 

90 

91 self.backend = BACKENDS[backend]() 

92 

93 self.event_messages = [] 

94 self.queryset = Article.objects.prefetch_related("extid_set").all() 

95 

96 if self.collection_pid: 

97 self.queryset = self.queryset.filter( 

98 my_container__my_collection__pid=self.collection_pid 

99 ) 

100 self.collection = Collection.objects.get(pid=self.collection_pid) 

101 

102 if self.source_id: 

103 source = Source.objects.get(domain=self.source_id) 

104 self.queryset = self.queryset.filter(my_container__origin__source=source) 

105 

106 if self.issue_pids: 

107 self.queryset = self.queryset.filter(my_container__pid__in=self.issue_pids) 

108 

109 def internal_do(self) -> None: 

110 _logger.info( 

111 f"Start augment [{self.backend.name}] " 

112 f"source={self.source_id} collection={self.collection_pid} " 

113 f"({self.queryset.count()} articles)" 

114 ) 

115 promises: "set[Future]" = set() 

116 with ThreadPoolExecutor(max_workers=1) as executor: 

117 for article in self.queryset.iterator(chunk_size=2000): 

118 xarticle, event = self._find_article(article) 

119 

120 if event: 

121 self.event_messages.append(event) 

122 

123 if not xarticle: 

124 self.callback() 

125 continue 

126 

127 # Handle asyncronous results and exceptions as soon as possible 

128 completed: "set[Future]" = set() 

129 for f in promises: 

130 if f.done(): 

131 completed.add(f) 

132 exc = f.exception() 

133 if exc: 

134 executor.shutdown(wait=False, cancel_futures=True) 

135 raise exc 

136 promises -= completed 

137 # Runs the database insertion in parallel 

138 promise = executor.submit( 

139 update_article_with_xarticle, 

140 article, 

141 xarticle, 

142 merge_titles=False, 

143 update_pdf_link=self.update_pdf_link, 

144 ) 

145 promise.add_done_callback(self.callback) 

146 promises.add(promise) 

147 _logger.info(f"Augment [{self.backend.name}] over") 

148 

149 def _find_article(self, article: Article) -> tuple: 

150 """ 

151 Calls ArticleData from backend. 

152 """ 

153 # With external ID 

154 xarticle = self.backend.find_by_extids(article=article) 

155 if xarticle: 

156 _logger.debug(f"[{self.backend.name}] {article.pid}: found via External ID") 

157 return xarticle, self._make_event(article, MatchingEventResult.ALREADY_PRESENT) 

158 

159 # Matching fuzzy 

160 xarticle = self.backend.find_by_matching(article) 

161 if xarticle: 

162 _logger.debug(f"[{self.backend.name}] {article.pid}: found via Matching") 

163 return xarticle, self._make_event(article, MatchingEventResult.ADDED) 

164 

165 _logger.info( 

166 f"[{self.backend.name}] {article.pid}: not found with External ID nor Matching" 

167 ) 

168 return None, self._make_event( 

169 article, MatchingEventResult.NOT_FOUND, status=HistoryEventStatus.WARNING 

170 ) 

171 

172 def _make_event( 

173 self, 

174 article: Article, 

175 status_message: str, 

176 status: HistoryEventStatus = HistoryEventStatus.OK, 

177 ): 

178 return { 

179 "pid": article.pid, 

180 "status": status, 

181 "status_message": status_message, 

182 "backend": self.backend.name, 

183 } 

184 

185 def insert_history_event(self): 

186 children = [] 

187 for msg in self.event_messages: 

188 resource = model_helpers.get_resource(msg["pid"]) 

189 if not resource: 

190 raise ValueError(f"Resource not found: {msg['pid']}") 

191 data: "HistoryChildDict" = { 

192 "resource": resource, 

193 "type": msg.get("backend", self.backend.name), 

194 "status": msg.get("status"), 

195 } 

196 if "score" in msg: 

197 data["score"] = msg["score"] 

198 if "message" in msg: 

199 data["message"] = msg["message"] 

200 if "status_message" in msg: 

201 data["status_message"] = msg["status_message"] 

202 children.append(data) 

203 

204 insert_history_event( 

205 { 

206 "pid": f"{self.source_id}_{self.collection_pid or ''}_{self.backend.name}_matching", 

207 "col": self.collection, 

208 "source": self.source_id, 

209 "status": HistoryEventStatus.OK, 

210 "type": "matching", 

211 "children": children, 

212 } 

213 )