Added n and qid to PULL (renamed from PULL_ALL)

Author: technige

neo4j/io/__init__.py

@@ -175,7 +175,7 @@ def run(self, statement, parameters=None, mode=None, bookmarks=None, metadata=No
         :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
+        :return: Response object
         """
 
     def discard(self, n=-1, qid=-1, **handlers):
@@ -184,11 +184,17 @@ def discard(self, n=-1, qid=-1, **handlers):
         :param n: number of records to discard, default = -1 (ALL)
         :param qid: query ID to discard for, default = -1 (last query)
         :param handlers: handler functions passed into the returned Response object
-        :return: :class:`neo4j.io._bolt3.Response` object
+        :return: Response object
         """
 
-    def pull_all(self, **handlers):
-        raise NotImplementedError
+    def pull(self, n=-1, qid=-1, **handlers):
+        """ Appends a PULL message to the output stream.
+
+        :param n: number of records to pull, default = -1 (ALL)
+        :param qid: query ID to pull for, default = -1 (last query)
+        :param handlers: handler functions passed into the returned Response object
+        :return: Response object
+        """
 
     def begin(self, mode=None, bookmarks=None, metadata=None, timeout=None, db=None, **handlers):
         """ Appends a BEGIN message to the output stream.
@@ -199,7 +205,7 @@ def begin(self, mode=None, bookmarks=None, metadata=None, timeout=None, db=None,
         :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
+        :return: Response object
         """
 
     def commit(self, **handlers):
@@ -492,7 +498,7 @@ def fail(md):
                 log.debug("[#%04X]  C: <ROUTING> query=%r", cx.local_port, self.routing_context or {})
                 cx.run("CALL dbms.cluster.routing.getRoutingTable($context)",
                        {"context": self.routing_context}, on_success=metadata.update, on_failure=fail)
-                cx.pull_all(on_success=metadata.update, on_records=records.extend)
+                cx.pull(on_success=metadata.update, on_records=records.extend)
                 cx.send_all()
                 cx.fetch_all()
                 routing_info = [dict(zip(metadata.get("fields", ()), values)) for values in records]

neo4j/io/_bolt3.py

@@ -181,7 +181,11 @@ def discard(self, n=-1, qid=-1, **handlers):
         log.debug("[#%04X]  C: DISCARD_ALL", self.local_port)
         self._append(b"\x2F", (), Response(self, **handlers))
 
-    def pull_all(self, **handlers):
+    def pull(self, n=-1, qid=-1, **handlers):
+        if n != -1:
+            raise ValueError("Incremental pull is not supported in Bolt 3")
+        if qid != -1:
+            raise ValueError("Query selection on pull is not supported in Bolt 3")
         log.debug("[#%04X]  C: PULL_ALL", self.local_port)
         self._append(b"\x3F", (), Response(self, **handlers))
 

neo4j/io/_bolt4x0.py

@@ -182,9 +182,14 @@ def discard(self, n=-1, qid=-1, **handlers):
         log.debug("[#%04X]  C: DISCARD %r", self.local_port, extra)
         self._append(b"\x2F", (extra,), Response(self, **handlers))
 
-    def pull_all(self, **handlers):
-        log.debug("[#%04X]  C: PULL_ALL", self.local_port)
-        self._append(b"\x3F", (), Response(self, **handlers))
+    def pull(self, n=-1, qid=-1, **handlers):
+        # TODO: find out whether qid is optional, and optimise if so
+        extra = {
+            "n": n,
+            "qid": qid,
+        }
+        log.debug("[#%04X]  C: PULL %r", self.local_port, extra)
+        self._append(b"\x3F", (extra,), Response(self, **handlers))
 
     def begin(self, mode=None, bookmarks=None, metadata=None, timeout=None,
               db=None, **handlers):

neo4j/work/pipelining.py

@@ -44,7 +44,7 @@ def __init__(self, pool, config):
 
     def push(self, statement, parameters=None):
         self._connection.run(statement, parameters)
-        self._connection.pull_all(on_records=self._data.extend)
+        self._connection.pull(on_records=self._data.extend)
         output_buffer_size = len(self._connection.outbox.view())
         if output_buffer_size >= self._flush_every:
             self._connection.send_all()

neo4j/work/simple.py

@@ -207,7 +207,7 @@ def done(summary_metadata):
         # TODO: capture ValueError and surface as SessionError/TransactionError if
         # TODO: explicit database selection has been made
         cx.run(statement_text, parameters, **run_metadata)
-        cx.pull_all(
+        cx.pull(
             on_records=lambda records: result._records.extend(
                 hydrant.hydrate_records(result.keys(), records)),
             on_success=done,

tests/unit/io/test_class_bolt3.py

@@ -80,3 +80,28 @@ def test_qid_extra_not_supported_in_discard(fake_socket):
     with pytest.raises(ValueError):
         connection.discard(qid=666)
 
+
+def test_simple_pull(fake_socket):
+    address = ("127.0.0.1", 7687)
+    socket = fake_socket(address)
+    connection = Bolt3(address, socket)
+    connection.pull()
+    connection.send_all()
+    tag, fields = socket.pop_message()
+    assert tag == b"\x3F"
+    assert len(fields) == 0
+
+
+def test_n_extra_not_supported_in_pull(fake_socket):
+    address = ("127.0.0.1", 7687)
+    connection = Bolt3(address, fake_socket(address))
+    with pytest.raises(ValueError):
+        connection.pull(n=666)
+
+
+def test_qid_extra_not_supported_in_pull(fake_socket):
+    address = ("127.0.0.1", 7687)
+    connection = Bolt3(address, fake_socket(address))
+    with pytest.raises(ValueError):
+        connection.pull(qid=666)
+

tests/unit/io/test_class_bolt4x0.py

@@ -100,3 +100,39 @@ def test_n_and_qid_extras_in_discard(fake_socket):
     assert tag == b"\x2F"
     assert len(fields) == 1
     assert fields[0] == {"n": 666, "qid": 777}
+
+
+def test_n_extra_in_pull(fake_socket):
+    address = ("127.0.0.1", 7687)
+    socket = fake_socket(address)
+    connection = Bolt4x0(address, socket)
+    connection.pull(n=666)
+    connection.send_all()
+    tag, fields = socket.pop_message()
+    assert tag == b"\x3F"
+    assert len(fields) == 1
+    assert fields[0] == {"n": 666, "qid": -1}
+
+
+def test_qid_extra_in_pull(fake_socket):
+    address = ("127.0.0.1", 7687)
+    socket = fake_socket(address)
+    connection = Bolt4x0(address, socket)
+    connection.pull(qid=666)
+    connection.send_all()
+    tag, fields = socket.pop_message()
+    assert tag == b"\x3F"
+    assert len(fields) == 1
+    assert fields[0] == {"n": -1, "qid": 666}
+
+
+def test_n_and_qid_extras_in_pull(fake_socket):
+    address = ("127.0.0.1", 7687)
+    socket = fake_socket(address)
+    connection = Bolt4x0(address, socket)
+    connection.pull(n=666, qid=777)
+    connection.send_all()
+    tag, fields = socket.pop_message()
+    assert tag == b"\x3F"
+    assert len(fields) == 1
+    assert fields[0] == {"n": 666, "qid": 777}