diff options
author | BrianEaton1 <eaton.brian@gmail.com> | 2019-05-07 09:38:58 -0400 |
---|---|---|
committer | Stephen SORRIAUX <stephen.sorriaux@gmail.com> | 2019-05-07 15:38:58 +0200 |
commit | 88b657a0977161f3815657878ba48f82a97a3846 (patch) | |
tree | 48d7cfcb8d11413a839d7f942345a7e41280d711 | |
parent | cd49b3fa01136848c5e6bfafb4c241b9704f249d (diff) | |
download | kazoo-88b657a0977161f3815657878ba48f82a97a3846.tar.gz |
feat(recipe): add support for curator SharedCount recipe (#559)
* feat(recipe): add support for curator SharedCount recipe
This feature allows Java clients using curator's SharedCount recipe
and python clients using kazoo's Counter recipe to read and
write from the same path without receiving type errors.
example use:
counter = zk.Counter("/curator", support_curator=True)
counter += 2
counter -= 1
counter.value == 1
counter.pre_value == 2
counter.post_value == 1
Closes #558
-rw-r--r-- | kazoo/recipe/counter.py | 35 | ||||
-rw-r--r-- | kazoo/tests/test_counter.py | 18 |
2 files changed, 48 insertions, 5 deletions
diff --git a/kazoo/recipe/counter.py b/kazoo/recipe/counter.py index b728bc2..9e68849 100644 --- a/kazoo/recipe/counter.py +++ b/kazoo/recipe/counter.py @@ -6,7 +6,7 @@ """ from kazoo.exceptions import BadVersionError from kazoo.retry import ForceRetryError - +import struct class Counter(object): """Kazoo Counter @@ -19,6 +19,12 @@ class Counter(object): `type(counter.default)(value)` both using an ascii encoding. As such other data types might be used for the counter value. + If you would like to support clients updating the same znode path using + either kazoo's counter recipe or curator's SharedCount recipe, you will + need to enable the support_curator flag. This flag limits + support to integers only and does not use ascii encoding as described + above. + Counter changes can raise :class:`~kazoo.exceptions.BadVersionError` if the retry policy wasn't able to apply a change. @@ -42,22 +48,35 @@ class Counter(object): counter.pre_value == 1.0 counter.post_value == 3.0 + counter = zk.Counter("/curator", support_curator=True) + counter += 2 + counter -= 1 + counter.value == 1 + counter.pre_value == 2 + counter.post_value == 1 + """ - def __init__(self, client, path, default=0): + def __init__(self, client, path, default=0, support_curator=False): """Create a Kazoo Counter :param client: A :class:`~kazoo.client.KazooClient` instance. :param path: The counter path to use. - :param default: The default value. + :param default: The default value to use for new counter paths. + :param support_curator: Enable if support for curator's SharedCount + recipe is desired. """ self.client = client self.path = path self.default = default self.default_type = type(default) + self.support_curator = support_curator self._ensured_path = False self.pre_value = None self.post_value = None + if self.support_curator and not isinstance(self.default, int): + raise TypeError("when support_curator is enabled the default " + "type must be an int") def _ensure_node(self): if not self._ensured_path: @@ -68,7 +87,10 @@ class Counter(object): def _value(self): self._ensure_node() old, stat = self.client.get(self.path) - old = old.decode('ascii') if old != b'' else self.default + if self.support_curator: + old = struct.unpack(">i", old)[0] if old != b'' else self.default + else: + old = old.decode('ascii') if old != b'' else self.default version = stat.version data = self.default_type(old) return data, version @@ -86,7 +108,10 @@ class Counter(object): def _inner_change(self, value): self.pre_value, version = self._value() post_value = self.pre_value + value - data = repr(post_value).encode('ascii') + if self.support_curator: + data = struct.pack(">i", post_value) + else: + data = repr(post_value).encode('ascii') try: self.client.set(self.path, data, version=version) except BadVersionError: # pragma: nocover diff --git a/kazoo/tests/test_counter.py b/kazoo/tests/test_counter.py index 50095a2..47e077e 100644 --- a/kazoo/tests/test_counter.py +++ b/kazoo/tests/test_counter.py @@ -21,6 +21,22 @@ class KazooCounterTests(KazooTestCase): counter - 1 eq_(counter.value, -1) + def test_int_curator_counter(self): + counter = self._makeOne(support_curator=True) + eq_(counter.value, 0) + counter += 2 + counter + 1 + eq_(counter.value, 3) + counter -= 3 + counter - 1 + eq_(counter.value, -1) + counter += 1 + counter += 2147483647 + eq_(counter.value, 2147483647) + counter -= 2147483647 + counter -= 2147483647 + eq_(counter.value, -2147483647) + def test_float_counter(self): counter = self._makeOne(default=0.0) eq_(counter.value, 0.0) @@ -33,6 +49,8 @@ class KazooCounterTests(KazooTestCase): counter = self._makeOne() self.assertRaises(TypeError, counter.__add__, 2.1) self.assertRaises(TypeError, counter.__add__, b"a") + with self.assertRaises(TypeError): + counter = self._makeOne(default=0.0, support_curator=True) def test_pre_post_values(self): counter = self._makeOne() |