Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/introspection/instance.py: 96.08%

1# Copyright (C) 2024 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 

15from google.protobuf.timestamp_pb2 import Timestamp 

16 

17from buildgrid._protos.buildgrid.v2.introspection_pb2 import ( 

18 DESCRIPTOR, 

19 ListWorkersRequest, 

20 ListWorkersResponse, 

21 OperationFilter, 

22 OperationFilters, 

23 RegisteredWorker, 

24 WorkerOperation, 

25) 

26from buildgrid.server.enums import BotStatus 

27from buildgrid.server.exceptions import InvalidArgumentError 

28from buildgrid.server.operations.filtering.interpreter import VALID_OPERATION_FILTERS, OperationFilterSpec 

29from buildgrid.server.operations.filtering.sanitizer import DatetimeValueSanitizer, SortKeyValueSanitizer 

30from buildgrid.server.scheduler.impl import Scheduler 

31from buildgrid.server.servicer import Instance 

32from buildgrid.server.settings import MAX_LIST_PAGE_SIZE 

33from buildgrid.server.utils.digests import parse_digest 

34 

35 

36class IntrospectionInstance(Instance): 

37 SERVICE_NAME = DESCRIPTOR.services_by_name["Introspection"].full_name 

38 

39 def __init__(self, sql: Scheduler) -> None: 

40 self._scheduler = sql 

41 

42 def list_workers(self, request: ListWorkersRequest) -> ListWorkersResponse: 

43 if request.page_size < 0: 

44 raise InvalidArgumentError("Page size must be a positive integer") 

45 if request.page < 0: 

46 raise InvalidArgumentError("Page number must be a positive integer") 

47 

48 page = request.page or 1 

49 page_size = min(request.page_size, MAX_LIST_PAGE_SIZE) or MAX_LIST_PAGE_SIZE 

50 bot_entries, count = self._scheduler.list_workers(request.worker_name, page, page_size) 

51 

52 workers = [] 

53 for bot in bot_entries: 

54 last_update = Timestamp() 

55 last_update.FromDatetime(bot.last_update_timestamp) 

56 expiry: Timestamp | None = None 

57 if bot.expiry_time is not None: 

58 expiry = Timestamp() 

59 expiry.FromDatetime(bot.expiry_time) 

60 worker = RegisteredWorker( 

61 worker_name=bot.bot_id, 

62 session_name=bot.name, 

63 last_updated=last_update, 

64 bot_status=BotStatus(bot.bot_status).value, 

65 expiry_time=expiry, 

66 ) 

67 if bot.job is not None: 

68 if digest := parse_digest(bot.job.action_digest): 

69 worker.action_digest.CopyFrom(digest) 

70 worker.operations.extend( 

71 [WorkerOperation(operation_name=operation.name) for operation in bot.job.operations] 

72 ) 

73 workers.append(worker) 

74 

75 return ListWorkersResponse(workers=workers, total=count, page=page, page_size=page_size) 

76 

77 def get_operation_filters(self) -> OperationFilters: 

78 def _generate_filter_spec(key: str, spec: OperationFilterSpec) -> OperationFilter: 

79 comparators = ["<", "<=", "=", "!=", ">=", ">"] 

80 filter_type = "text" 

81 if isinstance(spec.sanitizer, SortKeyValueSanitizer): 

82 comparators = ["="] 

83 elif isinstance(spec.sanitizer, DatetimeValueSanitizer): 

84 filter_type = "datetime" 

85 

86 try: 

87 values = spec.sanitizer.valid_values 

88 except NotImplementedError: 

89 values = [] 

90 

91 return OperationFilter( 

92 key=key, 

93 name=spec.name, 

94 type=filter_type, 

95 description=spec.description, 

96 comparators=comparators, 

97 values=values, 

98 ) 

99 

100 return OperationFilters( 

101 filters=[_generate_filter_spec(key, spec) for key, spec in VALID_OPERATION_FILTERS.items()] 

102 )