Add set_multi_cas for per-key CAS return from batched writes

set_multi's current return shape is a list of failed keys, which can't
carry a per-key CAS value. It also uses the quiet setq/addq opcodes,
which intentionally suppress successful responses -- so even if the
shape allowed it, the wire protocol wouldn't return a CAS per key.

Add a separate `set_multi_cas` method that uses the non-quiet set/add
opcodes (one response per key) and returns `{str_key: int | None}` for
every input key -- int on success, None on failure. The existing
`{(key, cas): value}` input syntax from set_multi is preserved; the
result dict is keyed by the string key regardless of which form was
passed.

For ReplicatingClient, set_multi_cas raises NotImplementedError when
the client has more than one server, for the same reason as the
single-key get_cas=True surfaces: each replica has its own CAS
counter, so any per-key value we returned would be unsafe to feed
back to cas(). Single-server ReplicatingClient and DistributedClient
work normally.

Author: alexmv

bmemcached/client/distributed.py

@@ -90,6 +90,36 @@ def set_multi(self, mappings, time=0, compress_level=-1):
 
         return list(returns)
 
+    def set_multi_cas(self, mappings, time=0, compress_level=-1):
+        """
+        Set multiple keys with their values on server, returning the new CAS
+        value for each successfully stored key.
+
+        :param mappings: A dict with keys/values. Keys may be (key, cas)
+            tuples as in set_multi.
+        :type mappings: dict
+        :param time: Time in seconds that your key will expire.
+        :type time: int
+        :param compress_level: How much to compress.
+            0 = no compression, 1 = fastest, 9 = slowest but best,
+            -1 = default compression level.
+        :type compress_level: int
+        :return: A dict keyed by the string key of every input mapping. The
+            value is the new CAS int on success or None on failure.
+        :rtype: dict
+        """
+        if not mappings:
+            return {}
+        result = {}
+        server_mappings = defaultdict(dict)
+        for key, value in mappings.items():
+            str_key = key[0] if isinstance(key, tuple) else key
+            server_key = self._get_server(str_key)
+            server_mappings[server_key][key] = value
+        for server, m in server_mappings.items():
+            result.update(server.set_multi_cas(m, time, compress_level))
+        return result
+
     def add(self, key, value, time=0, compress_level=-1, get_cas=False):
         """
         Add a key/value to server ony if it does not exist.

bmemcached/client/mixin.py

@@ -141,6 +141,9 @@ def cas(self, key, value, cas, time=0, compress_level=-1, get_cas=False):
     def set_multi(self, mappings, time=0, compress_level=-1):
         raise NotImplementedError()
 
+    def set_multi_cas(self, mappings, time=0, compress_level=-1):
+        raise NotImplementedError()
+
     def add(self, key, value, time=0, compress_level=-1, get_cas=False):
         raise NotImplementedError()
 

bmemcached/client/replicating.py

@@ -279,6 +279,38 @@ def set_multi(self, mappings, time=0, compress_level=-1):
 
         return list(returns)
 
+    def set_multi_cas(self, mappings, time=0, compress_level=-1):
+        """
+        Set multiple keys with their values on the server, returning the new
+        CAS value for each successfully stored key.
+
+        Only supported when the client is configured with a single server;
+        see the class docstring for why CAS and multi-server replication
+        don't mix.
+
+        :param mappings: A dict with keys/values. Keys may be (key, cas)
+            tuples as in set_multi.
+        :type mappings: dict
+        :param time: Time in seconds that your key will expire.
+        :type time: int
+        :param compress_level: How much to compress.
+            0 = no compression, 1 = fastest, 9 = slowest but best,
+            -1 = default compression level.
+        :type compress_level: int
+        :return: A dict keyed by the string key of every input mapping. The
+            value is the new CAS int on success or None on failure.
+        :rtype: dict
+        :raises NotImplementedError: if more than one server is configured.
+        """
+        if len(self._servers) > 1:
+            raise NotImplementedError(
+                "set_multi_cas is not supported on ReplicatingClient with "
+                "more than one server."
+            )
+        if not mappings:
+            return {}
+        return self._servers[0].set_multi_cas(mappings, time, compress_level=compress_level)
+
     def add(self, key, value, time=0, compress_level=-1, get_cas=False):
         """
         Add a key/value to server ony if it does not exist.

bmemcached/protocol.py

@@ -789,6 +789,72 @@ def set_multi(self, mappings, time=100, compress_level=-1):
 
         return failed
 
+    def set_multi_cas(self, mappings, time=100, compress_level=-1):
+        """
+        Set multiple keys with their values on server and return the new CAS
+        value for each successfully stored key.
+
+        If a key is a (key, cas) tuple, insert as if cas(key, value, cas) had
+        been called. A cas of 0 means add-if-not-exists.
+
+        Unlike set_multi, this uses the non-quiet set/add opcodes so that the
+        server responds to every request; this costs one response per key but
+        is what allows per-key CAS values to be returned.
+
+        :param mappings: A dict with keys/values
+        :type mappings: dict
+        :param time: Time in seconds that your key will expire.
+        :type time: int
+        :param compress_level: How much to compress.
+            0 = no compression, 1 = fastest, 9 = slowest but best,
+            -1 = default compression level.
+        :type compress_level: int
+        :return: A dict keyed by the string key of every input mapping. The
+            value is the new CAS int on success or None on failure.
+        :rtype: dict
+        """
+        mappings = list(mappings.items())
+        msg = bytearray()
+        result = {}
+
+        for opaque, (key, value) in enumerate(mappings):
+            if isinstance(key, tuple):
+                str_key, cas = key
+            else:
+                str_key, cas = key, None
+            result[str_key] = None
+
+            if cas == 0:
+                command = 'add'
+            else:
+                command = 'set'
+
+            keybytes = str_to_bytes(str_key)
+            flags, value = self.serialize(value, compress_level=compress_level)
+            msg += struct.pack(self.HEADER_STRUCT +
+                               self.COMMANDS[command]['struct'] % (len(keybytes), len(value)),
+                               self.MAGIC['request'],
+                               self.COMMANDS[command]['command'],
+                               len(keybytes),
+                               8, 0, 0, len(keybytes) + len(value) + 8, opaque, cas or 0,
+                               flags, time, keybytes, value)
+
+        self._send(msg)
+
+        # Non-quiet set/add return exactly one response per request, so we can
+        # read a fixed count rather than relying on a trailing noop sentinel.
+        for _ in range(len(mappings)):
+            (magic, opcode, keylen, extlen, datatype, status, bodylen, opaque,
+             cas, extra_content) = self._get_response()
+            if status == self.STATUS['server_disconnected']:
+                return result
+            if status == self.STATUS['success']:
+                key, value = mappings[opaque]
+                str_key = key[0] if isinstance(key, tuple) else key
+                result[str_key] = cas
+
+        return result
+
     def _incr_decr(self, command, key, value, default, time):
         """
         Function which increments and decrements.

test/test_simple_functions.py

@@ -28,6 +28,7 @@ def tearDown(self):
     def reset(self):
         self.client.delete('test_key')
         self.client.delete('test_key2')
+        self.client.delete('fresh_key')
 
     def testSet(self):
         self.assertTrue(self.client.set('test_key', 'test'))
@@ -121,6 +122,51 @@ def testMultiCas(self):
         }), [])
         self.assertEqual(self.client.get('test_key'), 'value4')
 
+    def testSetMultiCas(self):
+        # All-success plain keys: every input gets a non-None CAS, and each
+        # returned CAS matches what gets() reports afterwards.
+        result = self.client.set_multi_cas({
+            'test_key': 'value1',
+            'test_key2': 'value2',
+        })
+        self.assertEqual(set(result.keys()), {'test_key', 'test_key2'})
+        self.assertTrue(result['test_key'] is not None)
+        self.assertTrue(result['test_key2'] is not None)
+        _, cas1 = self.client.gets('test_key')
+        _, cas2 = self.client.gets('test_key2')
+        self.assertEqual(result['test_key'], cas1)
+        self.assertEqual(result['test_key2'], cas2)
+
+        # CAS failure: add-if-not-exists when the key already exists returns
+        # None for that key; unrelated keys still succeed.
+        result = self.client.set_multi_cas({
+            ('test_key', 0): 'shouldnt_store',
+            'fresh_key': 'fresh',
+        })
+        self.assertTrue(result['test_key'] is None)
+        self.assertTrue(result['fresh_key'] is not None)
+        self.assertEqual(self.client.get('test_key'), 'value1')
+        self.client.delete('fresh_key')
+
+        # Stale-CAS failure: capture cas, mutate out of band, then set_multi_cas
+        # with the stale cas must fail and leave the out-of-band value intact.
+        _, stale_cas = self.client.gets('test_key')
+        self.client.set('test_key', 'other')
+        result = self.client.set_multi_cas({
+            ('test_key', stale_cas): 'should_fail',
+        })
+        self.assertTrue(result['test_key'] is None)
+        self.assertEqual(self.client.get('test_key'), 'other')
+
+        # Returned CAS is usable directly in cas() without a gets() round-trip.
+        self.client.delete('test_key')
+        result = self.client.set_multi_cas({'test_key': 'v'})
+        self.assertTrue(self.client.cas('test_key', 'v2', result['test_key']))
+        self.assertEqual(self.client.get('test_key'), 'v2')
+
+    def testSetMultiCasEmpty(self):
+        self.assertEqual(self.client.set_multi_cas({}), {})
+
     def testGetMultiCas(self):
         self.client.set('test_key', 'value1')
         self.client.set('test_key2', 'value2')
@@ -374,6 +420,8 @@ def testGetCasMultiReplicaRaises(self):
                 client.replace('test_key', 'v', get_cas=True)
             with self.assertRaises(NotImplementedError):
                 client.cas('test_key', 'v', None, get_cas=True)
+            with self.assertRaises(NotImplementedError):
+                client.set_multi_cas({'test_key': 'v'})
 
             # get_cas=False (default) still works fine on multi-replica.
             self.assertTrue(client.set('test_key', 'v'))