Couple of session hooks

Author: technige

neo4j/bolt/connection.py

@@ -343,12 +343,18 @@ def send(self):
         self.channel.send()
 
     def fetch(self):
-        """ Receive exactly one message from the server.
+        """ Receive exactly one message from the server
+        (if one is available).
+
+        :return: number of messages fetched (zero or one)
         """
         if self.closed:
             raise ServiceUnavailable("Failed to read from closed connection %r" % (self.server.address,))
         if self.defunct:
             raise ServiceUnavailable("Failed to read from defunct connection %r" % (self.server.address,))
+        if not self.responses:
+            return 0
+
         try:
             message_data = self.buffering_socket.read_message()
         except ProtocolError:
@@ -388,6 +394,8 @@ def fetch(self):
         else:
             raise ProtocolError("Unexpected response message with signature %02X" % signature)
 
+        return 1
+
     def sync(self):
         """ Send and fetch all outstanding messages.
         """

neo4j/v1/session.py

@@ -33,6 +33,10 @@ class Session(object):
     method.
     """
 
+    response_class = Response
+
+    result_class = None
+
     transaction = None
 
     last_bookmark = None
@@ -66,9 +70,9 @@ def run(self, statement, parameters=None, **kwparameters):
         statement = _norm_statement(statement)
         parameters = _norm_parameters(parameters, **kwparameters)
 
-        run_response = Response(self.connection)
-        pull_all_response = Response(self.connection)
-        result = StatementResult(self, run_response, pull_all_response)
+        run_response = self.response_class(self.connection)
+        pull_all_response = self.response_class(self.connection)
+        result = self.result_class(self, run_response, pull_all_response)
         result.statement = statement
         result.parameters = parameters
 
@@ -79,6 +83,9 @@ def run(self, statement, parameters=None, **kwparameters):
         return result
 
     def fetch(self):
+        """ Fetch the next message if available and return
+        the number of messages fetched (one or zero).
+        """
         try:
             return self.connection.fetch()
         except ServiceUnavailable as cause:
@@ -331,6 +338,9 @@ def peek(self):
         raise ResultError("End of stream")
 
 
+Session.result_class = StatementResult
+
+
 class Record(object):
     """ Record is an ordered collection of fields.