From 4c72940009e91d5a9b94dbd7b5eb50047c216322 Mon Sep 17 00:00:00 2001 From: Charles-Henri de Boysson Date: Sun, 12 Feb 2023 10:55:44 -0500 Subject: [PATCH 1/2] feat(core): Add support for Create2 in transactions. --- kazoo/client.py | 9 ++++++++- kazoo/protocol/serialization.py | 9 +++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/kazoo/client.py b/kazoo/client.py index 3f2c3b94..554ff8a8 100644 --- a/kazoo/client.py +++ b/kazoo/client.py @@ -2000,6 +2000,7 @@ def create( acl: Sequence[ACL] | None = None, ephemeral: bool = False, sequence: bool = False, + include_data: bool = False, ) -> None: """Add a create ZNode to the transaction. Takes the same arguments as :meth:`KazooClient.create`, with the exception @@ -2023,6 +2024,8 @@ def create( raise TypeError("Invalid type for 'ephemeral' (bool expected)") if not isinstance(sequence, bool): raise TypeError("Invalid type for 'sequence' (bool expected)") + if not isinstance(include_data, bool): + raise TypeError("Invalid type for 'include_data' (bool expected)") flags = 0 if ephemeral: @@ -2031,9 +2034,13 @@ def create( flags |= 2 if acl is None: acl = OPEN_ACL_UNSAFE + if include_data: + opcode = Create2 + else: + opcode = Create self._add( - Create(_prefix_root(self.client.chroot, path), value, acl, flags), + opcode(_prefix_root(self.client.chroot, path), value, acl, flags), None, ) diff --git a/kazoo/protocol/serialization.py b/kazoo/protocol/serialization.py index 914540a8..72aa3c89 100644 --- a/kazoo/protocol/serialization.py +++ b/kazoo/protocol/serialization.py @@ -421,6 +421,11 @@ def deserialize( while not header.done: if header.type == Create.type: response, offset = read_string(bytes, offset) + elif header.type == Create2.type: + path, offset = read_string(bytes, offset) + stat = ZnodeStat._make(stat_struct.unpack_from(bytes, offset)) + offset += stat_struct.size + response = (path, stat) elif header.type == Delete.type: response = True elif header.type == SetData.type: @@ -445,6 +450,10 @@ def unchroot( for result in response: if isinstance(result, str): resp.append(client.unchroot(result)) + elif isinstance(result, ZnodeStat): # Need to test before tuple + resp.append(result) + elif isinstance(result, tuple): + resp.append((client.unchroot(result[0]), result[1])) else: resp.append(result) return resp From 1da625f200cae5b36f7a10e76dfe2c0f982c746c Mon Sep 17 00:00:00 2001 From: Charles-Henri de Boysson Date: Sun, 12 Feb 2023 10:55:58 -0500 Subject: [PATCH 2/2] feat(core): Add support for Container and TTL nodes Also add support through transactions. Closes #334, #496 --- kazoo/client.py | 247 ++++++++++++++++++-------------- kazoo/protocol/serialization.py | 49 +++++++ kazoo/testing/harness.py | 7 +- kazoo/tests/test_client.py | 28 ++++ 4 files changed, 222 insertions(+), 109 deletions(-) diff --git a/kazoo/client.py b/kazoo/client.py index 554ff8a8..0c7ac468 100644 --- a/kazoo/client.py +++ b/kazoo/client.py @@ -48,6 +48,8 @@ CloseInstance, Create, Create2, + CreateContainer, + CreateTTL, Delete, Exists, GetChildren, @@ -1116,6 +1118,8 @@ def create( sequence: bool = False, makepath: bool = False, include_data: Literal[False] = False, + container: bool = False, + ttl: int = 0, ) -> str: ... @@ -1129,6 +1133,8 @@ def create( sequence: bool = False, makepath: bool = False, include_data: Literal[True] = True, + container: bool = False, + ttl: int = 0, ) -> tuple[str, ZnodeStat]: ... @@ -1141,6 +1147,8 @@ def create( sequence: bool = False, makepath: bool = False, include_data: bool = False, + container: bool = False, + ttl: int = 0, ) -> str | tuple[str, ZnodeStat]: """Create a node with the given value as its data. Optionally set an ACL on the node. @@ -1218,6 +1226,9 @@ def create( The `makepath` option. .. versionadded:: 2.7 The `include_data` option. + .. versionadded:: 2.9 + The `container` and `ttl` options. + """ acl = acl or self.default_acl return cast( @@ -1230,6 +1241,8 @@ def create( sequence=sequence, makepath=makepath, include_data=include_data, + container=container, + ttl=ttl, ).get(), ) @@ -1242,6 +1255,8 @@ def create_async( sequence: bool = False, makepath: bool = False, include_data: bool = False, + container: bool = False, + ttl: int = 0, ) -> IAsyncResult: """Asynchronously create a ZNode. Takes the same arguments as :meth:`create`. @@ -1252,55 +1267,39 @@ def create_async( The makepath option. .. versionadded:: 2.7 The `include_data` option. + .. versionadded:: 2.9 + The `container` and `ttl` options. """ if acl is None and self.default_acl: acl = self.default_acl - if not isinstance(path, str): - raise TypeError("Invalid type for 'path' (string expected)") - if acl and ( - isinstance(acl, ACL) or not isinstance(acl, (tuple, list)) - ): - raise TypeError( - "Invalid type for 'acl' (acl must be a tuple/list" " of ACL's" - ) - if value is not None and not isinstance(value, bytes): - raise TypeError("Invalid type for 'value' (must be a byte string)") - if not isinstance(ephemeral, bool): - raise TypeError("Invalid type for 'ephemeral' (bool expected)") - if not isinstance(sequence, bool): - raise TypeError("Invalid type for 'sequence' (bool expected)") - if not isinstance(makepath, bool): - raise TypeError("Invalid type for 'makepath' (bool expected)") - if not isinstance(include_data, bool): - raise TypeError("Invalid type for 'include_data' (bool expected)") - - flags = 0 - if ephemeral: - flags |= 1 - if sequence: - flags |= 2 - if acl is None: - acl = OPEN_ACL_UNSAFE - + opcode = _create_opcode( + path, + value, + acl, + self.chroot, + ephemeral, + sequence, + include_data, + container, + ttl, + ) async_result = self.handler.async_result() @capture_exceptions(async_result) def do_create() -> None: - result = self._create_async_inner( - path, - value, - # The way acl is constructed ends up confusing mypy, which - # thinks that acl can be None here, even though the code - # above ensures that if acl is None, it gets set to - # OPEN_ACL_UNSAFE, so we ignore the type error here. - # behaves differently in python3.8 and python3.14, sigh. - acl, # type: ignore[arg-type] - flags, - trailing=sequence, - include_data=include_data, - ) - result.rawlink(create_completion) + inner_async_result = self.handler.async_result() + + call_result = self._call(opcode, inner_async_result) + if call_result is False: + # We hit a short-circuit exit on the _call. Because we are + # not using the original async_result here, we bubble the + # exception upwards to the do_create function in + # KazooClient.create so that it gets set on the correct + # async_result object + raise cast(Exception, inner_async_result.exception) + + inner_async_result.rawlink(create_completion) @capture_exceptions(async_result) def retry_completion(result: IAsyncResult) -> None: @@ -1312,11 +1311,11 @@ def create_completion( result: IAsyncResult, ) -> str | tuple[str, ZnodeStat] | None: try: - if include_data: + if opcode.type == Create.type: + return self.unchroot(result.get()) + else: new_path, stat = result.get() return self.unchroot(new_path), stat - else: - return self.unchroot(result.get()) except NoNodeError: if not makepath: raise @@ -1330,39 +1329,6 @@ def create_completion( do_create() return async_result - def _create_async_inner( - self, - path: str, - value: bytes | None, - acl: Sequence[ACL], - flags: int, - trailing: bool = False, - include_data: bool = False, - ) -> IAsyncResult: - async_result = self.handler.async_result() - opcode = Create2 if include_data else Create - - call_result = self._call( - opcode( - _prefix_root(self.chroot, path, trailing=trailing), - value, - acl, - flags, - ), - async_result, - ) - if call_result is False: - # We hit a short-circuit exit on the _call. Because we are - # not using the original async_result here, we bubble the - # exception upwards to the do_create function in - # KazooClient.create so that it gets set on the correct - # async_result object - # Note: Do we actually need call_result? It seems like we could - # just check the state of the exception, and avoid the typing - # stuff. - raise async_result.exception # type: ignore[misc] - return async_result - def ensure_path(self, path: str, acl: Sequence[ACL] | None = None) -> bool: """Recursively create a path if it doesn't exist. @@ -2001,6 +1967,8 @@ def create( ephemeral: bool = False, sequence: bool = False, include_data: bool = False, + container: bool = False, + ttl: int = 0, ) -> None: """Add a create ZNode to the transaction. Takes the same arguments as :meth:`KazooClient.create`, with the exception @@ -2008,41 +1976,24 @@ def create( :returns: None + .. versionadded:: 2.9 + The `include_data`, `container` and `ttl` options. """ if acl is None and self.client.default_acl: acl = self.client.default_acl - if not isinstance(path, str): - raise TypeError("Invalid type for 'path' (string expected)") - if acl and not isinstance(acl, (tuple, list)): - raise TypeError( - "Invalid type for 'acl' (acl must be a tuple/list" " of ACL's" - ) - if not isinstance(value, bytes): - raise TypeError("Invalid type for 'value' (must be a byte string)") - if not isinstance(ephemeral, bool): - raise TypeError("Invalid type for 'ephemeral' (bool expected)") - if not isinstance(sequence, bool): - raise TypeError("Invalid type for 'sequence' (bool expected)") - if not isinstance(include_data, bool): - raise TypeError("Invalid type for 'include_data' (bool expected)") - - flags = 0 - if ephemeral: - flags |= 1 - if sequence: - flags |= 2 - if acl is None: - acl = OPEN_ACL_UNSAFE - if include_data: - opcode = Create2 - else: - opcode = Create - - self._add( - opcode(_prefix_root(self.client.chroot, path), value, acl, flags), - None, + opcode = _create_opcode( + path, + value, + acl, + self.client.chroot, + ephemeral, + sequence, + include_data, + container, + ttl, ) + self._add(opcode, None) def delete(self, path: str, version: int = -1) -> None: """Add a delete ZNode to the transaction. Takes the same @@ -2133,3 +2084,85 @@ def _add( self._check_tx_state() self.client.logger.log(BLATHER, "Added %r to %r", request, self) self.operations.append(request) + + +def _create_opcode( + path: str, + value: bytes | None, + acl: Sequence[ACL] | None, + chroot: str | None, + ephemeral: bool, + sequence: bool, + include_data: bool, + container: bool, + ttl: int, +) -> Create | Create2 | CreateContainer | CreateTTL: + """Helper function. + Creates the create OpCode for regular `client.create()` operations as + well as in a `client.transaction()` context. + """ + if not isinstance(path, str): + raise TypeError("Invalid type for 'path' (string expected)") + if acl and (isinstance(acl, ACL) or not isinstance(acl, (tuple, list))): + raise TypeError( + "Invalid type for 'acl' (acl must be a tuple/list" " of ACL's" + ) + if value is not None and not isinstance(value, bytes): + raise TypeError("Invalid type for 'value' (must be a byte string)") + if not isinstance(ephemeral, bool): + raise TypeError("Invalid type for 'ephemeral' (bool expected)") + if not isinstance(sequence, bool): + raise TypeError("Invalid type for 'sequence' (bool expected)") + if not isinstance(include_data, bool): + raise TypeError("Invalid type for 'include_data' (bool expected)") + if not isinstance(container, bool): + raise TypeError("Invalid type for 'container' (bool expected)") + if not isinstance(ttl, int) or ttl < 0: + raise TypeError("Invalid 'ttl' (integer >= 0 expected)") + if ttl and ephemeral: + raise TypeError("Invalid node creation: ephemeral & ttl") + if container and (ephemeral or sequence or ttl): + raise TypeError( + "Invalid node creation: container & ephemeral/sequence/ttl" + ) + + # Should match Zookeeper's CreateMode fromFlag + # https://github.com/apache/zookeeper/blob/master/zookeeper-server/ + # src/main/java/org/apache/zookeeper/CreateMode.java#L112 + flags = 0 + if ephemeral: + flags |= 1 + if sequence: + flags |= 2 + if container: + flags = 4 + if ttl: + if sequence: + flags = 6 + else: + flags = 5 + + if acl is None: + acl = OPEN_ACL_UNSAFE + + # Figure out the OpCode we are going to send + if include_data: + return Create2( + _prefix_root(chroot, path, trailing=sequence), value, acl, flags + ) + elif container: + return CreateContainer( + _prefix_root(chroot, path, trailing=False), value, acl, flags + ) + elif ttl: + return CreateTTL( + _prefix_root(chroot, path, trailing=sequence), + value, + acl, + flags, + ttl, + ) + else: + return Create( + _prefix_root(chroot, path, trailing=sequence), value, acl, flags + ) diff --git a/kazoo/protocol/serialization.py b/kazoo/protocol/serialization.py index 72aa3c89..389b39f6 100644 --- a/kazoo/protocol/serialization.py +++ b/kazoo/protocol/serialization.py @@ -515,6 +515,55 @@ def deserialize( return data, stat +class CreateContainer(namedtuple("CreateContainer", "path data acl flags")): + type = 19 + + def serialize(self): + b = bytearray() + b.extend(write_string(self.path)) + b.extend(write_buffer(self.data)) + b.extend(int_struct.pack(len(self.acl))) + for acl in self.acl: + b.extend( + int_struct.pack(acl.perms) + + write_string(acl.id.scheme) + + write_string(acl.id.id) + ) + b.extend(int_struct.pack(self.flags)) + return b + + @classmethod + def deserialize(cls, bytes, offset): + path, offset = read_string(bytes, offset) + stat = ZnodeStat._make(stat_struct.unpack_from(bytes, offset)) + return path, stat + + +class CreateTTL(namedtuple("CreateTTL", "path data acl flags ttl")): + type = 21 + + def serialize(self): + b = bytearray() + b.extend(write_string(self.path)) + b.extend(write_buffer(self.data)) + b.extend(int_struct.pack(len(self.acl))) + for acl in self.acl: + b.extend( + int_struct.pack(acl.perms) + + write_string(acl.id.scheme) + + write_string(acl.id.id) + ) + b.extend(int_struct.pack(self.flags)) + b.extend(long_struct.pack(self.ttl)) + return b + + @classmethod + def deserialize(cls, bytes, offset): + path, offset = read_string(bytes, offset) + stat = ZnodeStat._make(stat_struct.unpack_from(bytes, offset)) + return path, stat + + class Auth(namedtuple("Auth", "auth_type scheme auth")): auth_type: int scheme: str diff --git a/kazoo/testing/harness.py b/kazoo/testing/harness.py index 5a3a55a0..55adc514 100644 --- a/kazoo/testing/harness.py +++ b/kazoo/testing/harness.py @@ -90,10 +90,13 @@ def get_global_cluster() -> ZookeeperCluster: "localSessionsEnabled=" + ZOOKEEPER_LOCAL_SESSION_RO, "localSessionsUpgradingEnabled=" + ZOOKEEPER_LOCAL_SESSION_RO, ] - # If defined, this sets the superuser password to "test" additional_java_system_properties = [ + # Enable extended types (container & ttl znodes) + "-Dzookeeper.extendedTypesEnabled=true", + "-Dznode.container.checkIntervalMs=100", + # If defined, this sets the superuser password to "test" "-Dzookeeper.DigestAuthenticationProvider.superDigest=" - "super:D/InIHSb7yEEbrWz8b9l71RjZJU=" + "super:D/InIHSb7yEEbrWz8b9l71RjZJU=", ] else: additional_configuration_entries = [] diff --git a/kazoo/tests/test_client.py b/kazoo/tests/test_client.py index a031b1ff..32591c5c 100644 --- a/kazoo/tests/test_client.py +++ b/kazoo/tests/test_client.py @@ -719,6 +719,34 @@ def test_create_stat(self) -> None: assert data == b"bytes" assert stat1 == stat2 + def test_create_container(self) -> None: + if CI_ZK_VERSION: + version = CI_ZK_VERSION + else: + version = self.client.server_version() + if not version or version < (3, 5): + pytest.skip("Must use Zookeeper 3.5 or above") + client = self.client + path, stat1 = client.create("/1_cnt", b"bytes", container=True) + data, stat2 = client.get(path) + assert path == "/1_cnt" + assert data == b"bytes" + assert stat1 == stat2 + + def test_create_ttl(self) -> None: + if CI_ZK_VERSION: + version = CI_ZK_VERSION + else: + version = self.client.server_version() + if not version or version < (3, 5): + pytest.skip("Must use Zookeeper 3.5 or above") + client = self.client + path, stat1 = client.create("/1_ttl", b"bytes", ttl=1) + data, stat2 = client.get(path) + assert path == "/1_ttl" + assert data == b"bytes" + assert stat1 == stat2 + def test_create_get_set(self) -> None: nodepath = "/" + uuid.uuid4().hex