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

1# Copyright (C) 2022 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 

15 

16import functools 

17from contextvars import ContextVar 

18from typing import Any, Callable, Optional, Tuple, TypeVar, Union, cast 

19 

20import grpc 

21 

22from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import RequestMetadata 

23from buildgrid.server.request_metadata_utils import extract_request_metadata 

24from buildgrid.settings import REQUEST_METADATA_HEADER_NAME, REQUEST_METADATA_TOOL_NAME, REQUEST_METADATA_TOOL_VERSION 

25 

26 

27def get_empty() -> RequestMetadata: 

28 """Function to create an empty request metadata structure, to use as 

29 the deafult for the ContextVar 

30 """ 

31 empty_metadata = RequestMetadata() 

32 empty_metadata.tool_details.tool_name = REQUEST_METADATA_TOOL_NAME 

33 empty_metadata.tool_details.tool_version = REQUEST_METADATA_TOOL_VERSION 

34 return empty_metadata 

35 

36 

37# ContextVar for request metadata 

38ctx_request_metadata: ContextVar[RequestMetadata] = ContextVar("ctx_request_metadata", default=get_empty()) 

39 

40 

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

42 

43 

44def metadatacontext() -> Callable[[Func], Func]: 

45 """Helper function to obtain metadata and set request metadata ContextVar, 

46 and then reset it on completion of method. 

47 

48 Note: 

49 args[2] of the method must be of type grpc.ServicerContext 

50 

51 This returns a decorator that extracts the invocation_metadata from the 

52 context argument and sets the ContextVar variable with it. Resetting the 

53 ContextVar variable after the method has completed. 

54 """ 

55 

56 def context_decorator(func: Func) -> Func: 

57 @functools.wraps(func) 

58 def context_wrapper(*args: Any, **kwargs: Any) -> Any: 

59 context = args[2] 

60 assert isinstance(context, grpc.ServicerContext) 

61 metadata = extract_request_metadata(context.invocation_metadata()) 

62 token = ctx_request_metadata.set(metadata) 

63 try: 

64 retval = func(*args, **kwargs) 

65 return retval 

66 finally: 

67 ctx_request_metadata.reset(token) 

68 

69 return cast(Func, context_wrapper) 

70 

71 return context_decorator 

72 

73 

74def metadata_list() -> Tuple[Tuple[str, Union[str, bytes]], ...]: 

75 """Helper function to construct the metadata list from the ContextVar.""" 

76 metadata = ctx_request_metadata.get() 

77 return ((REQUEST_METADATA_HEADER_NAME, metadata.SerializeToString()),) 

78 

79 

80ctx_grpc_request_id: ContextVar[Optional[str]] = ContextVar("grpc_request_id", default=None)