iSuch reactive! So stream!

Author: technige

neo4j/connection.py

@@ -34,6 +34,7 @@
 
 
 DEFAULT_PORT = 7687
+DEFAULT_USER_AGENT = "neo4j-python/0.0"
 
 # Signature bytes for each message type
 INIT = b"\x01"             # 0000 0001 // INIT <user_agent>
@@ -76,16 +77,20 @@ def hex2(x):
         return hex(x)[2:].upper()
 
 
-class ChunkWriter(object):
-    """ Writer for chunked data.
+class ChunkChannel(object):
+    """ Reader/writer for chunked data.
+
+    .. note:: logs at DEBUG level
     """
 
     max_chunk_size = 65535
 
-    def __init__(self):
+    def __init__(self, sock):
+        self.socket = sock
         self.raw = BytesIO()
         self.output_buffer = []
         self.output_size = 0
+        self._recv_buffer = b""
 
     def write(self, b):
         """ Write some bytes, splitting into chunks if necessary.
@@ -106,7 +111,7 @@ def write(self, b):
                 self.output_size = future_size
                 b = b""
 
-    def flush(self, zero_chunk=False):
+    def flush(self, end_of_message=False):
         """ Flush everything written since the last chunk to the
         stream, followed by a zero-chunk if required.
         """
@@ -115,28 +120,71 @@ def flush(self, zero_chunk=False):
             lines = [struct_pack(">H", self.output_size)] + output_buffer
         else:
             lines = []
-        if zero_chunk:
+        if end_of_message:
             lines.append(b"\x00\x00")
         if lines:
             self.raw.writelines(lines)
             self.raw.flush()
             del output_buffer[:]
             self.output_size = 0
 
-    def to_bytes(self):
-        """ Extract the written data as bytes.
+    def send(self):
+        """ Send all queued messages to the server.
         """
-        return self.raw.getvalue()
+        data = self.raw.getvalue()
+        if __debug__:
+            log_debug("C: %s", ":".join(map(hex2, data)))
+        self.socket.sendall(data)
 
-    def reset(self):
-        """ Reset the stream.
+        self.raw.seek(self.raw.truncate(0))
+
+    def _recv(self, size):
+        # If data is needed, keep reading until all bytes have been received
+        remaining = size - len(self._recv_buffer)
+        ready_to_read = None
+        while remaining > 0:
+            # Read up to the required amount remaining
+            b = self.socket.recv(8192)
+            if b:
+                if __debug__: log_debug("S: %s", ":".join(map(hex2, b)))
+            else:
+                if ready_to_read is not None:
+                    raise ProtocolError("Server closed connection")
+            remaining -= len(b)
+            self._recv_buffer += b
+
+            # If more is required, wait for available network data
+            if remaining > 0:
+                ready_to_read, _, _ = select((self.socket,), (), (), 0)
+                while not ready_to_read:
+                    ready_to_read, _, _ = select((self.socket,), (), (), 0)
+
+        # Split off the amount of data required and keep the rest in the buffer
+        data, self._recv_buffer = self._recv_buffer[:size], self._recv_buffer[size:]
+        return data
+
+    def chunk_reader(self):
+        chunk_size = -1
+        while chunk_size != 0:
+            chunk_header = self._recv(2)
+            chunk_size, = struct_unpack_from(">H", chunk_header)
+            if chunk_size > 0:
+                data = self._recv(chunk_size)
+                yield data
+
+    def close(self):
+        """ Shut down and close the connection.
         """
-        self.flush()
-        raw = self.raw
-        return raw.seek(raw.truncate(0))
+        if __debug__: log_info("~~ [CLOSE]")
+        socket = self.socket
+        socket.shutdown(SHUT_RDWR)
+        socket.close()
 
 
-class ResponseSubscriber(object):
+class Response(object):
+    """ Subscriber object for a full response (zero or
+    more detail messages followed by one summary message).
+    """
 
     def __init__(self, connection):
         self.connection = connection
@@ -160,151 +208,92 @@ def deliver(self, messages):
             if signature in SUMMARY:
                 self.__complete = True
             if signature == FAILURE:
-                self.connection.ack_failure()
+                self.ack_failure()
 
     def consume(self):
-        self.connection.receive_all(self)
+        self.connection.fetch_all(self)
+
+    def ack_failure(self):
+        """ Queue an acknowledgement for a previous failure.
+        """
+
+        def on_failure(metadata):
+            raise ProtocolError("Could not acknowledge failure")
+
+        subscriber = Response(self)
+        subscriber.on_failure = on_failure
+        self.connection.append(ACK_FAILURE, response=subscriber)
 
 
 class Connection(object):
     """ Server connection through which all protocol messages
     are sent and received. This class is designed for protocol
     version 1.
+
+    .. note:: logs at INFO level
     """
 
     def __init__(self, sock, **config):
-        self.socket = sock
-        self.raw = ChunkWriter()
-        self.packer = Packer(self.raw)
-        self.inbox = deque()
-        self.outbox = deque()
-        self._recv_buffer = b""
+        self.channel = ChunkChannel(sock)
+        self.packer = Packer(self.channel)
+        self.responses = deque()
 
-    def append(self, signature, fields=(), subscriber=None):
-        """ Add a message to the outgoing queue.
-        """
-        self.outbox.append((signature, fields))
-        self.inbox.append(subscriber)
+        # Determine the user agent and ensure it is a Unicode value
+        user_agent = config.get("user-agent", DEFAULT_USER_AGENT)
+        if isinstance(user_agent, bytes):
+            user_agent = user_agent.decode("UTF-8")
 
-    def send(self):
-        """ Send all queued messages to the server.
-        """
+        def on_failure(metadata):
+            raise ProtocolError("Initialisation failed")
 
-        # Shortcuts to avoid too many dots
-        raw = self.raw
-        packer = self.packer
-        pack_struct_header = packer.pack_struct_header
-        pack = packer.pack
-        flush = raw.flush
-
-        for signature, fields in self.outbox:
-            pack_struct_header(len(fields), signature)
-            for field in fields:
-                pack(field)
-            flush(zero_chunk=True)
-
-        data = raw.to_bytes()
-        if __debug__: log_debug("C: %s", ":".join(map(hex2, data)))
-        self.socket.sendall(data)
+        response = Response(self)
+        response.on_failure = on_failure
 
-        raw.reset()
-        self.outbox.clear()
+        self.append(INIT, (user_agent,), response=response)
+        self.send()
+        response.consume()
 
-    def _recv(self, size):
-        # If data is needed, keep reading until all bytes have been received
-        remaining = size - len(self._recv_buffer)
-        ready_to_read = None
-        while remaining > 0:
-            # Read up to the required amount remaining
-            b = self.socket.recv(8192)
-            if b:
-                if __debug__: log_debug("S: %s", ":".join(map(hex2, b)))
-            else:
-                if ready_to_read is not None:
-                    raise ProtocolError("Server closed connection")
-            remaining -= len(b)
-            self._recv_buffer += b
+    def append(self, signature, fields=(), response=None):
+        """ Add a message to the outgoing queue.
+        """
+        if __debug__:
+            log_info("C: %s %s", message_names[signature], " ".join(map(repr, fields)))
 
-            # If more is required, wait for available network data
-            if remaining > 0:
-                ready_to_read, _, _ = select((self.socket,), (), (), 0)
-                while not ready_to_read:
-                    ready_to_read, _, _ = select((self.socket,), (), (), 0)
+        self.packer.pack_struct_header(len(fields), signature)
+        for field in fields:
+            self.packer.pack(field)
+        self.channel.flush(end_of_message=True)
+        self.responses.append(response)
 
-        # Split off the amount of data required and keep the rest in the buffer
-        data, self._recv_buffer = self._recv_buffer[:size], self._recv_buffer[size:]
-        return data
+    def send(self):
+        """ Send all queued messages to the server.
+        """
+        self.channel.send()
 
-    def receive_next(self):
+    def fetch_next(self):
         """ Receive exactly one message from the server.
         """
         raw = BytesIO()
         unpack = Unpacker(raw).unpack
-
-        # Receive chunks of data until chunk_size == 0
-        chunk_size = None
-        while chunk_size != 0:
-            if chunk_size is None:
-                # Chunk header
-                data = self._recv(2)
-                chunk_size, = struct_unpack_from(">H", data)
-            else:
-                # Chunk content
-                data = self._recv(chunk_size)
-                raw.write(data)
-                chunk_size = None
+        raw.writelines(self.channel.chunk_reader())
 
         # Unpack the message from the raw byte stream into the inbox
         raw.seek(0)
-        response = self.inbox[0]
+        response = self.responses[0]
         response.deliver(unpack())
         if response.complete:
-            self.inbox.popleft()
+            self.responses.popleft()
         raw.close()
 
-    def receive_all(self, response):
-        receive_next = self.receive_next
+    def fetch_all(self, response):
+        fetch_next = self.fetch_next
         while not response.complete:
-            receive_next()
-
-    def init(self, user_agent):
-        """ Initialise a connection with a user agent string.
-        """
-
-        # Ensure the user agent is a Unicode value
-        if isinstance(user_agent, bytes):
-            user_agent = user_agent.decode("UTF-8")
-
-        if __debug__: log_info("C: INIT %r", user_agent)
-
-        def on_failure(metadata):
-            raise ProtocolError("Initialisation failed")
-
-        subscriber = ResponseSubscriber(self)
-        subscriber.on_failure = on_failure
-        self.append(INIT, (user_agent,), subscriber=subscriber)
-        self.send()
-        subscriber.consume()
-
-    def ack_failure(self):
-        """ Queue an acknowledgement for a previous failure.
-        """
-        if __debug__: log_info("C: ACK_FAILURE")
-
-        def on_failure(metadata):
-            raise ProtocolError("Could not acknowledge failure")
-
-        subscriber = ResponseSubscriber(self)
-        subscriber.on_failure = on_failure
-        self.append(ACK_FAILURE, subscriber=subscriber)
+            fetch_next()
 
     def close(self):
         """ Shut down and close the connection.
         """
-        if __debug__: log_info("~~ [CLOSE]")
-        socket = self.socket
-        socket.shutdown(SHUT_RDWR)
-        socket.close()
+        self.channel.close()
 
 
 def connect(host, port=None, **config):