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

1import logging 

2from concurrent.futures import Future, ThreadPoolExecutor 

3from typing import TYPE_CHECKING, NotRequired, TypedDict 

4 

5from history.model_data import HistoryChildDict, MatchingEventResult 

6from history.models import HistoryEventStatus 

7from history.utils import insert_history_event 

8from ptf import model_helpers 

9from ptf.cmds.base_cmds import baseCmd 

10from ptf.display.resolver import ExtIdName 

11from ptf.models import Article, Collection 

12from ptf_external.matching import get_matching_backend 

13 

14from crawler.cmds.augment import update_article_with_xarticle 

15from crawler.models.source import Source 

16 

17if TYPE_CHECKING: 

18 from collections.abc import Callable 

19 

20 from ptf.external.matching.base_matching import MatchingBackend 

21 from ptf.models import ArticleQuerySet 

22 

23 class AugmentArticlesParams(TypedDict): 

24 source_id: str 

25 collection_pid: NotRequired[str] 

26 issue_pids: NotRequired[list[str]] 

27 update_pdf_link: NotRequired[bool] 

28 

29 from typing import NotRequired, TypedDict 

30 

31 class MatchingOperationMessage(TypedDict): 

32 pid: str 

33 zbl_id: "NotRequired[str]" 

34 score: "NotRequired[int]" 

35 status: HistoryEventStatus 

36 message: "NotRequired[str]" 

37 status_message: "NotRequired[str]" 

38 

39 

40_logger = logging.getLogger(__name__) 

41 

42 

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

44 children = [] 

45 for msg in event_messages: 

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

47 if not resource: 

48 raise ValueError("Resource not found") 

49 data: HistoryChildDict = { 

50 "resource": resource, 

51 "type": "zbl-id", 

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

53 } 

54 

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

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

57 

58 if "score" in msg: 

59 data["score"] = msg["score"] 

60 

61 if "message" in msg: 

62 data["message"] = msg["message"] 

63 

64 if "status_message" in msg: 

65 data["status_message"] = msg["status_message"] 

66 

67 children.append(data) 

68 return children 

69 

70 

71class AugmentArticlesCmd(baseCmd[None]): 

72 """ 

73 Metadata augment command. 

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

75 """ 

76 

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

78 

79 backend: "MatchingBackend" 

80 source_id: str 

81 collection_pid: str | None = None 

82 issue_pids: list[str] = [] 

83 update_pdf_link: bool = False 

84 collection: "Collection | None" = None 

85 queryset: "ArticleQuerySet" 

86 event_messages: list[dict] 

87 callback: "Callable" = lambda _: _ 

88 

89 def __init__(self, params: "AugmentArticlesParams", backend: ExtIdName): 

90 super().__init__(params) 

91 

92 self.backend = get_matching_backend(backend) 

93 

94 self.event_messages = [] 

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

96 

97 if self.collection_pid: 

98 self.queryset = self.queryset.filter( 

99 my_container__my_collection__pid=self.collection_pid 

100 ) 

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

102 

103 if self.source_id: 

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

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

106 

107 if self.issue_pids: 

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

109 

110 def internal_do(self) -> None: 

111 _logger.info( 

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

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

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

115 ) 

116 promises: set[Future] = set() 

117 with ThreadPoolExecutor(max_workers=1) as executor: 

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

119 xarticle, event = self._find_article(article) 

120 

121 if event: 

122 self.event_messages.append(event) 

123 

124 if not xarticle: 

125 self.callback() 

126 continue 

127 

128 # Handle asyncronous results and exceptions as soon as possible 

129 completed: set[Future] = set() 

130 for f in promises: 

131 if f.done(): 

132 completed.add(f) 

133 exc = f.exception() 

134 if exc: 

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

136 raise exc 

137 promises -= completed 

138 # Runs the database insertion in parallel 

139 promise = executor.submit( 

140 update_article_with_xarticle, 

141 article, 

142 xarticle, 

143 merge_titles=False, 

144 update_pdf_link=self.update_pdf_link, 

145 ) 

146 promise.add_done_callback(self.callback) 

147 promises.add(promise) 

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

149 

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

151 """ 

152 Calls ArticleData from backend. 

153 """ 

154 # With external ID 

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

156 if xarticle: 

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

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

159 

160 # Matching fuzzy 

161 xarticle = self.backend.find_by_matching(article) 

162 if xarticle: 

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

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

165 

166 _logger.info( 

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

168 ) 

169 return None, self._make_event( 

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

171 ) 

172 

173 def _make_event( 

174 self, 

175 article: Article, 

176 status_message: str, 

177 status: HistoryEventStatus = HistoryEventStatus.OK, 

178 ): 

179 return { 

180 "pid": article.pid, 

181 "status": status, 

182 "status_message": status_message, 

183 "backend": self.backend.name, 

184 } 

185 

186 def insert_history_event(self): 

187 children = [] 

188 for msg in self.event_messages: 

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

190 if not resource: 

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

192 data: HistoryChildDict = { 

193 "resource": resource, 

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

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

196 } 

197 if "score" in msg: 

198 data["score"] = msg["score"] 

199 if "message" in msg: 

200 data["message"] = msg["message"] 

201 if "status_message" in msg: 

202 data["status_message"] = msg["status_message"] 

203 children.append(data) 

204 

205 insert_history_event( 

206 { 

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

208 "col": self.collection, 

209 "source": self.source_id, 

210 "status": HistoryEventStatus.OK, 

211 "type": "matching", 

212 "children": children, 

213 } 

214 )