diff --git a/aiomcache/client.py b/aiomcache/client.py index ef20418..c066117 100644 --- a/aiomcache/client.py +++ b/aiomcache/client.py @@ -77,7 +77,7 @@ def close(self): @asyncio.coroutine def _multi_get(self, conn, *keys): # req - get [ ...]\r\n - # resp - VALUE []\r\n + # resp - VALUE \r\n # \r\n (if exists) # [...] # END\r\n @@ -118,6 +118,51 @@ def _multi_get(self, conn, *keys): raise ClientException('received too many responses') return [received.get(k, None) for k in keys] + @asyncio.coroutine + def _multi_gets(self, conn, *keys): + # req - get [ ...]\r\n + # resp - VALUE \r\n + # \r\n (if exists) + # [...] + # END\r\n + if not keys: + return [] + + [self._validate_key(key) for key in keys] + if len(set(keys)) != len(keys): + raise ClientException('duplicate keys passed to multi_get') + + conn.writer.write(b'gets ' + b' '.join(keys) + b'\r\n') + + received = {} + line = yield from conn.reader.readline() + + while line != b'END\r\n': + terms = line.split() + + if len(terms) == 5 and terms[0] == b'VALUE': # exists + key = terms[1] + flags = int(terms[2]) + length = int(terms[3]) + cas = int(terms[4]) + + if flags != 0: + raise ClientException('received non zero flags') + + val = (yield from conn.reader.readexactly(length+2))[:-2] + if key in received: + raise ClientException('duplicate results from server') + + received[key] = val, cas + else: + raise ClientException('get failed', line) + + line = yield from conn.reader.readline() + + if len(received) > len(keys): + raise ClientException('received too many responses') + return [received.get(k, [None, None]) for k in keys] + @acquire def delete(self, conn, key): """Deletes a key/value pair from the server. @@ -146,6 +191,19 @@ def get(self, conn, key, default=None): result = yield from self._multi_get(conn, key) return (result[0] or default) if result else default + @acquire + def gets(self, conn, key, default=None): + """Gets a single value from the server. + + :param key: ``bytes``, is the key for the item being fetched + :param default: default value if there is no value. + :return: ``bytes``, is the data for this specified key. + """ + result = yield from self._multi_gets(conn, key) + if result and result[0][0] is not None: + return result[0] + return default, None + @acquire def multi_get(self, conn, *keys): """Takes a list of keys and returns a list of values. @@ -189,10 +247,13 @@ def stats(self, conn, args=None): @asyncio.coroutine def _storage_command(self, conn, command, key, value, - flags=0, exptime=0): + flags=0, exptime=0, cas=None): # req - set [noreply]\r\n # \r\n # resp - STORED\r\n (or others) + # req - set [noreply]\r\n + # \r\n + # resp - STORED\r\n (or others) # typically, if val is > 1024**2 bytes server returns: # SERVER_ERROR object too large for cache\r\n @@ -206,11 +267,14 @@ def _storage_command(self, conn, command, key, value, raise ValidationException('exptime negative', exptime) args = [str(a).encode('utf-8') for a in (flags, exptime, len(value))] - _cmd = b' '.join([command, key] + args) + b'\r\n' - cmd = _cmd + value + b'\r\n' + _cmd = b' '.join([command, key] + args) + if cas: + _cmd += b' ' + str(cas).encode('utf-8') + cmd = _cmd + b'\r\n' + value + b'\r\n' resp = yield from self._execute_simple_command(conn, cmd) - if resp not in (const.STORED, const.NOT_STORED): + if resp not in ( + const.STORED, const.NOT_STORED, const.EXISTS, const.NOT_FOUND): raise ClientException('stats {} failed'.format(command), resp) return resp == const.STORED @@ -230,6 +294,25 @@ def set(self, conn, key, value, exptime=0): conn, b'set', key, value, flags, exptime) return resp + @acquire + def cas(self, conn, key, value, cas_token, exptime=0): + """Sets a key to a value on the server + with an optional exptime (0 means don't auto-expire) + only if value hasn't change from first retrieval + + :param key: ``bytes``, is the key of the item. + :param value: ``bytes``, data to store. + :param exptime: ``int``, is expiration time. If it's 0, the + item never expires. + :param cas_token: ``int``, unique cas token retrieve from previous + ``gets`` + :return: ``bool``, True in case of success. + """ + flags = 0 # TODO: fix when exception removed + resp = yield from self._storage_command( + conn, b'cas', key, value, flags, exptime, cas=cas_token) + return resp + @acquire def add(self, conn, key, value, exptime=0): """Store this data, but only if the server *doesn't* already diff --git a/aiomcache/constants.py b/aiomcache/constants.py index fef65f4..bea058c 100644 --- a/aiomcache/constants.py +++ b/aiomcache/constants.py @@ -4,4 +4,5 @@ NOT_FOUND = b'NOT_FOUND' DELETED = b'DELETED' VERSION = b'VERSION' +EXISTS = b'EXISTS' OK = b'OK' diff --git a/tests/commands_test.py b/tests/commands_test.py index 9e83125..0154502 100644 --- a/tests/commands_test.py +++ b/tests/commands_test.py @@ -49,7 +49,7 @@ def test_set_get(mcache, loop): test_value = yield from mcache.get(b'not:' + key) assert test_value is None test_value = yield from mcache.get(b'not:' + key, default=value) - assert test_value is value + assert test_value == value with mock.patch.object(mcache, '_execute_simple_command') as patched: fut = asyncio.Future(loop=loop) @@ -59,6 +59,24 @@ def test_set_get(mcache, loop): yield from mcache.set(key, value) +@pytest.mark.run_loop +def test_gets(mcache, loop): + key, value = b'key:set', b'1' + yield from mcache.set(key, value) + + test_value, cas = yield from mcache.gets(key) + assert test_value == value + assert isinstance(cas, int) + + test_value, cas = yield from mcache.gets(b'not:' + key) + assert test_value is None + assert cas is None + + test_value, cas = yield from mcache.gets(b'not:' + key, default=value) + assert test_value == value + assert cas is None + + @pytest.mark.run_loop def test_multi_get(mcache): key1, value1 = b'key:multi_get:1', b'1' @@ -108,6 +126,27 @@ def test_set_errors(mcache): yield from mcache.set(key, value, exptime=3.14) +@pytest.mark.run_loop +def test_gets_cas(mcache, loop): + key, value = b'key:set', b'1' + yield from mcache.set(key, value) + + test_value, cas = yield from mcache.gets(key) + + stored = yield from mcache.cas(key, value, cas) + assert stored is True + + stored = yield from mcache.cas(key, value, cas) + assert stored is False + + +@pytest.mark.run_loop +def test_cas_missing(mcache, loop): + key, value = b'key:set', b'1' + stored = yield from mcache.cas(key, value, 123) + assert stored is False + + @pytest.mark.run_loop def test_add(mcache): key, value = b'key:add', b'1'