Result.consume

Author: technige

neo4j/connection.py

@@ -235,9 +235,8 @@ def on_failure(metadata):
 
         self.append(INIT, (user_agent,), response=response)
         self.send()
-        fetch_next = self.fetch_next
         while not response.complete:
-            fetch_next()
+            self.fetch_next()
 
     def append(self, signature, fields=(), response=None):
         """ Add a message to the outgoing queue.

neo4j/session.py

@@ -98,9 +98,11 @@ def session(self, **config):
 
 class Result(list):
 
-    def __init__(self):
+    def __init__(self, session):
+        self.session = session
         super(Result, self).__init__()
         self.keys = None
+        self.complete = False
         self.bench_test = None
 
     def on_header(self, metadata):
@@ -112,12 +114,18 @@ def on_record(self, values):
         self.append(Record(self.keys, tuple(map(hydrated, values))))
 
     def on_footer(self, metadata):
+        self.complete = True
         if self.bench_test:
             self.bench_test.end_recv = perf_counter()
 
     def on_failure(self, metadata):
         raise CypherError(metadata)
 
+    def consume(self):
+        fetch_next = self.session.connection.fetch_next
+        while not self.complete:
+            fetch_next()
+
 
 class Session(object):
     """ Logical session carried out over an established TCP connection.
@@ -154,7 +162,7 @@ def run(self, statement, parameters=None):
         t = BenchTest()
         t.init = perf_counter()
 
-        result = Result()
+        result = Result(self)
         result.bench_test = t
 
         run_response = Response(self.connection)
@@ -172,9 +180,7 @@ def run(self, statement, parameters=None):
         self.connection.send()
         t.end_send = perf_counter()
 
-        fetch_next = self.connection.fetch_next
-        while not pull_all_response.complete:
-            fetch_next()
+        result.consume()
 
         t.done = perf_counter()
         self.bench_tests.append(t)