Stream results instead of fetching the entire set

README.md

@@ -103,6 +103,12 @@ variables](https://www.postgresql.org/docs/current/libpq-envars.html), such as
 - **Queries** are sent to AGE as `SELECT * FROM cypher('<graph>', $$ <query> $$) AS (...)`.
   boltage works out the result columns from the final `RETURN` clause and inlines
   query parameters as Cypher literals.
+- **Results** of read queries are streamed through a PostgreSQL cursor as clients pull
+  them, so boltage holds at most a thousand rows of a result at a time. The query
+  keeps a transaction open until its result has been read. Queries that can write
+  (`CREATE`, `MERGE`, `SET`, `DELETE`, `REMOVE`, or `CALL` of a procedure) run to
+  completion first and are buffered, because AGE performs writes while their rows are
+  fetched.
 - **Neo4j procedures** that neo4j-browser relies on (`dbms.components`, `SHOW DATABASES`,
   `db.labels`, `db.schema.visualization`, `SHOW PROCEDURES`, ...) are emulated from the
   AGE catalog.
@@ -522,8 +528,8 @@ addressing it in [`plans/`](plans/README.md):
    ([plan](plans/04-users-transactions-sysinfo.md)).
 5. [TLS](#tls) certificates are only read at start-up, so renewing one needs a
    restart ([plan](plans/05-tls.md)).
-6. Each query's full result is fetched from PostgreSQL before it is streamed to the
-   client, so very large results use memory in boltage
+6. The full result of a query that can write is fetched from PostgreSQL before it is
+   sent to the client, so very large results of such queries use memory in boltage
    ([plan](plans/06-streaming-results.md)).
 
 ## Development
@@ -544,6 +550,13 @@ pip install neo4j
 python tests/smoke_test.py neo4j://localhost:7687 postgres secret
 ```
 
+[`tests/streaming_test.py`](tests/streaming_test.py) checks result streaming by speaking
+Bolt directly, which gives it control over `PULL`:
+
+```sh
+python tests/streaming_test.py localhost 7687 postgres secret neo4j
+```
+
 ## License
 
 boltage is free software, licensed under the GNU General Public License v3.0 or

boltage/backend.py

@@ -7,6 +7,7 @@ only answers the administrative queries that neo4j-browser issues.
 """
 
 import logging
+import itertools
 import re
 import secrets
 import ssl
@@ -18,7 +19,7 @@ from psycopg import errors as pg_errors
 from psycopg import sql
 
 from . import agtype
-from .cypher import (CypherError, expand_id_in_lists, first_keyword, inline_parameters,
+from .cypher import (CypherError, expand_id_in_lists, first_keyword, has_keyword, inline_parameters,
                      return_columns, rewrite_for_age, strip_statement)
 from .packstream import Node, Relationship
 
@@ -27,6 +28,7 @@ log = logging.getLogger("boltage.backend")
 NEO4J_VERSION = "5.26.0"
 SYSTEM_DB = "system"
 SYSTEM_LABELS = ("_ag_label_vertex", "_ag_label_edge")
+FETCH_SIZE = 1000  # rows fetched from a PostgreSQL cursor at a time
 GRAPH_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.\-]{2,62}$")
 
 # Procedures emulated by boltage, reported through SHOW PROCEDURES.
@@ -66,13 +68,68 @@ class Config:
     tls_context: ssl.SSLContext | None = None
 
 
+class Cursor:
+    """A PostgreSQL cursor. It runs in the Bolt transaction, or else in a transaction of its own."""
+
+    def __init__(self, session, name, own_transaction):
+        self.session, self.name, self.own_transaction = session, name, own_transaction
+
+    async def fetch(self):
+        return await self.session._sql(f"FETCH FORWARD {FETCH_SIZE} FROM {self.name}")
+
+    async def close(self):
+        await self.session._sql(f"CLOSE {self.name}" + ("; COMMIT" if self.own_transaction else ""))
+
+
 @dataclass
 class Result:
+    """The rows of a query that have not been sent yet.
+
+    Rows come from ``rows``, and from ``cursor`` when that is set. At most one batch
+    of a cursor's rows is held here at a time.
+    """
+
     fields: list
     rows: list = field(default_factory=list)
     db: str = ""
     query_type: str = "r"
     t_first: int = 0
+    cursor: Cursor | None = None
+    position: int = 0  # the next row to send
+
+    @property
+    def buffered(self):
+        return self.position < len(self.rows)
+
+    @property
+    def done(self):
+        return not self.buffered and self.cursor is None
+
+    async def add_batch(self, batch):
+        """Take a batch fetched from the cursor; a short batch means the cursor is finished."""
+        self.rows, self.position = [[agtype.parse(v) for v in row] for row in batch], 0
+        if len(batch) < FETCH_SIZE:
+            await self.close_cursor()
+
+    def take(self, limit: int) -> list:
+        """The next ``limit`` buffered rows or fewer."""
+        rows = self.rows[self.position:self.position + limit]
+        self.position += len(rows)
+        return rows
+
+    async def refill(self):
+        """Fetch more rows if none are buffered, so ``done`` tells whether any are left."""
+        if not self.buffered and self.cursor is not None:
+            await self.add_batch(await self.cursor.fetch())
+
+    async def close_cursor(self):
+        if self.cursor is not None:
+            cursor, self.cursor = self.cursor, None
+            await cursor.close()
+
+    async def close(self):
+        self.rows, self.position = [], 0
+        await self.close_cursor()
 
 
 class BoltFailure(Exception):
@@ -120,6 +177,7 @@ class Session:
         self.user = user
         self.conn: psycopg.AsyncConnection | None = None
         self.tx_db: str | None = None
+        self.cursor_ids = itertools.count(1)
 
     async def connect(self, password: str):
         try:
@@ -227,14 +285,40 @@ class Session:
         stmt = (f"SELECT * FROM ag_catalog.cypher({sql.Literal(graph).as_string(self.conn)}, "
                 f"${tag}$ {text} ${tag}$) AS ({col_defs})")
         log.debug("SQL: %s", stmt)
+        writes = has_keyword(text, "CREATE|MERGE|SET|DELETE|REMOVE")
+        # Only reads go through a cursor. AGE performs writes as their rows are fetched, so a
+        # cursor would delay them past the next query, and leave them half done if a client
+        # stops reading. A procedure called with CALL may write too.
+        if columns and not writes and not has_keyword(text, "CALL"):
+            return await self._stream(stmt, columns)
         rows = await self._sql(stmt)
         if columns:
             rows = [[agtype.parse(v) for v in row] for row in rows]
         else:
             rows = []
-        writes = re.search(r"(?<![\w.])(CREATE|MERGE|SET|DELETE|REMOVE)(?!\w)", text, re.IGNORECASE)
         return Result(columns, rows, query_type="rw" if writes else "r")
 
+    async def _stream(self, stmt, columns) -> Result:
+        """Open a cursor on a read query and fetch the first batch, which reports most errors."""
+        name = f"boltage_c{next(self.cursor_ids)}"
+        own_transaction = not self.in_transaction
+        # PostgreSQL plans cursors for fetching only a part of the rows by default; clients
+        # usually read everything, so plan like for a plain query.
+        script = (f"{'BEGIN; ' if own_transaction else ''}SET LOCAL cursor_tuple_fraction = 1; "
+                  f"DECLARE {name} NO SCROLL CURSOR FOR {stmt}; FETCH FORWARD {FETCH_SIZE} FROM {name}")
+        try:
+            cur = await self.conn.execute(script)
+            while cur.nextset():
+                pass
+            batch = await cur.fetchall()
+        except psycopg.Error as e:
+            if own_transaction:
+                await self.rollback()
+            raise failure_from_pg(e) from e
+        result = Result(columns, cursor=Cursor(self, name, own_transaction))
+        await result.add_batch(batch)
+        return result
+
     # -- emulated Neo4j procedures ---------------------------------------------
 
     async def _run_system(self, query, db) -> Result | None:

boltage/bolt.py

@@ -16,7 +16,7 @@ from websockets.protocol import State
 from websockets.server import ServerProtocol
 
 from . import packstream
-from .backend import BoltFailure, Config, Session
+from .backend import FETCH_SIZE, BoltFailure, Config, Session
 from .discovery import discovery_document
 from .packstream import Structure
 from .streams import TLS_HANDSHAKE, ConnectionClosed, TcpStream, TlsStream
@@ -165,7 +165,7 @@ class BoltConnection:
         self.session: Session | None = None
         self.failed = False
         self.routing = {}
-        self.results = {}  # qid -> [Result, next row index]
+        self.results = {}  # qid -> Result
         self.last_qid = -1
         self.qids = itertools.count()
 
@@ -200,14 +200,17 @@ class BoltConnection:
                 return packstream.unpack(bytes(data))
             data += await self.t.read(size)
 
-    async def send(self, tag, *fields):
+    def encode(self, tag, *fields) -> bytes:
         data = self.packer.pack(Structure(tag, list(fields)))
         out = bytearray()
         for i in range(0, len(data), 0xFFFF):
             chunk = data[i:i + 0xFFFF]
             out += struct.pack(">H", len(chunk)) + chunk
         out += b"\x00\x00"
-        await self.t.write(bytes(out))
+        return bytes(out)
+
+    async def send(self, tag, *fields):
+        await self.t.write(self.encode(tag, *fields))
 
     async def success(self, metadata=None):
         await self.send(SUCCESS, metadata or {})
@@ -253,6 +256,7 @@ class BoltConnection:
             await self.logon(f[0])
             await self.success()
         elif tag == LOGOFF:
+            self.results.clear()
             if self.session:
                 await self.session.close()
             self.session = None
@@ -267,18 +271,25 @@ class BoltConnection:
             await self.pull(f[0] if f else {}, discard=tag == DISCARD)
         elif tag == BEGIN:
             extra = f[0] if f else {}
+            await self.close_results()
             await self.require_session().begin(extra.get("db"))
-            self.results.clear()
             await self.success()
         elif tag == COMMIT:
+            self.results.clear()  # the transaction's cursors end with it
             await self.require_session().commit()
             await self.success({"bookmark": "boltage:bookmark"})
         elif tag == ROLLBACK:
+            self.results.clear()
             await self.require_session().rollback()
             await self.success()
         else:
             raise BoltFailure("Neo.ClientError.Request.Invalid", f"Unsupported Bolt message 0x{tag:02X}")
 
+    async def close_results(self):
+        results, self.results = self.results, {}
+        for result in results.values():
+            await result.close()
+
     def require_session(self) -> Session:
         if self.session is None:
             raise BoltFailure("Neo.ClientError.Security.Unauthorized", "Not authenticated")
@@ -308,10 +319,10 @@ class BoltConnection:
         extra = f[2] if len(f) > 2 else {}
         session = self.require_session()
         if not session.in_transaction:
-            self.results.clear()
+            await self.close_results()
         result = await session.run(query, params, extra.get("db"))
         qid = next(self.qids)
-        self.results[qid] = [result, 0]
+        self.results[qid] = result
         self.last_qid = qid
         metadata = {"fields": result.fields, "t_first": result.t_first}
         if session.in_transaction:
@@ -323,14 +334,19 @@ class BoltConnection:
         qid = self.last_qid if qid == -1 else qid
         if qid not in self.results:
             raise BoltFailure("Neo.ClientError.Request.Invalid", "No result available to pull from")
-        result, index = self.results[qid]
-        n = extra.get("n", -1)
-        end = len(result.rows) if n == -1 or discard else min(index + n, len(result.rows))
-        if not discard:
-            for row in result.rows[index:end]:
-                await self.send(RECORD, row)
-        self.results[qid][1] = end
-        if end < len(result.rows):
+        result = self.results[qid]
+        if discard:
+            await result.close()
+        else:
+            remaining = extra.get("n", -1)
+            while remaining != 0 and not result.done:
+                rows = result.take(FETCH_SIZE if remaining < 0 else remaining)
+                await self.t.write(b"".join(self.encode(RECORD, row) for row in rows))
+                if remaining > 0:
+                    remaining -= len(rows)
+                # After sending, so an error in the next batch does not lose these rows.
+                await result.refill()
+        if not result.done:
             await self.success({"has_more": True})
             return
         del self.results[qid]

boltage/cypher.py

@@ -236,3 +236,8 @@ def expand_id_in_lists(text: str) -> str:
 def first_keyword(text: str) -> str:
     m = re.match(r"\s*(\w+)", scan(text).masked)
     return m.group(1).upper() if m else ""
+
+
+def has_keyword(text: str, keywords: str) -> bool:
+    """Whether a keyword from a ``|``-separated list occurs outside strings and comments."""
+    return re.search(rf"(?<![\w.])({keywords})(?!\w)", scan(text).masked, re.IGNORECASE) is not None

plans/02-write-statistics.md

@@ -77,8 +77,8 @@ counters for the graph's schema before and after the query and take the differen
    `query_type` (`r`, `w` or `rw`, depending on whether there are columns and writes)
    instead of the regex.
 
-With streaming (plan 06) the "after" snapshot must be taken once the cursor is
-exhausted or drained, just before the final summary.
+Streaming (plan 06) does not affect this: only read queries go through a cursor, and
+queries that can write have run to completion when `RUN` returns.
 
 ## Alternatives considered
 

plans/06-streaming-results.md

@@ -1,5 +1,24 @@
 # Streaming results instead of fetching them fully
 
+## Status
+
+Done for read queries: they run through a cursor, in a transaction of their own
+outside an explicit Bolt transaction. On `UNWIND range(1, 3000000) AS i RETURN i,
+'some text in every row' AS t`, boltage's peak memory went from 1419 MB to 44 MB, the
+first record from 10.8 s to 0.5 s, and the whole query from 44 s to 33 s.
+
+Queries that can write are deliberately still buffered, unlike what this plan
+proposes below. The spike found that AGE performs writes as rows are fetched: fetching
+2 of 5 rows of `UNWIND ... CREATE ... RETURN` and committing creates 2 nodes. Writes in
+a cursor would therefore not be visible to the next query in the transaction, and
+their errors would move from `RUN` to `PULL`. Reads have none of this, and have no
+need for `MOVE FORWARD ALL`. Any query containing `CREATE`, `MERGE`, `SET`, `DELETE`,
+`REMOVE` or `CALL` is buffered, and runs in autocommit as before. The sections on
+draining below only apply if write results are ever streamed; that needs a way to
+keep write ordering, and write queries rarely return much.
+
+Tests: `tests/streaming_test.py`, plus additions to `tests/smoke_test.py`.
+
 ## Problem
 
 `Session._sql` calls `fetchall()` (`boltage/backend.py:178`), `_run_cypher` parses

plans/README.md

@@ -28,7 +28,9 @@ plans small:
    transaction-local counters, streaming (06) needs a transaction to keep a server-side
    cursor open, and `PROFILE` (01) needs a savepoint. All three want `Session.run` to
    wrap auto-commit queries in `BEGIN ... COMMIT` instead of relying on the
-   connection's autocommit mode. Do this once, in 06 or 02, whichever comes first.
+   connection's autocommit mode. Plan 06 only did this for read queries, which run
+   in a cursor's transaction; queries that can write still use autocommit, so 02 still
+   has to add it for them.
 3. **Unit tests.** The only test is the end-to-end `tests/smoke_test.py`. Most of
    this work is text rewriting in `cypher.py`, which is cheap to unit test without a
    database. Add `tests/test_cypher.py` (pytest) alongside the first plan that

tests/smoke_test.py

@@ -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()
 

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)