Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/servicer.py: 100.00%

53 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-11-01 16:30 +0000

1# Copyright (C) 2023 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 abc import ABC 

16from contextlib import ExitStack 

17from typing import Any, Callable, ClassVar, Dict, Generic, Optional, TypeVar, cast 

18 

19import grpc 

20 

21from buildgrid.server.context import current_instance 

22from buildgrid.server.exceptions import InvalidArgumentError 

23 

24_Instance = TypeVar("_Instance", bound="Instance") 

25 

26 

27class Instance(ABC): 

28 """ 

29 An Instance is the underlying implementation of a given Servicer. 

30 """ 

31 

32 SERVICE_NAME: ClassVar[str] 

33 """ 

34 The expected FULL_NAME of the Service which will wrap this instance. 

35 This value should be declared on the class of any Instance implementations. 

36 """ 

37 

38 def __enter__(self: _Instance) -> _Instance: 

39 self.start() 

40 return self 

41 

42 def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: 

43 self.stop() 

44 

45 def start(self) -> None: 

46 """ 

47 A method called when the grpc service is starting. 

48 

49 This method may be overriden if startup logic is required. 

50 """ 

51 

52 def stop(self) -> None: 

53 """ 

54 A method called when the grpc service is shutting down. 

55 

56 This method may be overriden if shutdown logic is required. 

57 """ 

58 

59 

60_InstancedServicer = TypeVar("_InstancedServicer", bound="InstancedServicer[Any]") 

61 

62 

63class InstancedServicer(ABC, Generic[_Instance]): 

64 REGISTER_METHOD: ClassVar[Callable[[Any, grpc.Server], None]] 

65 """ 

66 The method to be invoked when attaching the service to a grpc.Server instance. 

67 This value should be declared on the class of any Servicer implementations. 

68 """ 

69 

70 FULL_NAME: ClassVar[str] 

71 """ 

72 The full name of the servicer, used to match instances to the servicer and configure reflection. 

73 This value should be declared on the class of any Servicer implementations. 

74 """ 

75 

76 def __init__(self) -> None: 

77 """ 

78 The InstancedServicer base class allows easily creating implementations for services 

79 which require delegating logic to distinct instance implementations. 

80 

81 The base class provides logic for registering new instances with the service. 

82 """ 

83 

84 self._stack = ExitStack() 

85 self.instances: Dict[str, _Instance] = {} 

86 

87 def setup_grpc(self, server: grpc.Server) -> None: 

88 """ 

89 A method called when the Service is being attached to a grpc server. 

90 """ 

91 

92 if self.enabled: 

93 self.REGISTER_METHOD(server) 

94 

95 def __enter__(self: _InstancedServicer) -> _InstancedServicer: 

96 self.start() 

97 return self 

98 

99 def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: 

100 self.stop() 

101 

102 def start(self) -> None: 

103 for instance in self.instances.values(): 

104 self._stack.enter_context(instance) 

105 

106 def stop(self) -> None: 

107 self._stack.close() 

108 

109 @property 

110 def enabled(self) -> bool: 

111 """ 

112 By default, a servicer is disabled if there are no registered instances. 

113 If a servicer is not enabled, it will not be attached to the grpc.Server instance. 

114 

115 This property may be overriden if servicer enablement follows other rules. 

116 """ 

117 

118 return len(self.instances) > 0 

119 

120 def add_instance(self, name: str, instance: _Instance) -> None: 

121 """ 

122 Adds an instance to the servicer. 

123 

124 This method may be overriden if adding an instance requires additional setup. 

125 

126 Args: 

127 name (str): The name of the instance. 

128 

129 instance (_Instance): The instance implementation. 

130 """ 

131 

132 self.instances[name] = instance 

133 

134 @property 

135 def current_instance(self) -> _Instance: 

136 return self.get_instance(current_instance()) 

137 

138 def get_instance(self, instance_name: str) -> _Instance: 

139 """ 

140 Provides a wrapper to access the instance, throwing a InvalidArgumentError 

141 if the instance requested does not exist. 

142 

143 This method may be overriden if you wish to create a custom error message. 

144 

145 Args: 

146 instance_name (str): The name of the instance. 

147 

148 Returns: 

149 _Instance: The requested instance. 

150 """ 

151 

152 try: 

153 return self.instances[instance_name] 

154 except KeyError: 

155 raise InvalidArgumentError(f"Invalid instance name: [{instance_name}]") 

156 

157 def cast(self, instance: Instance) -> Optional[_Instance]: 

158 """ 

159 A helper tool used by the BuildGrid Server startup logic to determine the correct 

160 servicer to attach an instance to. This method will also cast the instance to the 

161 correct type required by the servicer implementation. 

162 

163 Args: 

164 instance (Instance): The instance to check. 

165 

166 Returns: 

167 Optional[_Instance]: The validated instance or None if invalid. 

168 """ 

169 

170 if instance.SERVICE_NAME == self.FULL_NAME: 

171 return cast(_Instance, instance) 

172 return None