Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/decorators/errors.py: 82.72%

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 

15import contextlib 

16import inspect 

17import itertools 

18from functools import wraps 

19from typing import Any, Callable, Iterator, TypeVar, cast 

20 

21import grpc 

22 

23from buildgrid.server.exceptions import ( 

24 BotSessionCancelledError, 

25 BotSessionClosedError, 

26 BotSessionMismatchError, 

27 DuplicateBotSessionError, 

28 FailedPreconditionError, 

29 IncompleteReadError, 

30 InvalidArgumentError, 

31 NotFoundError, 

32 PermissionDeniedError, 

33 ResourceExhaustedError, 

34 RetriableError, 

35 UnknownBotSessionError, 

36) 

37from buildgrid.server.logging import buildgrid_logger 

38from buildgrid.server.sentry import send_exception_to_sentry 

39 

40Func = TypeVar("Func", bound=Callable) # type: ignore[type-arg] 

41 

42 

43LOGGER = buildgrid_logger(__name__) 

44 

45 

46@contextlib.contextmanager 

47def error_context(context: grpc.ServicerContext, unhandled_message: str) -> Iterator[None]: 

48 try: 

49 yield 

50 

51 except BotSessionCancelledError as e: 

52 LOGGER.info(e) 

53 context.abort(grpc.StatusCode.CANCELLED, str(e)) 

54 

55 except BotSessionClosedError as e: 

56 LOGGER.debug(e) 

57 context.abort(grpc.StatusCode.DATA_LOSS, str(e)) 

58 

59 except IncompleteReadError as e: 

60 LOGGER.exception(e) 

61 context.abort(grpc.StatusCode.DATA_LOSS, str(e)) 

62 

63 except ConnectionError as e: 

64 LOGGER.exception(e) 

65 context.abort(grpc.StatusCode.UNAVAILABLE, str(e)) 

66 

67 except DuplicateBotSessionError as e: 

68 LOGGER.info(e) 

69 context.abort(grpc.StatusCode.ABORTED, str(e)) 

70 

71 except FailedPreconditionError as e: 

72 LOGGER.error(e) 

73 context.abort(grpc.StatusCode.FAILED_PRECONDITION, str(e)) 

74 

75 except (InvalidArgumentError, BotSessionMismatchError, UnknownBotSessionError) as e: 

76 LOGGER.info(e) 

77 context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(e)) 

78 

79 except NotFoundError as e: 

80 LOGGER.debug(e) 

81 context.abort(grpc.StatusCode.NOT_FOUND, str(e)) 

82 

83 except NotImplementedError as e: 

84 LOGGER.info(e) 

85 context.abort(grpc.StatusCode.UNIMPLEMENTED, str(e)) 

86 

87 except PermissionDeniedError as e: 

88 LOGGER.exception(e) 

89 context.abort(grpc.StatusCode.PERMISSION_DENIED, str(e)) 

90 

91 except RetriableError as e: 

92 LOGGER.info("Retriable error.", tags=dict(client_retry_delay=e.retry_info.retry_delay)) 

93 context.abort_with_status(e.error_status) 

94 

95 except ResourceExhaustedError as e: 

96 LOGGER.exception(e) 

97 context.abort(grpc.StatusCode.RESOURCE_EXHAUSTED, str(e)) 

98 

99 except Exception as e: 

100 if context.code() is None: # type: ignore[attr-defined] 

101 send_exception_to_sentry(e) 

102 LOGGER.exception(unhandled_message) 

103 context.abort(grpc.StatusCode.INTERNAL, str(e)) 

104 raise 

105 

106 

107def handle_errors(get_printable_request: Callable[[Any], Any] = lambda r: str(r)) -> Callable[[Func], Func]: 

108 def decorator(f: Func) -> Func: 

109 @wraps(f) 

110 def return_wrapper(self: Any, request: Any, context: grpc.ServicerContext) -> Any: 

111 if isinstance(request, Iterator): 

112 # Pop the message out to get the instance from it, then and recreate the iterator. 

113 initial_request = next(request) 

114 printed_request = get_printable_request(initial_request) 

115 request = itertools.chain([initial_request], request) 

116 else: 

117 printed_request = get_printable_request(request) 

118 

119 with error_context(context, f"Unexpected error in {f.__name__}; request=[{printed_request}]"): 

120 return f(self, request, context) 

121 

122 @wraps(f) 

123 def yield_wrapper(self: Any, request: Any, context: grpc.ServicerContext) -> Any: 

124 if isinstance(request, Iterator): 

125 # Pop the message out to get the instance from it, then and recreate the iterator. 

126 initial_request = next(request) 

127 printed_request = get_printable_request(initial_request) 

128 request = itertools.chain([initial_request], request) 

129 else: 

130 printed_request = get_printable_request(request) 

131 

132 with error_context(context, f"Unexpected error in {f.__name__}; request=[{printed_request}]"): 

133 yield from f(self, request, context) 

134 

135 if inspect.isgeneratorfunction(f): 

136 return cast(Func, yield_wrapper) 

137 return cast(Func, return_wrapper) 

138 

139 return decorator