Coverage for api/controllers/elasticsearch/helpers.py: 82%

1from __future__ import annotations 

2 

3import functools 

4import pprint 

5import time 

6from itertools import accumulate 

7from math import ceil 

8 

9from django.conf import settings 

10 

11import structlog 

12from elasticsearch import BadRequestError, NotFoundError 

13from elasticsearch_dsl import Search 

14 

15from api.utils.dead_link_mask import get_query_hash, get_query_mask 

16 

17 

18logger = structlog.get_logger(__name__) 

19 

20 

21def log_timing_info(func): 

22 @functools.wraps(func) 

23 def wrapper(*args, es_query, **kwargs): 

24 start_time = time.time() 

25 

26 # Call the original function 

27 result = func(*args, **kwargs) 

28 

29 response_time_in_ms = int((time.time() - start_time) * 1000) 

30 if hasattr(result, "took"): 

31 es_time_in_ms = result.took 

32 else: 

33 es_time_in_ms = result.get("took") 

34 logger.info( 

35 "Performed ES query", 

36 func=func.__name__, 

37 response_time=response_time_in_ms, 

38 es_time=es_time_in_ms, 

39 es_query=es_query, 

40 ) 

41 

42 return result 

43 

44 return wrapper 

45 

46 

47@log_timing_info 

48def get_es_response(s, *args, **kwargs): 

49 if settings.VERBOSE_ES_RESPONSE: 49 ↛ 50line 49 didn't jump to line 50 because the condition on line 49 was never true

50 logger.info(pprint.pprint(s.to_dict())) 

51 

52 try: 

53 search_response = s.execute() 

54 

55 if settings.VERBOSE_ES_RESPONSE: 55 ↛ 56line 55 didn't jump to line 56 because the condition on line 55 was never true

56 logger.info(pprint.pprint(search_response.to_dict())) 

57 except (BadRequestError, NotFoundError) as e: 

58 raise ValueError(e) 

59 

60 return search_response 

61 

62 

63@log_timing_info 

64def get_raw_es_response(index, body, *args, **kwargs): 

65 return settings.ES.search(index=index, body=body, *args, **kwargs) 

66 

67 

68ELASTICSEARCH_MAX_RESULT_WINDOW = 10000 

69DEAD_LINK_RATIO = 1 / 2 

70DEEP_PAGINATION_ERROR = "Deep pagination is not allowed." 

71 

72 

73def _unmasked_query_end(page_size, page): 

74 """ 

75 Calculate the upper index of results to retrieve from Elasticsearch. 

76 

77 Used to retrieve the upper index of results to retrieve from Elasticsearch under the 

78 following conditions: 

79 1. There is no query mask 

80 2. The lower index is beyond the scope of the existing query mask 

81 3. The lower index is within the scope of the existing query mask 

82 but the upper index exceeds it 

83 

84 In all these cases, the query mask is not used to calculate the upper index. 

85 """ 

86 return ceil(page_size * page / (1 - DEAD_LINK_RATIO)) 

87 

88 

89def _paginate_with_dead_link_mask( 

90 s: Search, page_size: int, page: int 

91) -> tuple[int, int]: 

92 """ 

93 Return the start and end of the results slice, given the query, page and page size. 

94 

95 In almost all cases the ``DEAD_LINK_RATIO`` will effectively double 

96 the page size (given the current configuration of 0.5). 

97 

98 The "branch X" labels are for cross-referencing with the tests. 

99 

100 :param s: The elasticsearch Search object 

101 :param page_size: How big the page should be. 

102 :param page: The page number. 

103 :return: Tuple of start and end. 

104 """ 

105 query_hash = get_query_hash(s) 

106 query_mask = get_query_mask(query_hash) 

107 if not query_mask: # branch 1 

108 start = 0 

109 end = _unmasked_query_end(page_size, page) 

110 elif page_size * (page - 1) > sum(query_mask): # branch 2 110 ↛ 111line 110 didn't jump to line 111 because the condition on line 110 was never true

111 start = len(query_mask) 

112 end = _unmasked_query_end(page_size, page) 

113 else: # branch 3 

114 # query_mask is a list of 0 and 1 where 0 indicates the result position 

115 # for the given query will be an invalid link. If we accumulate a query 

116 # mask you end up, at each index, with the number of live results you 

117 # will get back when you query that deeply. 

118 # We then query for the start and end index _of the results_ in ES based 

119 # on the number of results that we think will be valid based on the query mask. 

120 # If we're requesting `page=2 page_size=3` and the mask is [0, 1, 0, 1, 0, 1], 

121 # then we know that we have to _start_ with at least the sixth result of the 

122 # overall query to skip the first page of 3 valid results. The "end" of the 

123 # query will then follow the same pattern to reach the number of valid results 

124 # required to fill the requested page. If the mask is not deep enough to 

125 # account for the entire range, then we follow the typical assumption when 

126 # a mask is not available that the end should be `page * page_size / 0.5` 

127 # (i.e., double the page size) 

128 accu_query_mask = list(accumulate(query_mask)) 

129 start = 0 

130 if page > 1: 130 ↛ 131line 130 didn't jump to line 131 because the condition on line 130 was never true

131 try: # branch 3_start_A 

132 # find the index at which we can skip N valid results where N = all 

133 # the results that would be skipped to arrive at the start of the 

134 # requested page 

135 # This will effectively be the index at which we have the number of 

136 # previous valid results + 1 because we don't want to include the 

137 # last valid result from the previous page 

138 start = accu_query_mask.index(page_size * (page - 1) + 1) 

139 except ValueError: # branch 3_start_B 

140 # Cannot fail because of the check on branch 2 which verifies that 

141 # the query mask already includes at least enough masked valid 

142 # results to fulfill the requested page size 

143 start = accu_query_mask.index(page_size * (page - 1)) + 1 

144 # else: branch 3_start_C 

145 # Always start page=1 queries at 0 

146 

147 if page_size * page > sum(query_mask): # branch 3_end_A 

148 end = _unmasked_query_end(page_size, page) 

149 else: # branch 3_end_B 

150 end = accu_query_mask.index(page_size * page) + 1 

151 return start, end 

152 

153 

154def get_query_slice( 

155 s: Search, page_size: int, page: int, filter_dead: bool | None = False 

156) -> tuple[int, int]: 

157 """Select the start and end of the search results for this query.""" 

158 

159 if filter_dead: 

160 start_slice, end_slice = _paginate_with_dead_link_mask(s, page_size, page) 

161 else: 

162 # Paginate search query. 

163 start_slice = page_size * (page - 1) 

164 end_slice = page_size * page 

165 if start_slice + end_slice > ELASTICSEARCH_MAX_RESULT_WINDOW: 165 ↛ 166line 165 didn't jump to line 166 because the condition on line 165 was never true

166 raise ValueError(DEEP_PAGINATION_ERROR) 

167 return start_slice, end_slice