Added n and qid to DISCARD (renamed from DISCARD_ALL)

Author: technige

neo4j/io/__init__.py

@@ -178,8 +178,14 @@ def run(self, statement, parameters=None, mode=None, bookmarks=None, metadata=No
         :return: :class:`neo4j.io._bolt3.Response` object
         """
 
-    def discard_all(self, **handlers):
-        raise NotImplementedError
+    def discard(self, n=-1, qid=-1, **handlers):
+        """ Appends a DISCARD message to the output stream.
+
+        :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
+        """
 
     def pull_all(self, **handlers):
         raise NotImplementedError

neo4j/io/_bolt3.py

@@ -173,7 +173,11 @@ def run(self, statement, parameters=None, mode=None, bookmarks=None, metadata=No
         else:
             self._append(b"\x10", fields, Response(self, **handlers))
 
-    def discard_all(self, **handlers):
+    def discard(self, n=-1, qid=-1, **handlers):
+        if n != -1:
+            raise ValueError("Incremental discard is not supported in Bolt 3")
+        if qid != -1:
+            raise ValueError("Query selection on discard is not supported in Bolt 3")
         log.debug("[#%04X]  C: DISCARD_ALL", self.local_port)
         self._append(b"\x2F", (), Response(self, **handlers))
 

neo4j/io/_bolt4x0.py

@@ -173,9 +173,14 @@ def run(self, statement, parameters=None, mode=None, bookmarks=None, metadata=No
         else:
             self._append(b"\x10", fields, Response(self, **handlers))
 
-    def discard_all(self, **handlers):
-        log.debug("[#%04X]  C: DISCARD_ALL", self.local_port)
-        self._append(b"\x2F", (), Response(self, **handlers))
+    def discard(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: 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)

tests/unit/io/test_class_bolt3.py

@@ -54,3 +54,29 @@ def test_db_extra_not_supported_in_run(fake_socket):
     connection = Bolt3(address, fake_socket(address))
     with pytest.raises(ValueError):
         connection.run("", db="something")
+
+
+def test_simple_discard(fake_socket):
+    address = ("127.0.0.1", 7687)
+    socket = fake_socket(address)
+    connection = Bolt3(address, socket)
+    connection.discard()
+    connection.send_all()
+    tag, fields = socket.pop_message()
+    assert tag == b"\x2F"
+    assert len(fields) == 0
+
+
+def test_n_extra_not_supported_in_discard(fake_socket):
+    address = ("127.0.0.1", 7687)
+    connection = Bolt3(address, fake_socket(address))
+    with pytest.raises(ValueError):
+        connection.discard(n=666)
+
+
+def test_qid_extra_not_supported_in_discard(fake_socket):
+    address = ("127.0.0.1", 7687)
+    connection = Bolt3(address, fake_socket(address))
+    with pytest.raises(ValueError):
+        connection.discard(qid=666)
+

tests/unit/io/test_class_bolt4x0.py

@@ -64,3 +64,39 @@ def test_db_extra_in_run(fake_socket):
     assert fields[0] == ""
     assert fields[1] == {}
     assert fields[2] == {"db": "something"}
+
+
+def test_n_extra_in_discard(fake_socket):
+    address = ("127.0.0.1", 7687)
+    socket = fake_socket(address)
+    connection = Bolt4x0(address, socket)
+    connection.discard(n=666)
+    connection.send_all()
+    tag, fields = socket.pop_message()
+    assert tag == b"\x2F"
+    assert len(fields) == 1
+    assert fields[0] == {"n": 666, "qid": -1}
+
+
+def test_qid_extra_in_discard(fake_socket):
+    address = ("127.0.0.1", 7687)
+    socket = fake_socket(address)
+    connection = Bolt4x0(address, socket)
+    connection.discard(qid=666)
+    connection.send_all()
+    tag, fields = socket.pop_message()
+    assert tag == b"\x2F"
+    assert len(fields) == 1
+    assert fields[0] == {"n": -1, "qid": 666}
+
+
+def test_n_and_qid_extras_in_discard(fake_socket):
+    address = ("127.0.0.1", 7687)
+    socket = fake_socket(address)
+    connection = Bolt4x0(address, socket)
+    connection.discard(n=666, qid=777)
+    connection.send_all()
+    tag, fields = socket.pop_message()
+    assert tag == b"\x2F"
+    assert len(fields) == 1
+    assert fields[0] == {"n": 666, "qid": 777}