Added db to BEGIN and introduced unit test tool

Author: technige

neo4j/io/__init__.py

@@ -172,8 +172,17 @@ def discard_all(self, **handlers):
     def pull_all(self, **handlers):
         raise NotImplementedError
 
-    def begin(self, mode=None, bookmarks=None, metadata=None, timeout=None, **handlers):
-        raise NotImplementedError
+    def begin(self, mode=None, bookmarks=None, metadata=None, timeout=None, db=None, **handlers):
+        """ Appends a BEGIN message to the output stream.
+
+        :param mode: access mode for routing - "READ" or "WRITE" (default)
+        :param bookmarks: iterable of bookmark values after which this transaction should begin
+        :param metadata: custom metadata dictionary to attach to the transaction
+        :param timeout: timeout for transaction execution (seconds)
+        :param db: name of the database against which to begin the transaction
+        :param handlers: handler functions passed into the returned Response object
+        :return: :class:`neo4j.io._bolt3.Response` object
+        """
 
     def commit(self, **handlers):
         raise NotImplementedError

neo4j/io/_bolt3.py

@@ -26,6 +26,7 @@
 from neo4j.api import (
     Version,
 )
+from neo4j.io._courier import MessageInbox
 from neo4j.meta import get_user_agent
 from neo4j.exceptions import (
     ProtocolError,
@@ -39,7 +40,6 @@
     SessionExpired,
 )
 from neo4j.packstream import (
-    UnpackableBuffer,
     Unpacker,
     Packer,
 )
@@ -178,7 +178,10 @@ def pull_all(self, **handlers):
         log.debug("[#%04X]  C: PULL_ALL", self.local_port)
         self._append(b"\x3F", (), Response(self, **handlers))
 
-    def begin(self, mode=None, bookmarks=None, metadata=None, timeout=None, **handlers):
+    def begin(self, mode=None, bookmarks=None, metadata=None, timeout=None,
+              db=None, **handlers):
+        if db is not None:
+            raise ValueError("Database selection is not supported in Bolt 3")
         extra = {}
         if mode:
             extra["mode"] = mode
@@ -538,56 +541,16 @@ def recv_into(self, buffer, n_bytes=0, flags=0):
         return n_bytes
 
 
-class Inbox:
-
-    def __init__(self, s, on_error):
-        super(Inbox, self).__init__()
-        self.on_error = on_error
-        self._messages = self._yield_messages(s)
-
-    def __iter__(self):
-        return self
+class Inbox(MessageInbox):
 
     def __next__(self):
-        return next(self._messages)
-
-    @classmethod
-    def _load_chunks(cls, sock, buffer):
-        chunk_size = 0
-        while True:
-            if chunk_size == 0:
-                buffer.receive(sock, 2)
-            chunk_size = buffer.pop_u16()
-            if chunk_size > 0:
-                buffer.receive(sock, chunk_size + 2)
-            yield chunk_size
-
-    def _yield_messages(self, sock):
-        try:
-            buffer = UnpackableBuffer()
-            chunk_loader = self._load_chunks(sock, buffer)
-            unpacker = Unpacker(buffer)
-            details = []
-            while True:
-                unpacker.reset()
-                details[:] = ()
-                chunk_size = -1
-                while chunk_size != 0:
-                    chunk_size = next(chunk_loader)
-                summary_signature = None
-                summary_metadata = None
-                size, signature = unpacker.unpack_structure_header()
-                if size > 1:
-                    raise ProtocolError("Expected one field")
-                if signature == b"\x71":
-                    data = unpacker.unpack()
-                    details.append(data)
-                else:
-                    summary_signature = signature
-                    summary_metadata = unpacker.unpack_map()
-                yield details, summary_signature, summary_metadata
-        except OSError as error:
-            self.on_error(error)
+        tag, fields = self.pop()
+        if tag == b"\x71":
+            return fields, None, None
+        elif fields:
+            return [], tag, fields[0]
+        else:
+            return [], tag, None
 
 
 class Response:

neo4j/io/_bolt4x0.py

@@ -26,6 +26,7 @@
 from neo4j.api import (
     Version,
 )
+from neo4j.io._courier import MessageInbox
 from neo4j.meta import get_user_agent
 from neo4j.exceptions import (
     ProtocolError,
@@ -39,7 +40,6 @@
     SessionExpired,
 )
 from neo4j.packstream import (
-    UnpackableBuffer,
     Unpacker,
     Packer,
 )
@@ -178,10 +178,13 @@ def pull_all(self, **handlers):
         log.debug("[#%04X]  C: PULL_ALL", self.local_port)
         self._append(b"\x3F", (), Response(self, **handlers))
 
-    def begin(self, mode=None, bookmarks=None, metadata=None, timeout=None, **handlers):
+    def begin(self, mode=None, bookmarks=None, metadata=None, timeout=None,
+              db=None, **handlers):
         extra = {}
         if mode:
             extra["mode"] = mode
+        if db:
+            extra["db"] = db
         if bookmarks:
             try:
                 extra["bookmarks"] = list(bookmarks)
@@ -538,56 +541,16 @@ def recv_into(self, buffer, n_bytes=0, flags=0):
         return n_bytes
 
 
-class Inbox:
-
-    def __init__(self, s, on_error):
-        super(Inbox, self).__init__()
-        self.on_error = on_error
-        self._messages = self._yield_messages(s)
-
-    def __iter__(self):
-        return self
+class Inbox(MessageInbox):
 
     def __next__(self):
-        return next(self._messages)
-
-    @classmethod
-    def _load_chunks(cls, sock, buffer):
-        chunk_size = 0
-        while True:
-            if chunk_size == 0:
-                buffer.receive(sock, 2)
-            chunk_size = buffer.pop_u16()
-            if chunk_size > 0:
-                buffer.receive(sock, chunk_size + 2)
-            yield chunk_size
-
-    def _yield_messages(self, sock):
-        try:
-            buffer = UnpackableBuffer()
-            chunk_loader = self._load_chunks(sock, buffer)
-            unpacker = Unpacker(buffer)
-            details = []
-            while True:
-                unpacker.reset()
-                details[:] = ()
-                chunk_size = -1
-                while chunk_size != 0:
-                    chunk_size = next(chunk_loader)
-                summary_signature = None
-                summary_metadata = None
-                size, signature = unpacker.unpack_structure_header()
-                if size > 1:
-                    raise ProtocolError("Expected one field")
-                if signature == b"\x71":
-                    data = unpacker.unpack()
-                    details.append(data)
-                else:
-                    summary_signature = signature
-                    summary_metadata = unpacker.unpack_map()
-                yield details, summary_signature, summary_metadata
-        except OSError as error:
-            self.on_error(error)
+        tag, fields = self.pop()
+        if tag == b"\x71":
+            return fields, None, None
+        elif fields:
+            return [], tag, fields[0]
+        else:
+            return [], tag, None
 
 
 class Response:

neo4j/io/_courier.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+# Copyright (c) 2002-2020 "Neo4j,"
+# Neo4j Sweden AB [http://neo4j.com]
+#
+# This file is part of Neo4j.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from neo4j.packstream import (
+    UnpackableBuffer,
+    Unpacker,
+)
+
+
+class MessageInbox:
+
+    def __init__(self, s, on_error):
+        self.on_error = on_error
+        self._messages = self._yield_messages(s)
+
+    @classmethod
+    def _load_chunks(cls, sock, buffer):
+        chunk_size = 0
+        while True:
+            if chunk_size == 0:
+                buffer.receive(sock, 2)
+            chunk_size = buffer.pop_u16()
+            if chunk_size > 0:
+                buffer.receive(sock, chunk_size + 2)
+            yield chunk_size
+
+    def _yield_messages(self, sock):
+        try:
+            buffer = UnpackableBuffer()
+            chunk_loader = self._load_chunks(sock, buffer)
+            unpacker = Unpacker(buffer)
+            while True:
+                unpacker.reset()
+                chunk_size = -1
+                while chunk_size != 0:
+                    chunk_size = next(chunk_loader)
+                size, tag = unpacker.unpack_structure_header()
+                fields = [unpacker.unpack() for _ in range(size)]
+                yield tag, fields
+        except OSError as error:
+            self.on_error(error)
+
+    def pop(self):
+        return next(self._messages)

neo4j/work/simple.py

@@ -308,6 +308,8 @@ def begin_transaction(self, bookmark=None, metadata=None, timeout=None):
     def _open_transaction(self, access_mode=None, metadata=None, timeout=None):
         self._transaction = Transaction(self, on_close=self._close_transaction)
         self._connect(access_mode)
+        # TODO: capture ValueError and surface as SessionError/TransactionError if
+        # TODO: explicit database selection has been made
         self._connection.begin(bookmarks=self._bookmarks_in, metadata=metadata, timeout=timeout)
 
     def commit_transaction(self):

tests/unit/io/conftest.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+# Copyright (c) 2002-2020 "Neo4j,"
+# Neo4j Sweden AB [http://neo4j.com]
+#
+# This file is part of Neo4j.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytest
+
+from neo4j.io._courier import MessageInbox
+
+
+class FakeSocket:
+
+    def __init__(self, address):
+        self.address = address
+        self.captured = b""
+        self.messages = MessageInbox(self, on_error=print)
+
+    def setblocking(self, flag):
+        pass
+
+    def getsockname(self):
+        return "127.0.0.1", 0xFFFF
+
+    def getpeername(self):
+        return self.address
+
+    def recv_into(self, buffer, nbytes):
+        data = self.captured[:nbytes]
+        actual = len(data)
+        buffer[:actual] = data
+        self.captured = self.captured[actual:]
+        return actual
+
+    def sendall(self, data):
+        self.captured += data
+
+    def close(self):
+        return
+
+    def pop_message(self):
+        return self.messages.pop()
+
+
+@pytest.fixture
+def fake_socket():
+    return FakeSocket