element id tests

boltstub/bolt_protocol.py

@@ -16,8 +16,7 @@ def get_bolt_protocol(version):
     if version is None:
         raise BoltMissingVersionError()
     for sub in recursive_subclasses(BoltProtocol):
-        if (version == sub.protocol_version
-                or version in sub.version_aliases):
+        if version == sub.protocol_version or version in sub.version_aliases:
             return sub
     raise BoltUnknownVersionError(
         "unsupported bolt version {}".format(version)
@@ -371,3 +370,13 @@ class Bolt4x4Protocol(Bolt4x3Protocol):
     equivalent_versions = set()
 
     server_agent = "Neo4j/4.4.0"
+
+
+class Bolt5x0Protocol(Bolt4x4Protocol):
+
+    protocol_version = (5, 0)
+    version_aliases = set()
+    # allow the server to negotiate other bolt versions
+    equivalent_versions = set()
+
+    server_agent = "Neo4j/5.0.0"

boltstub/channel.py

@@ -10,8 +10,9 @@
 class Channel:
     def __init__(self, wire, bolt_version, log_cb=None, handshake_data=None):
         self.wire = wire
-        self.stream = PackStream(wire)
         self.bolt_protocol = get_bolt_protocol(bolt_version)
+        self.stream = PackStream(wire,
+                                 self.bolt_protocol.protocol_version[0] < 5)
         self.log = log_cb
         self.handshake_data = handshake_data
         self._buffered_msg = None

boltstub/packstream.py

@@ -17,7 +17,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
+import inspect
 from codecs import decode
 from io import BytesIO
 from struct import pack as struct_pack
@@ -62,10 +62,15 @@ class StructTag:
 
 class Structure:
 
-    def __init__(self, tag, *fields, verified=True):
+    def __init__(self, tag, *fields, verified=True,
+                 validate_old_id=None):
         self.tag = tag
         self.fields = list(fields)
         self._verified = verified
+        if validate_old_id is None:
+            self.validate_either = True
+        self.validate_v3 = validate_old_id
+
         if verified:
             self._verify_fields()
 
@@ -76,58 +81,124 @@ def verified(self):
     def _verify_fields(self):
         tag, fields = self.tag, self.fields
 
-        def verify_node():
-            if (len(fields) != 3
-                    or not isinstance(fields[0], int)
-                    or not isinstance(fields[1], list)
-                    or not all(isinstance(label, str) for label in fields[1])
-                    or not isinstance(fields[2], dict)
-                    or not all(isinstance(k, str) for k in fields[2].keys())):
-                raise ValueError("Invalid Node struct received %r" % self)
-
-        def verify_relationship():
-            if (len(fields) != 5
-                    or not isinstance(fields[0], int)
-                    or not isinstance(fields[1], int)
-                    or not isinstance(fields[2], int)
-                    or not isinstance(fields[3], str)
-                    or not isinstance(fields[4], dict)
-                    or not all(isinstance(k, str) for k in fields[4].keys())):
-                raise ValueError(
-                    "Invalid Relationship struct received %r" % self
-                )
-
-        def verify_unbound_relationship():
-            if (len(fields) != 3
-                    or not isinstance(fields[0], int)
-                    or not isinstance(fields[1], str)
-                    or not isinstance(fields[2], dict)
-                    or not all(isinstance(k, str) for k in fields[2].keys())):
-                raise ValueError(
-                    "Invalid UnboundRelationship struct received %r" % self
-                )
+        def validate_validations(validations, type_name):
+            for validation in validations:
+                if not validation(fields):
+                    raise ValueError(
+                        f"Invalid {type_name} struct received.\n"
+                        f"validation failed: {inspect.getsource(validation)}"
+                        f" {self}")
+
+        def verify_node_3_up():
+            validations = [
+                lambda f: len(f) == 3,
+                lambda f: isinstance(f[0], int),
+                lambda f: isinstance(f[1], list),
+                lambda f: all(isinstance(label, str) for label in f[1]),
+                lambda f: isinstance(f[2], dict),
+                lambda f: all(isinstance(k, str) for k in f[2].keys())
+            ]
+            validate_validations(validations, "V3_node")
+
+        def verify_node_5_up():
+            if self.validate_either:
+                try:
+                    verify_node_3_up()
+                    return
+                except ValueError as _:
+                    pass
+
+            validations = [
+                lambda f: len(f) == 4,
+                lambda f: isinstance(f[0], int) or f[0] is None,
+                lambda f: isinstance(f[1], list),
+                lambda f: all(isinstance(label, str) for label in f[1]),
+                lambda f: isinstance(f[2], dict),
+                lambda f: all(isinstance(k, str) for k in f[2].keys()),
+                lambda f: isinstance(f[3], str)
+            ]
+            validate_validations(validations, "V5_node")
+
+        def verify_unbound_relationship_3_up():
+            validations = [
+                lambda f: len(f) == 3,
+                lambda f: isinstance(f[0], int),
+                lambda f: isinstance(f[1], str),
+                lambda f: isinstance(f[2], dict),
+                lambda f: all(isinstance(k, str) for k in f[2].keys())
+            ]
+            validate_validations(validations, "V3_UnboundRelationship")
+
+        def verify_unbound_relationship_5_up():
+            if self.validate_either:
+                try:
+                    verify_unbound_relationship_3_up()
+                    return
+                except ValueError as _:
+                    pass
+            validations = [
+                lambda f: len(f) == 4,
+                lambda f: isinstance(f[0], int) or f[0] is None,
+                lambda f: isinstance(f[1], str),
+                lambda f: isinstance(f[2], dict),
+                lambda f: all(isinstance(k, str) for k in f[2].keys()),
+                lambda f: isinstance(f[3], str)
+            ]
+            validate_validations(validations, "V5_UnboundRelationship")
+
+        def verify_relationship_3_up():
+            validations = [
+                lambda f: len(f) == 5,
+                lambda f: isinstance(f[0], int),
+                lambda f: isinstance(f[1], int),
+                lambda f: isinstance(f[2], int),
+                lambda f: isinstance(f[3], str),
+                lambda f: isinstance(f[4], dict),
+                lambda f: all(isinstance(k, str) for k in f[4].keys())
+            ]
+            validate_validations(validations, "V3_Relationship")
+
+        def verify_relationship_5_up():
+            if self.validate_either:
+                try:
+                    verify_relationship_3_up()
+                    return
+                except ValueError as _:
+                    pass
+
+            validations = [
+                lambda f: len(f) == 8,
+                lambda f: (isinstance(f[0], int) or f[0] is None),
+                lambda f: (isinstance(f[1], int) or f[1] is None),
+                lambda f: (isinstance(f[2], int) or f[2] is None),
+                lambda f: isinstance(f[3], str),
+                lambda f: isinstance(f[4], dict),
+                lambda f: all(isinstance(k, str) for k in f[4].keys()),
+                lambda f: isinstance(f[5], str),
+                lambda f: isinstance(f[6], str),
+                lambda f: isinstance(f[7], str),
+            ]
+            validate_validations(validations, "V5_Relationship")
 
         def verify_path():
-            if (len(fields) != 3
-                    or not isinstance(fields[0], list)
-                    or not all(isinstance(n, Structure)
-                               and n.tag == StructTag.node
-                               and n.fields[0] in fields[2]  # id is used
-                               for n in fields[0])
-                    or not isinstance(fields[1], list)
-                    or not all(isinstance(rel, Structure)
-                               and rel.tag == StructTag.unbound_relationship
-                               and rel.fields[0] in fields[2]  # id is used
-                               for rel in fields[1])
-                    or not isinstance(fields[2], list)
-                    or not all(isinstance(id_, int)
-                               # id exists in nodes or relationships
-                               and id_ in {s.fields[0]
-                                           for s in fields[0] + fields[1]}
-                               for id_ in fields[2])):
-                raise ValueError(
-                    "Invalid Path struct received %r" % self
-                )
+            validations = [
+                lambda f: len(f) == 3,
+                lambda f: all(isinstance(n, Structure)
+                              and n.tag == StructTag.node
+                              for n in f[0]),
+                lambda f: isinstance(f[1], list),
+                lambda f: all(isinstance(rel, Structure)
+                              and rel.tag == StructTag.unbound_relationship
+                              for rel in f[1]),
+                lambda f: isinstance(f[2], list),
+                # index is less than respective array length
+                # rels indexed from 1, nodes 0
+                lambda f: all(isinstance(id_, int)
+                              and abs(id_) <= len(f[1])
+                              if i % 2 == 0 else abs(id_) < len(f[0])
+                              for i, id_ in enumerate(f[2]))
+            ]
+            validate_validations(validations, "Path")
 
         def build_generic_verifier(types, name):
             def verify():
@@ -142,11 +213,14 @@ def verify():
 
         field_validator = {
             StructTag.node:
-                verify_node,
+                verify_node_3_up if self.validate_v3
+                else verify_node_5_up,
             StructTag.relationship:
-                verify_relationship,
+                verify_relationship_3_up if self.validate_v3
+                else verify_relationship_5_up,
             StructTag.unbound_relationship:
-                verify_unbound_relationship,
+                verify_unbound_relationship_3_up if self.validate_v3
+                else verify_unbound_relationship_5_up,
             StructTag.path:
                 verify_path,
             StructTag.date:
@@ -264,27 +338,80 @@ def from_jolt_type(cls, jolt: jolt_types.JoltType):
                 return cls(StructTag.point_3d, jolt.srid, jolt.x, jolt.y,
                            jolt.z)
         if isinstance(jolt, jolt_types.JoltNode):
+            if jolt.element_id is not None:
+                return cls(StructTag.node, jolt.id, jolt.labels,
+                           jolt.properties, jolt.element_id)
             return cls(StructTag.node, jolt.id, jolt.labels, jolt.properties)
         if isinstance(jolt, jolt_types.JoltRelationship):
+            if jolt.element_id is None:
+                return cls(StructTag.relationship, jolt.id, jolt.start_node_id,
+                           jolt.end_node_id, jolt.rel_type, jolt.properties)
             return cls(StructTag.relationship, jolt.id, jolt.start_node_id,
-                       jolt.end_node_id, jolt.rel_type, jolt.properties)
+                       jolt.end_node_id, jolt.rel_type, jolt.properties,
+                       jolt.element_id, jolt.start_node_element_id,
+                       jolt.end_node_element_id)
         if isinstance(jolt, jolt_types.JoltPath):
-            # Node structs
+            uniq_nodes = []
+            uniq_rels = []
+            ids = []
             nodes = []
+            rels = []
+            # Node structs
+            node_idxs = {}
+            node_idx = 0
             for node in jolt.path[::2]:
-                node = cls(StructTag.node, node.id, node.labels,
-                           node.properties)
-                if node not in nodes:
-                    nodes.append(node)
+                nodes.append(node)
+                map_node = Structure.from_jolt_type(node)
+                if map_node not in uniq_nodes:
+                    node_idxs[str(node.id)
+                              if node.element_id is None
+                              else node.element_id] = node_idx
+                    node_idx = node_idx + 1
+                    # nodesDict[str(node.id)] = node
+                    uniq_nodes.append(map_node)
+
             # UnboundRelationship structs
-            rels = []
+
+            rel_idxs = {}
+            rel_idx = 1
             for rel in jolt.path[1::2]:
-                rel = cls(StructTag.unbound_relationship, rel.id, rel.rel_type,
-                          rel.properties)
-                if rel not in rels:
-                    rels.append(rel)
-            ids = [e.id for e in jolt.path]
-            return cls(StructTag.path, nodes, rels, ids)
+                rels.append(rel)
+
+                ub_rel = (cls(StructTag.unbound_relationship, rel.id,
+                              rel.rel_type,
+                              rel.properties, rel.element_id)
+                          if rel.element_id is not None else
+                          cls(StructTag.unbound_relationship, rel.id,
+                              rel.rel_type,
+                              rel.properties))
+                if ub_rel not in uniq_rels:
+                    rel_idxs[str(rel.id)
+                             if rel.element_id is None
+                             else rel.element_id] = rel_idx
+                    rel_idx = rel_idx + 1
+                    # relsDict[rel.id_] = rel
+                    uniq_rels.append(ub_rel)
+
+            last_node = nodes[0]
+            for i in range(1, (len(rels)*2)+1):
+                if i % 2 == 0:
+                    last_node = nodes[int(i/2)]
+                    index = node_idxs[str(last_node.id)
+                                      if last_node.element_id is None
+                                      else last_node.element_id]
+                    ids.append(index)
+                else:
+                    rel = rels[int(i/2)]
+                    index = rel_idxs[str(rel.id)
+                                     if rel.element_id is None
+                                     else rel.element_id]
+                    # print(f"{rel}")
+                    if last_node.id == rel.start_node_id:
+                        ids.append(index)
+                    else:
+                        ids.append(-index)
+
+            return cls(StructTag.path, uniq_nodes, uniq_rels, ids)
         raise TypeError("Unsupported jolt type: {}".format(type(jolt)))
 
     def to_jolt_type(self):
@@ -633,8 +760,9 @@ def pack_end_of_stream(self):
 
 class Unpacker:
 
-    def __init__(self, unpackable):
+    def __init__(self, unpackable, validate_old_id):
         self.unpackable = unpackable
+        self.validate_old_id = validate_old_id
 
     def reset(self):
         self.unpackable.reset()
@@ -731,7 +859,8 @@ def _unpack(self, verify_struct=True):
                 fields = [None] * size
                 for i in range(len(fields)):
                     fields[i] = self._unpack(verify_struct=True)
-                return Structure(tag, *fields, verified=verify_struct)
+                return Structure(tag, *fields, verified=verify_struct,
+                                 validate_old_id=self.validate_old_id)
 
             elif marker == 0xDF:  # END_OF_STREAM:
                 return EndOfStream
@@ -888,8 +1017,9 @@ def receive(self, sock, n_bytes):
 class PackStream:
     """Chunked message reader/writer for PackStream messaging."""
 
-    def __init__(self, wire):
+    def __init__(self, wire, validate_old_id):
         self.wire = wire
+        self.validate_old_id = validate_old_id
         self.data_buffer = []
         self.next_chunk_size = None
 
@@ -911,7 +1041,7 @@ def read_message(self):
                 break
         buffer = UnpackableBuffer(b"".join(self.data_buffer))
         self.data_buffer = []
-        unpacker = Unpacker(buffer)
+        unpacker = Unpacker(buffer, self.validate_old_id)
         return unpacker.unpack_message()
 
     def write_message(self, message):