diff --git a/README.md b/README.md index 0c2fd8a..1eea0d1 100644 --- a/README.md +++ b/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('', $$ $$) 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 diff --git a/boltage/backend.py b/boltage/backend.py index 72bf0e1..be8360c 100644 --- a/boltage/backend.py +++ b/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"(? 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: diff --git a/boltage/bolt.py b/boltage/bolt.py index e54b8cb..ebebb10 100644 --- a/boltage/bolt.py +++ b/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] diff --git a/boltage/cypher.py b/boltage/cypher.py index e85e824..18e39ab 100644 --- a/boltage/cypher.py +++ b/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"(?() 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() diff --git a/tests/streaming_test.py b/tests/streaming_test.py new file mode 100644 index 0000000..a2dfdaf --- /dev/null +++ b/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)