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
« 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.
15from abc import ABC
16from contextlib import ExitStack
17from typing import Any, Callable, ClassVar, Dict, Generic, Optional, TypeVar, cast
19import grpc
21from buildgrid.server.context import current_instance
22from buildgrid.server.exceptions import InvalidArgumentError
24_Instance = TypeVar("_Instance", bound="Instance")
27class Instance(ABC):
28 """
29 An Instance is the underlying implementation of a given Servicer.
30 """
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 """
38 def __enter__(self: _Instance) -> _Instance:
39 self.start()
40 return self
42 def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
43 self.stop()
45 def start(self) -> None:
46 """
47 A method called when the grpc service is starting.
49 This method may be overriden if startup logic is required.
50 """
52 def stop(self) -> None:
53 """
54 A method called when the grpc service is shutting down.
56 This method may be overriden if shutdown logic is required.
57 """
60_InstancedServicer = TypeVar("_InstancedServicer", bound="InstancedServicer[Any]")
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 """
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 """
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.
81 The base class provides logic for registering new instances with the service.
82 """
84 self._stack = ExitStack()
85 self.instances: Dict[str, _Instance] = {}
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 """
92 if self.enabled:
93 self.REGISTER_METHOD(server)
95 def __enter__(self: _InstancedServicer) -> _InstancedServicer:
96 self.start()
97 return self
99 def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
100 self.stop()
102 def start(self) -> None:
103 for instance in self.instances.values():
104 self._stack.enter_context(instance)
106 def stop(self) -> None:
107 self._stack.close()
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.
115 This property may be overriden if servicer enablement follows other rules.
116 """
118 return len(self.instances) > 0
120 def add_instance(self, name: str, instance: _Instance) -> None:
121 """
122 Adds an instance to the servicer.
124 This method may be overriden if adding an instance requires additional setup.
126 Args:
127 name (str): The name of the instance.
129 instance (_Instance): The instance implementation.
130 """
132 self.instances[name] = instance
134 @property
135 def current_instance(self) -> _Instance:
136 return self.get_instance(current_instance())
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.
143 This method may be overriden if you wish to create a custom error message.
145 Args:
146 instance_name (str): The name of the instance.
148 Returns:
149 _Instance: The requested instance.
150 """
152 try:
153 return self.instances[instance_name]
154 except KeyError:
155 raise InvalidArgumentError(f"Invalid instance name: [{instance_name}]")
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.
163 Args:
164 instance (Instance): The instance to check.
166 Returns:
167 Optional[_Instance]: The validated instance or None if invalid.
168 """
170 if instance.SERVICE_NAME == self.FULL_NAME:
171 return cast(_Instance, instance)
172 return None