Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/request_metadata_utils.py: 96.15%

1# Copyright (C) 2020 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 buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import RequestMetadata, ToolDetails 

16from buildgrid.settings import REQUEST_METADATA_HEADER_NAME 

17 

18 

19def printable_request_metadata(metadata_entries) -> str: 

20 """Given a metadata object, return a human-readable representation 

21 of its `RequestMetadata` entry. 

22 

23 Args: 

24 metadata_entries: tuple of entries obtained from a gRPC context 

25 with, for example, `context.invocation_metadata()`. 

26 

27 Returns: 

28 A string with the metadata contents. 

29 """ 

30 metadata = extract_request_metadata(metadata_entries) 

31 return request_metadata_to_string(metadata) 

32 

33 

34def extract_request_metadata(metadata_entries) -> RequestMetadata: 

35 """Given a list of string tuples, extract the RequestMetadata 

36 header values if they are present. If they were not provided, 

37 returns an empty message. 

38 

39 Args: 

40 metadata_entries: tuple of entries obtained from a gRPC context 

41 with, for example, `context.invocation_metadata()`. 

42 

43 Returns: 

44 A `RequestMetadata` proto. If the metadata is not defined in the 

45 request, the message will be empty. 

46 """ 

47 request_metadata_entry = next((entry for entry in metadata_entries 

48 if entry.key == REQUEST_METADATA_HEADER_NAME), 

49 None) 

50 

51 request_metadata = RequestMetadata() 

52 if request_metadata_entry: 

53 request_metadata.ParseFromString(request_metadata_entry.value) 

54 return request_metadata 

55 

56 

57def request_metadata_to_string(request_metadata): 

58 if request_metadata.tool_details: 

59 tool_name = request_metadata.tool_details.tool_name 

60 tool_version = request_metadata.tool_details.tool_version 

61 else: 

62 tool_name = tool_version = '' 

63 

64 return ( 

65 f'tool_name="{tool_name}", tool_version="{tool_version}", ' 

66 f'action_id="{request_metadata.action_id}", ' 

67 f'tool_invocation_id="{request_metadata.tool_invocation_id}", ' 

68 f'correlated_invocations_id="{request_metadata.correlated_invocations_id}"' 

69 ) 

70 

71 

72def request_metadata_from_scheduler_dict(scheduler_request_metadata: dict) -> RequestMetadata: 

73 tool_details = ToolDetails() 

74 tool_details.tool_name = scheduler_request_metadata['tool-name'] 

75 tool_details.tool_version = scheduler_request_metadata['tool-version'] 

76 

77 request_metadata = RequestMetadata() 

78 request_metadata.tool_details.CopyFrom(tool_details) 

79 request_metadata.tool_invocation_id = scheduler_request_metadata['invocation-id'] 

80 request_metadata.correlated_invocations_id = scheduler_request_metadata['correlated-invocations-id'] 

81 

82 return request_metadata