Stream results instead of fetching the entire set

tests/streaming_test.py

@@ -0,0 +1,137 @@
+# SPDX-FileCopyrightText: 2026 Bart Broere
+# SPDX-License-Identifier: GPL-3.0-or-later
+"""End-to-end test of result streaming, speaking Bolt directly for control over PULL.
+
+Usage: python tests/streaming_test.py [host] [port] [user] [password] [graph]
+The graph must exist. Requires an importable boltage.
+"""
+
+import socket
+import struct
+import sys
+
+from boltage.packstream import Packer, Structure, unpack
+
+host = sys.argv[1] if len(sys.argv) > 1 else "localhost"
+port = int(sys.argv[2]) if len(sys.argv) > 2 else 7687
+user = sys.argv[3] if len(sys.argv) > 3 else "postgres"
+password = sys.argv[4] if len(sys.argv) > 4 else "postgres"
+graph = sys.argv[5] if len(sys.argv) > 5 else "neo4j"
+
+HELLO, RESET, RUN, BEGIN, COMMIT, DISCARD, PULL = 0x01, 0x0F, 0x10, 0x11, 0x12, 0x2F, 0x3F
+SUCCESS, RECORD, FAILURE = 0x70, 0x71, 0x7F
+RANGE = "UNWIND range(1, 2500) AS i RETURN i"
+
+
+def check(label, condition):
+    print(("ok   " if condition else "FAIL ") + label)
+    if not condition:
+        check.failed = True
+
+
+check.failed = False
+
+
+class Client:
+    def __init__(self):
+        self.sock = socket.create_connection((host, port), timeout=30)
+        self.sock.sendall(b"\x60\x60\xB0\x17" + bytes([0, 0, 4, 4]) + bytes(12))
+        assert self.read(4) == bytes([0, 0, 4, 4])
+        self.packer = Packer(element_ids=False)
+        self.request(HELLO, {"user_agent": "boltage-test", "scheme": "basic", "principal": user,
+                             "credentials": password})
+
+    def read(self, n):
+        data = b""
+        while len(data) < n:
+            chunk = self.sock.recv(n - len(data))
+            if not chunk:
+                raise ConnectionError("closed by server")
+            data += chunk
+        return data
+
+    def message(self):
+        data = b""
+        while True:
+            size = struct.unpack(">H", self.read(2))[0]
+            if size == 0:
+                return unpack(data)
+            data += self.read(size)
+
+    def request(self, tag, *fields):
+        """Send a message; return its records and the final SUCCESS or FAILURE."""
+        data = self.packer.pack(Structure(tag, list(fields)))
+        self.sock.sendall(struct.pack(">H", len(data)) + data + b"\x00\x00")
+        records = []
+        while True:
+            reply = self.message()
+            if reply.tag != RECORD:
+                return records, reply
+            records.append(reply.fields[0])
+
+    def run(self, query):
+        return self.request(RUN, query, {}, {"db": graph})[1]
+
+
+
+c = Client()
+
+check("RUN reports the columns", c.run(RANGE).fields[0]["fields"] == ["i"])
+records, end = c.request(PULL, {"n": 10})
+check("PULL n sends n records", [r[0] for r in records] == list(range(1, 11)))
+check("PULL n reports has_more", end.tag == SUCCESS and end.fields[0].get("has_more") is True)
+records, end = c.request(PULL, {"n": 1500})
+check("PULL n crosses a batch boundary", [r[0] for r in records] == list(range(11, 1511)))
+check("has_more while rows remain", end.fields[0].get("has_more") is True)
+records, end = c.request(PULL, {"n": -1})
+check("PULL all sends the rest", [r[0] for r in records] == list(range(1511, 2501)))
+check("final summary", "has_more" not in end.fields[0] and "bookmark" in end.fields[0])
+
+c.run("UNWIND range(1, 2000) AS i RETURN i")
+records, end = c.request(PULL, {"n": 1000})
+check("exactly one batch leaves has_more", len(records) == 1000 and end.fields[0].get("has_more") is True)
+records, end = c.request(PULL, {"n": 1000})
+check("exactly two batches end the result", len(records) == 1000 and "has_more" not in end.fields[0])
+
+c.run("UNWIND range(1, 1000) AS i RETURN i")
+records, end = c.request(PULL, {"n": 1000})
+check("a result of exactly one batch ends", len(records) == 1000 and "has_more" not in end.fields[0])
+
+c.run(RANGE)
+c.request(PULL, {"n": 10})
+records, end = c.request(DISCARD, {"n": -1})
+check("DISCARD mid-result ends it", end.tag == SUCCESS and not records and "has_more" not in end.fields[0])
+
+c.run(RANGE)
+c.request(PULL, {"n": 10})
+check("RUN replaces an unfinished result", c.run("RETURN 2 AS x").tag == SUCCESS)
+records, _ = c.request(PULL, {"n": -1})
+check("the new result is read", records == [[2]])
+
+c.run(RANGE)
+c.request(PULL, {"n": 10})
+check("BEGIN after an unfinished result", c.request(BEGIN, {"db": graph})[1].tag == SUCCESS)
+first = c.run("UNWIND range(1, 2500) AS i RETURN i").fields[0]["qid"]
+second = c.run("UNWIND range(5001, 7500) AS i RETURN i").fields[0]["qid"]
+values = []
+for _ in range(1500):
+    values.append(c.request(PULL, {"n": 1, "qid": first})[0][0][0])
+    values.append(c.request(PULL, {"n": 1, "qid": second})[0][0][0])
+check("two results in a transaction, pulled alternately",
+      values[:2] == [1, 5001] and values[-2:] == [1500, 6500])
+check("COMMIT with results open", c.request(COMMIT, {})[1].tag == SUCCESS)
+
+c.run("UNWIND range(1, 1500) AS i RETURN 1 / (1500 - i) AS v")
+records, end = c.request(PULL, {"n": -1})
+check("error in a later batch fails the PULL", end.tag == FAILURE and len(records) >= 1000)
+check("RESET recovers", c.request(RESET)[1].tag == SUCCESS)
+check("no transaction left aborted", c.run("RETURN 1 AS x").tag == SUCCESS)
+c.request(PULL, {"n": -1})
+
+c.request(BEGIN, {"db": graph})
+c.run("RETURN 1 / 0 AS v")
+check("error in a transaction's first batch keeps it failed until RESET", c.request(RESET)[1].tag == SUCCESS)
+c.run("RETURN 1 AS x")
+check("usable after that RESET", c.request(PULL, {"n": -1})[0] == [[1]])
+
+sys.exit(1 if check.failed else 0)