valkey-glide 1.3.5__cp39-cp39-macosx_11_0_arm64.whl → 2.2.2__cp39-cp39-macosx_11_0_arm64.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.
Files changed (47) hide show
  1. glide/__init__.py +165 -107
  2. glide/async_commands/cluster_commands.py +318 -136
  3. glide/async_commands/core.py +1770 -992
  4. glide/async_commands/{server_modules/ft.py → ft.py} +91 -21
  5. glide/async_commands/{server_modules/glide_json.py → glide_json.py} +148 -134
  6. glide/async_commands/standalone_commands.py +203 -137
  7. glide/glide.cpython-39-darwin.so +0 -0
  8. glide/glide.pyi +26 -1
  9. glide/glide_client.py +352 -135
  10. glide/logger.py +34 -22
  11. glide/opentelemetry.py +185 -0
  12. glide_shared/__init__.py +330 -0
  13. glide_shared/commands/__init__.py +0 -0
  14. glide/async_commands/transaction.py → glide_shared/commands/batch.py +1839 -1017
  15. glide_shared/commands/batch_options.py +261 -0
  16. {glide/async_commands → glide_shared/commands}/bitmap.py +94 -85
  17. {glide/async_commands → glide_shared/commands}/command_args.py +7 -6
  18. glide_shared/commands/core_options.py +407 -0
  19. {glide/async_commands → glide_shared/commands}/server_modules/ft_options/ft_aggregate_options.py +18 -11
  20. {glide/async_commands → glide_shared/commands}/server_modules/ft_options/ft_create_options.py +27 -13
  21. {glide/async_commands → glide_shared/commands}/server_modules/ft_options/ft_profile_options.py +16 -11
  22. {glide/async_commands → glide_shared/commands}/server_modules/ft_options/ft_search_options.py +16 -8
  23. {glide/async_commands → glide_shared/commands}/server_modules/json_batch.py +160 -130
  24. glide_shared/commands/server_modules/json_options.py +93 -0
  25. {glide/async_commands → glide_shared/commands}/sorted_set.py +41 -31
  26. {glide/async_commands → glide_shared/commands}/stream.py +95 -88
  27. glide_shared/config.py +975 -0
  28. {glide → glide_shared}/constants.py +11 -7
  29. {glide → glide_shared}/exceptions.py +27 -1
  30. glide_shared/protobuf/command_request_pb2.py +56 -0
  31. glide_shared/protobuf/connection_request_pb2.py +56 -0
  32. {glide → glide_shared}/protobuf/response_pb2.py +6 -6
  33. {glide → glide_shared}/protobuf_codec.py +7 -6
  34. glide_shared/routes.py +161 -0
  35. valkey_glide-2.2.2.dist-info/METADATA +211 -0
  36. valkey_glide-2.2.2.dist-info/RECORD +40 -0
  37. glide/config.py +0 -590
  38. glide/protobuf/command_request_pb2.py +0 -54
  39. glide/protobuf/command_request_pb2.pyi +0 -1164
  40. glide/protobuf/connection_request_pb2.py +0 -52
  41. glide/protobuf/connection_request_pb2.pyi +0 -292
  42. glide/protobuf/response_pb2.pyi +0 -101
  43. glide/routes.py +0 -114
  44. valkey_glide-1.3.5.dist-info/METADATA +0 -125
  45. valkey_glide-1.3.5.dist-info/RECORD +0 -37
  46. {glide/async_commands → glide_shared/commands}/server_modules/ft_options/ft_constants.py +0 -0
  47. {valkey_glide-1.3.5.dist-info → valkey_glide-2.2.2.dist-info}/WHEEL +0 -0
glide/glide_client.py CHANGED
@@ -1,62 +1,121 @@
1
1
  # Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0
2
2
 
3
- import asyncio
4
3
  import sys
5
4
  import threading
6
- from typing import Any, Dict, List, Optional, Tuple, Type, Union, cast
7
-
8
- from glide.async_commands.cluster_commands import ClusterCommands
9
- from glide.async_commands.command_args import ObjectType
10
- from glide.async_commands.core import CoreCommands
11
- from glide.async_commands.standalone_commands import StandaloneCommands
12
- from glide.config import BaseClientConfiguration, ServerCredentials
13
- from glide.constants import DEFAULT_READ_BYTES_SIZE, OK, TEncodable, TRequest, TResult
14
- from glide.exceptions import (
15
- ClosingError,
16
- ConfigurationError,
17
- ConnectionError,
18
- ExecAbortError,
19
- RequestError,
20
- TimeoutError,
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,
21
16
  )
22
- from glide.logger import Level as LogLevel
23
- from glide.logger import Logger as ClientLogger
24
- from glide.protobuf.command_request_pb2 import Command, CommandRequest, RequestType
25
- from glide.protobuf.connection_request_pb2 import ConnectionRequest
26
- from glide.protobuf.response_pb2 import RequestErrorType, Response
27
- from glide.protobuf_codec import PartialMessageException, ProtobufCodec
28
- from glide.routes import Route, set_protobuf_route
29
-
30
- from .glide import (
17
+
18
+ import anyio
19
+ import sniffio
20
+ from anyio import to_thread
21
+ from glide.glide import (
31
22
  DEFAULT_TIMEOUT_IN_MILLISECONDS,
32
23
  MAX_REQUEST_ARGS_LEN,
33
24
  ClusterScanCursor,
34
25
  create_leaked_bytes_vec,
26
+ create_otel_span,
27
+ drop_otel_span,
35
28
  get_statistics,
36
29
  start_socket_listener_external,
37
30
  value_from_pointer,
38
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
39
65
 
40
66
  if sys.version_info >= (3, 11):
41
- import asyncio as async_timeout
42
67
  from typing import Self
43
68
  else:
44
- import async_timeout
45
69
  from typing_extensions import Self
46
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__()
47
101
 
48
- def get_request_error_class(
49
- error_type: Optional[RequestErrorType.ValueType],
50
- ) -> Type[RequestError]:
51
- if error_type == RequestErrorType.Disconnect:
52
- return ConnectionError
53
- if error_type == RequestErrorType.ExecAbort:
54
- return ExecAbortError
55
- if error_type == RequestErrorType.Timeout:
56
- return TimeoutError
57
- if error_type == RequestErrorType.Unspecified:
58
- return RequestError
59
- return RequestError
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()
60
119
 
61
120
 
62
121
  class BaseClient(CoreCommands):
@@ -65,32 +124,121 @@ class BaseClient(CoreCommands):
65
124
  To create a new client, use the `create` classmethod
66
125
  """
67
126
  self.config: BaseClientConfiguration = config
68
- self._available_futures: Dict[int, asyncio.Future] = {}
127
+ self._available_futures: Dict[int, "TFuture"] = {}
69
128
  self._available_callback_indexes: List[int] = list()
70
129
  self._buffered_requests: List[TRequest] = list()
71
130
  self._writer_lock = threading.Lock()
72
131
  self.socket_path: Optional[str] = None
73
- self._reader_task: Optional[asyncio.Task] = None
132
+ self._reader_task: Optional["TTask"] = None
74
133
  self._is_closed: bool = False
75
- self._pubsub_futures: List[asyncio.Future] = []
134
+ self._pubsub_futures: List["TFuture"] = []
76
135
  self._pubsub_lock = threading.Lock()
77
136
  self._pending_push_notifications: List[Response] = list()
78
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
+
79
168
  @classmethod
80
169
  async def create(cls, config: BaseClientConfiguration) -> Self:
81
170
  """Creates a Glide client.
82
171
 
83
172
  Args:
84
- config (ClientConfiguration): The client configurations.
85
- If no configuration is provided, a default client to "localhost":6379 will be created.
173
+ config (ClientConfiguration): The configuration options for the client, including cluster addresses,
174
+ authentication credentials, TLS settings, periodic checks, and Pub/Sub subscriptions.
86
175
 
87
176
  Returns:
88
- Self: a Glide Client instance.
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
+
89
237
  """
90
238
  config = config
91
239
  self = cls(config)
92
- init_future: asyncio.Future = asyncio.Future()
93
- loop = asyncio.get_event_loop()
240
+
241
+ init_event: threading.Event = threading.Event()
94
242
 
95
243
  def init_callback(socket_path: Optional[str], err: Optional[str]):
96
244
  if err is not None:
@@ -102,7 +250,7 @@ class BaseClient(CoreCommands):
102
250
  else:
103
251
  # Received socket path
104
252
  self.socket_path = socket_path
105
- loop.call_soon_threadsafe(init_future.set_result, True)
253
+ init_event.set()
106
254
 
107
255
  start_socket_listener_external(init_callback=init_callback)
108
256
 
@@ -110,36 +258,27 @@ class BaseClient(CoreCommands):
110
258
  # level or higher
111
259
  ClientLogger.log(LogLevel.INFO, "connection info", "new connection established")
112
260
  # Wait for the socket listener to complete its initialization
113
- await init_future
261
+ await to_thread.run_sync(init_event.wait)
114
262
  # Create UDS connection
115
263
  await self._create_uds_connection()
264
+
116
265
  # Start the reader loop as a background task
117
- self._reader_task = asyncio.create_task(self._reader_loop())
266
+ self._reader_task = self._create_task(self._reader_loop)
267
+
118
268
  # Set the client configurations
119
269
  await self._set_connection_configurations()
270
+
120
271
  return self
121
272
 
122
273
  async def _create_uds_connection(self) -> None:
123
274
  try:
124
275
  # Open an UDS connection
125
- async with async_timeout.timeout(DEFAULT_TIMEOUT_IN_MILLISECONDS):
126
- reader, writer = await asyncio.open_unix_connection(
127
- path=self.socket_path
276
+ with anyio.fail_after(DEFAULT_TIMEOUT_IN_MILLISECONDS):
277
+ self._stream = await anyio.connect_unix(
278
+ path=cast(str, self.socket_path)
128
279
  )
129
- self._reader = reader
130
- self._writer = writer
131
280
  except Exception as e:
132
- await self.close(f"Failed to create UDS connection: {e}")
133
- raise
134
-
135
- def __del__(self) -> None:
136
- try:
137
- if self._reader_task:
138
- self._reader_task.cancel()
139
- except RuntimeError as e:
140
- if "no running event loop" in str(e):
141
- # event loop already closed
142
- pass
281
+ raise ClosingError("Failed to create UDS connection") from e
143
282
 
144
283
  async def close(self, err_message: Optional[str] = None) -> None:
145
284
  """
@@ -147,28 +286,28 @@ class BaseClient(CoreCommands):
147
286
  All open futures will be closed with an exception.
148
287
 
149
288
  Args:
150
- err_message (Optional[str]): If not None, this error message will be passed along with the exceptions when closing all open futures.
289
+ err_message (Optional[str]): If not None, this error message will be passed along with the exceptions when
290
+ closing all open futures.
151
291
  Defaults to None.
152
292
  """
153
- self._is_closed = True
154
- for response_future in self._available_futures.values():
155
- if not response_future.done():
156
- err_message = "" if err_message is None else err_message
157
- response_future.set_exception(ClosingError(err_message))
158
- try:
159
- self._pubsub_lock.acquire()
160
- for pubsub_future in self._pubsub_futures:
161
- if not pubsub_future.done() and not pubsub_future.cancelled():
162
- pubsub_future.set_exception(ClosingError(""))
163
- finally:
164
- self._pubsub_lock.release()
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()
165
306
 
166
- self._writer.close()
167
- await self._writer.wait_closed()
168
- self.__del__()
307
+ await self._stream.aclose()
169
308
 
170
- def _get_future(self, callback_idx: int) -> asyncio.Future:
171
- response_future: asyncio.Future = asyncio.Future()
309
+ def _get_future(self, callback_idx: int) -> "TFuture":
310
+ response_future: "TFuture" = _get_new_future_instance()
172
311
  self._available_futures.update({callback_idx: response_future})
173
312
  return response_future
174
313
 
@@ -177,14 +316,15 @@ class BaseClient(CoreCommands):
177
316
 
178
317
  async def _set_connection_configurations(self) -> None:
179
318
  conn_request = self._get_protobuf_conn_request()
180
- response_future: asyncio.Future = self._get_future(0)
181
- await self._write_or_buffer_request(conn_request)
319
+ response_future: "TFuture" = self._get_future(0)
320
+ self._create_write_task(conn_request)
182
321
  await response_future
183
- if response_future.result() is not OK:
184
- raise ClosingError(response_future.result())
322
+ res = response_future.result()
323
+ if res is not OK:
324
+ raise ClosingError(res)
185
325
 
186
326
  def _create_write_task(self, request: TRequest):
187
- asyncio.create_task(self._write_or_buffer_request(request))
327
+ self._create_task(self._write_or_buffer_request, request)
188
328
 
189
329
  async def _write_or_buffer_request(self, request: TRequest):
190
330
  self._buffered_requests.append(request)
@@ -192,7 +332,21 @@ class BaseClient(CoreCommands):
192
332
  try:
193
333
  while len(self._buffered_requests) > 0:
194
334
  await self._write_buffered_requests_to_socket()
195
-
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
+ )
196
350
  finally:
197
351
  self._writer_lock.release()
198
352
 
@@ -202,8 +356,10 @@ class BaseClient(CoreCommands):
202
356
  b_arr = bytearray()
203
357
  for request in requests:
204
358
  ProtobufCodec.encode_delimited(b_arr, request)
205
- self._writer.write(b_arr)
206
- await self._writer.drain()
359
+ try:
360
+ await self._stream.send(b_arr)
361
+ except (anyio.ClosedResourceError, anyio.EndOfStream):
362
+ raise ClosingError("The communication layer was unexpectedly closed.")
207
363
 
208
364
  def _encode_arg(self, arg: TEncodable) -> bytes:
209
365
  """
@@ -254,6 +410,13 @@ class BaseClient(CoreCommands):
254
410
  raise ClosingError(
255
411
  "Unable to execute requests; the client is closed. Please create a new client."
256
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
+
257
420
  request = CommandRequest()
258
421
  request.callback_idx = self._get_callback_index()
259
422
  request.single_command.request_type = request_type
@@ -268,21 +431,39 @@ class BaseClient(CoreCommands):
268
431
  request.single_command.args_vec_pointer = create_leaked_bytes_vec(
269
432
  encoded_args
270
433
  )
434
+
435
+ # Add span pointer to request if span was created
436
+ if span:
437
+ request.root_span_ptr = span
438
+
271
439
  set_protobuf_route(request, route)
272
440
  return await self._write_request_await_response(request)
273
441
 
274
- async def _execute_transaction(
442
+ async def _execute_batch(
275
443
  self,
276
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,
277
449
  route: Optional[Route] = None,
450
+ timeout: Optional[int] = None,
278
451
  ) -> List[TResult]:
279
452
  if self._is_closed:
280
453
  raise ClosingError(
281
454
  "Unable to execute requests; the client is closed. Please create a new client."
282
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
+
283
464
  request = CommandRequest()
284
465
  request.callback_idx = self._get_callback_index()
285
- transaction_commands = []
466
+ batch_commands = []
286
467
  for requst_type, args in commands:
287
468
  command = Command()
288
469
  command.request_type = requst_type
@@ -293,8 +474,19 @@ class BaseClient(CoreCommands):
293
474
  command.args_array.args[:] = encoded_args
294
475
  else:
295
476
  command.args_vec_pointer = create_leaked_bytes_vec(encoded_args)
296
- transaction_commands.append(command)
297
- request.transaction.commands.extend(transaction_commands)
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
+
298
490
  set_protobuf_route(request, route)
299
491
  return await self._write_request_await_response(request)
300
492
 
@@ -329,7 +521,7 @@ class BaseClient(CoreCommands):
329
521
  set_protobuf_route(request, route)
330
522
  return await self._write_request_await_response(request)
331
523
 
332
- async def get_pubsub_message(self) -> CoreCommands.PubSubMsg:
524
+ async def get_pubsub_message(self) -> PubSubMsg:
333
525
  if self._is_closed:
334
526
  raise ClosingError(
335
527
  "Unable to execute requests; the client is closed. Please create a new client."
@@ -346,16 +538,17 @@ class BaseClient(CoreCommands):
346
538
  )
347
539
 
348
540
  # locking might not be required
349
- response_future: asyncio.Future = asyncio.Future()
541
+ response_future: "TFuture" = _get_new_future_instance()
350
542
  try:
351
543
  self._pubsub_lock.acquire()
352
544
  self._pubsub_futures.append(response_future)
353
545
  self._complete_pubsub_futures_safe()
354
546
  finally:
355
547
  self._pubsub_lock.release()
356
- return await response_future
548
+ await response_future
549
+ return response_future.result()
357
550
 
358
- def try_get_pubsub_message(self) -> Optional[CoreCommands.PubSubMsg]:
551
+ def try_get_pubsub_message(self) -> Optional[PubSubMsg]:
359
552
  if self._is_closed:
360
553
  raise ClosingError(
361
554
  "Unable to execute requests; the client is closed. Please create a new client."
@@ -372,7 +565,7 @@ class BaseClient(CoreCommands):
372
565
  )
373
566
 
374
567
  # locking might not be required
375
- msg: Optional[CoreCommands.PubSubMsg] = None
568
+ msg: Optional[PubSubMsg] = None
376
569
  try:
377
570
  self._pubsub_lock.acquire()
378
571
  self._complete_pubsub_futures_safe()
@@ -386,12 +579,11 @@ class BaseClient(CoreCommands):
386
579
  def _cancel_pubsub_futures_with_exception_safe(self, exception: ConnectionError):
387
580
  while len(self._pubsub_futures):
388
581
  next_future = self._pubsub_futures.pop(0)
389
- if not next_future.cancelled():
390
- next_future.set_exception(exception)
582
+ next_future.set_exception(exception)
391
583
 
392
584
  def _notification_to_pubsub_message_safe(
393
585
  self, response: Response
394
- ) -> Optional[CoreCommands.PubSubMsg]:
586
+ ) -> Optional[PubSubMsg]:
395
587
  pubsub_message = None
396
588
  push_notification = cast(
397
589
  Dict[str, Any], value_from_pointer(response.resp_pointer)
@@ -410,11 +602,11 @@ class BaseClient(CoreCommands):
410
602
  ):
411
603
  values: List = push_notification["values"]
412
604
  if message_kind == "PMessage":
413
- pubsub_message = BaseClient.PubSubMsg(
605
+ pubsub_message = PubSubMsg(
414
606
  message=values[2], channel=values[1], pattern=values[0]
415
607
  )
416
608
  else:
417
- pubsub_message = BaseClient.PubSubMsg(
609
+ pubsub_message = PubSubMsg(
418
610
  message=values[1], channel=values[0], pattern=None
419
611
  )
420
612
  elif (
@@ -467,10 +659,16 @@ class BaseClient(CoreCommands):
467
659
  if response.HasField("closing_error")
468
660
  else f"Client Error - closing due to unknown error. callback index: {response.callback_idx}"
469
661
  )
662
+ exc = ClosingError(err_msg)
470
663
  if res_future is not None:
471
- res_future.set_exception(ClosingError(err_msg))
472
- await self.close(err_msg)
473
- raise ClosingError(err_msg)
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
474
672
  else:
475
673
  self._available_callback_indexes.append(response.callback_idx)
476
674
  if response.HasField("request_error"):
@@ -483,6 +681,10 @@ class BaseClient(CoreCommands):
483
681
  else:
484
682
  res_future.set_result(None)
485
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
+
486
688
  async def _process_push(self, response: Response) -> None:
487
689
  if response.HasField("closing_error") or not response.HasField("resp_pointer"):
488
690
  err_msg = (
@@ -490,9 +692,7 @@ class BaseClient(CoreCommands):
490
692
  if response.HasField("closing_error")
491
693
  else "Client Error - push notification without resp_pointer"
492
694
  )
493
- await self.close(err_msg)
494
695
  raise ClosingError(err_msg)
495
-
496
696
  try:
497
697
  self._pubsub_lock.acquire()
498
698
  callback, context = self.config._get_pubsub_callback_and_context()
@@ -508,30 +708,36 @@ class BaseClient(CoreCommands):
508
708
 
509
709
  async def _reader_loop(self) -> None:
510
710
  # Socket reader loop
511
- remaining_read_bytes = bytearray()
512
- while True:
513
- read_bytes = await self._reader.read(DEFAULT_READ_BYTES_SIZE)
514
- if len(read_bytes) == 0:
515
- err_msg = "The communication layer was unexpectedly closed."
516
- await self.close(err_msg)
517
- raise ClosingError(err_msg)
518
- read_bytes = remaining_read_bytes + bytearray(read_bytes)
519
- read_bytes_view = memoryview(read_bytes)
520
- offset = 0
521
- while offset <= len(read_bytes):
711
+ try:
712
+ remaining_read_bytes = bytearray()
713
+ while True:
522
714
  try:
523
- response, offset = ProtobufCodec.decode_delimited(
524
- read_bytes, read_bytes_view, offset, Response
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."
525
719
  )
526
- except PartialMessageException:
527
- # Received only partial response, break the inner loop
528
- remaining_read_bytes = read_bytes[offset:]
529
- break
530
- response = cast(Response, response)
531
- if response.is_push:
532
- await self._process_push(response=response)
533
- else:
534
- await self._process_response(response=response)
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))
535
741
 
536
742
  async def get_statistics(self) -> dict:
537
743
  return get_statistics()
@@ -552,12 +758,22 @@ class BaseClient(CoreCommands):
552
758
  self.config.credentials.password = password or ""
553
759
  return response
554
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
+
555
770
 
556
771
  class GlideClusterClient(BaseClient, ClusterCommands):
557
772
  """
558
773
  Client used for connection to cluster servers.
774
+ Use :func:`~BaseClient.create` to request a client.
559
775
  For full documentation, see
560
- https://github.com/valkey-io/valkey-glide/wiki/Python-wrapper#cluster
776
+ [Valkey GLIDE Wiki](https://github.com/valkey-io/valkey-glide/wiki/Python-wrapper#cluster)
561
777
  """
562
778
 
563
779
  async def _cluster_scan(
@@ -596,8 +812,9 @@ class GlideClusterClient(BaseClient, ClusterCommands):
596
812
  class GlideClient(BaseClient, StandaloneCommands):
597
813
  """
598
814
  Client used for connection to standalone servers.
815
+ Use :func:`~BaseClient.create` to request a client.
599
816
  For full documentation, see
600
- https://github.com/valkey-io/valkey-glide/wiki/Python-wrapper#standalone
817
+ [Valkey GLIDE Wiki](https://github.com/valkey-io/valkey-glide/wiki/Python-wrapper#standalone)
601
818
  """
602
819
 
603
820