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.
- ni_grpc_extensions-0.1.0.dev0/PKG-INFO +52 -0
- ni_grpc_extensions-0.1.0.dev0/README.md +20 -0
- ni_grpc_extensions-0.1.0.dev0/pyproject.toml +99 -0
- ni_grpc_extensions-0.1.0.dev0/src/ni_grpc_extensions/__init__.py +1 -0
- ni_grpc_extensions-0.1.0.dev0/src/ni_grpc_extensions/_tracelogging.py +210 -0
- ni_grpc_extensions-0.1.0.dev0/src/ni_grpc_extensions/channelpool.py +116 -0
- ni_grpc_extensions-0.1.0.dev0/src/ni_grpc_extensions/loggers.py +506 -0
|
@@ -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: ...
|