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

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 contextvars 

17import functools 

18 

19import grpc 

20 

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

22from buildgrid.server.request_metadata_utils import extract_request_metadata 

23from buildgrid.settings import ( 

24 REQUEST_METADATA_TOOL_NAME, 

25 REQUEST_METADATA_TOOL_VERSION, 

26 REQUEST_METADATA_HEADER_NAME 

27) 

28 

29 

30def get_empty(): 

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

32 the deafult for the ContextVar 

33 """ 

34 empty_metadata = RequestMetadata() 

35 empty_metadata.tool_details.tool_name = REQUEST_METADATA_TOOL_NAME 

36 empty_metadata.tool_details.tool_version = REQUEST_METADATA_TOOL_VERSION 

37 return empty_metadata 

38 

39 

40# ContextVar for request metadata 

41ctx_request_metadata: contextvars.ContextVar = (contextvars.ContextVar('ctx_request_metadata', 

42 default=get_empty())) 

43 

44 

45def metadatacontext(): 

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

47 and then reset it on completion of method. 

48 

49 Note: 

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

51 

52 This returns a decorator that extracts the invocation_metadata from the 

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

54 ContextVar variable after the method has completed. 

55 """ 

56 def context_decorator(func): 

57 @functools.wraps(func) 

58 def context_wrapper(*args, **kwargs): 

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 return context_wrapper 

69 return context_decorator 

70 

71 

72def metadata_list(): 

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

74 metadata = ctx_request_metadata.get() 

75 return [(REQUEST_METADATA_HEADER_NAME, metadata.SerializeToString())]