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

57 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-06-11 15:37 +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._exceptions import InvalidArgumentError 

22 

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

24 

25 

26class Instance(ABC): 

27 """ 

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

29 """ 

30 

31 SERVICE_NAME: ClassVar[str] 

32 """ 

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

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

35 """ 

36 

37 # Special case due to delayed setup methods 

38 _instance_name: str = None # type: ignore[assignment] 

39 """ 

40 The name of the instance 

41 """ 

42 

43 def set_instance_name(self, instance_name: str) -> None: 

44 """ 

45 Set the instance name associated with this instance. 

46 This method may be overriden if extra referencing is required. 

47 

48 Args: 

49 instance_name (str): The new instance's name. 

50 

51 Raises: 

52 AssertionError: if the instance is already registered. 

53 """ 

54 

55 if self._instance_name is None: 

56 self._instance_name = instance_name 

57 else: 

58 raise AssertionError("Instance already registered") 

59 

60 @property 

61 def instance_name(self) -> str: 

62 """The name of the instance""" 

63 return self._instance_name 

64 

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

66 self.start() 

67 return self 

68 

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

70 self.stop() 

71 

72 def start(self) -> None: 

73 """ 

74 A method called when the grpc service is starting. 

75 

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

77 """ 

78 

79 def stop(self) -> None: 

80 """ 

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

82 

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

84 """ 

85 

86 

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

88 

89 

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

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

92 """ 

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

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

95 """ 

96 

97 FULL_NAME: ClassVar[str] 

98 """ 

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

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

101 """ 

102 

103 def __init__(self) -> None: 

104 """ 

105 The InstancedServicer base class allows easily creating implementations for services 

106 which require delegating logic to distinct instance implementations. 

107 

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

109 """ 

110 

111 self._stack = ExitStack() 

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

113 

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

115 """ 

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

117 """ 

118 

119 if self.enabled: 

120 self.REGISTER_METHOD(server) 

121 

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

123 self.start() 

124 return self 

125 

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

127 self.stop() 

128 

129 def start(self) -> None: 

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

131 self._stack.enter_context(instance) 

132 

133 def stop(self) -> None: 

134 self._stack.close() 

135 

136 @property 

137 def enabled(self) -> bool: 

138 """ 

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

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

141 

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

143 """ 

144 

145 return len(self.instances) > 0 

146 

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

148 """ 

149 Adds an instance to the servicer. 

150 

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

152 

153 Args: 

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

155 

156 instance (_Instance): The instance implementation. 

157 """ 

158 

159 self.instances[name] = instance 

160 

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

162 """ 

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

164 if the instance requested does not exist. 

165 

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

167 

168 Args: 

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

170 

171 Returns: 

172 _Instance: The requested instance. 

173 """ 

174 

175 try: 

176 return self.instances[instance_name] 

177 except KeyError: 

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

179 

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

181 """ 

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

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

184 correct type required by the servicer implementation. 

185 

186 Args: 

187 instance (Instance): The instance to check. 

188 

189 Returns: 

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

191 """ 

192 

193 if instance.SERVICE_NAME == self.FULL_NAME: 

194 return cast(_Instance, instance) 

195 return None