Stream results instead of fetching the entire set
This commit is contained in:
1 parent
eb989c9155
commit
dc978fd0fd
9 files changed
+324
-22
No files matched your search
@@ -91,6 +91,26 @@ MATCH ()-[]->() RETURN { name:'relationships', data: count(*)} AS result
|
||||
rows = list(s.run("UNWIND range(1, 2500) AS i RETURN i"))
|
||||
check("large result over several PULLs", len(rows) == 2500 and rows[-1]["i"] == 2500)
|
||||
|
||||
rows = [r["i"] for r in s.run("UNWIND range(1, 2500) AS i RETURN i ORDER BY i")]
|
||||
check("rows in order across batches", rows == list(range(1, 2501)))
|
||||
|
||||
s.run("UNWIND range(1, 5) AS i CREATE (n:Discarded {i: i}) RETURN n.i").consume()
|
||||
check("discarded result of a write still writes",
|
||||
s.run("MATCH (n:Discarded) RETURN count(n) AS c").single()["c"] == 5)
|
||||
s.run("MATCH (n:Discarded) DELETE n RETURN 1").consume()
|
||||
|
||||
check("write keyword in a string is a read",
|
||||
s.run("RETURN 'SET' AS v").consume().query_type == "r")
|
||||
|
||||
rows = []
|
||||
try:
|
||||
for r in s.run("UNWIND range(1, 1500) AS i RETURN 1 / (1500 - i) AS v"):
|
||||
rows.append(r["v"])
|
||||
check("error after the first batch is reported", False)
|
||||
except neo4j.exceptions.Neo4jError:
|
||||
check("error after the first batch is reported", len(rows) >= 1000)
|
||||
check("recovers after error in a later batch", s.run("RETURN 1 AS x").single()["x"] == 1)
|
||||
|
||||
check("string escaping roundtrip",
|
||||
s.run("RETURN $v AS v", v="a'b\\c\n\"d\" $x").single()["v"] == "a'b\\c\n\"d\" $x")
|
||||
check("map/list params",
|
||||
@@ -117,6 +137,12 @@ MATCH ()-[]->() RETURN { name:'relationships', data: count(*)} AS result
|
||||
|
||||
check("managed write transaction", s.execute_write(work) == 1)
|
||||
|
||||
def read_your_writes(tx):
|
||||
tx.run("CREATE (:Temp {n: 3}) RETURN 1")
|
||||
return tx.run("UNWIND range(1, 1500) AS i MATCH (n:Temp {n: 3}) RETURN i").value()
|
||||
|
||||
check("query sees an earlier write in its transaction", len(s.execute_write(read_your_writes)) == 1500)
|
||||
|
||||
with driver.session(database="system") as s:
|
||||
s.run(f"DROP DATABASE {GRAPH}").consume()
|
||||
|
||||
|
||||
@@ -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)
|
||||
Reference in new issue
Block a user