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

39 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-06-11 15:37 +0000

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 typing import Any, Dict, Iterable, Optional, Tuple 

16 

17from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import RequestMetadata, ToolDetails 

18from buildgrid.server.persistence.sql.models import ClientIdentityEntry 

19from buildgrid.settings import REQUEST_METADATA_HEADER_NAME 

20 

21 

22def printable_request_metadata(metadata_entries: Any) -> str: 

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

24 of its `RequestMetadata` entry. 

25 

26 Args: 

27 metadata_entries: tuple of entries obtained from a gRPC context 

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

29 

30 Returns: 

31 A string with the metadata contents. 

32 """ 

33 metadata = extract_request_metadata(metadata_entries) 

34 return request_metadata_to_string(metadata) 

35 

36 

37def extract_request_metadata(metadata_entries: Any) -> RequestMetadata: 

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

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

40 returns an empty message. 

41 

42 Args: 

43 metadata_entries: tuple of entries obtained from a gRPC context 

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

45 

46 Returns: 

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

48 request, the message will be empty. 

49 """ 

50 request_metadata_entry = next( 

51 (entry for entry in metadata_entries if entry.key == REQUEST_METADATA_HEADER_NAME), None 

52 ) 

53 

54 request_metadata = RequestMetadata() 

55 if request_metadata_entry: 

56 request_metadata.ParseFromString(request_metadata_entry.value) 

57 return request_metadata 

58 

59 

60def request_metadata_to_string(request_metadata: RequestMetadata) -> str: 

61 if request_metadata.tool_details: 

62 tool_name = request_metadata.tool_details.tool_name 

63 tool_version = request_metadata.tool_details.tool_version 

64 else: 

65 tool_name = tool_version = "" 

66 

67 return ( 

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

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

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

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

72 ) 

73 

74 

75def request_metadata_from_scheduler_dict(scheduler_request_metadata: Dict[str, str]) -> RequestMetadata: 

76 tool_details = ToolDetails() 

77 tool_details.tool_name = scheduler_request_metadata["tool-name"] 

78 tool_details.tool_version = scheduler_request_metadata["tool-version"] 

79 

80 request_metadata = RequestMetadata() 

81 request_metadata.tool_details.CopyFrom(tool_details) 

82 request_metadata.tool_invocation_id = scheduler_request_metadata["invocation-id"] 

83 request_metadata.correlated_invocations_id = scheduler_request_metadata["correlated-invocations-id"] 

84 

85 return request_metadata 

86 

87 

88def extract_client_identity( 

89 instance: str, invocation_metadata: Iterable[Tuple[str, Any]] 

90) -> Optional[ClientIdentityEntry]: 

91 """Extract the ClientIdentity from gRPC metadata 

92 

93 Args: 

94 invocation_metadata (List[Tuple[str, str]]): grpc metadata 

95 

96 Returns: 

97 Optional[ClientIdentityEntry]: identity of the client if exists 

98 """ 

99 metadata_dict = dict(invocation_metadata) 

100 workflow = metadata_dict.get("x-request-workflow", None) 

101 actor = metadata_dict.get("x-request-actor", None) 

102 subject = metadata_dict.get("x-request-subject", None) 

103 

104 # None or empty strings are invalid 

105 if workflow and actor and subject: 

106 return ClientIdentityEntry(instance=instance, workflow=workflow, actor=actor, subject=subject) 

107 

108 return None 

109 

110 

111def printable_client_identity(instance: str, invocation_metadata: Iterable[Tuple[str, Any]]) -> str: 

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

113 of its `ClientIdentity` entry. 

114 

115 Args: 

116 instance: REAPI instance name 

117 invocation_metadata: tuple of entries obtained from a gRPC context 

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

119 

120 Returns: 

121 A string with the ClientIdentity contents. 

122 """ 

123 client_id = extract_client_identity(instance, invocation_metadata) 

124 return str(client_id)