Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/persistence/sql/utils.py: 95.65%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

92 statements  

1# Copyright (C) 2020 Bloomberg LP 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# <http://www.apache.org/licenses/LICENSE-2.0> 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15""" Holds constants and utility functions for the SQL scheduler. """ 

16 

17 

18from datetime import datetime 

19from distutils.util import strtobool 

20import operator 

21from typing import Any, List, Tuple 

22 

23from sqlalchemy.sql.expression import and_, or_, ClauseElement 

24 

25from buildgrid._exceptions import InvalidArgumentError 

26from buildgrid.server.operations.filtering import OperationFilter, SortKey 

27from buildgrid.server.persistence.sql.models import Job, Operation 

28 

29 

30DATETIME_FORMAT = "%Y-%m-%d-%H-%M-%S-%f" 

31 

32 

33LIST_OPERATIONS_PARAMETER_MODEL_MAP = { 

34 "stage": Job.stage, 

35 "name": Operation.name, 

36 "queued_time": Job.queued_timestamp, 

37 "start_time": Job.worker_start_timestamp, 

38 "completed_time": Job.worker_completed_timestamp, 

39 "invocation_id": Operation.invocation_id, 

40 "correlated_invocations_id": Operation.correlated_invocations_id, 

41 "tool_name": Operation.tool_name, 

42 "tool_version": Operation.tool_version 

43} 

44 

45 

46def parse_list_operations_sort_value(value: str, column) -> Any: 

47 """ Convert the string representation of a value to the proper Python type. """ 

48 python_type = column.property.columns[0].type.python_type 

49 if python_type == datetime: 

50 return datetime.strptime(value, DATETIME_FORMAT) 

51 elif python_type == bool: 

52 # Using this distutils function to cover a few different bool representations 

53 return bool(strtobool(value)) 

54 else: 

55 return python_type(value) 

56 

57 

58def dump_list_operations_token_value(token_value) -> str: 

59 """ Convert a value to a string for use in the page_token. """ 

60 if isinstance(token_value, datetime): 

61 return datetime.strftime(token_value, DATETIME_FORMAT) 

62 else: 

63 return str(token_value) 

64 

65 

66def build_pagination_clause_for_sort_key( 

67 sort_value: Any, previous_sort_values: List[Any], sort_keys: List[SortKey]) -> ClauseElement: 

68 """ Build part of a filter clause to figure out the starting point of the page given 

69 by the page_token. See the docstring of build_page_filter for more details. """ 

70 if len(sort_keys) <= len(previous_sort_values): 

71 raise ValueError("Not enough keys to unpack") 

72 

73 filter_clause_list = [] 

74 for i, previous_sort_value in enumerate(previous_sort_values): 

75 previous_sort_col = LIST_OPERATIONS_PARAMETER_MODEL_MAP[sort_keys[i].name] 

76 filter_clause_list.append(previous_sort_col == previous_sort_value) 

77 sort_key = sort_keys[len(previous_sort_values)] 

78 sort_col = LIST_OPERATIONS_PARAMETER_MODEL_MAP[sort_key.name] 

79 if sort_key.descending: 

80 filter_clause_list.append(sort_col < sort_value) 

81 else: 

82 filter_clause_list.append(sort_col > sort_value) 

83 return and_(*filter_clause_list) 

84 

85 

86def build_page_filter(page_token, sort_keys: List[SortKey]) -> ClauseElement: 

87 """ Build a filter to determine the starting point of the rows to fetch, based 

88 on the page_token. 

89 

90 The page_token is directly related to the sort order, and in this way it acts as a 

91 "cursor." It is given in the format Xval|Yval|Zval|..., where each element is a value 

92 corresponding to an orderable column in the database. If the corresponding rows are 

93 X, Y, and Z, then X is the primary sort key, with Y breaking ties between X, and Z 

94 breaking ties between X and Y. The corresponding filter clause is then: 

95 

96 (X > Xval) OR (X == XVal AND Y > Yval) OR (X == Xval AND Y == Yval AND Z > Zval) ... 

97 """ 

98 # The page token is used as a "cursor" to determine the starting point 

99 # of the rows to fetch. It is derived from the sort keys. 

100 token_elements = page_token.split("|") 

101 if len(token_elements) != len(sort_keys): 

102 # It is possible that an extra "|" was in the input 

103 # TODO: Handle extra "|"s somehow? Or at least allow escaping them 

104 raise InvalidArgumentError( 

105 f"Wrong number of \"|\"-separated elements in page token [{page_token}]. " 

106 f"Expected {len(sort_keys)}, got {len(token_elements)}.") 

107 

108 sort_key_clause_list = [] 

109 previous_sort_values: List[Any] = [] 

110 # Build the compound clause for each sort key in the token 

111 for i, sort_key in enumerate(sort_keys): 

112 col = LIST_OPERATIONS_PARAMETER_MODEL_MAP[sort_key.name] 

113 sort_value = parse_list_operations_sort_value(token_elements[i], col) 

114 filter_clause = build_pagination_clause_for_sort_key( 

115 sort_value, previous_sort_values, sort_keys) 

116 sort_key_clause_list.append(filter_clause) 

117 previous_sort_values.append(sort_value) 

118 

119 return or_(*sort_key_clause_list) 

120 

121 

122def build_page_token(operation, sort_keys): 

123 """ Use the sort keys to build a page token from the given operation. """ 

124 token_values = [] 

125 for sort_key in sort_keys: 

126 col = LIST_OPERATIONS_PARAMETER_MODEL_MAP[sort_key.name] 

127 col_properties = col.property.columns[0] 

128 column_name = col_properties.name 

129 table_name = col_properties.table.name 

130 if table_name == "operations": 

131 token_value = getattr(operation, column_name) 

132 elif table_name == "jobs": 

133 token_value = getattr(operation.job, column_name) 

134 else: 

135 raise ValueError("Got invalid table f{table_name} for sort key {sort_key.name} while building page_token") 

136 

137 token_values.append(dump_list_operations_token_value(token_value)) 

138 

139 next_page_token = "|".join(token_values) 

140 return next_page_token 

141 

142 

143def extract_sort_keys(operation_filters: List[OperationFilter]) -> Tuple[List[SortKey], List[OperationFilter]]: 

144 """ Splits the operation filters into sort keys and non-sort filters, returning both as 

145 separate lists. 

146 

147 Sort keys are specified with the "sort_order" parameter in the filter string. Multiple 

148 "sort_order"s can appear in the filter string, and all are extracted and returned. """ 

149 # pylint: disable=comparison-with-callable 

150 sort_keys = [] 

151 non_sort_filters = [] 

152 for op_filter in operation_filters: 

153 if op_filter.parameter == "sort_order": 

154 if op_filter.operator != operator.eq: 

155 raise InvalidArgumentError("sort_order must be specified with the \"=\" operator.") 

156 sort_keys.append(op_filter.value) 

157 else: 

158 non_sort_filters.append(op_filter) 

159 

160 return sort_keys, non_sort_filters 

161 

162 

163def build_sort_column_list(sort_keys): 

164 """ Convert the list of sort keys into a list of columns that can be 

165 passed to an order_by. 

166 

167 This function checks the sort keys to ensure that they are in the 

168 parameter-model map and raises an InvalidArgumentError if they are not. """ 

169 sort_columns = [] 

170 for sort_key in sort_keys: 

171 try: 

172 col = LIST_OPERATIONS_PARAMETER_MODEL_MAP[sort_key.name] 

173 if sort_key.descending: 

174 sort_columns.append(col.desc()) 

175 else: 

176 sort_columns.append(col.asc()) 

177 except KeyError: 

178 raise InvalidArgumentError(f"[{sort_key.name}] is not a valid sort key.") 

179 return sort_columns 

180 

181 

182def convert_filter_to_sql_filter(operation_filter: OperationFilter) -> ClauseElement: 

183 """ Convert the internal representation of a filter to a representation that SQLAlchemy 

184 can understand. The return type is a "ColumnElement," per the end of this section in 

185 the SQLAlchemy docs: https://docs.sqlalchemy.org/en/13/core/tutorial.html#selecting-specific-columns 

186 

187 This function assumes that the parser has appropriately converted the filter 

188 value to a Python type that can be compared to the parameter. """ 

189 try: 

190 param = LIST_OPERATIONS_PARAMETER_MODEL_MAP[operation_filter.parameter] 

191 except KeyError: 

192 raise InvalidArgumentError(f"Invalid parameter: [{operation_filter.parameter}]") 

193 

194 return operation_filter.operator(param, operation_filter.value) 

195 

196 

197def build_custom_filters(operation_filters: List[OperationFilter]) -> List[ClauseElement]: 

198 return [ 

199 convert_filter_to_sql_filter(operation_filter) 

200 for operation_filter in operation_filters 

201 ]