Adopt multi get

Author: bisho

pyproject.toml

@@ -3,7 +3,7 @@ dependencies = [
     "uhashring >= 2.3, < 3",
     "marisa-trie >= 1.2.1, < 1.3",
     "zstandard >= 0.25.0, < 0.26",
-    "meta-memcache-socket >= 0.2.0, < 0.3.0",
+    "meta-memcache-socket >= 0.2.1, < 0.3.0",
 ]
 authors = [{ name = "Guillermo Perez", email = "bisho@revenuecat.com" }]
 license = { text = "MIT License" }

src/meta_memcache/executors/default.py

@@ -159,6 +159,37 @@ def exec_on_pool(
             else:
                 return NotStored()
 
+    def _exec_multi_get_on_pool(
+        self,
+        pool: ConnectionPool,
+        keys: List[Union[bytes, str]],
+        flags: Optional[RequestFlags],
+        raise_on_server_error: Optional[bool] = None,
+    ) -> Dict[Key, MemcacheResponse]:
+        try:
+            conn = pool.pop_connection()
+            error = False
+            try:
+                responses = conn.meta_multiget(keys, flags)
+            except Exception as e:
+                error = True
+                raise MemcacheServerError(pool.server, "Memcache error") from e
+            finally:
+                pool.release_connection(conn, error=error)
+            return {
+                key: self._process_response(response)
+                for key, response in zip(keys, responses, strict=True)
+            }
+        except MemcacheServerError:
+            raise_on_server_error = (
+                raise_on_server_error
+                if raise_on_server_error is not None
+                else self._raise_on_server_error
+            )
+            if raise_on_server_error:
+                raise
+            return {key: Miss() for key in keys}
+
     def exec_multi_on_pool(  # noqa: C901
         self,
         pool: ConnectionPool,
@@ -168,12 +199,16 @@ def exec_multi_on_pool(  # noqa: C901
         track_write_failures: bool,
         raise_on_server_error: Optional[bool] = None,
     ) -> Dict[Key, MemcacheResponse]:
+        if command == MetaCommand.META_GET:
+            return self._exec_multi_get_on_pool(
+                pool, [key.key for key, _ in key_values], flags, raise_on_server_error
+            )
+
         results: Dict[Key, MemcacheResponse] = {}
         try:
             conn = pool.pop_connection()
             error = False
             try:
-                # with pool.get_connection() as conn:
                 for key, value in key_values:
                     cmd_value, flags = (
                         (None, flags)

tests/commands_test.py

@@ -1131,7 +1131,7 @@ def test_multi_get_with_values(
     int_encoded = serializer.serialize(Key("k2"), 42)
     list_encoded = serializer.serialize(Key("k3"), [1, 2])
 
-    memcache_socket.get_response.side_effect = [
+    memcache_socket.meta_multiget.return_value = [
         Value(
             size=len(str_encoded.data),
             value=str_encoded.data,
@@ -1151,12 +1151,9 @@ def test_multi_get_with_values(
 
     results = cache_client.multi_get(keys=[Key("k1"), Key("k2"), Key("k3")])
 
-    calls = memcache_socket.send_meta_get.call_args_list
-    assert len(calls) == 3
-    assert calls[0][0][0] == "k1"
-    assert calls[1][0][0] == "k2"
-    assert calls[2][0][0] == "k3"
-    assert memcache_socket.get_response.call_count == 3
+    memcache_socket.meta_multiget.assert_called_once()
+    keys_sent = memcache_socket.meta_multiget.call_args.args[0]
+    assert keys_sent == ["k1", "k2", "k3"]
     assert results == {
         Key("k1"): "hello",
         Key("k2"): 42,
@@ -1167,8 +1164,8 @@ def test_multi_get_with_values(
 def test_multi_get_with_touch_ttl(
     memcache_socket: MemcacheSocket, cache_client: CacheClient
 ) -> None:
-    """Test multi_get passes touch_ttl through to the pipelining path."""
-    memcache_socket.get_response.side_effect = [
+    """Test multi_get passes touch_ttl through to meta_multiget."""
+    memcache_socket.meta_multiget.return_value = [
         Miss(),
         Miss(),
     ]
@@ -1178,54 +1175,49 @@ def test_multi_get_with_touch_ttl(
         touch_ttl=300,
     )
 
-    calls = memcache_socket.send_meta_get.call_args_list
-    assert len(calls) == 2
-    assert calls[0][0][0] == "a"
-    assert calls[1][0][0] == "b"
-    # Verify the flags contain the touch TTL
-    for c in calls:
-        flags = c[0][1]
-        assert flags is not None
-        assert flags.cache_ttl == 300
+    memcache_socket.meta_multiget.assert_called_once()
+    keys_sent = memcache_socket.meta_multiget.call_args.args[0]
+    flags = memcache_socket.meta_multiget.call_args.kwargs["request_flags"]
+    assert keys_sent == ["a", "b"]
+    assert flags is not None
+    assert flags.cache_ttl == 300
     assert results == {Key("a"): None, Key("b"): None}
 
 
 def test_multi_get_with_recache(
     memcache_socket: MemcacheSocket, cache_client: CacheClient
 ) -> None:
-    """Test multi_get with recache_policy passes recache_ttl in flags."""
-    memcache_socket.get_response.side_effect = [Miss()]
+    """Test multi_get with recache_policy passes recache_ttl in request flags."""
+    memcache_socket.meta_multiget.return_value = [Miss()]
 
     cache_client.multi_get(
         keys=[Key("x")],
         recache_policy=RecachePolicy(ttl=60),
     )
 
-    calls = memcache_socket.send_meta_get.call_args_list
-    assert len(calls) == 1
-    flags = calls[0][0][1]
+    memcache_socket.meta_multiget.assert_called_once()
+    flags = memcache_socket.meta_multiget.call_args.kwargs["request_flags"]
     assert flags is not None
     assert flags.recache_ttl == 60
 
 
 def test_meta_multiget_no_reply(
     memcache_socket: MemcacheSocket, cache_client: CacheClient
 ) -> None:
-    """Test meta_multiget with no_reply flag uses pipelining path correctly.
+    """Test meta_multiget with no_reply uses socket batch API."""
+    memcache_socket.meta_multiget.return_value = [
+        Success(flags=ResponseFlags()),
+        Success(flags=ResponseFlags()),
+    ]
 
-    With no_reply on mg, the executor skips get_response and returns Success.
-    """
     results = cache_client.meta_multiget(
         keys=[Key("a"), Key("b")],
         flags=RequestFlags(no_reply=True, cache_ttl=60),
     )
 
-    calls = memcache_socket.send_meta_get.call_args_list
-    assert len(calls) == 2
-    assert calls[0][0][0] == "a"
-    assert calls[1][0][0] == "b"
-    # no_reply: get_response not called, executor returns Success directly
-    memcache_socket.get_response.assert_not_called()
+    memcache_socket.meta_multiget.assert_called_once()
+    keys_sent = memcache_socket.meta_multiget.call_args.args[0]
+    assert keys_sent == ["a", "b"]
     assert all(isinstance(r, Success) for r in results.values())
 
 

tests/conftest.py

@@ -39,6 +39,9 @@ def mock_memcache_socket(mocker, mock_raw_socket) -> MagicMock:
     mock.get_version.return_value = ServerVersion.STABLE
     _success = Success(flags=ResponseFlags())
     mock.meta_get.return_value = Miss()
+    mock.meta_multiget.side_effect = lambda keys, request_flags=None: [
+        Miss() for _ in keys
+    ]
     mock.meta_set.return_value = _success
     mock.meta_delete.return_value = _success
     mock.meta_arithmetic.return_value = _success