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

22 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2025-02-11 15:07 +0000

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 inspect 

16import uuid 

17from functools import wraps 

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

19 

20from buildgrid.server.metadata import ctx_grpc_request_id 

21 

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

23 

24 

25def track_request_id(f: Func) -> Func: 

26 """Decorator to set the request ID ContextVar. 

27 

28 This decorator sets the ``ctx_grpc_request_id`` ContextVar to a UUID 

29 for the duration of the decorated function. This ContextVar is used 

30 in logging output to allow log lines for the same request to be 

31 identified. 

32 

33 """ 

34 

35 @wraps(f) 

36 def return_wrapper(*args: Any, **kwargs: Any) -> Any: 

37 ctx_grpc_request_id.set(str(uuid.uuid4())) 

38 try: 

39 return f(*args, **kwargs) 

40 finally: 

41 ctx_grpc_request_id.set(None) 

42 

43 @wraps(f) 

44 def yield_wrapper(*args: Any, **kwargs: Any) -> Iterator[Any]: 

45 ctx_grpc_request_id.set(str(uuid.uuid4())) 

46 try: 

47 yield from f(*args, **kwargs) 

48 finally: 

49 ctx_grpc_request_id.set(None) 

50 

51 if inspect.isgeneratorfunction(f): 

52 return cast(Func, yield_wrapper) 

53 return cast(Func, return_wrapper)