valkey-glide 2.2.3__pp39-pypy39_pp73-macosx_10_7_x86_64.whl
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.
- glide/__init__.py +388 -0
- glide/async_commands/__init__.py +5 -0
- glide/async_commands/cluster_commands.py +1476 -0
- glide/async_commands/core.py +7818 -0
- glide/async_commands/ft.py +465 -0
- glide/async_commands/glide_json.py +1269 -0
- glide/async_commands/standalone_commands.py +1001 -0
- glide/glide.pyi +61 -0
- glide/glide.pypy39-pp73-darwin.so +0 -0
- glide/glide_client.py +821 -0
- glide/logger.py +97 -0
- glide/opentelemetry.py +185 -0
- glide/py.typed +0 -0
- glide_shared/__init__.py +330 -0
- glide_shared/commands/__init__.py +0 -0
- glide_shared/commands/batch.py +5997 -0
- glide_shared/commands/batch_options.py +261 -0
- glide_shared/commands/bitmap.py +320 -0
- glide_shared/commands/command_args.py +103 -0
- glide_shared/commands/core_options.py +407 -0
- glide_shared/commands/server_modules/ft_options/ft_aggregate_options.py +300 -0
- glide_shared/commands/server_modules/ft_options/ft_constants.py +84 -0
- glide_shared/commands/server_modules/ft_options/ft_create_options.py +423 -0
- glide_shared/commands/server_modules/ft_options/ft_profile_options.py +113 -0
- glide_shared/commands/server_modules/ft_options/ft_search_options.py +139 -0
- glide_shared/commands/server_modules/json_batch.py +820 -0
- glide_shared/commands/server_modules/json_options.py +93 -0
- glide_shared/commands/sorted_set.py +412 -0
- glide_shared/commands/stream.py +449 -0
- glide_shared/config.py +975 -0
- glide_shared/constants.py +124 -0
- glide_shared/exceptions.py +88 -0
- glide_shared/protobuf/command_request_pb2.py +56 -0
- glide_shared/protobuf/connection_request_pb2.py +56 -0
- glide_shared/protobuf/response_pb2.py +32 -0
- glide_shared/protobuf_codec.py +110 -0
- glide_shared/routes.py +161 -0
- valkey_glide-2.2.3.dist-info/METADATA +211 -0
- valkey_glide-2.2.3.dist-info/RECORD +40 -0
- valkey_glide-2.2.3.dist-info/WHEEL +4 -0
glide/glide_client.py
ADDED
|
@@ -0,0 +1,821 @@
|
|
|
1
|
+
# Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0
|
|
2
|
+
|
|
3
|
+
import sys
|
|
4
|
+
import threading
|
|
5
|
+
from typing import (
|
|
6
|
+
TYPE_CHECKING,
|
|
7
|
+
Any,
|
|
8
|
+
Awaitable,
|
|
9
|
+
Dict,
|
|
10
|
+
List,
|
|
11
|
+
Optional,
|
|
12
|
+
Set,
|
|
13
|
+
Tuple,
|
|
14
|
+
Union,
|
|
15
|
+
cast,
|
|
16
|
+
)
|
|
17
|
+
|
|
18
|
+
import anyio
|
|
19
|
+
import sniffio
|
|
20
|
+
from anyio import to_thread
|
|
21
|
+
from glide.glide import (
|
|
22
|
+
DEFAULT_TIMEOUT_IN_MILLISECONDS,
|
|
23
|
+
MAX_REQUEST_ARGS_LEN,
|
|
24
|
+
ClusterScanCursor,
|
|
25
|
+
create_leaked_bytes_vec,
|
|
26
|
+
create_otel_span,
|
|
27
|
+
drop_otel_span,
|
|
28
|
+
get_statistics,
|
|
29
|
+
start_socket_listener_external,
|
|
30
|
+
value_from_pointer,
|
|
31
|
+
)
|
|
32
|
+
from glide_shared.commands.command_args import ObjectType
|
|
33
|
+
from glide_shared.commands.core_options import PubSubMsg
|
|
34
|
+
from glide_shared.config import BaseClientConfiguration, ServerCredentials
|
|
35
|
+
from glide_shared.constants import (
|
|
36
|
+
DEFAULT_READ_BYTES_SIZE,
|
|
37
|
+
OK,
|
|
38
|
+
TEncodable,
|
|
39
|
+
TRequest,
|
|
40
|
+
TResult,
|
|
41
|
+
)
|
|
42
|
+
from glide_shared.exceptions import (
|
|
43
|
+
ClosingError,
|
|
44
|
+
ConfigurationError,
|
|
45
|
+
ConnectionError,
|
|
46
|
+
get_request_error_class,
|
|
47
|
+
)
|
|
48
|
+
from glide_shared.protobuf.command_request_pb2 import (
|
|
49
|
+
Command,
|
|
50
|
+
CommandRequest,
|
|
51
|
+
RefreshIamToken,
|
|
52
|
+
RequestType,
|
|
53
|
+
)
|
|
54
|
+
from glide_shared.protobuf.connection_request_pb2 import ConnectionRequest
|
|
55
|
+
from glide_shared.protobuf.response_pb2 import Response
|
|
56
|
+
from glide_shared.protobuf_codec import PartialMessageException, ProtobufCodec
|
|
57
|
+
from glide_shared.routes import Route, set_protobuf_route
|
|
58
|
+
|
|
59
|
+
from .async_commands.cluster_commands import ClusterCommands
|
|
60
|
+
from .async_commands.core import CoreCommands
|
|
61
|
+
from .async_commands.standalone_commands import StandaloneCommands
|
|
62
|
+
from .logger import Level as LogLevel
|
|
63
|
+
from .logger import Logger as ClientLogger
|
|
64
|
+
from .opentelemetry import OpenTelemetry
|
|
65
|
+
|
|
66
|
+
if sys.version_info >= (3, 11):
|
|
67
|
+
from typing import Self
|
|
68
|
+
else:
|
|
69
|
+
from typing_extensions import Self
|
|
70
|
+
|
|
71
|
+
if TYPE_CHECKING:
|
|
72
|
+
import asyncio
|
|
73
|
+
|
|
74
|
+
import trio
|
|
75
|
+
|
|
76
|
+
TTask = Union[asyncio.Task[None], trio.lowlevel.Task]
|
|
77
|
+
TFuture = Union[asyncio.Future[Any], "_CompatFuture"]
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
class _CompatFuture:
|
|
81
|
+
"""anyio shim for asyncio.Future-like functionality"""
|
|
82
|
+
|
|
83
|
+
def __init__(self) -> None:
|
|
84
|
+
self._is_done = anyio.Event()
|
|
85
|
+
self._result: Any = None
|
|
86
|
+
self._exception: Optional[Exception] = None
|
|
87
|
+
|
|
88
|
+
def set_result(self, result: Any) -> None:
|
|
89
|
+
self._result = result
|
|
90
|
+
self._is_done.set()
|
|
91
|
+
|
|
92
|
+
def set_exception(self, exception: Exception) -> None:
|
|
93
|
+
self._exception = exception
|
|
94
|
+
self._is_done.set()
|
|
95
|
+
|
|
96
|
+
def done(self) -> bool:
|
|
97
|
+
return self._is_done.is_set()
|
|
98
|
+
|
|
99
|
+
def __await__(self):
|
|
100
|
+
return self._is_done.wait().__await__()
|
|
101
|
+
|
|
102
|
+
def result(self) -> Any:
|
|
103
|
+
if self._exception:
|
|
104
|
+
raise self._exception
|
|
105
|
+
|
|
106
|
+
return self._result
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
def _get_new_future_instance() -> "TFuture":
|
|
110
|
+
if sniffio.current_async_library() == "asyncio":
|
|
111
|
+
import asyncio
|
|
112
|
+
|
|
113
|
+
return asyncio.get_running_loop().create_future()
|
|
114
|
+
|
|
115
|
+
# _CompatFuture is also compatible with asyncio, but is not as closely integrated
|
|
116
|
+
# into the asyncio event loop and thus introduces a noticeable performance
|
|
117
|
+
# degradation. so we only use it for trio
|
|
118
|
+
return _CompatFuture()
|
|
119
|
+
|
|
120
|
+
|
|
121
|
+
class BaseClient(CoreCommands):
|
|
122
|
+
def __init__(self, config: BaseClientConfiguration):
|
|
123
|
+
"""
|
|
124
|
+
To create a new client, use the `create` classmethod
|
|
125
|
+
"""
|
|
126
|
+
self.config: BaseClientConfiguration = config
|
|
127
|
+
self._available_futures: Dict[int, "TFuture"] = {}
|
|
128
|
+
self._available_callback_indexes: List[int] = list()
|
|
129
|
+
self._buffered_requests: List[TRequest] = list()
|
|
130
|
+
self._writer_lock = threading.Lock()
|
|
131
|
+
self.socket_path: Optional[str] = None
|
|
132
|
+
self._reader_task: Optional["TTask"] = None
|
|
133
|
+
self._is_closed: bool = False
|
|
134
|
+
self._pubsub_futures: List["TFuture"] = []
|
|
135
|
+
self._pubsub_lock = threading.Lock()
|
|
136
|
+
self._pending_push_notifications: List[Response] = list()
|
|
137
|
+
|
|
138
|
+
self._pending_tasks: Optional[Set[Awaitable[None]]] = None
|
|
139
|
+
"""asyncio-only to avoid gc on pending write tasks"""
|
|
140
|
+
|
|
141
|
+
def _create_task(self, task, *args, **kwargs):
|
|
142
|
+
"""framework agnostic free-floating task shim"""
|
|
143
|
+
framework = sniffio.current_async_library()
|
|
144
|
+
if framework == "trio":
|
|
145
|
+
from functools import partial
|
|
146
|
+
|
|
147
|
+
import trio
|
|
148
|
+
|
|
149
|
+
return trio.lowlevel.spawn_system_task(partial(task, **kwargs), *args)
|
|
150
|
+
elif framework == "asyncio":
|
|
151
|
+
import asyncio
|
|
152
|
+
|
|
153
|
+
# the asyncio event loop holds weak refs to tasks, so it's recommended to
|
|
154
|
+
# hold strong refs to them during their lifetime to prevent garbage
|
|
155
|
+
# collection
|
|
156
|
+
t = asyncio.create_task(task(*args, **kwargs))
|
|
157
|
+
|
|
158
|
+
if self._pending_tasks is None:
|
|
159
|
+
self._pending_tasks = set()
|
|
160
|
+
|
|
161
|
+
self._pending_tasks.add(t)
|
|
162
|
+
t.add_done_callback(self._pending_tasks.discard)
|
|
163
|
+
|
|
164
|
+
return t
|
|
165
|
+
|
|
166
|
+
raise RuntimeError(f"Unsupported async framework {framework}")
|
|
167
|
+
|
|
168
|
+
@classmethod
|
|
169
|
+
async def create(cls, config: BaseClientConfiguration) -> Self:
|
|
170
|
+
"""Creates a Glide client.
|
|
171
|
+
|
|
172
|
+
Args:
|
|
173
|
+
config (ClientConfiguration): The configuration options for the client, including cluster addresses,
|
|
174
|
+
authentication credentials, TLS settings, periodic checks, and Pub/Sub subscriptions.
|
|
175
|
+
|
|
176
|
+
Returns:
|
|
177
|
+
Self: A promise that resolves to a connected client instance.
|
|
178
|
+
|
|
179
|
+
Examples:
|
|
180
|
+
# Connecting to a Standalone Server
|
|
181
|
+
>>> from glide import GlideClientConfiguration, NodeAddress, GlideClient, ServerCredentials, BackoffStrategy
|
|
182
|
+
>>> config = GlideClientConfiguration(
|
|
183
|
+
... [
|
|
184
|
+
... NodeAddress('primary.example.com', 6379),
|
|
185
|
+
... NodeAddress('replica1.example.com', 6379),
|
|
186
|
+
... ],
|
|
187
|
+
... use_tls = True,
|
|
188
|
+
... database_id = 1,
|
|
189
|
+
... credentials = ServerCredentials(username = 'user1', password = 'passwordA'),
|
|
190
|
+
... reconnect_strategy = BackoffStrategy(num_of_retries = 5, factor = 1000, exponent_base = 2),
|
|
191
|
+
... pubsub_subscriptions = GlideClientConfiguration.PubSubSubscriptions(
|
|
192
|
+
... channels_and_patterns = {GlideClientConfiguration.PubSubChannelModes.Exact: {'updates'}},
|
|
193
|
+
... callback = lambda message,context : print(message),
|
|
194
|
+
... ),
|
|
195
|
+
... )
|
|
196
|
+
>>> client = await GlideClient.create(config)
|
|
197
|
+
|
|
198
|
+
# Connecting to a Cluster
|
|
199
|
+
>>> from glide import GlideClusterClientConfiguration, NodeAddress, GlideClusterClient,
|
|
200
|
+
... PeriodicChecksManualInterval
|
|
201
|
+
>>> config = GlideClusterClientConfiguration(
|
|
202
|
+
... [
|
|
203
|
+
... NodeAddress('address1.example.com', 6379),
|
|
204
|
+
... NodeAddress('address2.example.com', 6379),
|
|
205
|
+
... ],
|
|
206
|
+
... use_tls = True,
|
|
207
|
+
... periodic_checks = PeriodicChecksManualInterval(duration_in_sec = 30),
|
|
208
|
+
... credentials = ServerCredentials(username = 'user1', password = 'passwordA'),
|
|
209
|
+
... reconnect_strategy = BackoffStrategy(num_of_retries = 5, factor = 1000, exponent_base = 2),
|
|
210
|
+
... pubsub_subscriptions = GlideClusterClientConfiguration.PubSubSubscriptions(
|
|
211
|
+
... channels_and_patterns = {
|
|
212
|
+
... GlideClusterClientConfiguration.PubSubChannelModes.Exact: {'updates'},
|
|
213
|
+
... GlideClusterClientConfiguration.PubSubChannelModes.Sharded: {'sharded_channel'},
|
|
214
|
+
... },
|
|
215
|
+
... callback = lambda message,context : print(message),
|
|
216
|
+
... ),
|
|
217
|
+
... )
|
|
218
|
+
>>> client = await GlideClusterClient.create(config)
|
|
219
|
+
|
|
220
|
+
Remarks:
|
|
221
|
+
Use this static method to create and connect a client to a Valkey server.
|
|
222
|
+
The client will automatically handle connection establishment, including cluster topology discovery and
|
|
223
|
+
handling of authentication and TLS configurations.
|
|
224
|
+
|
|
225
|
+
- **Cluster Topology Discovery**: The client will automatically discover the cluster topology based
|
|
226
|
+
on the seed addresses provided.
|
|
227
|
+
- **Authentication**: If `ServerCredentials` are provided, the client will attempt to authenticate
|
|
228
|
+
using the specified username and password.
|
|
229
|
+
- **TLS**: If `use_tls` is set to `true`, the client will establish secure connections using TLS.
|
|
230
|
+
- **Periodic Checks**: The `periodic_checks` setting allows you to configure how often the client
|
|
231
|
+
checks for cluster topology changes.
|
|
232
|
+
- **Reconnection Strategy**: The `BackoffStrategy` settings define how the client will attempt to
|
|
233
|
+
reconnect in case of disconnections.
|
|
234
|
+
- **Pub/Sub Subscriptions**: Any channels or patterns specified in `PubSubSubscriptions` will be
|
|
235
|
+
subscribed to upon connection.
|
|
236
|
+
|
|
237
|
+
"""
|
|
238
|
+
config = config
|
|
239
|
+
self = cls(config)
|
|
240
|
+
|
|
241
|
+
init_event: threading.Event = threading.Event()
|
|
242
|
+
|
|
243
|
+
def init_callback(socket_path: Optional[str], err: Optional[str]):
|
|
244
|
+
if err is not None:
|
|
245
|
+
raise ClosingError(err)
|
|
246
|
+
elif socket_path is None:
|
|
247
|
+
raise ClosingError(
|
|
248
|
+
"Socket initialization error: Missing valid socket path."
|
|
249
|
+
)
|
|
250
|
+
else:
|
|
251
|
+
# Received socket path
|
|
252
|
+
self.socket_path = socket_path
|
|
253
|
+
init_event.set()
|
|
254
|
+
|
|
255
|
+
start_socket_listener_external(init_callback=init_callback)
|
|
256
|
+
|
|
257
|
+
# will log if the logger was created (wrapper or costumer) on info
|
|
258
|
+
# level or higher
|
|
259
|
+
ClientLogger.log(LogLevel.INFO, "connection info", "new connection established")
|
|
260
|
+
# Wait for the socket listener to complete its initialization
|
|
261
|
+
await to_thread.run_sync(init_event.wait)
|
|
262
|
+
# Create UDS connection
|
|
263
|
+
await self._create_uds_connection()
|
|
264
|
+
|
|
265
|
+
# Start the reader loop as a background task
|
|
266
|
+
self._reader_task = self._create_task(self._reader_loop)
|
|
267
|
+
|
|
268
|
+
# Set the client configurations
|
|
269
|
+
await self._set_connection_configurations()
|
|
270
|
+
|
|
271
|
+
return self
|
|
272
|
+
|
|
273
|
+
async def _create_uds_connection(self) -> None:
|
|
274
|
+
try:
|
|
275
|
+
# Open an UDS connection
|
|
276
|
+
with anyio.fail_after(DEFAULT_TIMEOUT_IN_MILLISECONDS):
|
|
277
|
+
self._stream = await anyio.connect_unix(
|
|
278
|
+
path=cast(str, self.socket_path)
|
|
279
|
+
)
|
|
280
|
+
except Exception as e:
|
|
281
|
+
raise ClosingError("Failed to create UDS connection") from e
|
|
282
|
+
|
|
283
|
+
async def close(self, err_message: Optional[str] = None) -> None:
|
|
284
|
+
"""
|
|
285
|
+
Terminate the client by closing all associated resources, including the socket and any active futures.
|
|
286
|
+
All open futures will be closed with an exception.
|
|
287
|
+
|
|
288
|
+
Args:
|
|
289
|
+
err_message (Optional[str]): If not None, this error message will be passed along with the exceptions when
|
|
290
|
+
closing all open futures.
|
|
291
|
+
Defaults to None.
|
|
292
|
+
"""
|
|
293
|
+
if not self._is_closed:
|
|
294
|
+
self._is_closed = True
|
|
295
|
+
err_message = "" if err_message is None else err_message
|
|
296
|
+
for response_future in self._available_futures.values():
|
|
297
|
+
if not response_future.done():
|
|
298
|
+
response_future.set_exception(ClosingError(err_message))
|
|
299
|
+
try:
|
|
300
|
+
self._pubsub_lock.acquire()
|
|
301
|
+
for pubsub_future in self._pubsub_futures:
|
|
302
|
+
if not pubsub_future.done():
|
|
303
|
+
pubsub_future.set_exception(ClosingError(err_message))
|
|
304
|
+
finally:
|
|
305
|
+
self._pubsub_lock.release()
|
|
306
|
+
|
|
307
|
+
await self._stream.aclose()
|
|
308
|
+
|
|
309
|
+
def _get_future(self, callback_idx: int) -> "TFuture":
|
|
310
|
+
response_future: "TFuture" = _get_new_future_instance()
|
|
311
|
+
self._available_futures.update({callback_idx: response_future})
|
|
312
|
+
return response_future
|
|
313
|
+
|
|
314
|
+
def _get_protobuf_conn_request(self) -> ConnectionRequest:
|
|
315
|
+
return self.config._create_a_protobuf_conn_request()
|
|
316
|
+
|
|
317
|
+
async def _set_connection_configurations(self) -> None:
|
|
318
|
+
conn_request = self._get_protobuf_conn_request()
|
|
319
|
+
response_future: "TFuture" = self._get_future(0)
|
|
320
|
+
self._create_write_task(conn_request)
|
|
321
|
+
await response_future
|
|
322
|
+
res = response_future.result()
|
|
323
|
+
if res is not OK:
|
|
324
|
+
raise ClosingError(res)
|
|
325
|
+
|
|
326
|
+
def _create_write_task(self, request: TRequest):
|
|
327
|
+
self._create_task(self._write_or_buffer_request, request)
|
|
328
|
+
|
|
329
|
+
async def _write_or_buffer_request(self, request: TRequest):
|
|
330
|
+
self._buffered_requests.append(request)
|
|
331
|
+
if self._writer_lock.acquire(False):
|
|
332
|
+
try:
|
|
333
|
+
while len(self._buffered_requests) > 0:
|
|
334
|
+
await self._write_buffered_requests_to_socket()
|
|
335
|
+
except Exception as e:
|
|
336
|
+
# trio system tasks cannot raise exceptions, so gracefully propagate
|
|
337
|
+
# any error to the pending future instead
|
|
338
|
+
callback_idx = (
|
|
339
|
+
request.callback_idx if isinstance(request, CommandRequest) else 0
|
|
340
|
+
)
|
|
341
|
+
res_future = self._available_futures.pop(callback_idx, None)
|
|
342
|
+
if res_future and not res_future.done():
|
|
343
|
+
res_future.set_exception(e)
|
|
344
|
+
else:
|
|
345
|
+
ClientLogger.log(
|
|
346
|
+
LogLevel.WARN,
|
|
347
|
+
"unhandled response error",
|
|
348
|
+
f"Unhandled response error for unknown request: {callback_idx}",
|
|
349
|
+
)
|
|
350
|
+
finally:
|
|
351
|
+
self._writer_lock.release()
|
|
352
|
+
|
|
353
|
+
async def _write_buffered_requests_to_socket(self) -> None:
|
|
354
|
+
requests = self._buffered_requests
|
|
355
|
+
self._buffered_requests = list()
|
|
356
|
+
b_arr = bytearray()
|
|
357
|
+
for request in requests:
|
|
358
|
+
ProtobufCodec.encode_delimited(b_arr, request)
|
|
359
|
+
try:
|
|
360
|
+
await self._stream.send(b_arr)
|
|
361
|
+
except (anyio.ClosedResourceError, anyio.EndOfStream):
|
|
362
|
+
raise ClosingError("The communication layer was unexpectedly closed.")
|
|
363
|
+
|
|
364
|
+
def _encode_arg(self, arg: TEncodable) -> bytes:
|
|
365
|
+
"""
|
|
366
|
+
Converts a string argument to bytes.
|
|
367
|
+
|
|
368
|
+
Args:
|
|
369
|
+
arg (str): An encodable argument.
|
|
370
|
+
|
|
371
|
+
Returns:
|
|
372
|
+
bytes: The encoded argument as bytes.
|
|
373
|
+
"""
|
|
374
|
+
if isinstance(arg, str):
|
|
375
|
+
# TODO: Allow passing different encoding options
|
|
376
|
+
return bytes(arg, encoding="utf8")
|
|
377
|
+
return arg
|
|
378
|
+
|
|
379
|
+
def _encode_and_sum_size(
|
|
380
|
+
self,
|
|
381
|
+
args_list: Optional[List[TEncodable]],
|
|
382
|
+
) -> Tuple[List[bytes], int]:
|
|
383
|
+
"""
|
|
384
|
+
Encodes the list and calculates the total memory size.
|
|
385
|
+
|
|
386
|
+
Args:
|
|
387
|
+
args_list (Optional[List[TEncodable]]): A list of strings to be converted to bytes.
|
|
388
|
+
If None or empty, returns ([], 0).
|
|
389
|
+
|
|
390
|
+
Returns:
|
|
391
|
+
int: The total memory size of the encoded arguments in bytes.
|
|
392
|
+
"""
|
|
393
|
+
args_size = 0
|
|
394
|
+
encoded_args_list: List[bytes] = []
|
|
395
|
+
if not args_list:
|
|
396
|
+
return (encoded_args_list, args_size)
|
|
397
|
+
for arg in args_list:
|
|
398
|
+
encoded_arg = self._encode_arg(arg) if isinstance(arg, str) else arg
|
|
399
|
+
encoded_args_list.append(encoded_arg)
|
|
400
|
+
args_size += len(encoded_arg)
|
|
401
|
+
return (encoded_args_list, args_size)
|
|
402
|
+
|
|
403
|
+
async def _execute_command(
|
|
404
|
+
self,
|
|
405
|
+
request_type: RequestType.ValueType,
|
|
406
|
+
args: List[TEncodable],
|
|
407
|
+
route: Optional[Route] = None,
|
|
408
|
+
) -> TResult:
|
|
409
|
+
if self._is_closed:
|
|
410
|
+
raise ClosingError(
|
|
411
|
+
"Unable to execute requests; the client is closed. Please create a new client."
|
|
412
|
+
)
|
|
413
|
+
|
|
414
|
+
# Create span if OpenTelemetry is configured and sampling indicates we should trace
|
|
415
|
+
span = None
|
|
416
|
+
if OpenTelemetry.should_sample():
|
|
417
|
+
command_name = RequestType.Name(request_type)
|
|
418
|
+
span = create_otel_span(command_name)
|
|
419
|
+
|
|
420
|
+
request = CommandRequest()
|
|
421
|
+
request.callback_idx = self._get_callback_index()
|
|
422
|
+
request.single_command.request_type = request_type
|
|
423
|
+
request.single_command.args_array.args[:] = [
|
|
424
|
+
bytes(elem, encoding="utf8") if isinstance(elem, str) else elem
|
|
425
|
+
for elem in args
|
|
426
|
+
]
|
|
427
|
+
(encoded_args, args_size) = self._encode_and_sum_size(args)
|
|
428
|
+
if args_size < MAX_REQUEST_ARGS_LEN:
|
|
429
|
+
request.single_command.args_array.args[:] = encoded_args
|
|
430
|
+
else:
|
|
431
|
+
request.single_command.args_vec_pointer = create_leaked_bytes_vec(
|
|
432
|
+
encoded_args
|
|
433
|
+
)
|
|
434
|
+
|
|
435
|
+
# Add span pointer to request if span was created
|
|
436
|
+
if span:
|
|
437
|
+
request.root_span_ptr = span
|
|
438
|
+
|
|
439
|
+
set_protobuf_route(request, route)
|
|
440
|
+
return await self._write_request_await_response(request)
|
|
441
|
+
|
|
442
|
+
async def _execute_batch(
|
|
443
|
+
self,
|
|
444
|
+
commands: List[Tuple[RequestType.ValueType, List[TEncodable]]],
|
|
445
|
+
is_atomic: bool,
|
|
446
|
+
raise_on_error: bool = False,
|
|
447
|
+
retry_server_error: bool = False,
|
|
448
|
+
retry_connection_error: bool = False,
|
|
449
|
+
route: Optional[Route] = None,
|
|
450
|
+
timeout: Optional[int] = None,
|
|
451
|
+
) -> List[TResult]:
|
|
452
|
+
if self._is_closed:
|
|
453
|
+
raise ClosingError(
|
|
454
|
+
"Unable to execute requests; the client is closed. Please create a new client."
|
|
455
|
+
)
|
|
456
|
+
|
|
457
|
+
# Create span if OpenTelemetry is configured and sampling indicates we should trace
|
|
458
|
+
span = None
|
|
459
|
+
|
|
460
|
+
if OpenTelemetry.should_sample():
|
|
461
|
+
# Use "Batch" as span name for batches
|
|
462
|
+
span = create_otel_span("Batch")
|
|
463
|
+
|
|
464
|
+
request = CommandRequest()
|
|
465
|
+
request.callback_idx = self._get_callback_index()
|
|
466
|
+
batch_commands = []
|
|
467
|
+
for requst_type, args in commands:
|
|
468
|
+
command = Command()
|
|
469
|
+
command.request_type = requst_type
|
|
470
|
+
# For now, we allow the user to pass the command as array of strings
|
|
471
|
+
# we convert them here into bytes (the datatype that our rust core expects)
|
|
472
|
+
(encoded_args, args_size) = self._encode_and_sum_size(args)
|
|
473
|
+
if args_size < MAX_REQUEST_ARGS_LEN:
|
|
474
|
+
command.args_array.args[:] = encoded_args
|
|
475
|
+
else:
|
|
476
|
+
command.args_vec_pointer = create_leaked_bytes_vec(encoded_args)
|
|
477
|
+
batch_commands.append(command)
|
|
478
|
+
request.batch.commands.extend(batch_commands)
|
|
479
|
+
request.batch.is_atomic = is_atomic
|
|
480
|
+
request.batch.raise_on_error = raise_on_error
|
|
481
|
+
if timeout is not None:
|
|
482
|
+
request.batch.timeout = timeout
|
|
483
|
+
request.batch.retry_server_error = retry_server_error
|
|
484
|
+
request.batch.retry_connection_error = retry_connection_error
|
|
485
|
+
|
|
486
|
+
# Add span pointer to request if span was created
|
|
487
|
+
if span:
|
|
488
|
+
request.root_span_ptr = span
|
|
489
|
+
|
|
490
|
+
set_protobuf_route(request, route)
|
|
491
|
+
return await self._write_request_await_response(request)
|
|
492
|
+
|
|
493
|
+
async def _execute_script(
|
|
494
|
+
self,
|
|
495
|
+
hash: str,
|
|
496
|
+
keys: Optional[List[Union[str, bytes]]] = None,
|
|
497
|
+
args: Optional[List[Union[str, bytes]]] = None,
|
|
498
|
+
route: Optional[Route] = None,
|
|
499
|
+
) -> TResult:
|
|
500
|
+
if self._is_closed:
|
|
501
|
+
raise ClosingError(
|
|
502
|
+
"Unable to execute requests; the client is closed. Please create a new client."
|
|
503
|
+
)
|
|
504
|
+
request = CommandRequest()
|
|
505
|
+
request.callback_idx = self._get_callback_index()
|
|
506
|
+
(encoded_keys, keys_size) = self._encode_and_sum_size(keys)
|
|
507
|
+
(encoded_args, args_size) = self._encode_and_sum_size(args)
|
|
508
|
+
if (keys_size + args_size) < MAX_REQUEST_ARGS_LEN:
|
|
509
|
+
request.script_invocation.hash = hash
|
|
510
|
+
request.script_invocation.keys[:] = encoded_keys
|
|
511
|
+
request.script_invocation.args[:] = encoded_args
|
|
512
|
+
|
|
513
|
+
else:
|
|
514
|
+
request.script_invocation_pointers.hash = hash
|
|
515
|
+
request.script_invocation_pointers.keys_pointer = create_leaked_bytes_vec(
|
|
516
|
+
encoded_keys
|
|
517
|
+
)
|
|
518
|
+
request.script_invocation_pointers.args_pointer = create_leaked_bytes_vec(
|
|
519
|
+
encoded_args
|
|
520
|
+
)
|
|
521
|
+
set_protobuf_route(request, route)
|
|
522
|
+
return await self._write_request_await_response(request)
|
|
523
|
+
|
|
524
|
+
async def get_pubsub_message(self) -> PubSubMsg:
|
|
525
|
+
if self._is_closed:
|
|
526
|
+
raise ClosingError(
|
|
527
|
+
"Unable to execute requests; the client is closed. Please create a new client."
|
|
528
|
+
)
|
|
529
|
+
|
|
530
|
+
if not self.config._is_pubsub_configured():
|
|
531
|
+
raise ConfigurationError(
|
|
532
|
+
"The operation will never complete since there was no pubsub subscriptions applied to the client."
|
|
533
|
+
)
|
|
534
|
+
|
|
535
|
+
if self.config._get_pubsub_callback_and_context()[0] is not None:
|
|
536
|
+
raise ConfigurationError(
|
|
537
|
+
"The operation will never complete since messages will be passed to the configured callback."
|
|
538
|
+
)
|
|
539
|
+
|
|
540
|
+
# locking might not be required
|
|
541
|
+
response_future: "TFuture" = _get_new_future_instance()
|
|
542
|
+
try:
|
|
543
|
+
self._pubsub_lock.acquire()
|
|
544
|
+
self._pubsub_futures.append(response_future)
|
|
545
|
+
self._complete_pubsub_futures_safe()
|
|
546
|
+
finally:
|
|
547
|
+
self._pubsub_lock.release()
|
|
548
|
+
await response_future
|
|
549
|
+
return response_future.result()
|
|
550
|
+
|
|
551
|
+
def try_get_pubsub_message(self) -> Optional[PubSubMsg]:
|
|
552
|
+
if self._is_closed:
|
|
553
|
+
raise ClosingError(
|
|
554
|
+
"Unable to execute requests; the client is closed. Please create a new client."
|
|
555
|
+
)
|
|
556
|
+
|
|
557
|
+
if not self.config._is_pubsub_configured():
|
|
558
|
+
raise ConfigurationError(
|
|
559
|
+
"The operation will never succeed since there was no pubsbub subscriptions applied to the client."
|
|
560
|
+
)
|
|
561
|
+
|
|
562
|
+
if self.config._get_pubsub_callback_and_context()[0] is not None:
|
|
563
|
+
raise ConfigurationError(
|
|
564
|
+
"The operation will never succeed since messages will be passed to the configured callback."
|
|
565
|
+
)
|
|
566
|
+
|
|
567
|
+
# locking might not be required
|
|
568
|
+
msg: Optional[PubSubMsg] = None
|
|
569
|
+
try:
|
|
570
|
+
self._pubsub_lock.acquire()
|
|
571
|
+
self._complete_pubsub_futures_safe()
|
|
572
|
+
while len(self._pending_push_notifications) and not msg:
|
|
573
|
+
push_notification = self._pending_push_notifications.pop(0)
|
|
574
|
+
msg = self._notification_to_pubsub_message_safe(push_notification)
|
|
575
|
+
finally:
|
|
576
|
+
self._pubsub_lock.release()
|
|
577
|
+
return msg
|
|
578
|
+
|
|
579
|
+
def _cancel_pubsub_futures_with_exception_safe(self, exception: ConnectionError):
|
|
580
|
+
while len(self._pubsub_futures):
|
|
581
|
+
next_future = self._pubsub_futures.pop(0)
|
|
582
|
+
next_future.set_exception(exception)
|
|
583
|
+
|
|
584
|
+
def _notification_to_pubsub_message_safe(
|
|
585
|
+
self, response: Response
|
|
586
|
+
) -> Optional[PubSubMsg]:
|
|
587
|
+
pubsub_message = None
|
|
588
|
+
push_notification = cast(
|
|
589
|
+
Dict[str, Any], value_from_pointer(response.resp_pointer)
|
|
590
|
+
)
|
|
591
|
+
message_kind = push_notification["kind"]
|
|
592
|
+
if message_kind == "Disconnection":
|
|
593
|
+
ClientLogger.log(
|
|
594
|
+
LogLevel.WARN,
|
|
595
|
+
"disconnect notification",
|
|
596
|
+
"Transport disconnected, messages might be lost",
|
|
597
|
+
)
|
|
598
|
+
elif (
|
|
599
|
+
message_kind == "Message"
|
|
600
|
+
or message_kind == "PMessage"
|
|
601
|
+
or message_kind == "SMessage"
|
|
602
|
+
):
|
|
603
|
+
values: List = push_notification["values"]
|
|
604
|
+
if message_kind == "PMessage":
|
|
605
|
+
pubsub_message = PubSubMsg(
|
|
606
|
+
message=values[2], channel=values[1], pattern=values[0]
|
|
607
|
+
)
|
|
608
|
+
else:
|
|
609
|
+
pubsub_message = PubSubMsg(
|
|
610
|
+
message=values[1], channel=values[0], pattern=None
|
|
611
|
+
)
|
|
612
|
+
elif (
|
|
613
|
+
message_kind == "PSubscribe"
|
|
614
|
+
or message_kind == "Subscribe"
|
|
615
|
+
or message_kind == "SSubscribe"
|
|
616
|
+
or message_kind == "Unsubscribe"
|
|
617
|
+
or message_kind == "PUnsubscribe"
|
|
618
|
+
or message_kind == "SUnsubscribe"
|
|
619
|
+
):
|
|
620
|
+
pass
|
|
621
|
+
else:
|
|
622
|
+
ClientLogger.log(
|
|
623
|
+
LogLevel.WARN,
|
|
624
|
+
"unknown notification",
|
|
625
|
+
f"Unknown notification message: '{message_kind}'",
|
|
626
|
+
)
|
|
627
|
+
|
|
628
|
+
return pubsub_message
|
|
629
|
+
|
|
630
|
+
def _complete_pubsub_futures_safe(self):
|
|
631
|
+
while len(self._pending_push_notifications) and len(self._pubsub_futures):
|
|
632
|
+
next_push_notification = self._pending_push_notifications.pop(0)
|
|
633
|
+
pubsub_message = self._notification_to_pubsub_message_safe(
|
|
634
|
+
next_push_notification
|
|
635
|
+
)
|
|
636
|
+
if pubsub_message:
|
|
637
|
+
self._pubsub_futures.pop(0).set_result(pubsub_message)
|
|
638
|
+
|
|
639
|
+
async def _write_request_await_response(self, request: CommandRequest):
|
|
640
|
+
# Create a response future for this request and add it to the available
|
|
641
|
+
# futures map
|
|
642
|
+
response_future = self._get_future(request.callback_idx)
|
|
643
|
+
self._create_write_task(request)
|
|
644
|
+
await response_future
|
|
645
|
+
return response_future.result()
|
|
646
|
+
|
|
647
|
+
def _get_callback_index(self) -> int:
|
|
648
|
+
try:
|
|
649
|
+
return self._available_callback_indexes.pop()
|
|
650
|
+
except IndexError:
|
|
651
|
+
# The list is empty
|
|
652
|
+
return len(self._available_futures)
|
|
653
|
+
|
|
654
|
+
async def _process_response(self, response: Response) -> None:
|
|
655
|
+
res_future = self._available_futures.pop(response.callback_idx, None)
|
|
656
|
+
if not res_future or response.HasField("closing_error"):
|
|
657
|
+
err_msg = (
|
|
658
|
+
response.closing_error
|
|
659
|
+
if response.HasField("closing_error")
|
|
660
|
+
else f"Client Error - closing due to unknown error. callback index: {response.callback_idx}"
|
|
661
|
+
)
|
|
662
|
+
exc = ClosingError(err_msg)
|
|
663
|
+
if res_future is not None:
|
|
664
|
+
res_future.set_exception(exc)
|
|
665
|
+
else:
|
|
666
|
+
ClientLogger.log(
|
|
667
|
+
LogLevel.WARN,
|
|
668
|
+
"unhandled response error",
|
|
669
|
+
f"Unhandled response error for unknown request: {response.callback_idx}",
|
|
670
|
+
)
|
|
671
|
+
raise exc
|
|
672
|
+
else:
|
|
673
|
+
self._available_callback_indexes.append(response.callback_idx)
|
|
674
|
+
if response.HasField("request_error"):
|
|
675
|
+
error_type = get_request_error_class(response.request_error.type)
|
|
676
|
+
res_future.set_exception(error_type(response.request_error.message))
|
|
677
|
+
elif response.HasField("resp_pointer"):
|
|
678
|
+
res_future.set_result(value_from_pointer(response.resp_pointer))
|
|
679
|
+
elif response.HasField("constant_response"):
|
|
680
|
+
res_future.set_result(OK)
|
|
681
|
+
else:
|
|
682
|
+
res_future.set_result(None)
|
|
683
|
+
|
|
684
|
+
# Clean up span if it was created
|
|
685
|
+
if response.HasField("root_span_ptr"):
|
|
686
|
+
drop_otel_span(response.root_span_ptr)
|
|
687
|
+
|
|
688
|
+
async def _process_push(self, response: Response) -> None:
|
|
689
|
+
if response.HasField("closing_error") or not response.HasField("resp_pointer"):
|
|
690
|
+
err_msg = (
|
|
691
|
+
response.closing_error
|
|
692
|
+
if response.HasField("closing_error")
|
|
693
|
+
else "Client Error - push notification without resp_pointer"
|
|
694
|
+
)
|
|
695
|
+
raise ClosingError(err_msg)
|
|
696
|
+
try:
|
|
697
|
+
self._pubsub_lock.acquire()
|
|
698
|
+
callback, context = self.config._get_pubsub_callback_and_context()
|
|
699
|
+
if callback:
|
|
700
|
+
pubsub_message = self._notification_to_pubsub_message_safe(response)
|
|
701
|
+
if pubsub_message:
|
|
702
|
+
callback(pubsub_message, context)
|
|
703
|
+
else:
|
|
704
|
+
self._pending_push_notifications.append(response)
|
|
705
|
+
self._complete_pubsub_futures_safe()
|
|
706
|
+
finally:
|
|
707
|
+
self._pubsub_lock.release()
|
|
708
|
+
|
|
709
|
+
async def _reader_loop(self) -> None:
|
|
710
|
+
# Socket reader loop
|
|
711
|
+
try:
|
|
712
|
+
remaining_read_bytes = bytearray()
|
|
713
|
+
while True:
|
|
714
|
+
try:
|
|
715
|
+
read_bytes = await self._stream.receive(DEFAULT_READ_BYTES_SIZE)
|
|
716
|
+
except (anyio.ClosedResourceError, anyio.EndOfStream):
|
|
717
|
+
raise ClosingError(
|
|
718
|
+
"The communication layer was unexpectedly closed."
|
|
719
|
+
)
|
|
720
|
+
read_bytes = remaining_read_bytes + bytearray(read_bytes)
|
|
721
|
+
read_bytes_view = memoryview(read_bytes)
|
|
722
|
+
offset = 0
|
|
723
|
+
while offset <= len(read_bytes):
|
|
724
|
+
try:
|
|
725
|
+
response, offset = ProtobufCodec.decode_delimited(
|
|
726
|
+
read_bytes, read_bytes_view, offset, Response
|
|
727
|
+
)
|
|
728
|
+
except PartialMessageException:
|
|
729
|
+
# Received only partial response, break the inner loop
|
|
730
|
+
remaining_read_bytes = read_bytes[offset:]
|
|
731
|
+
break
|
|
732
|
+
response = cast(Response, response)
|
|
733
|
+
if response.is_push:
|
|
734
|
+
await self._process_push(response=response)
|
|
735
|
+
else:
|
|
736
|
+
await self._process_response(response=response)
|
|
737
|
+
except Exception as e:
|
|
738
|
+
# close and stop reading at terminal exceptions from incoming responses or
|
|
739
|
+
# stream closures
|
|
740
|
+
await self.close(str(e))
|
|
741
|
+
|
|
742
|
+
async def get_statistics(self) -> dict:
|
|
743
|
+
return get_statistics()
|
|
744
|
+
|
|
745
|
+
async def _update_connection_password(
|
|
746
|
+
self, password: Optional[str], immediate_auth: bool
|
|
747
|
+
) -> TResult:
|
|
748
|
+
request = CommandRequest()
|
|
749
|
+
request.callback_idx = self._get_callback_index()
|
|
750
|
+
if password is not None:
|
|
751
|
+
request.update_connection_password.password = password
|
|
752
|
+
request.update_connection_password.immediate_auth = immediate_auth
|
|
753
|
+
response = await self._write_request_await_response(request)
|
|
754
|
+
# Update the client binding side password if managed to change core configuration password
|
|
755
|
+
if response is OK:
|
|
756
|
+
if self.config.credentials is None:
|
|
757
|
+
self.config.credentials = ServerCredentials(password=password or "")
|
|
758
|
+
self.config.credentials.password = password or ""
|
|
759
|
+
return response
|
|
760
|
+
|
|
761
|
+
async def _refresh_iam_token(self) -> TResult:
|
|
762
|
+
request = CommandRequest()
|
|
763
|
+
request.callback_idx = self._get_callback_index()
|
|
764
|
+
request.refresh_iam_token.CopyFrom(
|
|
765
|
+
RefreshIamToken()
|
|
766
|
+
) # Empty message, just triggers the refresh
|
|
767
|
+
response = await self._write_request_await_response(request)
|
|
768
|
+
return response
|
|
769
|
+
|
|
770
|
+
|
|
771
|
+
class GlideClusterClient(BaseClient, ClusterCommands):
|
|
772
|
+
"""
|
|
773
|
+
Client used for connection to cluster servers.
|
|
774
|
+
Use :func:`~BaseClient.create` to request a client.
|
|
775
|
+
For full documentation, see
|
|
776
|
+
[Valkey GLIDE Wiki](https://github.com/valkey-io/valkey-glide/wiki/Python-wrapper#cluster)
|
|
777
|
+
"""
|
|
778
|
+
|
|
779
|
+
async def _cluster_scan(
|
|
780
|
+
self,
|
|
781
|
+
cursor: ClusterScanCursor,
|
|
782
|
+
match: Optional[TEncodable] = None,
|
|
783
|
+
count: Optional[int] = None,
|
|
784
|
+
type: Optional[ObjectType] = None,
|
|
785
|
+
allow_non_covered_slots: bool = False,
|
|
786
|
+
) -> List[Union[ClusterScanCursor, List[bytes]]]:
|
|
787
|
+
if self._is_closed:
|
|
788
|
+
raise ClosingError(
|
|
789
|
+
"Unable to execute requests; the client is closed. Please create a new client."
|
|
790
|
+
)
|
|
791
|
+
request = CommandRequest()
|
|
792
|
+
request.callback_idx = self._get_callback_index()
|
|
793
|
+
# Take out the id string from the wrapping object
|
|
794
|
+
cursor_string = cursor.get_cursor()
|
|
795
|
+
request.cluster_scan.cursor = cursor_string
|
|
796
|
+
request.cluster_scan.allow_non_covered_slots = allow_non_covered_slots
|
|
797
|
+
if match is not None:
|
|
798
|
+
request.cluster_scan.match_pattern = (
|
|
799
|
+
self._encode_arg(match) if isinstance(match, str) else match
|
|
800
|
+
)
|
|
801
|
+
if count is not None:
|
|
802
|
+
request.cluster_scan.count = count
|
|
803
|
+
if type is not None:
|
|
804
|
+
request.cluster_scan.object_type = type.value
|
|
805
|
+
response = await self._write_request_await_response(request)
|
|
806
|
+
return [ClusterScanCursor(bytes(response[0]).decode()), response[1]]
|
|
807
|
+
|
|
808
|
+
def _get_protobuf_conn_request(self) -> ConnectionRequest:
|
|
809
|
+
return self.config._create_a_protobuf_conn_request(cluster_mode=True)
|
|
810
|
+
|
|
811
|
+
|
|
812
|
+
class GlideClient(BaseClient, StandaloneCommands):
|
|
813
|
+
"""
|
|
814
|
+
Client used for connection to standalone servers.
|
|
815
|
+
Use :func:`~BaseClient.create` to request a client.
|
|
816
|
+
For full documentation, see
|
|
817
|
+
[Valkey GLIDE Wiki](https://github.com/valkey-io/valkey-glide/wiki/Python-wrapper#standalone)
|
|
818
|
+
"""
|
|
819
|
+
|
|
820
|
+
|
|
821
|
+
TGlideClient = Union[GlideClient, GlideClusterClient]
|