ni-grpc-extensions 0.1.0.dev0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,52 @@
1
+ Metadata-Version: 2.1
2
+ Name: ni-grpc-extensions
3
+ Version: 0.1.0.dev0
4
+ Summary: gRPC Extensions
5
+ Home-page: https://github.com/ni/ni-apis-python
6
+ License: MIT
7
+ Keywords: grpc,extensions,channel_pool
8
+ Author: NI
9
+ Author-email: opensource@ni.com
10
+ Maintainer: Brad Keryan
11
+ Maintainer-email: brad.keryan@emerson.com
12
+ Requires-Python: >=3.9,<4.0
13
+ Classifier: Development Status :: 3 - Alpha
14
+ Classifier: Intended Audience :: Developers
15
+ Classifier: Intended Audience :: Manufacturing
16
+ Classifier: Intended Audience :: Science/Research
17
+ Classifier: License :: OSI Approved :: MIT License
18
+ Classifier: Operating System :: Microsoft :: Windows
19
+ Classifier: Operating System :: POSIX
20
+ Classifier: Programming Language :: Python :: 3
21
+ Classifier: Programming Language :: Python :: 3.9
22
+ Classifier: Programming Language :: Python :: 3.10
23
+ Classifier: Programming Language :: Python :: 3.11
24
+ Classifier: Programming Language :: Python :: 3.12
25
+ Classifier: Programming Language :: Python :: 3.13
26
+ Classifier: Programming Language :: Python :: Implementation :: CPython
27
+ Requires-Dist: grpcio (>=1.49.0,<2.0)
28
+ Requires-Dist: traceloggingdynamic (>=1.0) ; sys_platform == "win32"
29
+ Project-URL: Repository, https://github.com/ni/ni-apis-python
30
+ Description-Content-Type: text/markdown
31
+
32
+ # Table of Contents
33
+
34
+ - [Table of Contents](#table-of-contents)
35
+ - [About](#about)
36
+ - [Operating System Support](#operating-system-support)
37
+ - [Python Version Support](#python-version-support)
38
+
39
+ # About
40
+
41
+ `ni-grpc-extensions` is a Python package that provides channel pool handling for gRPC communication.
42
+
43
+ NI created and supports this package.
44
+
45
+ ## Operating System Support
46
+
47
+ ``ni-grpc-extensions` supports Windows and Linux operating systems.
48
+
49
+ ## Python Version Support
50
+
51
+ ``ni-grpc-extensions` supports CPython 3.9+.
52
+
@@ -0,0 +1,20 @@
1
+ # Table of Contents
2
+
3
+ - [Table of Contents](#table-of-contents)
4
+ - [About](#about)
5
+ - [Operating System Support](#operating-system-support)
6
+ - [Python Version Support](#python-version-support)
7
+
8
+ # About
9
+
10
+ `ni-grpc-extensions` is a Python package that provides channel pool handling for gRPC communication.
11
+
12
+ NI created and supports this package.
13
+
14
+ ## Operating System Support
15
+
16
+ ``ni-grpc-extensions` supports Windows and Linux operating systems.
17
+
18
+ ## Python Version Support
19
+
20
+ ``ni-grpc-extensions` supports CPython 3.9+.
@@ -0,0 +1,99 @@
1
+ [tool.poetry]
2
+ name = "ni-grpc-extensions"
3
+ version = "0.1.0-dev0"
4
+ license = "MIT"
5
+ description = "gRPC Extensions"
6
+ authors = ["NI <opensource@ni.com>"]
7
+ maintainers = [
8
+ "Brad Keryan <brad.keryan@emerson.com>",
9
+ "Joe Friedrichsen <joe.friedrichsen@emerson.com>",
10
+ ]
11
+ readme = "README.md"
12
+ repository = "https://github.com/ni/ni-apis-python"
13
+ keywords = ["grpc", "extensions", "channel_pool"]
14
+ classifiers = [
15
+ "Development Status :: 3 - Alpha",
16
+ "Intended Audience :: Developers",
17
+ "Intended Audience :: Manufacturing",
18
+ "Intended Audience :: Science/Research",
19
+ "License :: OSI Approved :: MIT License",
20
+ "Operating System :: Microsoft :: Windows",
21
+ "Operating System :: POSIX",
22
+ "Programming Language :: Python :: 3",
23
+ "Programming Language :: Python :: 3.9",
24
+ "Programming Language :: Python :: 3.10",
25
+ "Programming Language :: Python :: 3.11",
26
+ "Programming Language :: Python :: 3.12",
27
+ "Programming Language :: Python :: 3.13",
28
+ "Programming Language :: Python :: Implementation :: CPython",
29
+ ]
30
+ packages = [{include = "ni_grpc_extensions", from = "src"}]
31
+
32
+ [tool.poetry.dependencies]
33
+ python = "^3.9"
34
+ grpcio = {version=">=1.49.0,<2.0"}
35
+ traceloggingdynamic = { version = ">=1.0", platform = "win32" }
36
+
37
+ [tool.poetry.group.dev.dependencies]
38
+ types-grpcio = ">=1.0"
39
+ # Install traceloggingdynamic on Linux for type checking.
40
+ traceloggingdynamic = { version = ">=1.0", platform = "linux" }
41
+
42
+ [tool.poetry.group.lint.dependencies]
43
+ bandit = { version = ">=1.7", extras = ["toml"] }
44
+ ni-python-styleguide = ">=0.4.1"
45
+ mypy = ">=1.0"
46
+ pyright = { version = ">=1.1.400", extras = ["nodejs"] }
47
+
48
+ [tool.poetry.group.test.dependencies]
49
+ pytest = ">=7.2"
50
+ pytest-cov = ">=4.0"
51
+ pytest-doctestplus = ">=1.4"
52
+ pytest-mock = ">=3.0"
53
+ ni-measurementlink-discovery-v1-proto = { path = "../ni.measurementlink.discovery.v1.proto" }
54
+
55
+ [tool.poetry.group.docs]
56
+ optional = true
57
+
58
+ [tool.poetry.group.docs.dependencies]
59
+ # The latest Sphinx requires a recent Python version.
60
+ Sphinx = { version = ">=8.2", python = "^3.11" }
61
+ sphinx-rtd-theme = ">=1.0.0"
62
+ sphinx-autoapi = ">=1.8.4"
63
+ m2r2 = ">=0.3.2"
64
+ toml = ">=0.10.2"
65
+
66
+ [build-system]
67
+ requires = ["poetry-core>=1.8"]
68
+ build-backend = "poetry.core.masonry.api"
69
+
70
+
71
+ [tool.bandit]
72
+ skips = [
73
+ "B101", # assert_used
74
+ ]
75
+
76
+ [tool.ni-python-styleguide]
77
+ extend_exclude = "docs"
78
+
79
+ [tool.black]
80
+ extend-exclude = 'docs/'
81
+ line-length = 100
82
+
83
+ [tool.mypy]
84
+ mypy_path = "src"
85
+ files = "."
86
+ namespace_packages = true
87
+ strict = true
88
+ explicit_package_bases = true
89
+ exclude = ["docs"]
90
+
91
+ [tool.pyright]
92
+ include = ["src/", "tests/"]
93
+
94
+ [tool.pytest.ini_options]
95
+ addopts = "--doctest-modules --doctest-plus --strict-markers"
96
+ testpaths = ["tests"]
97
+
98
+ [tool.isort]
99
+ known_first_party = ["ni_grpc_extensions"]
@@ -0,0 +1 @@
1
+ """gRPC extensions."""
@@ -0,0 +1,210 @@
1
+ from __future__ import annotations
2
+
3
+ import ctypes
4
+ import sys
5
+ import uuid
6
+ from typing import TYPE_CHECKING
7
+
8
+ if sys.platform == "win32":
9
+ try:
10
+ import traceloggingdynamic
11
+
12
+ _event_provider: traceloggingdynamic.Provider | None = traceloggingdynamic.Provider(
13
+ b"NI-Grpc-Python"
14
+ )
15
+ except ImportError:
16
+ _event_provider = None
17
+ else:
18
+ if TYPE_CHECKING:
19
+ import traceloggingdynamic
20
+
21
+ _event_provider = None
22
+
23
+ _LEVEL_LOG_ALWAYS = 0
24
+ _LEVEL_CRITICAL = 1
25
+ _LEVEL_ERROR = 3
26
+ _LEVEL_WARNING = 3
27
+ _LEVEL_INFO = 4
28
+ _LEVEL_VERBOSE = 5
29
+
30
+ _OPCODE_INFO = 0
31
+ _OPCODE_START = 1
32
+ _OPCODE_STOP = 2
33
+
34
+ _KEYWORD_NONE = 0
35
+ _KEYWORD_GRPC = 1 << 0
36
+
37
+ _TASK_GRPC_CLIENT_CALL = 1
38
+ _TASK_GRPC_SERVER_CALL = 2
39
+
40
+
41
+ if sys.platform == "win32":
42
+ # 0x00000800 = LOAD_LIBRARY_SEARCH_SYSTEM32 (Win8 or later)
43
+ _eventing_dll = ctypes.WinDLL("api-ms-win-eventing-provider-l1-1-0.dll", mode=0x00000800)
44
+
45
+ _EventActivityIdControl = _eventing_dll.EventActivityIdControl
46
+ _EventActivityIdControl.restype = ctypes.c_uint32
47
+ _EventActivityIdControl.argtypes = (ctypes.c_uint32, ctypes.c_void_p)
48
+
49
+ _EVENT_ACTIVITY_CTRL_GET_ID = 1
50
+ _EVENT_ACTIVITY_CTRL_SET_ID = 2
51
+ _EVENT_ACTIVITY_CTRL_CREATE_ID = 3
52
+ _EVENT_ACTIVITY_CTRL_GET_SET_ID = 4
53
+ _EVENT_ACTIVITY_CTRL_CREATE_SET_ID = 5
54
+
55
+ def _create_activity_id() -> uuid.UUID:
56
+ activity_bytes = (ctypes.c_byte * 16)()
57
+ status = _EventActivityIdControl(
58
+ _EVENT_ACTIVITY_CTRL_CREATE_ID, ctypes.pointer(activity_bytes)
59
+ )
60
+ if status != 0:
61
+ raise OSError("EventActivityIdControl error", status)
62
+ return uuid.UUID(bytes_le=bytes(activity_bytes))
63
+
64
+ def _get_current_thread_activity_id() -> uuid.UUID:
65
+ activity_bytes = (ctypes.c_byte * 16)()
66
+ status = _EventActivityIdControl(
67
+ _EVENT_ACTIVITY_CTRL_GET_ID, ctypes.pointer(activity_bytes)
68
+ )
69
+ if status != 0:
70
+ raise OSError("EventActivityIdControl error", status)
71
+ return uuid.UUID(bytes_le=bytes(activity_bytes))
72
+
73
+ else:
74
+
75
+ def _create_activity_id() -> uuid.UUID:
76
+ return uuid.uuid4()
77
+
78
+ def _get_current_thread_activity_id() -> uuid.UUID:
79
+ return uuid.UUID()
80
+
81
+
82
+ def is_enabled() -> bool:
83
+ """Queries whether the event provider is enabled."""
84
+ return _event_provider is not None and _event_provider.is_enabled()
85
+
86
+
87
+ def log_grpc_client_call_start(method_name: str) -> uuid.UUID | None:
88
+ """Log when starting a gRPC client call."""
89
+ if _event_provider and _event_provider.is_enabled(level=_LEVEL_INFO, keyword=_KEYWORD_GRPC):
90
+ eb = traceloggingdynamic.EventBuilder() # pyright: ignore[reportPossiblyUnboundVariable]
91
+ eb.reset(
92
+ b"GrpcClientCall",
93
+ level=_LEVEL_INFO,
94
+ keyword=_KEYWORD_GRPC,
95
+ opcode=_OPCODE_START,
96
+ task=_TASK_GRPC_CLIENT_CALL,
97
+ )
98
+ eb.add_str8(b"FormattedMessage", "gRPC client call starting: " + method_name)
99
+ activity_id = _create_activity_id()
100
+ related_activity_id = _get_current_thread_activity_id()
101
+ _event_provider.write(eb, activity_id, related_activity_id)
102
+ return activity_id
103
+ else:
104
+ return None
105
+
106
+
107
+ def log_grpc_client_call_stop(method_name: str, activity_id: uuid.UUID | None = None) -> None:
108
+ """Log when a gRPC client call has completed."""
109
+ if _event_provider and _event_provider.is_enabled(level=_LEVEL_INFO, keyword=_KEYWORD_GRPC):
110
+ eb = traceloggingdynamic.EventBuilder() # pyright: ignore[reportPossiblyUnboundVariable]
111
+ eb.reset(
112
+ b"GrpcClientCall",
113
+ level=_LEVEL_INFO,
114
+ keyword=_KEYWORD_GRPC,
115
+ opcode=_OPCODE_STOP,
116
+ task=_TASK_GRPC_CLIENT_CALL,
117
+ )
118
+ eb.add_str8(b"FormattedMessage", "gRPC client call complete: " + method_name)
119
+ _event_provider.write(eb, activity_id)
120
+
121
+
122
+ def log_grpc_client_call_streaming_request(method_name: str) -> None:
123
+ """Log when a gRPC client call is sending a client-streaming request."""
124
+ if _event_provider and _event_provider.is_enabled(level=_LEVEL_INFO, keyword=_KEYWORD_GRPC):
125
+ eb = traceloggingdynamic.EventBuilder() # pyright: ignore[reportPossiblyUnboundVariable]
126
+ eb.reset(
127
+ b"GrpcClientCallStreamingRequest",
128
+ level=_LEVEL_INFO,
129
+ keyword=_KEYWORD_GRPC,
130
+ opcode=_OPCODE_INFO,
131
+ )
132
+ eb.add_str8(b"FormattedMessage", "gRPC client call streaming request: " + method_name)
133
+ _event_provider.write(eb)
134
+
135
+
136
+ def log_grpc_client_call_streaming_response(method_name: str) -> None:
137
+ """Log when a gRPC client call has received a server-streaming response."""
138
+ if _event_provider and _event_provider.is_enabled(level=_LEVEL_INFO, keyword=_KEYWORD_GRPC):
139
+ eb = traceloggingdynamic.EventBuilder() # pyright: ignore[reportPossiblyUnboundVariable]
140
+ eb.reset(
141
+ b"GrpcClientCallStreamingResponse",
142
+ level=_LEVEL_INFO,
143
+ keyword=_KEYWORD_GRPC,
144
+ opcode=_OPCODE_INFO,
145
+ )
146
+ eb.add_str8(b"FormattedMessage", "gRPC client call streaming response: " + method_name)
147
+ _event_provider.write(eb)
148
+
149
+
150
+ def log_grpc_server_call_start(method_name: str) -> uuid.UUID | None:
151
+ """Log when starting a gRPC server call."""
152
+ if _event_provider and _event_provider.is_enabled(level=_LEVEL_INFO, keyword=_KEYWORD_GRPC):
153
+ eb = traceloggingdynamic.EventBuilder() # pyright: ignore[reportPossiblyUnboundVariable]
154
+ eb.reset(
155
+ b"GrpcServerCall",
156
+ level=_LEVEL_INFO,
157
+ keyword=_KEYWORD_GRPC,
158
+ opcode=_OPCODE_START,
159
+ task=_TASK_GRPC_SERVER_CALL,
160
+ )
161
+ eb.add_str8(b"FormattedMessage", "gRPC server call starting: " + method_name)
162
+ activity_id = _create_activity_id()
163
+ related_activity_id = _get_current_thread_activity_id()
164
+ _event_provider.write(eb, activity_id, related_activity_id)
165
+ return activity_id
166
+ else:
167
+ return None
168
+
169
+
170
+ def log_grpc_server_call_stop(method_name: str, activity_id: uuid.UUID | None = None) -> None:
171
+ """Log when a gRPC server call has completed."""
172
+ if _event_provider and _event_provider.is_enabled(level=_LEVEL_INFO, keyword=_KEYWORD_GRPC):
173
+ eb = traceloggingdynamic.EventBuilder() # pyright: ignore[reportPossiblyUnboundVariable]
174
+ eb.reset(
175
+ b"GrpcServerCall",
176
+ level=_LEVEL_INFO,
177
+ keyword=_KEYWORD_GRPC,
178
+ opcode=_OPCODE_STOP,
179
+ task=_TASK_GRPC_SERVER_CALL,
180
+ )
181
+ eb.add_str8(b"FormattedMessage", "gRPC server call complete: " + method_name)
182
+ _event_provider.write(eb, activity_id)
183
+
184
+
185
+ def log_grpc_server_call_streaming_request(method_name: str) -> None:
186
+ """Log when a gRPC server call is sending a server-streaming request."""
187
+ if _event_provider and _event_provider.is_enabled(level=_LEVEL_INFO, keyword=_KEYWORD_GRPC):
188
+ eb = traceloggingdynamic.EventBuilder() # pyright: ignore[reportPossiblyUnboundVariable]
189
+ eb.reset(
190
+ b"GrpcServerCallStreamingRequest",
191
+ level=_LEVEL_INFO,
192
+ keyword=_KEYWORD_GRPC,
193
+ opcode=_OPCODE_INFO,
194
+ )
195
+ eb.add_str8(b"FormattedMessage", "gRPC server call streaming request: " + method_name)
196
+ _event_provider.write(eb)
197
+
198
+
199
+ def log_grpc_server_call_streaming_response(method_name: str) -> None:
200
+ """Log when a gRPC server call has received a server-streaming response."""
201
+ if _event_provider and _event_provider.is_enabled(level=_LEVEL_INFO, keyword=_KEYWORD_GRPC):
202
+ eb = traceloggingdynamic.EventBuilder() # pyright: ignore[reportPossiblyUnboundVariable]
203
+ eb.reset(
204
+ b"GrpcServerCallStreamingResponse",
205
+ level=_LEVEL_INFO,
206
+ keyword=_KEYWORD_GRPC,
207
+ opcode=_OPCODE_INFO,
208
+ )
209
+ eb.add_str8(b"FormattedMessage", "gRPC server call streaming response: " + method_name)
210
+ _event_provider.write(eb)
@@ -0,0 +1,116 @@
1
+ """gRPC channel pool."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import ipaddress
6
+ import re
7
+ import sys
8
+ from threading import Lock
9
+ from types import TracebackType
10
+ from typing import TYPE_CHECKING, Literal
11
+ from urllib.parse import urlparse
12
+
13
+ import grpc
14
+
15
+ from ni_grpc_extensions.loggers import ClientLogger
16
+
17
+ if TYPE_CHECKING:
18
+ if sys.version_info >= (3, 11):
19
+ from typing import Self
20
+ else:
21
+ from typing_extensions import Self
22
+
23
+
24
+ class GrpcChannelPool:
25
+ """Class that manages :any:`grpc.Channel` lifetimes."""
26
+
27
+ def __init__(self) -> None:
28
+ """Initialize the :class:`GrpcChannelPool` object."""
29
+ self._lock: Lock = Lock()
30
+ self._channel_cache: dict[str, grpc.Channel] = {}
31
+
32
+ def __enter__(self: Self) -> Self:
33
+ """Enter the runtime context of the GrpcChannelPool."""
34
+ return self
35
+
36
+ def __exit__(
37
+ self,
38
+ exc_type: type[BaseException] | None,
39
+ exc_val: BaseException | None,
40
+ traceback: TracebackType | None,
41
+ ) -> Literal[False]:
42
+ """Exit the runtime context of the GrpcChannelPool."""
43
+ self.close()
44
+ return False
45
+
46
+ def get_channel(self, target: str) -> grpc.Channel:
47
+ """Return a gRPC channel.
48
+
49
+ Args:
50
+ target (str): The server address
51
+
52
+ """
53
+ new_channel = None
54
+ with self._lock:
55
+ if target not in self._channel_cache:
56
+ self._lock.release()
57
+ new_channel = self._create_channel(target)
58
+ self._lock.acquire()
59
+ if target not in self._channel_cache:
60
+ self._channel_cache[target] = new_channel
61
+ new_channel = None
62
+ channel = self._channel_cache[target]
63
+
64
+ # Close new_channel if it was not stored in _channel_cache.
65
+ if new_channel is not None:
66
+ new_channel.close()
67
+
68
+ return channel
69
+
70
+ def close(self) -> None:
71
+ """Close channels opened by get_channel()."""
72
+ with self._lock:
73
+ for channel in self._channel_cache.values():
74
+ channel.close()
75
+ self._channel_cache.clear()
76
+
77
+ def _create_channel(self, target: str) -> grpc.Channel:
78
+ options = [
79
+ ("grpc.max_receive_message_length", -1),
80
+ ("grpc.max_send_message_length", -1),
81
+ ]
82
+ if self._is_local(target):
83
+ options.append(("grpc.enable_http_proxy", 0))
84
+ channel = grpc.insecure_channel(target, options)
85
+ if ClientLogger.is_enabled():
86
+ channel = grpc.intercept_channel(channel, ClientLogger())
87
+ return channel
88
+
89
+ def _is_local(self, target: str) -> bool:
90
+ hostname = ""
91
+ # First, check if the target string is in URL format
92
+ parse_result = urlparse(target)
93
+ if parse_result.scheme and parse_result.hostname and parse_result.port:
94
+ hostname = parse_result.hostname
95
+ else:
96
+ # Next, check for target string in <host_name>:<port> format
97
+ match = re.match(r"^(.*):(\d+)$", target)
98
+ if match:
99
+ hostname = match.group(1)
100
+
101
+ if not hostname:
102
+ return False
103
+ if hostname == "localhost" or hostname == "LOCALHOST":
104
+ return True
105
+
106
+ # IPv6 addresses don't support parsing with leading/trailing brackets
107
+ # so we need to remove them.
108
+ match = re.match(r"^\[(.*)\]$", hostname)
109
+ if match:
110
+ hostname = match.group(1)
111
+
112
+ try:
113
+ address = ipaddress.ip_address(hostname)
114
+ return address.is_loopback
115
+ except ValueError:
116
+ return False
@@ -0,0 +1,506 @@
1
+ """gRPC logging interceptors."""
2
+
3
+ # Suppress mypy errors related to interceptor base classes.
4
+ # https://github.com/grpc/grpc/issues/40550
5
+ # https://github.com/python/typeshed/issues/14641
6
+ # mypy: allow-any-generics
7
+
8
+ from __future__ import annotations
9
+
10
+ import abc
11
+ import functools
12
+ import logging
13
+ import sys
14
+ import threading
15
+ import time
16
+ from collections.abc import Iterator
17
+ from types import TracebackType
18
+ from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar
19
+
20
+ import grpc
21
+
22
+ from ni_grpc_extensions import _tracelogging
23
+
24
+ if TYPE_CHECKING:
25
+ if sys.version_info >= (3, 11):
26
+ from typing import Self
27
+ else:
28
+ from typing_extensions import Self
29
+
30
+ _logger = logging.getLogger(__name__)
31
+
32
+
33
+ class ClientLogger(
34
+ grpc.UnaryUnaryClientInterceptor,
35
+ grpc.UnaryStreamClientInterceptor,
36
+ grpc.StreamUnaryClientInterceptor,
37
+ grpc.StreamStreamClientInterceptor,
38
+ ):
39
+ """Intercepts gRPC client calls and logs them for debugging."""
40
+
41
+ @classmethod
42
+ def is_enabled(cls) -> bool:
43
+ """Indicates whether gRPC client call logging is enabled for the current log level."""
44
+ return _ClientCallLogger.is_enabled()
45
+
46
+ def intercept_unary_unary(
47
+ self,
48
+ continuation: Callable[
49
+ [grpc.ClientCallDetails, grpc._TRequest], grpc._CallFuture[grpc._TResponse]
50
+ ],
51
+ client_call_details: grpc.ClientCallDetails,
52
+ request: grpc._TRequest,
53
+ ) -> grpc._CallFuture[grpc._TResponse]:
54
+ """Intercept and log a unary call."""
55
+ if _ClientCallLogger.is_enabled():
56
+ call_logger = _ClientCallLogger(client_call_details.method)
57
+ try:
58
+ call = continuation(client_call_details, request)
59
+ return _LoggingResponseCallFuture(call_logger, call)
60
+ except Exception as e:
61
+ call_logger.close(e)
62
+ raise
63
+ else:
64
+ return continuation(client_call_details, request)
65
+
66
+ def intercept_unary_stream(
67
+ self,
68
+ continuation: Callable[
69
+ [grpc.ClientCallDetails, grpc._TRequest], grpc._CallIterator[grpc._TResponse]
70
+ ],
71
+ client_call_details: grpc.ClientCallDetails,
72
+ request: grpc._TRequest,
73
+ ) -> grpc._CallIterator[grpc._TResponse]:
74
+ """Intercept and log a server-streaming call."""
75
+ if _ClientCallLogger.is_enabled():
76
+ call_logger = _ClientCallLogger(client_call_details.method)
77
+ try:
78
+ call_iterator = continuation(client_call_details, request)
79
+ return _LoggingResponseCallIterator(call_logger, call_iterator)
80
+ except Exception as e:
81
+ call_logger.close(e)
82
+ raise
83
+ else:
84
+ return continuation(client_call_details, request)
85
+
86
+ def intercept_stream_unary(
87
+ self,
88
+ continuation: Callable[
89
+ [grpc.ClientCallDetails, Iterator[grpc._TRequest]], grpc._CallFuture[grpc._TResponse]
90
+ ],
91
+ client_call_details: grpc.ClientCallDetails,
92
+ request_iterator: Iterator[grpc._TRequest],
93
+ ) -> grpc._CallFuture[grpc._TResponse]:
94
+ """Intercept and log a client-streaming call."""
95
+ if _ClientCallLogger.is_enabled():
96
+ call_logger = _ClientCallLogger(client_call_details.method)
97
+ try:
98
+ call = continuation(
99
+ client_call_details, _LoggingRequestIterator(call_logger, request_iterator)
100
+ )
101
+ return _LoggingResponseCallFuture(call_logger, call)
102
+ except Exception as e:
103
+ call_logger.close(e)
104
+ raise
105
+ else:
106
+ return continuation(client_call_details, request_iterator)
107
+
108
+ def intercept_stream_stream(
109
+ self,
110
+ continuation: Callable[
111
+ [grpc.ClientCallDetails, Iterator[grpc._TRequest]], grpc._CallIterator[grpc._TResponse]
112
+ ],
113
+ client_call_details: grpc.ClientCallDetails,
114
+ request_iterator: Iterator[grpc._TRequest],
115
+ ) -> grpc._CallIterator[grpc._TResponse]:
116
+ """Intercept and log a bidirectional streaming call."""
117
+ if _ClientCallLogger.is_enabled():
118
+ call_logger = _ClientCallLogger(client_call_details.method)
119
+ try:
120
+ call_iterator = continuation(
121
+ client_call_details, _LoggingRequestIterator(call_logger, request_iterator)
122
+ )
123
+ return _LoggingResponseCallIterator(call_logger, call_iterator)
124
+ except Exception as e:
125
+ call_logger.close(e)
126
+ raise
127
+ else:
128
+ return continuation(client_call_details, request_iterator)
129
+
130
+
131
+ class ServerLogger(grpc.ServerInterceptor):
132
+ """Intercepts gRPC server calls and logs them for debugging."""
133
+
134
+ @classmethod
135
+ def is_enabled(cls) -> bool:
136
+ """Indicates whether gRPC client call logging is enabled for the current log level."""
137
+ return _ServerCallLogger.is_enabled()
138
+
139
+ def intercept_service(
140
+ self,
141
+ continuation: Callable[
142
+ [grpc.HandlerCallDetails], grpc.RpcMethodHandler[grpc._TRequest, grpc._TResponse] | None
143
+ ],
144
+ handler_call_details: grpc.HandlerCallDetails,
145
+ ) -> grpc.RpcMethodHandler[grpc._TRequest, grpc._TResponse] | None:
146
+ """Intercept and log a server call."""
147
+ if _ServerCallLogger.is_enabled():
148
+ call_logger = _ServerCallLogger(handler_call_details.method)
149
+ handler = continuation(handler_call_details)
150
+ if handler is None:
151
+ return handler
152
+ elif handler.unary_unary:
153
+ return grpc.unary_unary_rpc_method_handler(
154
+ functools.partial(self._log_unary_unary, call_logger, handler.unary_unary),
155
+ handler.request_deserializer,
156
+ handler.response_serializer,
157
+ )
158
+ elif handler.unary_stream:
159
+ return grpc.unary_stream_rpc_method_handler(
160
+ functools.partial(self._log_unary_stream, call_logger, handler.unary_stream),
161
+ handler.request_deserializer,
162
+ handler.response_serializer,
163
+ )
164
+ elif handler.stream_unary:
165
+ return grpc.stream_unary_rpc_method_handler(
166
+ functools.partial(self._log_stream_unary, call_logger, handler.stream_unary),
167
+ handler.request_deserializer,
168
+ handler.response_serializer,
169
+ )
170
+ elif handler.stream_stream:
171
+ return grpc.stream_stream_rpc_method_handler(
172
+ functools.partial(self._log_stream_stream, call_logger, handler.stream_stream),
173
+ handler.request_deserializer,
174
+ handler.response_serializer,
175
+ )
176
+ else:
177
+ raise RuntimeError("Invalid RpcMethodHandler")
178
+ else:
179
+ return continuation(handler_call_details)
180
+
181
+ def _log_unary_unary(
182
+ self,
183
+ call_logger: _CallLogger,
184
+ handler_function: Callable[[grpc._TRequest, grpc.ServicerContext], grpc._TResponse],
185
+ request: grpc._TRequest,
186
+ context: grpc.ServicerContext,
187
+ ) -> grpc._TResponse:
188
+ with call_logger:
189
+ return handler_function(request, context)
190
+
191
+ def _log_unary_stream(
192
+ self,
193
+ call_logger: _CallLogger,
194
+ handler_function: Callable[
195
+ [grpc._TRequest, grpc.ServicerContext], Iterator[grpc._TResponse]
196
+ ],
197
+ request: grpc._TRequest,
198
+ context: grpc.ServicerContext,
199
+ ) -> Iterator[grpc._TResponse]:
200
+ try:
201
+ return _LoggingResponseIterator(call_logger, handler_function(request, context))
202
+ except Exception as e:
203
+ call_logger.close(e)
204
+ raise
205
+
206
+ def _log_stream_unary(
207
+ self,
208
+ call_logger: _CallLogger,
209
+ handler_function: Callable[
210
+ [Iterator[grpc._TRequest], grpc.ServicerContext], grpc._TResponse
211
+ ],
212
+ request_iterator: Iterator[grpc._TRequest],
213
+ context: grpc.ServicerContext,
214
+ ) -> grpc._TResponse:
215
+ with call_logger:
216
+ return handler_function(_LoggingRequestIterator(call_logger, request_iterator), context)
217
+
218
+ def _log_stream_stream(
219
+ self,
220
+ call_logger: _CallLogger,
221
+ handler_function: Callable[
222
+ [Iterator[grpc._TRequest], grpc.ServicerContext], Iterator[grpc._TResponse]
223
+ ],
224
+ request_iterator: Iterator[grpc._TRequest],
225
+ context: grpc.ServicerContext,
226
+ ) -> Iterator[grpc._TResponse]:
227
+ try:
228
+ return _LoggingResponseIterator(
229
+ call_logger,
230
+ handler_function(_LoggingRequestIterator(call_logger, request_iterator), context),
231
+ )
232
+ except Exception as e:
233
+ call_logger.close(e)
234
+ raise
235
+
236
+
237
+ class _CallLogger(abc.ABC):
238
+ """Logs a single call."""
239
+
240
+ __slots__ = ["_closed"]
241
+
242
+ # As of 2023, the Python stdlib doesn't support atomic operations, so use a global lock to
243
+ # atomically update self._closed.
244
+ _lock = threading.Lock()
245
+
246
+ def __init__(self) -> None:
247
+ self._closed = False
248
+
249
+ def close(self, exception: BaseException | None = None) -> None:
250
+ # close() is idempotent to avoid duplicate logs.
251
+ with _CallLogger._lock:
252
+ if self._closed:
253
+ return
254
+ self._closed = True
255
+
256
+ self._close(exception)
257
+
258
+ @abc.abstractmethod
259
+ def _close(self, exception: BaseException | None) -> None:
260
+ raise NotImplementedError()
261
+
262
+ def __enter__(self: Self) -> Self:
263
+ return self
264
+
265
+ def __exit__(
266
+ self,
267
+ exc_type: type[BaseException] | None,
268
+ exc_val: BaseException | None,
269
+ traceback: TracebackType | None,
270
+ ) -> None:
271
+ self.close(exc_val)
272
+
273
+ @abc.abstractmethod
274
+ def log_streaming_request(self) -> None:
275
+ raise NotImplementedError()
276
+
277
+ @abc.abstractmethod
278
+ def log_streaming_response(self) -> None:
279
+ raise NotImplementedError()
280
+
281
+
282
+ class _ClientCallLogger(_CallLogger):
283
+ __slots__ = ["_method_name", "_activity_id"]
284
+
285
+ @classmethod
286
+ def is_enabled(cls) -> bool:
287
+ return _logger.isEnabledFor(logging.DEBUG) or _tracelogging.is_enabled()
288
+
289
+ def __init__(self, method_name: str) -> None:
290
+ super().__init__()
291
+ self._method_name = method_name
292
+ _logger.debug("gRPC client call starting: %s", self._method_name)
293
+ self._activity_id = _tracelogging.log_grpc_client_call_start(self._method_name)
294
+
295
+ def _close(self, exception: BaseException | None = None) -> None:
296
+ _logger.debug("gRPC client call complete: %s", self._method_name)
297
+ _tracelogging.log_grpc_client_call_stop(self._method_name, self._activity_id)
298
+
299
+ def log_streaming_request(self) -> None:
300
+ _logger.debug("gRPC client call streaming request: %s", self._method_name)
301
+ _tracelogging.log_grpc_client_call_streaming_request(self._method_name)
302
+
303
+ def log_streaming_response(self) -> None:
304
+ _logger.debug("gRPC client call streaming response: %s", self._method_name)
305
+ _tracelogging.log_grpc_client_call_streaming_response(self._method_name)
306
+
307
+
308
+ class _ServerCallLogger(_CallLogger):
309
+ __slots__ = ["_method_name", "_start_time", "_activity_id"]
310
+
311
+ @classmethod
312
+ def is_enabled(cls) -> bool:
313
+ return _logger.isEnabledFor(logging.INFO) or _tracelogging.is_enabled()
314
+
315
+ def __init__(self, method_name: str) -> None:
316
+ super().__init__()
317
+ self._method_name = method_name
318
+ self._start_time = time.perf_counter()
319
+ _logger.debug("gRPC server call starting: %s", self._method_name)
320
+ self._activity_id = _tracelogging.log_grpc_server_call_start(self._method_name)
321
+
322
+ def _close(self, exception: BaseException | None = None) -> None:
323
+ if _logger.isEnabledFor(logging.INFO):
324
+ # For production usage, log a concise one-line summary of the call, similar to what
325
+ # Serilog provides for ASP.NET Core services. Don't log exception details because
326
+ # grpcio's logging_pool already handles this.
327
+ elapsed_time = time.perf_counter() - self._start_time
328
+ _logger.info(
329
+ "gRPC server call %s responded %s in %.4f ms",
330
+ self._method_name,
331
+ str(_get_status_code(exception)).replace("StatusCode.", ""),
332
+ elapsed_time * 1000.0,
333
+ )
334
+ _logger.debug("gRPC server call complete: %s", self._method_name)
335
+ _tracelogging.log_grpc_server_call_stop(self._method_name, self._activity_id)
336
+
337
+ def log_streaming_request(self) -> None:
338
+ _logger.debug("gRPC server call streaming request: %s", self._method_name)
339
+ _tracelogging.log_grpc_server_call_streaming_request(self._method_name)
340
+
341
+ def log_streaming_response(self) -> None:
342
+ _logger.debug("gRPC server call streaming response: %s", self._method_name)
343
+ _tracelogging.log_grpc_server_call_streaming_response(self._method_name)
344
+
345
+
346
+ def _get_status_code(exception: BaseException | None) -> grpc.StatusCode:
347
+ if exception is None:
348
+ return grpc.StatusCode.OK
349
+ elif isinstance(exception, grpc.RpcError):
350
+ return exception.code()
351
+ else:
352
+ return grpc.StatusCode.UNKNOWN
353
+
354
+
355
+ _T = TypeVar("_T")
356
+
357
+
358
+ class _LoggingRequestIterator(Generic[_T]):
359
+ __slots__ = ["_call_logger", "_inner_iterator"]
360
+
361
+ def __init__(self, call_logger: _CallLogger, inner_iterator: Iterator[_T]) -> None:
362
+ self._call_logger = call_logger
363
+ self._inner_iterator = inner_iterator
364
+
365
+ def __iter__(self) -> Iterator[_T]:
366
+ return self
367
+
368
+ def __next__(self) -> _T:
369
+ request = next(self._inner_iterator)
370
+ self._call_logger.log_streaming_request()
371
+ return request
372
+
373
+
374
+ class _LoggingResponseIterator(Generic[_T]):
375
+ __slots__ = ["_call_logger", "_inner_iterator"]
376
+
377
+ def __init__(self, call_logger: _CallLogger, inner_iterator: Iterator[_T]) -> None:
378
+ self._call_logger = call_logger
379
+ self._inner_iterator = inner_iterator
380
+
381
+ def __iter__(self) -> Iterator[_T]:
382
+ return self
383
+
384
+ def __next__(self) -> _T:
385
+ # For server-streaming and bidirectional RPCs, the call is complete when the response
386
+ # stream is closed or it throws an exception.
387
+ try:
388
+ response = next(self._inner_iterator)
389
+ self._call_logger.log_streaming_response()
390
+ return response
391
+ except StopIteration:
392
+ self._call_logger.close()
393
+ raise
394
+ except Exception as e:
395
+ self._call_logger.close(e)
396
+ raise
397
+
398
+
399
+ if TYPE_CHECKING:
400
+ # These types are marked as @type_check_only and can be used only in type annotations,
401
+ # so we redefine them as empty classes at run time. Pyright still warns due to the behavior
402
+ # described in https://github.com/microsoft/pyright/issues/9934
403
+ _CallFuture = grpc._CallFuture # pyright: ignore[reportGeneralTypeIssues]
404
+ _CallIterator = grpc._CallIterator # pyright: ignore[reportGeneralTypeIssues]
405
+ else:
406
+
407
+ class _CallFuture(Generic[_T]):
408
+ pass
409
+
410
+ class _CallIterator(Generic[_T]):
411
+ pass
412
+
413
+
414
+ @grpc.Call.register
415
+ @grpc.Future.register
416
+ class _LoggingResponseCallFuture(_CallFuture[_T]): # pyright: ignore[reportGeneralTypeIssues]
417
+ __slots__ = ["_call_logger", "_inner_call_future"]
418
+
419
+ def __init__(self, call_logger: _CallLogger, inner_call_future: grpc._CallFuture[_T]) -> None:
420
+ self._call_logger = call_logger
421
+ self._inner_call_future = inner_call_future
422
+
423
+ def __getattr__(self, name: str) -> Any:
424
+ return getattr(self._inner_call_future, name)
425
+
426
+ def result(self, timeout: float | None = None) -> _T:
427
+ # For unary and client-streaming RPCs, the call is complete when it returns a response or
428
+ # throws an exception.
429
+ with self._call_logger:
430
+ return self._inner_call_future.result(timeout)
431
+
432
+ def exception(self, timeout: float | None = None) -> Exception | None:
433
+ with self._call_logger:
434
+ return self._inner_call_future.exception(timeout)
435
+
436
+ def traceback(self, timeout: float | None = None) -> Any:
437
+ with self._call_logger:
438
+ return self._inner_call_future.traceback(timeout)
439
+
440
+ # At run time, __getattr__ forwards these to the inner CallFuture.
441
+ if TYPE_CHECKING:
442
+
443
+ def add_callback(self, callback: Callable[[], None]) -> bool: ...
444
+ def add_done_callback(
445
+ self, fn: Callable[[grpc.Future[grpc._TFutureValue]], None]
446
+ ) -> None: ...
447
+ def cancel(self) -> bool: ...
448
+ def cancelled(self) -> bool: ...
449
+ def code(self) -> grpc.StatusCode: ...
450
+ def details(self) -> str: ...
451
+ def done(self) -> bool: ...
452
+ def initial_metadata(self) -> grpc._Metadata: ...
453
+ def is_active(self) -> bool: ...
454
+ def running(self) -> bool: ...
455
+ def time_remaining(self) -> float: ...
456
+ def trailing_metadata(self) -> grpc._Metadata: ...
457
+
458
+
459
+ @grpc.Call.register
460
+ @grpc.Future.register
461
+ class _LoggingResponseCallIterator(_CallIterator[_T]): # pyright: ignore
462
+ __slots__ = ["_call_logger", "_inner_call_iterator"]
463
+
464
+ def __init__(
465
+ self, call_logger: _CallLogger, inner_call_iterator: grpc._CallIterator[_T]
466
+ ) -> None:
467
+ self._call_logger = call_logger
468
+ self._inner_call_iterator = inner_call_iterator
469
+
470
+ def __getattr__(self, name: str) -> Any:
471
+ return getattr(self._inner_call_iterator, name)
472
+
473
+ def __iter__(self) -> Iterator[_T]:
474
+ return self
475
+
476
+ def __next__(self) -> _T:
477
+ # For server-streaming and bidirectional RPCs, the call is complete when the response
478
+ # stream is closed or throws an exception.
479
+ try:
480
+ response = next(self._inner_call_iterator)
481
+ self._call_logger.log_streaming_response()
482
+ return response
483
+ except StopIteration:
484
+ self._call_logger.close()
485
+ raise
486
+ except Exception as e:
487
+ self._call_logger.close(e)
488
+ raise
489
+
490
+ # At run time, __getattr__ forwards these to the inner CallIterator.
491
+ if TYPE_CHECKING:
492
+
493
+ def add_callback(self, callback: Callable[[], None]) -> bool: ...
494
+ def add_done_callback(
495
+ self, fn: Callable[[grpc.Future[grpc._TFutureValue]], None]
496
+ ) -> None: ...
497
+ def cancel(self) -> bool: ...
498
+ def cancelled(self) -> bool: ...
499
+ def code(self) -> grpc.StatusCode: ...
500
+ def details(self) -> str: ...
501
+ def done(self) -> bool: ...
502
+ def initial_metadata(self) -> grpc._Metadata: ...
503
+ def is_active(self) -> bool: ...
504
+ def running(self) -> bool: ...
505
+ def time_remaining(self) -> float: ...
506
+ def trailing_metadata(self) -> grpc._Metadata: ...