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

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 functools 

16import inspect 

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

18 

19import grpc 

20 

21from buildgrid.server.metadata import ctx_request_metadata, extract_request_metadata 

22 

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

24 

25 

26def metadatacontext(f: Func) -> Func: 

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

28 and then reset it on completion of method. 

29 

30 Note: 

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

32 

33 This returns a decorator that extracts the invocation_metadata from the 

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

35 ContextVar variable after the method has completed. 

36 """ 

37 

38 @functools.wraps(f) 

39 def server_stream_wrapper(self: Any, message: Any, context: grpc.ServicerContext) -> Iterator[Any]: 

40 metadata = extract_request_metadata(context.invocation_metadata()) 

41 token = ctx_request_metadata.set(metadata) 

42 try: 

43 yield from f(self, message, context) 

44 finally: 

45 ctx_request_metadata.reset(token) 

46 

47 @functools.wraps(f) 

48 def server_unary_wrapper(self: Any, message: Any, context: grpc.ServicerContext) -> Any: 

49 metadata = extract_request_metadata(context.invocation_metadata()) 

50 token = ctx_request_metadata.set(metadata) 

51 try: 

52 return f(self, message, context) 

53 finally: 

54 ctx_request_metadata.reset(token) 

55 

56 if inspect.isgeneratorfunction(f): 

57 return cast(Func, server_stream_wrapper) 

58 return cast(Func, server_unary_wrapper)