feat(core): Add support for Create2 in transactions.

ceache · Feb 12, 2023 · b7b2b15 · b7b2b15
1 parent 5225b3e
commit b7b2b15
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/kazoo/client.py b/kazoo/client.py
@@ -1673,7 +1673,13 @@ def __init__(self, client):
         self.committed = False
 
     def create(
-        self, path, value=b"", acl=None, ephemeral=False, sequence=False
+        self,
+        path,
+        value=b"",
+        acl=None,
+        ephemeral=False,
+        sequence=False,
+        include_data=False,
     ):
         """Add a create ZNode to the transaction. Takes the same
         arguments as :meth:`KazooClient.create`, with the exception
@@ -1697,6 +1703,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:
@@ -1705,9 +1713,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
@@ -343,6 +343,11 @@ 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:
@@ -367,6 +372,10 @@ def unchroot(client, response):
         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