Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/context.py: 93.48%

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 

15from contextlib import contextmanager 

16from contextvars import ContextVar 

17from typing import Iterator, Optional 

18 

19from buildgrid.server.exceptions import InvalidArgumentError 

20 

21_instance_name: "ContextVar[Optional[str]]" = ContextVar("_instance_name", default=None) 

22_service_name: "ContextVar[Optional[str]]" = ContextVar("_service_name", default=None) 

23_method_name: "ContextVar[Optional[str]]" = ContextVar("_method_name", default=None) 

24 

25 

26@contextmanager 

27def instance_context(instance_name: Optional[str]) -> Iterator[None]: 

28 token = _instance_name.set(instance_name) 

29 try: 

30 yield 

31 finally: 

32 _instance_name.reset(token) 

33 

34 

35def try_current_instance() -> Optional[str]: 

36 return _instance_name.get() 

37 

38 

39def current_instance() -> str: 

40 instance_name = try_current_instance() 

41 if instance_name is None: 

42 raise InvalidArgumentError("Instance name not set") 

43 return instance_name 

44 

45 

46@contextmanager 

47def service_context(service_name: Optional[str]) -> Iterator[None]: 

48 token = _service_name.set(service_name) 

49 try: 

50 yield 

51 finally: 

52 _service_name.reset(token) 

53 

54 

55def try_current_service() -> Optional[str]: 

56 return _service_name.get() 

57 

58 

59def current_service() -> str: 

60 service_name = try_current_service() 

61 if service_name is None: 

62 raise InvalidArgumentError("Service name not set") 

63 return service_name 

64 

65 

66@contextmanager 

67def method_context(method_name: Optional[str]) -> Iterator[None]: 

68 token = _method_name.set(method_name) 

69 try: 

70 yield 

71 finally: 

72 _method_name.reset(token) 

73 

74 

75def try_current_method() -> Optional[str]: 

76 return _method_name.get() 

77 

78 

79def current_method() -> str: 

80 method_name = try_current_method() 

81 if method_name is None: 

82 raise InvalidArgumentError("Method name not set") 

83 return method_name