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
« 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.
15import inspect
16import uuid
17from functools import wraps
18from typing import Any, Callable, Iterator, TypeVar, cast
20from buildgrid.server.metadata import ctx_grpc_request_id
22Func = TypeVar("Func", bound=Callable) # type: ignore[type-arg]
25def track_request_id(f: Func) -> Func:
26 """Decorator to set the request ID ContextVar.
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.
33 """
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)
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)
51 if inspect.isgeneratorfunction(f):
52 return cast(Func, yield_wrapper)
53 return cast(Func, return_wrapper)