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 get_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_back.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", "crossred"]): 

89 self.backends = get_backends(backend) 

90 super().__init__(params) 

91 self.event_messages = [] 

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

93 

94 if self.collection_pid: 

95 self.queryset = self.queryset.filter( 

96 my_container__my_collection__pid=self.collection_pid 

97 ) 

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

99 

100 if self.source_id: 

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

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

103 

104 if self.issue_pids: 

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

106 

107 def internal_do(self) -> None: 

108 for self.backend in self.backends: 

109 _logger.info( 

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

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

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

113 ) 

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

115 with ThreadPoolExecutor(max_workers=1) as executor: 

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

117 xarticle, event = self._find_article(article) 

118 

119 if event: 

120 self.event_messages.append(event) 

121 

122 if not xarticle: 

123 self.callback() 

124 continue 

125 

126 # Handle asyncronous results and exceptions as soon as possible 

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

128 for f in promises: 

129 if f.done(): 

130 completed.add(f) 

131 exc = f.exception() 

132 if exc: 

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

134 raise exc 

135 promises -= completed 

136 # Runs the database insertion in parallel 

137 promise = executor.submit( 

138 update_article_with_xarticle, 

139 article, 

140 xarticle, 

141 merge_titles=False, 

142 update_pdf_link=self.update_pdf_link, 

143 ) 

144 promise.add_done_callback(self.callback) 

145 promises.add(promise) 

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

147 

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

149 """ 

150 Calls ArticleData from backend. 

151 """ 

152 # With external ID 

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

154 if xarticle: 

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

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

157 

158 # Matching fuzzy 

159 xarticle = self.backend.find_by_matching(article) 

160 if xarticle: 

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

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

163 

164 _logger.info( 

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

166 ) 

167 return None, self._make_event( 

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

169 ) 

170 

171 def _make_event( 

172 self, 

173 article: Article, 

174 status_message: str, 

175 status: HistoryEventStatus = HistoryEventStatus.OK, 

176 ) -> dict: 

177 return { 

178 "pid": article.pid, 

179 "status": status, 

180 "status_message": status_message, 

181 "backend": self.backend.name, 

182 } 

183 

184 def insert_history_event(self): 

185 children = [] 

186 for msg in self.event_messages: 

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

188 if not resource: 

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

190 data: "HistoryChildDict" = { 

191 "resource": resource, 

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

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

194 } 

195 if "score" in msg: 

196 data["score"] = msg["score"] 

197 if "message" in msg: 

198 data["message"] = msg["message"] 

199 if "status_message" in msg: 

200 data["status_message"] = msg["status_message"] 

201 children.append(data) 

202 

203 insert_history_event( 

204 { 

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

206 "col": self.collection, 

207 "source": self.source_id, 

208 "status": HistoryEventStatus.OK, 

209 "type": "matching", 

210 "children": children, 

211 } 

212 )