adjust herd and shard with mset, mget

Author: amirreza8002

django_valkey/client/herd.py

@@ -1,15 +1,14 @@
 import random
 import socket
 import time
-from collections import OrderedDict
-from typing import Tuple, Any
+from typing import Tuple, Any, Iterable
 
 from django.conf import settings
 from valkey import Valkey
 from valkey.exceptions import ConnectionError, ResponseError, TimeoutError
 from valkey.typing import KeyT, EncodableT
 
-from django_valkey.base_client import DEFAULT_TIMEOUT
+from django_valkey.base_client import DEFAULT_TIMEOUT, Backend
 from django_valkey.client.default import DefaultClient
 from django_valkey.exceptions import ConnectionInterrupted
 
@@ -113,13 +112,16 @@ def get_many(self, keys, version=None, client=None):
         if not keys:
             return {}
 
-        recovered_data = OrderedDict()
+        recovered_data = {}
 
         new_keys = [self.make_key(key, version=version) for key in keys]
         map_keys = dict(zip(new_keys, keys))
 
         try:
-            results = client.mget(*new_keys)
+            pipeline = client.pipeline()
+            for key in new_keys:
+                pipeline.get(key)
+            results = pipeline.execute()
         except _main_exceptions as e:
             raise ConnectionInterrupted(connection=client) from e
 
@@ -132,6 +134,32 @@ def get_many(self, keys, version=None, client=None):
 
         return recovered_data
 
+    def mget(
+        self,
+        keys: Iterable[KeyT],
+        version: int | None = None,
+        client: Backend | Any | None = None,
+    ) -> dict:
+        client = self._get_client(write=False, client=client)
+        if not keys:
+            return {}
+
+        recovered_data = {}
+
+        new_keys = [self.make_key(key, version=version) for key in keys]
+
+        try:
+            results = client.mget(new_keys)
+        except _main_exceptions as e:
+            raise ConnectionInterrupted(connection=client) from e
+
+        for key, value in zip(keys, results):
+            if value is None:
+                continue
+            val, refresh = self._unpack(self.decode(value))
+            recovered_data[key] = None if refresh else val
+        return recovered_data
+
     def set_many(
         self, data, timeout=DEFAULT_TIMEOUT, version=None, client=None, herd=True
     ):

django_valkey/client/sharded.py

@@ -98,6 +98,9 @@ def get_many(
             recovered_data[map_keys[key]] = value
         return recovered_data
 
+    def mget(self, *args, **kwargs):
+        raise NotImplementedError
+
     def set(
         self,
         key: KeyT,
@@ -142,6 +145,9 @@ def set_many(
         for key, value in data.items():
             self.set(key, value, timeout, version=version, client=client)
 
+    def mset(self, *args, **kwargs):
+        raise NotImplementedError
+
     def has_key(
         self, key: KeyT, version: int | None = None, client: Valkey | Any | None = None
     ) -> bool: