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

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 typing import Optional 

16 

17from google.protobuf.timestamp_pb2 import Timestamp 

18 

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

20 DESCRIPTOR, 

21 ListWorkersRequest, 

22 ListWorkersResponse, 

23 OperationFilter, 

24 OperationFilters, 

25 RegisteredWorker, 

26 WorkerOperation, 

27) 

28from buildgrid.server.enums import BotStatus 

29from buildgrid.server.exceptions import InvalidArgumentError 

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

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

32from buildgrid.server.scheduler.impl import Scheduler 

33from buildgrid.server.servicer import Instance 

34from buildgrid.server.settings import MAX_LIST_PAGE_SIZE 

35from buildgrid.server.utils.digests import parse_digest 

36 

37 

38class IntrospectionInstance(Instance): 

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

40 

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

42 self._scheduler = sql 

43 

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

45 if request.page_size < 0: 

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

47 if request.page < 0: 

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

49 

50 page = request.page or 1 

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

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

53 

54 workers = [] 

55 for bot in bot_entries: 

56 last_update = Timestamp() 

57 last_update.FromDatetime(bot.last_update_timestamp) 

58 expiry: Optional[Timestamp] = None 

59 if bot.expiry_time is not None: 

60 expiry = Timestamp() 

61 expiry.FromDatetime(bot.expiry_time) 

62 worker = RegisteredWorker( 

63 worker_name=bot.bot_id, 

64 session_name=bot.name, 

65 last_updated=last_update, 

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

67 expiry_time=expiry, 

68 ) 

69 if bot.job is not None: 

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

71 worker.action_digest.CopyFrom(digest) 

72 worker.operations.extend( 

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

74 ) 

75 workers.append(worker) 

76 

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

78 

79 def get_operation_filters(self) -> OperationFilters: 

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

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

82 filter_type = "text" 

83 if isinstance(spec.sanitizer, SortKeyValueSanitizer): 

84 comparators = ["="] 

85 elif isinstance(spec.sanitizer, DatetimeValueSanitizer): 

86 filter_type = "datetime" 

87 

88 try: 

89 values = spec.sanitizer.valid_values 

90 except NotImplementedError: 

91 values = [] 

92 

93 return OperationFilter( 

94 key=key, 

95 name=spec.name, 

96 type=filter_type, 

97 description=spec.description, 

98 comparators=comparators, 

99 values=values, 

100 ) 

101 

102 return OperationFilters( 

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

104 )