Skip to content

Commit

Permalink
feat(core): Add support for Container and TTL nodes
Browse files Browse the repository at this point in the history 
Also add support through transations.

Closes #334, #496
  • Loading branch information
@ceache
ceache committed Jun 2, 2022
1 parent 9bb8499 commit 2bb9471
Show file tree
Hide file tree
Showing 2 changed files with 177 additions and 80 deletions.
202 changes: 123 additions & 79 deletions kazoo/client.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@
CloseInstance,
Create,
Create2,
CreateContainer,
CreateTTL,
Delete,
Exists,
GetChildren,
Expand Down Expand Up @@ -873,7 +875,8 @@ def sync(self, path):
return self.sync_async(path).get()

def create(self, path, value=b"", acl=None, ephemeral=False,
sequence=False, makepath=False, include_data=False):
sequence=False, makepath=False, include_data=False,
container=False, ttl=0):
"""Create a node with the given value as its data. Optionally
set an ACL on the node.
Expand Down Expand Up @@ -950,15 +953,19 @@ def create(self, path, value=b"", acl=None, ephemeral=False,
The `makepath` option.
.. versionadded:: 2.7
The `include_data` option.
.. versionadded:: 2.8
The container and ttl options.
"""
acl = acl or self.default_acl
return self.create_async(
path, value, acl=acl, ephemeral=ephemeral,
sequence=sequence, makepath=makepath, include_data=include_data
).get()
sequence=sequence, makepath=makepath, include_data=include_data,
container=container, ttl=ttl).get()

def create_async(self, path, value=b"", acl=None, ephemeral=False,
sequence=False, makepath=False, include_data=False):
sequence=False, makepath=False, include_data=False,
container=False, ttl=0):
"""Asynchronously create a ZNode. Takes the same arguments as
:meth:`create`.
Expand All @@ -967,45 +974,35 @@ def create_async(self, path, value=b"", acl=None, ephemeral=False,
.. versionadded:: 1.1
The makepath option.
.. versionadded:: 2.7
The `include_data` option.
The include_data option.
.. versionadded:: 2.8
The container and ttl options.
"""
if acl is None and self.default_acl:
acl = self.default_acl

if not isinstance(path, string_types):
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_types):
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():
result = self._create_async_inner(
path, value, acl, 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 inner_async_result.exception

inner_async_result.rawlink(create_completion)

@capture_exceptions(async_result)
def retry_completion(result):
Expand All @@ -1015,11 +1012,11 @@ def retry_completion(result):
@wrap(async_result)
def create_completion(result):
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
Expand All @@ -1032,26 +1029,6 @@ def create_completion(result):
do_create()
return async_result

def _create_async_inner(self, path, value, acl, flags,
trailing=False, include_data=False):
async_result = self.handler.async_result()
if include_data:
opcode = Create2
else:
opcode = 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
raise async_result.exception
return async_result

def ensure_path(self, path, acl=None):
"""Recursively create a path if it doesn't exist.
Expand Down Expand Up @@ -1590,39 +1567,24 @@ def __init__(self, client):
self.committed = False

def create(self, path, value=b"", acl=None, ephemeral=False,
sequence=False):
sequence=False, include_data=False, container=False, ttl=0):
"""Add a create ZNode to the transaction. Takes the same
arguments as :meth:`KazooClient.create`, with the exception
of `makepath`.
:returns: None
.. versionadded:: 2.8
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, string_types):
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_types):
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)")

flags = 0
if ephemeral:
flags |= 1
if sequence:
flags |= 2
if acl is None:
acl = OPEN_ACL_UNSAFE

self._add(Create(_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, version=-1):
"""Add a delete ZNode to the transaction. Takes the same
Expand Down Expand Up @@ -1701,3 +1663,85 @@ def _add(self, request, post_processor=None):
self._check_tx_state()
self.client.logger.log(BLATHER, 'Added %r to %r', request, self)
self.operations.append(request)


def _create_opcode(path, value, acl, chroot,
ephemeral, sequence, include_data, container, ttl):
"""Helper function.
Creates the create OpCode for regular `client.create()` operations as
well as in a `client.transaction()` context.
"""
if not isinstance(path, string_types):
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_types):
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
)
55 changes: 54 additions & 1 deletion kazoo/protocol/serialization.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -323,11 +323,17 @@ def deserialize(cls, bytes, offset):
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:
response = ZnodeStat._make(
stat_struct.unpack_from(bytes, offset))
stat_struct.unpack_from(bytes, offset)
)
offset += stat_struct.size
elif header.type == CheckVersion.type:
response = True
Expand All @@ -346,6 +352,10 @@ def unchroot(client, response):
for result in response:
if isinstance(result, six.string_types):
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
Expand Down Expand Up @@ -391,6 +401,49 @@ def deserialize(cls, bytes, offset):
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')):
type = 100

Expand Down

0 comments on commit 2bb9471

Please sign in to comment.