From a769713a3bde2a806378fff7c5e0f2b40baa041e Mon Sep 17 00:00:00 2001 From: Tam Nguyen Duc <1218621+tamnd@users.noreply.github.com> Date: Wed, 19 Aug 2026 16:39:10 +0700 Subject: [PATCH] Prepare a statement, and see what one would do `prepare` compiles a statement once and gives back an object that runs it, `explain` says what the engine would run without running it, and `profile` runs it and says what it really did. The honest part is in the README. There is no round trip to save here, and the engine already caches the plan it compiled under the text of the statement, so a prepared statement bound per run costs 30.1 microseconds and the same text handed to `execute` costs 30.1 as well. What preparing buys is the compile happening at the line that asked for it, and `params` saying which names the statement wants, and an object a program can hold and close rather than a string it passes around. A plan and a profile each answer twice: `print` gives the engine's own listing, rendered by the engine so the two cannot drift apart, and `root` gives the same plan as objects for a program to walk. An operator carries `op` for what it is and `name` for what the listing calls it, which differ inside a bracket, and a program should match on `op` because that is the one that does not change with the company an operator keeps. Profile counts are `int`, exact at any size, and an operator the optimizer had nothing to say about carries `None` rather than a zero somebody would read as an estimate. A prepared statement runs through the same call `execute` does, so it gets the GIL released, the interrupt handling and the catalog names by being that call rather than a second path beside it. `zudb.aio` mirrors all three, and `statement`, `params` and `closed` are properties there because the names were read at the compile and nothing waits to say them again. Forty-two tests: nineteen for the lifetime of a prepared statement and twenty-three for the shape of a plan and a profile, including the listing asserted against the tree it was rendered from rather than against a string written out here. --- README.md | 59 ++++++- python/zudb/__init__.py | 14 ++ python/zudb/_zudb.pyi | 208 ++++++++++++++++++++++ python/zudb/aio.py | 105 ++++++++++- src/conn.rs | 128 +++++++++++++- src/lib.rs | 9 + src/plan.rs | 373 ++++++++++++++++++++++++++++++++++++++++ src/prepared.rs | 158 +++++++++++++++++ tests/test_plan.py | 275 +++++++++++++++++++++++++++++ tests/test_prepared.py | 165 ++++++++++++++++++ 10 files changed, 1487 insertions(+), 7 deletions(-) create mode 100644 src/plan.rs create mode 100644 src/prepared.rs create mode 100644 tests/test_plan.py create mode 100644 tests/test_prepared.py diff --git a/README.md b/README.md index 9ddbb00..8d13e30 100644 --- a/README.md +++ b/README.md @@ -126,6 +126,63 @@ result.record_batches() # a reader, for a result larger than memory `Result` implements `__arrow_c_stream__`, so anything that reads the protocol reads a result directly and none of the four methods above is needed: `pyarrow.table(result)` and `polars.DataFrame(result)` both work. Batches are 65,536 rows. A column holds one type, which the values decide, and integers beside floats are the one mixture that widens rather than being refused. Nodes, rels and paths go across as structs. The copy runs with the GIL released, and on this machine 300,000 rows across three columns take 44 ms as Arrow against 67 ms as Python objects, and a single integer column takes 13.8 ms against 44.5 ms. +## Preparing a statement + +A statement a program runs many times with different values can be compiled once and kept. + +```python +find = conn.prepare("MATCH (p:person) WHERE p.name = $name RETURN p.uid AS uid") +find.params # ['name'] +find.execute({"name": "ada"}).fetchall() # [(10,)] +find.close() +``` + +```python +with conn.prepare("INSERT (p:person {uid: $uid, name: $name})") as add: + for uid, name in people: + add.execute({"uid": uid, "name": name}) +``` + +What it saves here is smaller than the word usually promises, and the number is worth printing rather than hiding. The database is in this process, so there is no round trip to skip, and the engine already keeps the plan it compiled for a statement under the text of that statement, so running the same text twice compiles it once whether anybody prepared it or not. On this machine, over two thousand rows, a prepared statement bound per run costs 30.1 microseconds, the same text passed to `execute` and bound per run costs 30.1, and a statement whose text differs every time costs 35.8. A statement of 1.2 KB with sixty predicates in it says the same thing: 34.9 microseconds either way. If you came here expecting the first number to be half the second, the honest answer is that `execute` was already doing what preparing does. + +What preparing does buy is worth having for other reasons. The compile happens at the line that asked for it, so a statement that will not compile fails where it was written rather than in the middle of a loop at three in the morning, and `params` tells the program which names the statement wants, which is the difference between binding what it asked for and finding out at the run that a key was misspelled. It also gives a name to the intent: a prepared statement is an object a program can pass around, hold on a class and close, and a program built around one is a program whose statements live in one place. + +A prepared statement belongs to the connection that made it. Closing it gives the statement back and is safe to do twice, closing the connection closes all of them, and a closed one refuses to run with a message that says so rather than a segfault. The `with` block closes it at the end, including the end an exception makes. `zudb.aio` has the same thing under `async with`, where `prepare`, `execute` and `close` are awaited and `statement`, `params` and `closed` are properties, since the names were read at the compile and nothing has to wait to say them again. + +## Seeing what a statement will do + +`explain` answers what the engine would run, and `profile` runs it and answers what it really did. + +```python +plan = conn.explain("MATCH (p:person) WHERE p.name = $name RETURN p.uid AS uid") +print(plan) +# Project p.uid AS uid +# Filter p.name = $name +# ScanNodes p: person + +plan.columns # ['uid'] +plan.params # ['name'] +plan.root.op # 'Project' +plan.root.children[0].children[0].tables # ['person'] +``` + +Both calls answer twice over and that is deliberate. `print` gives the engine's own listing, which is what a person reads, and it is rendered by the engine rather than assembled here so that the listing and the plan cannot drift apart from one release to the next. `root` gives the same plan as objects, which is what a program walks: a test that wants to know a scan became a seek asks the tree, rather than matching on a string that was written to be read. An operator carries `op` for what it is, `name` for what the listing calls it, `detail` for what it works on, `binds` for the variables it introduces, `tables` for what it reads and `children` for what it pulls from. The two names differ where an operator sits inside a bracket, so an expand under an `OPTIONAL MATCH` has `op` `Expand`, `name` `OptionalExpand` and `bracket` `Optional`, and `op` is the one to match on because it is the one that does not change with the company the operator keeps. + +`explain` takes no parameters, which is not an oversight. A plan is chosen from the shape of the statement and the values are bound when it runs, so a plan asked for with values would suggest the values had changed it. It also does not run the statement, so explaining an `INSERT` inserts nothing. A query written where a value belongs gets a plan of its own in `scalars`, along with `reads`, which is which variables of the query around it that one reads and therefore whether the executor runs it once or once a row. + +```python +run = conn.profile("MATCH (p:person) WHERE p.score > 40.0 RETURN p.name AS name") +print(run) +# stage 1: Project [2 rows, 247.8 us] +# Filter p.score > 40 pulls 1 rows 2 flat 2 est 1 q 2.0 ... +# Scan p: person pulls 1 rows 3 flat 3 est 3 q 1.0 ... +# Source pulls 1 rows 1 flat 1 est - q - ... +``` + +A profile takes parameters, because it is a run. `est` is what the optimizer thought an operator would produce and `rows` is what it did, so `qerror` is the larger of the two over the smaller: one where the estimate was right, ten where it was out by an order of magnitude either way, and that column is where a plan that went wrong announces itself. All three are `None` on an operator the optimizer had nothing to say about, which is honest rather than a zero somebody would read as an estimate of none. `stages[i].ops` runs from the operator that read to the one that fed the sink, which is the order they ran in and the reverse of the order the listing prints them in. Every count is an `int`, exact however large it gets. + +A statement that writes is refused rather than profiled, because a measurement that also inserted two rows changed the thing it was measuring, and because a write runs as the parts it was split at rather than as the one plan a profile describes. Explaining costs 1.3 microseconds, since the plan was already compiled and cached; profiling costs what the statement costs plus the counters. In a notebook both draw themselves as their listing, preformatted, since the indentation is what says which operator pulls from which. + ## In a notebook A result in a cell draws itself as a table, because Jupyter asks an object for `_repr_html_` before it falls back to `repr` and a line saying how many rows there are is a strictly worse answer than the rows. Nodes, rels and paths draw themselves too, a path as the walk it is: `(person #0) -[knows]-> (person #1)`. @@ -229,7 +286,7 @@ Half of this package is compiled, which is the one thing an inspection cannot se ## What works today -The list above is what this client is for. What it does so far is the core of it: `connect`, `execute` and `sql` with named parameters, results that iterate and fetch, values as Python objects both ways including dates, times, datetimes and durations, `Node`, `Rel` and `Path` as classes, `load` for building a graph with edges in it, an appender for growing one, transactions as a context manager that commits at the end of a block and rolls back when it raises, every condition as an exception class carrying its code, its position and its documentation link, results as Arrow columns and as pandas and polars frames, `register` for putting a frame under a name a statement can match on and reading it where it lies, stubs inside the wheel with a gate that keeps them true, the GIL released around every statement, every load and every copy out, `Ctrl-C` and `interrupt()` stopping a statement without touching the connection under it, `zudb.aio` for the same calls awaited on an event loop, results, nodes, rels and paths that draw themselves in a notebook with `%gql` and `%%gql` to run statements in one, and `zudb.dbapi` for code written against PEP 249. Each one landed with the tests that say it works. +The list above is what this client is for. What it does so far is the core of it: `connect`, `execute` and `sql` with named parameters, results that iterate and fetch, values as Python objects both ways including dates, times, datetimes and durations, `Node`, `Rel` and `Path` as classes, `load` for building a graph with edges in it, an appender for growing one, transactions as a context manager that commits at the end of a block and rolls back when it raises, every condition as an exception class carrying its code, its position and its documentation link, results as Arrow columns and as pandas and polars frames, `register` for putting a frame under a name a statement can match on and reading it where it lies, stubs inside the wheel with a gate that keeps them true, the GIL released around every statement, every load and every copy out, `Ctrl-C` and `interrupt()` stopping a statement without touching the connection under it, `zudb.aio` for the same calls awaited on an event loop, results, nodes, rels and paths that draw themselves in a notebook with `%gql` and `%%gql` to run statements in one, `zudb.dbapi` for code written against PEP 249, and `prepare`, `explain` and `profile` for a statement compiled once, the plan it would run and the plan it did. Each one landed with the tests that say it works. ## Wheels diff --git a/python/zudb/__init__.py b/python/zudb/__init__.py index 7f54920..d04594e 100644 --- a/python/zudb/__init__.py +++ b/python/zudb/__init__.py @@ -25,8 +25,15 @@ Duration, Node, Path, + Plan, + PlanNode, + Prepared, + Profile, + ProfileOp, + ProfileStage, Rel, Result, + ScalarPlan, Transaction, __abi_version__, connect, @@ -52,7 +59,14 @@ "Connection", "Transaction", "Appender", + "Prepared", "Result", + "Plan", + "PlanNode", + "ScalarPlan", + "Profile", + "ProfileStage", + "ProfileOp", "Node", "Rel", "Path", diff --git a/python/zudb/_zudb.pyi b/python/zudb/_zudb.pyi index 324cd75..f43d95b 100644 --- a/python/zudb/_zudb.pyi +++ b/python/zudb/_zudb.pyi @@ -76,6 +76,15 @@ class Connection: def sql(self, statement: str, params: Mapping[str, Value] | None = None) -> Result: """The same call, named for the way it reads in a notebook.""" + def prepare(self, statement: str) -> Prepared: + """Compiles a statement now and hands back something that runs it later.""" + + def explain(self, statement: str) -> Plan: + """What the statement would do, without doing it.""" + + def profile(self, statement: str, params: Mapping[str, Value] | None = None) -> Profile: + """Runs the statement with the counters on and answers what its operators really did.""" + def transaction(self, *, read_only: bool = False) -> Transaction: """Starts a transaction and hands it back for a `with` block.""" @@ -162,6 +171,205 @@ class Appender: def __exit__(self, *_exception: object) -> bool: ... def __repr__(self) -> str: ... +class Prepared: + """A statement the engine has compiled and is holding for you.""" + + @property + def statement(self) -> str: + """The text it was compiled from.""" + + @property + def params(self) -> list[str]: + """The names this statement wants bound, in the order it uses them.""" + + @property + def closed(self) -> bool: + """Whether this prepared statement has been closed.""" + + def execute(self, params: Mapping[str, Value] | None = None) -> Result: + """Runs it with these parameters and gives back its rows.""" + + def sql(self, params: Mapping[str, Value] | None = None) -> Result: + """The same call, named for the way it reads in a notebook.""" + + def close(self) -> None: + """Closes it and gives the statement back to the connection.""" + + def __enter__(self) -> Prepared: ... + def __exit__(self, *_exception: object) -> bool: ... + def __repr__(self) -> str: ... + +class PlanNode: + """One operator of a plan.""" + + @property + def op(self) -> str: + """The kind of operator this is.""" + + @property + def name(self) -> str: + """What the listing calls it, bracket and all.""" + + @property + def bracket(self) -> str | None: + """The bracket it sits inside, if it sits inside one.""" + + @property + def detail(self) -> str: + """What it works on, in the words the listing prints.""" + + @property + def binds(self) -> list[str]: + """The variables it introduces.""" + + @property + def tables(self) -> list[str]: + """The tables it reads.""" + + @property + def children(self) -> list[PlanNode]: + """What it pulls from, in the order the listing prints them.""" + + def __repr__(self) -> str: ... + +class ScalarPlan: + """A query written where a value belongs, and the plan it gets.""" + + @property + def reads(self) -> list[str]: + """The variables it reads from the query it is written inside.""" + + @property + def exists(self) -> bool: + """Whether it is asking whether there is a row rather than for one.""" + + @property + def plan(self) -> Plan: + """The plan itself.""" + + def __repr__(self) -> str: ... + +class Plan: + """What a statement would do, without doing it.""" + + @property + def root(self) -> PlanNode | None: + """The operator everything else feeds.""" + + @property + def columns(self) -> list[str]: + """The columns the statement projects, in order.""" + + @property + def params(self) -> list[str]: + """The parameters it wants bound.""" + + @property + def notes(self) -> list[str]: + """What the planner has to say about it, if anything.""" + + @property + def scalars(self) -> list[ScalarPlan]: + """The plans of the queries written where values belong.""" + + @property + def text(self) -> str: + """The engine's own listing, which is what `print` gives.""" + + def _repr_html_(self) -> str: + """The listing as a notebook shows it, which is the listing.""" + + def __str__(self) -> str: ... + def __repr__(self) -> str: ... + +class ProfileOp: + """One operator of a statement that ran, and what it really did.""" + + @property + def op(self) -> str: + """The kind of operator this is.""" + + @property + def detail(self) -> str: + """What it worked on, in the words the listing prints.""" + + @property + def pulls(self) -> int: + """How many times the operator above it asked for rows.""" + + @property + def rows(self) -> int: + """How many rows it answered with.""" + + @property + def flat(self) -> int: + """The same count with the vectors unpacked.""" + + @property + def estimate(self) -> float | None: + """What the optimizer thought it would answer.""" + + @property + def bound(self) -> float | None: + """The upper bound the optimizer had for it.""" + + @property + def nanos(self) -> int: + """How long it spent, in nanoseconds.""" + + @property + def qerror(self) -> float | None: + """The estimate over the truth, or the truth over the estimate, + whichever is the larger.""" + + def __repr__(self) -> str: ... + +class ProfileStage: + """One stage of a statement that ran.""" + + @property + def sink(self) -> str: + """What the stage feeds.""" + + @property + def rows(self) -> int: + """How many rows came out of it.""" + + @property + def nanos(self) -> int: + """How long it took, in nanoseconds.""" + + @property + def ops(self) -> list[ProfileOp]: + """Its operators, from the one that read to the one that fed the + sink, which is the order they ran in and the reverse of the order + the listing prints them. + """ + + def __repr__(self) -> str: ... + +class Profile: + """What a statement did, measured while it did it.""" + + @property + def stages(self) -> list[ProfileStage]: + """The stages, in the order they ran.""" + + @property + def nanos(self) -> int: + """Every stage added up, in nanoseconds.""" + + @property + def text(self) -> str: + """The engine's own listing, which is what `print` gives.""" + + def _repr_html_(self) -> str: + """The listing as a notebook shows it, preformatted for the reason + a plan's is.""" + + def __str__(self) -> str: ... + def __repr__(self) -> str: ... + class Result: """The rows a statement gave back.""" diff --git a/python/zudb/aio.py b/python/zudb/aio.py index c0f07d2..23d0dd8 100644 --- a/python/zudb/aio.py +++ b/python/zudb/aio.py @@ -49,10 +49,16 @@ from typing import Any, Generic, TypeVar from . import _zudb -from ._zudb import Appender, Connection, Result, Transaction +from ._zudb import Appender, Connection, Plan, Prepared, Profile, Result, Transaction from .types import Value -__all__ = ["connect", "AsyncConnection", "AsyncTransaction", "AsyncAppender"] +__all__ = [ + "connect", + "AsyncConnection", + "AsyncTransaction", + "AsyncAppender", + "AsyncPrepared", +] T = TypeVar("T") @@ -217,6 +223,35 @@ async def sql(self, statement: str, params: Mapping[str, Value] | None = None) - """The same call, named for the way it reads in a notebook.""" return await self._call(functools.partial(self._conn.sql, statement, params)) + def prepare(self, statement: str) -> _Opening[AsyncPrepared]: + """Compiles a statement now and hands back something that runs + it later, as often as you like. + + statement = "MATCH (p:person) WHERE p.name = $name RETURN p.id AS id" + async with conn.prepare(statement) as find: + rows = await find.execute({"name": "ada"}) + + Compiling reaches the engine, so it is awaited like everything + else here, and what it buys is what the sync client's docstring + says it buys: the compile at this line rather than at the first + request, and the names the statement wants. + """ + return _Opening(functools.partial(self._prepare, statement)) + + async def _prepare(self, statement: str) -> AsyncPrepared: + compiled = await self._call(functools.partial(self._conn.prepare, statement)) + return AsyncPrepared(self, compiled) + + async def explain(self, statement: str) -> Plan: + """What the statement would do, without doing it.""" + return await self._call(functools.partial(self._conn.explain, statement)) + + async def profile(self, statement: str, params: Mapping[str, Value] | None = None) -> Profile: + """Runs the statement with the counters on and answers what its + operators really did. + """ + return await self._call(functools.partial(self._conn.profile, statement, params)) + def transaction(self, *, read_only: bool = False) -> _Opening[AsyncTransaction]: """Starts a transaction and hands it back for an `async with` block. @@ -496,3 +531,69 @@ async def __aexit__(self, *_exception: Any) -> bool: def __repr__(self) -> str: return f"" + + +class AsyncPrepared: + """A statement the engine has compiled and is holding for you. + + The prepared statement underneath is `zudb.Prepared` and the rules + are its rules: running one that has been closed is refused, closing + twice does nothing, and a run that binds no value the statement + wants fails at the run rather than at the compile. + """ + + __slots__ = ("_conn", "_prepared") + + def __init__(self, conn: AsyncConnection, prepared: Prepared) -> None: + self._conn = conn + self._prepared = prepared + + @property + def statement(self) -> str: + """The text it was compiled from.""" + return self._prepared.statement + + @property + def params(self) -> list[str]: + """The names this statement wants bound, in the order it uses + them. + + A property rather than a coroutine, unlike most of what is here: + the names were read at the compile and are held beside the id, + so asking for them reaches nothing that could wait. + """ + return self._prepared.params + + @property + def closed(self) -> bool: + """Whether this prepared statement has been closed.""" + return self._prepared.closed + + async def execute(self, params: Mapping[str, Value] | None = None) -> Result: + """Runs it with these parameters and gives back its rows.""" + return await self._conn._call(functools.partial(self._prepared.execute, params)) + + async def sql(self, params: Mapping[str, Value] | None = None) -> Result: + """The same call, named for the way it reads in a notebook.""" + return await self._conn._call(functools.partial(self._prepared.sql, params)) + + async def close(self) -> None: + """Closes it and gives the statement back to the connection.""" + await self._conn._call(self._prepared.close) + + async def __aenter__(self) -> AsyncPrepared: + return self + + async def __aexit__(self, *_exception: Any) -> bool: + """Closes on the way out, whether the block ended well or badly. + + There is nothing to undo and nothing to write, so unlike a + transaction's exit this one has only the one thing it could + mean. + """ + await self.close() + return False + + def __repr__(self) -> str: + closed = ", closed" if self.closed else "" + return f"" diff --git a/src/conn.rs b/src/conn.rs index 8e52472..54e1600 100644 --- a/src/conn.rs +++ b/src/conn.rs @@ -22,6 +22,8 @@ use crate::columns; use crate::error::{closed, programming, to_py_err}; use crate::html; use crate::interrupt; +use crate::plan; +use crate::prepared::Prepared; use crate::register; use crate::txn::Transaction; use crate::value::{Names, from_py, to_py}; @@ -31,6 +33,17 @@ use crate::value::{Names, from_py, to_py}; /// and a capsule named anything else is not one of these. const STREAM: &CStr = c"arrow_array_stream"; +/// What a run is: the text of a statement, or the id the session pinned +/// a prepared one under. +/// +/// One enum rather than two paths, so that a prepared statement gets +/// the interrupt handling, the GIL release and the catalog names that +/// every other statement gets, by being the same call. +pub(crate) enum Source { + Text(String), + Prepared(u64), +} + /// One connection to one database. /// /// Statements run on it in order, one at a time. It reads the database @@ -104,6 +117,76 @@ impl Connection { self.execute(py, statement, params) } + /// Compiles a statement now and hands back something that runs it + /// later, as often as you like, with different values bound each + /// time. + /// + /// ```python + /// with conn.prepare("MATCH (p:person) WHERE p.name = $name RETURN p.id AS id") as find: + /// rows = find.execute({"name": "ada"}) + /// ``` + /// + /// It is not the speedup the word usually promises, and the class + /// says why at length: there is no round trip to save here, and the + /// engine already caches a plan by the text of the statement. What + /// it buys is the compile happening at this line, at startup, and + /// `params` coming back. + fn prepare(slf: Py, py: Python<'_>, statement: &str) -> PyResult { + Prepared::compile(py, slf, statement) + } + + /// What the statement would do, without doing it. + /// + /// ```python + /// print(conn.explain("MATCH (p:person) RETURN p.name AS name")) + /// ``` + /// + /// The plan comes back as the engine's own listing, which is what + /// `print` gives, and as a tree of operators, which is what + /// `plan.root` walks. It takes no parameters: a plan is chosen from + /// the shape of the statement rather than from the values bound to + /// it, and one asked for with values would suggest otherwise. + fn explain(&self, py: Python<'_>, statement: &str) -> PyResult { + let out = self + .engine(py, |conn| conn.explain_plan(statement))? + .map_err(|err| to_py_err(py, err))?; + plan::planned(py, &out) + } + + /// Runs the statement with the counters on and answers what its + /// operators really did. + /// + /// ```python + /// run = conn.profile("MATCH (p:person)-[:knows]->(q:person) RETURN q.name AS name") + /// print(run) + /// ``` + /// + /// Beside every operator's rows is what the optimizer thought it + /// would produce, so the operator that surprised it is the one to + /// look at. A statement that writes is refused rather than + /// profiled, because a measurement that also inserted two rows + /// changed the thing it was measuring. + #[pyo3(signature = (statement, params = None))] + fn profile( + &self, + py: Python<'_>, + statement: &str, + params: Option<&Bound<'_, PyDict>>, + ) -> PyResult { + let params = bind(params)?; + let statement = statement.to_string(); + let out = interrupt::watched(py, &self.runner, &self.inner, &self.stop, move |conn| { + let borrowed: Vec<(&str, Value)> = params + .iter() + .map(|(name, value)| (name.as_str(), value.clone())) + .collect(); + conn.profile(&statement, &borrowed) + }) + .map_err(|stopped| stopped.raise(py))? + .map_err(|err| to_py_err(py, err))?; + plan::profiled(py, &out) + } + /// Starts a transaction and hands it back for a `with` block. /// /// Several statements as one unit of work: the block commits when @@ -358,7 +441,41 @@ impl Connection { // statement is not this one and the borrow would have to say // so. It is one allocation against a statement, which is // nothing beside the parse it is about to have. - let statement = statement.to_string(); + self.run_source(py, Source::Text(statement.to_string()), params) + } + + /// Runs a statement the session has already compiled, under the id + /// it pinned it with. + /// + /// The same path as every other statement, so a prepared one + /// releases the GIL, queues for the connection's lock and feels a + /// `Ctrl-C` exactly as one written out does. + pub(crate) fn query_prepared( + &self, + py: Python<'_>, + id: u64, + params: Vec<(String, Value)>, + ) -> PyResult { + self.run_source(py, Source::Prepared(id), params) + } + + /// Gives a prepared statement's id back to the session. + /// + /// Every reason the connection could not be had is ignored, since + /// they all mean the same thing here: a session that is gone is not + /// holding the statement either. + pub(crate) fn release(&self, py: Python<'_>, id: u64) { + let _ = self.engine(py, |conn| -> std::result::Result { + Ok(conn.close_prepared(id)) + }); + } + + fn run_source( + &self, + py: Python<'_>, + source: Source, + params: Vec<(String, Value)>, + ) -> PyResult { // The GIL goes down for the whole statement, waiting for the // connection's own lock included. That is the point of a // compiled engine in a Python process: another thread runs @@ -373,8 +490,11 @@ impl Connection { .map(|(name, value)| (name.as_str(), value.clone())) .collect(); let names = Names::of(conn.session_mut().catalog()); - conn.query_with(&statement, &borrowed) - .map(|result| (result, names)) + let out = match &source { + Source::Text(statement) => conn.query_with(statement, &borrowed), + Source::Prepared(id) => conn.execute_prepared(*id, &borrowed), + }; + out.map(|result| (result, names)) }) .map_err(|stopped| stopped.raise(py))? .map_err(|err| to_py_err(py, err))?; @@ -649,7 +769,7 @@ fn needed<'py>(py: Python<'py>, module: &str, extra: &str) -> PyResult>) -> PyResult> { +pub(crate) fn bind(params: Option<&Bound<'_, PyDict>>) -> PyResult> { let Some(params) = params else { return Ok(Vec::new()); }; diff --git a/src/lib.rs b/src/lib.rs index 60eeb47..6c5bd75 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -21,6 +21,8 @@ mod frame; mod html; mod interrupt; mod load; +mod plan; +mod prepared; mod register; mod txn; mod value; @@ -59,6 +61,13 @@ fn _zudb(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_class::()?; module.add_class::()?; module.add_class::()?; + module.add_class::()?; + module.add_class::()?; + module.add_class::()?; + module.add_class::()?; + module.add_class::()?; + module.add_class::()?; + module.add_class::()?; module.add_class::()?; module.add_class::()?; module.add_class::()?; diff --git a/src/plan.rs b/src/plan.rs new file mode 100644 index 0000000..6570f6a --- /dev/null +++ b/src/plan.rs @@ -0,0 +1,373 @@ +//! What a statement would do, and what it did. +//! +//! ```python +//! plan = conn.explain("MATCH (p:person) WHERE p.name = $name RETURN p.id AS id") +//! print(plan) +//! # Project p.id AS id +//! # Filter p.name = $name +//! # ScanNodes p: person +//! ``` +//! +//! Both calls answer twice over, and that is on purpose. `print(plan)` +//! gives the engine's own listing, which is what a person reads and is +//! rendered by the engine rather than here so the two cannot drift +//! apart from one release to the next. `plan.root` gives the same plan +//! as objects, which is what a program walks: a test that wants to know +//! a scan became a seek asks the tree rather than matching on a string +//! that was written to be read. +//! +//! `explain` takes no parameters, which is not an oversight. A plan is +//! chosen from the shape of the statement and the values are bound when +//! it runs, so a plan asked for with values would suggest the values +//! had changed it. `profile` does take them, because it is a run. +//! +//! Every count here is an `int`, which in Python is exact however large +//! it gets, so nothing is lost and nothing has to be converted before +//! it is added up. + +use pyo3::prelude::*; +use zudb::{OpProfile, Profile as EngineProfile, QueryPlan, StageProfile}; + +use crate::html; + +/// One operator of a plan. +/// +/// `op` is what the operator is and `name` is what the listing calls +/// it, and they differ where an operator sits inside a bracket: an +/// expand inside an `OPTIONAL MATCH` has `op` `Expand` and `name` +/// `OptionalExpand`. Matching on `op` is what a program should do, +/// since it is the one of the two that does not change with the +/// company an operator keeps. +#[pyclass(module = "zudb", frozen)] +pub struct PlanNode { + /// The kind of operator this is. + #[pyo3(get)] + op: String, + /// What the listing calls it, bracket and all. + #[pyo3(get)] + name: String, + /// The bracket it sits inside, if it sits inside one: `"Optional"`, + /// `"Semi"`, `"Anti"` or `"Mark"`. + #[pyo3(get)] + bracket: Option, + /// What it works on, in the words the listing prints. + #[pyo3(get)] + detail: String, + /// The variables it introduces. + #[pyo3(get)] + binds: Vec, + /// The tables it reads. + #[pyo3(get)] + tables: Vec, + /// What it pulls from, in the order the listing prints them. + #[pyo3(get)] + children: Vec>, +} + +#[pymethods] +impl PlanNode { + fn __repr__(&self) -> String { + format!("", self.name, self.detail) + } +} + +/// A query written where a value belongs, and the plan it gets. +/// +/// `reads` is which variables of the query around it this one reads, +/// and it is the whole of the difference between a subquery that runs +/// once and one that runs once a row: a plan that reads nothing is a +/// plan the executor can run once and keep. +#[pyclass(module = "zudb", frozen)] +pub struct ScalarPlan { + /// The variables it reads from the query it is written inside. + #[pyo3(get)] + reads: Vec, + /// Whether it is asking whether there is a row rather than for one. + #[pyo3(get)] + exists: bool, + /// The plan itself. + #[pyo3(get)] + plan: Py, +} + +#[pymethods] +impl ScalarPlan { + fn __repr__(&self) -> String { + let reads = if self.reads.is_empty() { + "once".to_string() + } else { + format!("reads {}", self.reads.join(", ")) + }; + format!("") + } +} + +/// What a statement would do, without doing it. +/// +/// Take one with `Connection.explain`. `print` it for the listing and +/// walk `root` for the tree, which are the same plan said twice. +#[pyclass(module = "zudb", frozen)] +pub struct Plan { + /// The operator everything else feeds, or `None` for a statement + /// that compiled to no operator at all. + #[pyo3(get)] + root: Option>, + /// The columns the statement projects, in order. + #[pyo3(get)] + columns: Vec, + /// The parameters it wants bound. + #[pyo3(get)] + params: Vec, + /// What the planner has to say about it, if anything. + #[pyo3(get)] + notes: Vec, + /// The plans of the queries written where values belong. + #[pyo3(get)] + scalars: Vec>, + /// The engine's own listing, which is what `print` gives. + #[pyo3(get)] + text: String, +} + +#[pymethods] +impl Plan { + fn __str__(&self) -> String { + self.text.clone() + } + + fn __repr__(&self) -> String { + format!("", self.text.lines().next().unwrap_or("")) + } + + /// The listing as a notebook shows it, which is the listing. + /// + /// Preformatted rather than a table, because the indentation is + /// what says which operator pulls from which and a browser that + /// collapsed the spaces would take the plan's shape away. + fn _repr_html_(&self) -> String { + html::wrap(&format!( + "
{}
", + html::escape(&self.text) + )) + } +} + +/// One operator of a statement that ran, and what it really did. +/// +/// `estimate` is what the optimizer thought the operator would produce +/// and `rows` is what it did, so `qerror` is the larger over the +/// smaller: one where the estimate was right, ten where it was out by +/// an order of magnitude either way. All three are `None` on an +/// operator the optimizer had nothing to say about. +#[pyclass(module = "zudb", frozen)] +pub struct ProfileOp { + /// The kind of operator this is. + #[pyo3(get)] + op: String, + /// What it worked on, in the words the listing prints. + #[pyo3(get)] + detail: String, + /// How many times the operator above it asked for rows. + #[pyo3(get)] + pulls: u64, + /// How many rows it answered with. + #[pyo3(get)] + rows: u64, + /// The same count with the vectors unpacked. + #[pyo3(get)] + flat: u64, + /// What the optimizer thought it would answer. + #[pyo3(get)] + estimate: Option, + /// The upper bound the optimizer had for it. + #[pyo3(get)] + bound: Option, + /// How long it spent, in nanoseconds. + #[pyo3(get)] + nanos: u64, + /// The estimate over the truth, or the truth over the estimate, + /// whichever is the larger. + #[pyo3(get)] + qerror: Option, +} + +#[pymethods] +impl ProfileOp { + fn __repr__(&self) -> String { + format!("", self.op, self.rows) + } +} + +/// One stage of a statement that ran. +/// +/// A stage is the run of operators between two points the executor has +/// to gather at, and `sink` is what it gathers into. +#[pyclass(module = "zudb", frozen)] +pub struct ProfileStage { + /// What the stage feeds. + #[pyo3(get)] + sink: String, + /// How many rows came out of it. + #[pyo3(get)] + rows: u64, + /// How long it took, in nanoseconds. + #[pyo3(get)] + nanos: u64, + /// Its operators, from the one that read to the one that fed the + /// sink, which is the order they ran in and the reverse of the + /// order the listing prints them. + #[pyo3(get)] + ops: Vec>, +} + +#[pymethods] +impl ProfileStage { + fn __repr__(&self) -> String { + format!("", self.sink, self.rows) + } +} + +/// What a statement did, measured while it did it. +/// +/// Take one with `Connection.profile`. It runs the statement, so a +/// statement that writes is refused rather than profiled: a measurement +/// that also inserted two rows changed the thing it was measuring. +#[pyclass(module = "zudb", frozen)] +pub struct Profile { + /// The stages, in the order they ran. + #[pyo3(get)] + stages: Vec>, + /// Every stage added up, in nanoseconds. + #[pyo3(get)] + nanos: u64, + /// The engine's own listing, which is what `print` gives. + #[pyo3(get)] + text: String, +} + +#[pymethods] +impl Profile { + fn __str__(&self) -> String { + self.text.clone() + } + + fn __repr__(&self) -> String { + format!( + "", + self.stages.len(), + spelled(self.nanos) + ) + } + + /// The listing as a notebook shows it, preformatted for the reason + /// a plan's is. + fn _repr_html_(&self) -> String { + html::wrap(&format!( + "
{}
", + html::escape(&self.text) + )) + } +} + +/// A count of nanoseconds in the unit a person would have said it in. +/// +/// A `repr` reporting `0 ms` for a statement that took a quarter of a +/// millisecond is a `repr` that reads as though nothing was measured. +fn spelled(nanos: u64) -> String { + if nanos < 1_000 { + format!("{nanos} ns") + } else if nanos < 1_000_000 { + format!("{:.1} us", nanos as f64 / 1e3) + } else if nanos < 1_000_000_000 { + format!("{:.1} ms", nanos as f64 / 1e6) + } else { + format!("{:.2} s", nanos as f64 / 1e9) + } +} + +/// A whole plan as the objects a caller reads. +pub fn planned(py: Python<'_>, plan: &QueryPlan) -> PyResult { + let root = match &plan.root { + Some(node) => Some(Py::new(py, operator(py, node)?)?), + None => None, + }; + let mut scalars = Vec::with_capacity(plan.scalars.len()); + for scalar in &plan.scalars { + scalars.push(Py::new( + py, + ScalarPlan { + reads: scalar.reads.clone(), + exists: scalar.exists, + plan: Py::new(py, planned(py, &scalar.plan)?)?, + }, + )?); + } + Ok(Plan { + root, + columns: plan.columns.clone(), + params: plan.params.clone(), + notes: plan.notes.clone(), + scalars, + text: plan.render(), + }) +} + +/// One operator and everything it pulls from. +fn operator(py: Python<'_>, node: &zudb::PlanNode) -> PyResult { + let mut children = Vec::with_capacity(node.children.len()); + for child in &node.children { + children.push(Py::new(py, operator(py, child)?)?); + } + Ok(PlanNode { + op: node.op.to_string(), + name: node.name(), + bracket: node + .bracket + .as_ref() + .map(|bracket| bracket.prefix().to_string()), + detail: node.detail.clone(), + binds: node.binds.clone(), + tables: node.tables.clone(), + children, + }) +} + +/// A whole profile as the objects a caller reads. +pub fn profiled(py: Python<'_>, profile: &EngineProfile) -> PyResult { + let mut stages = Vec::with_capacity(profile.stages.len()); + for stage in &profile.stages { + stages.push(Py::new(py, staged(py, stage)?)?); + } + Ok(Profile { + stages, + nanos: profile.stages.iter().map(|stage| stage.nanos).sum(), + text: profile.render(), + }) +} + +fn staged(py: Python<'_>, stage: &StageProfile) -> PyResult { + let mut ops = Vec::with_capacity(stage.ops.len()); + for op in &stage.ops { + ops.push(Py::new(py, counted(op))?); + } + Ok(ProfileStage { + sink: stage.sink.clone(), + rows: stage.out_rows, + nanos: stage.nanos, + ops, + }) +} + +fn counted(op: &OpProfile) -> ProfileOp { + ProfileOp { + op: op.kind.to_string(), + detail: op.detail.clone(), + pulls: op.pulls, + rows: op.rows, + flat: op.flat, + estimate: op.est, + bound: op.bnd, + nanos: op.nanos, + qerror: op.qerror(), + } +} diff --git a/src/prepared.rs b/src/prepared.rs new file mode 100644 index 0000000..416842e --- /dev/null +++ b/src/prepared.rs @@ -0,0 +1,158 @@ +//! A statement compiled once and run many times. +//! +//! ```python +//! with conn.prepare("MATCH (p:person) WHERE p.name = $name RETURN p.id AS id") as find: +//! for name in names: +//! print(find.execute({"name": name}).fetchall()) +//! ``` +//! +//! What this buys is not what a driver's `prepare` buys, and it is +//! worth saying so here rather than letting a reader assume it. A +//! driver prepares to save a round trip to a server, and there is no +//! server and no round trip: the engine is in this process. It caches a +//! plan by the text of the statement, so the second `conn.execute` of +//! the same string is not compiled a second time either, and a loop +//! that prepares and a loop that repeats the same string run at the +//! same speed. +//! +//! Two things it does buy. The compile happens at the line that asked +//! for it, so a program that prepares its statements at startup finds a +//! statement that does not compile there, rather than on the first +//! request that needed it. And `params` comes back, which is what the +//! statement wants bound, so a layer binding from a record knows what +//! to look for without reading the text. +//! +//! The one thing that is a speedup is the case this is written beside: +//! a statement whose text is different every run, which is what a +//! program that formats its values into the string is writing. That one +//! pays the compile every time and no cache can help it, and binding +//! parameters is what fixes it, prepared or not. + +use std::sync::atomic::{AtomicBool, Ordering}; + +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple}; + +use crate::conn::{Connection, Result, bind}; +use crate::error::{programming, to_py_err}; + +/// A statement the engine has compiled and is holding for you. +/// +/// Take one with `Connection.prepare`, run it with `execute`, and close +/// it, which gives the statement back to the connection. A `with` block +/// closes it at the end, and running one that has been closed is +/// refused rather than quietly recompiled. +#[pyclass(module = "zudb")] +pub struct Prepared { + /// The connection it was compiled on, held rather than borrowed: a + /// prepared statement whose connection was collected would be an id + /// for a session that is gone. + conn: Py, + /// The text it was compiled from, kept so that a caller holding one + /// in a dictionary can see which it is. + #[pyo3(get)] + statement: String, + params: Vec, + /// What the session pinned it under. + id: u64, + open: AtomicBool, +} + +#[pymethods] +impl Prepared { + /// The names this statement wants bound, in the order it uses them. + #[getter] + fn params(&self) -> Vec { + self.params.clone() + } + + /// Whether this prepared statement has been closed. + #[getter] + fn closed(&self) -> bool { + !self.open.load(Ordering::Acquire) + } + + /// Runs it with these parameters and gives back its rows. + /// + /// The same answer `Connection.execute` gives, because it is the + /// same statement: a `zudb.Result`, which is rows in memory and + /// knows how to become an Arrow table or a DataFrame. A name the + /// statement wants and the caller did not bind is an error from the + /// engine at this call, not at the prepare, since a missing value + /// is nothing to do with the statement. + #[pyo3(signature = (params = None))] + fn execute( + &self, + py: Python<'_>, + params: Option<&Bound<'_, PyDict>>, + ) -> PyResult { + self.alive(py)?; + let params = bind(params)?; + self.conn.borrow(py).query_prepared(py, self.id, params) + } + + /// The same call, named for the way it reads in a notebook. + #[pyo3(signature = (params = None))] + fn sql(&self, py: Python<'_>, params: Option<&Bound<'_, PyDict>>) -> PyResult { + self.execute(py, params) + } + + /// Closes it and gives the statement back to the connection. + /// + /// Doing it twice does nothing, which is what a `with` block around + /// a caller who closed it themselves needs. A prepared statement + /// whose connection has already closed is closed too, since the + /// session that was holding the id went with it, and closing that + /// one is not an error either. + fn close(&self, py: Python<'_>) { + if self.open.swap(false, Ordering::AcqRel) { + self.conn.borrow(py).release(py, self.id); + } + } + + fn __enter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + #[pyo3(signature = (*_exception))] + fn __exit__(&self, py: Python<'_>, _exception: &Bound<'_, PyTuple>) -> bool { + self.close(py); + // False, so an exception raised inside the block carries on out + // of it. + false + } + + fn __repr__(&self) -> String { + let closed = if self.closed() { ", closed" } else { "" }; + format!("", self.statement) + } +} + +impl Prepared { + /// Compiles one, which is what makes one. + pub fn compile(py: Python<'_>, conn: Py, statement: &str) -> PyResult { + let (id, params) = conn + .borrow(py) + .engine(py, |engine| engine.prepare(statement))? + .map_err(|err| to_py_err(py, err))?; + Ok(Prepared { + conn, + statement: statement.to_string(), + params, + id, + open: AtomicBool::new(true), + }) + } + + /// Refuses a run on one that has been closed. + fn alive(&self, py: Python<'_>) -> PyResult<()> { + if self.closed() { + return Err(programming( + py, + "this prepared statement is closed, and a closed one has given its \ + statement back to the connection", + )); + } + Ok(()) + } +} diff --git a/tests/test_plan.py b/tests/test_plan.py new file mode 100644 index 0000000..ec51a21 --- /dev/null +++ b/tests/test_plan.py @@ -0,0 +1,275 @@ +"""What a statement would do, and what it did. + +A plan belongs to the engine and this client only carries it, so what +these assert is that the carrying is faithful: every operator, in the +shape the tree had, with the fields that mean something and `None` +where the engine had nothing to say. The listing is asserted against the +tree it was rendered from rather than against a string written out here, +because a test that pinned the exact words would fail every time the +optimizer learned to print one better. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import zudb + +BY_NAME = "MATCH (p:person) WHERE p.name = $name RETURN p.uid AS uid" + + +def operators(node: zudb.PlanNode | None) -> list[zudb.PlanNode]: + """Every operator of a plan, depth first, which is the order the + listing prints them in. + """ + if node is None: + return [] + return [node, *[found for child in node.children for found in operators(child)]] + + +def test_a_plan_is_the_tree_of_operators_the_statement_would_run( + social: zudb.Connection, +) -> None: + plan = social.explain(BY_NAME) + + assert isinstance(plan, zudb.Plan) + assert [op.op for op in operators(plan.root)] == ["Project", "Filter", "ScanNodes"] + assert plan.columns == ["uid"] + assert plan.params == ["name"] + assert plan.notes == [] + assert plan.scalars == [] + + +def test_an_operator_carries_what_it_works_on_binds_and_touches(social: zudb.Connection) -> None: + project, filtered, scan = operators(social.explain(BY_NAME).root) + + assert project.detail == "p.uid AS uid" + assert project.binds == ["uid"] + assert project.tables == [] + + assert filtered.detail == "p.name = $name" + assert filtered.binds == [] + + assert scan.detail == "p: person" + assert scan.binds == ["p"] + assert scan.tables == ["person"] + assert scan.children == [] + + +def test_an_operator_inside_a_bracket_is_named_for_it_and_is_not_it( + loaded: zudb.Connection, +) -> None: + plan = loaded.explain( + "MATCH (a:person) OPTIONAL MATCH (a)-[:knows]->(b:person) RETURN a.name AS a, b.name AS b" + ) + expand = next(op for op in operators(plan.root) if op.op == "Expand") + + assert expand.name == "OptionalExpand" + assert expand.bracket == "Optional" + assert expand.tables == ["knows"] + + +def test_an_operator_outside_a_bracket_has_none_and_is_named_for_itself( + loaded: zudb.Connection, +) -> None: + plan = loaded.explain("MATCH (a:person)-[:knows]->(b:person) RETURN a.name AS a") + expand = next(op for op in operators(plan.root) if op.op == "Expand") + + assert expand.name == "Expand" + assert expand.bracket is None + + +def test_printing_a_plan_gives_the_listing_and_the_listing_is_the_tree( + social: zudb.Connection, +) -> None: + plan = social.explain(BY_NAME) + + assert str(plan) == plan.text + assert plan.text == "Project p.uid AS uid\n Filter p.name = $name\n ScanNodes p: person\n" + # Written twice on purpose: the listing is what a person reads and + # the tree is what a program walks, and this is the one assertion + # that says the two describe the same plan. + printed = [line.strip().split(" ")[0] for line in plan.text.rstrip("\n").split("\n")] + assert printed == [op.name for op in operators(plan.root)] + + +def test_a_query_written_where_a_value_belongs_is_a_plan_of_its_own( + social: zudb.Connection, +) -> None: + plan = social.explain( + "MATCH (p:person) RETURN VALUE " + "{ MATCH (q:person) WHERE q.name = p.name RETURN q.uid LIMIT 1 } AS v" + ) + + assert len(plan.scalars) == 1 + scalar = plan.scalars[0] + # It reads a name from the query around it, which is the whole test + # for whether it runs once or once a row. + assert scalar.reads == ["p"] + assert scalar.exists is False + assert scalar.plan.root.op == "Limit" + assert "ScanNodes q: person" in scalar.plan.text + + +def test_a_subquery_that_reads_nothing_runs_once_and_says_so(social: zudb.Connection) -> None: + plan = social.explain( + "MATCH (p:person) RETURN VALUE { MATCH (q:person) RETURN q.uid LIMIT 1 } AS v" + ) + + assert plan.scalars[0].reads == [] + assert "(once)" in plan.text + + +def test_explaining_does_not_run_the_statement(social: zudb.Connection) -> None: + social.explain("INSERT (p:person {uid: 40, name: 'hedy', score: 1.0})") + + assert social.execute("MATCH (p:person) RETURN count(*) AS n").fetchall() == [(3,)] + + +def test_a_statement_that_does_not_compile_fails_at_the_explain(social: zudb.Connection) -> None: + with pytest.raises(zudb.SyntaxError): + social.explain("MATCH (") + + +def test_explaining_on_a_closed_connection_is_refused(tmp_path: Path) -> None: + conn = zudb.connect(tmp_path / "shut.zu1") + conn.close() + + with pytest.raises(zudb.ProgrammingError, match="closed"): + conn.explain("MATCH (p:person) RETURN p.name AS name") + + +def test_a_plan_says_what_it_is(social: zudb.Connection) -> None: + plan = social.explain(BY_NAME) + + assert repr(plan) == "" + assert repr(plan.root) == "" + + +def test_a_notebook_gets_the_listing_preformatted(social: zudb.Connection) -> None: + plan = social.explain(BY_NAME) + markup = plan._repr_html_() + + assert "` is not markup. + assert "$name" in markup + + +def test_a_profile_is_what_the_operators_really_did(social: zudb.Connection) -> None: + run = social.profile("MATCH (p:person) RETURN p.name AS name") + + assert isinstance(run, zudb.Profile) + assert len(run.stages) == 1 + stage = run.stages[0] + assert stage.sink == "Project" + assert stage.rows == 3 + assert stage.nanos > 0 + assert [op.op for op in stage.ops] == ["Source", "Scan"] + + scan = next(op for op in stage.ops if op.op == "Scan") + assert scan.detail == "p: person" + assert scan.pulls == 1 + assert scan.rows == 3 + assert scan.flat == 3 + assert scan.estimate == 3 + # The optimizer was right about a table it has the statistics for, + # which is what a q-error of one means. + assert scan.qerror == 1 + assert scan.nanos > 0 + + +def test_an_operator_the_optimizer_had_nothing_to_say_about_carries_none( + social: zudb.Connection, +) -> None: + run = social.profile("MATCH (p:person) RETURN p.name AS name") + source = next(op for op in run.stages[0].ops if op.op == "Source") + + assert source.estimate is None + assert source.bound is None + assert source.qerror is None + + +def test_the_profile_totals_its_stages_and_prints_them(social: zudb.Connection) -> None: + run = social.profile("MATCH (p:person) RETURN p.name AS name") + + assert run.nanos == sum(stage.nanos for stage in run.stages) + assert str(run) == run.text + assert run.text.startswith("stage 1: Project") + assert "Scan p: person" in run.text + + +def test_the_operators_are_in_the_order_they_ran(social: zudb.Connection) -> None: + run = social.profile("MATCH (p:person) WHERE p.score > 40.0 RETURN p.name AS name") + stage = run.stages[0] + + assert [op.op for op in stage.ops] == ["Source", "Scan", "Filter"] + # The listing reads the other way, top down from the sink, and this + # is the assertion that says the reversal is the only difference. + lines = run.text.rstrip("\n").split("\n") + printed = [line.strip().split(" ")[0] for line in lines if line.startswith(" ")] + assert printed == [op.op for op in reversed(stage.ops)] + + +def test_the_counts_are_whole_numbers_python_can_hold(social: zudb.Connection) -> None: + run = social.profile("MATCH (p:person) RETURN p.name AS name") + + for stage in run.stages: + assert isinstance(stage.rows, int) + assert isinstance(stage.nanos, int) + for op in stage.ops: + assert isinstance(op.pulls, int) + assert isinstance(op.rows, int) + assert isinstance(op.flat, int) + assert isinstance(op.nanos, int) + + +def test_a_profile_binds_its_parameters(social: zudb.Connection) -> None: + run = social.profile(BY_NAME, {"name": "ada"}) + filtered = next(op for op in run.stages[0].ops if op.op == "Filter") + + assert filtered.detail == "p.name = $name" + assert run.stages[0].rows == 1 + + +def test_a_profile_that_binds_nothing_fails_the_way_the_statement_would( + social: zudb.Connection, +) -> None: + with pytest.raises(zudb.SyntaxError, match=r"\$name"): + social.profile(BY_NAME) + + +def test_an_expand_is_its_own_operator_with_the_rows_it_walked(loaded: zudb.Connection) -> None: + run = loaded.profile("MATCH (a:person)-[:knows]->(b:person) RETURN b.name AS name") + expand = next(op for op in run.stages[0].ops if op.op == "Expand") + + assert "knows" in expand.detail + assert run.stages[0].rows == 2 + + +def test_a_statement_that_writes_is_refused_rather_than_profiled(social: zudb.Connection) -> None: + with pytest.raises(zudb.Error, match="profiling a statement that writes"): + social.profile("INSERT (p:person {uid: 40, name: 'hedy', score: 1.0})") + + assert social.execute("MATCH (p:person) RETURN count(*) AS n").fetchall() == [(3,)] + + +def test_profiling_on_a_closed_connection_is_refused(tmp_path: Path) -> None: + conn = zudb.connect(tmp_path / "shut.zu1") + conn.close() + + with pytest.raises(zudb.ProgrammingError, match="closed"): + conn.profile("MATCH (p:person) RETURN p.name AS name") + + +def test_a_profile_says_how_long_it_took_in_words(social: zudb.Connection) -> None: + run = social.profile("MATCH (p:person) RETURN p.name AS name") + + # A repr that said `0 ms` for a statement measured in microseconds + # would read as though nothing had been measured at all. + assert repr(run).startswith("" in repr(run) or " us>" in repr(run) or " ns>" in repr(run) + assert repr(run.stages[0]) == "" + assert repr(run.stages[0].ops[0]).startswith(" None: + with social.prepare(BY_NAME) as find: + assert isinstance(find, zudb.Prepared) + assert find.statement == BY_NAME + assert find.params == ["name"] + assert find.closed is False + + +def test_a_statement_that_takes_no_parameters_reports_none(social: zudb.Connection) -> None: + with social.prepare("MATCH (p:person) RETURN p.name AS name") as everyone: + assert everyone.params == [] + + +def test_it_runs_as_often_as_it_is_asked_with_different_bindings(social: zudb.Connection) -> None: + with social.prepare(BY_NAME) as find: + assert find.execute({"name": "ada"}).fetchall() == [(10,)] + assert find.execute({"name": "grace"}).fetchall() == [(20,)] + assert find.execute({"name": "nobody"}).fetchall() == [] + assert find.execute({"name": "ada"}).fetchall() == [(10,)] + + +def test_what_comes_back_is_the_result_a_statement_gives(social: zudb.Connection) -> None: + with social.prepare(BY_NAME) as find: + rows = find.execute({"name": "ada"}) + assert isinstance(rows, zudb.Result) + assert rows.columns == ["uid"] + assert rows.notices == [] + assert len(rows) == 1 + + +def test_sql_is_the_same_call_under_the_notebook_name(social: zudb.Connection) -> None: + with social.prepare(BY_NAME) as find: + assert find.sql({"name": "grace"}).fetchall() == [(20,)] + + +def test_a_prepared_write_writes(social: zudb.Connection) -> None: + with social.prepare("INSERT (p:person {uid: $uid, name: $name, score: $score})") as add: + add.execute({"uid": 40, "name": "hedy", "score": 51.0}) + add.execute({"uid": 50, "name": "edith", "score": 33.5}) + + rows = social.execute("MATCH (p:person) RETURN count(*) AS n") + assert rows.fetchall() == [(5,)] + + +def test_a_statement_that_does_not_compile_fails_at_the_prepare(social: zudb.Connection) -> None: + with pytest.raises(zudb.SyntaxError): + social.prepare("MATCH (") + + +def test_a_name_the_caller_did_not_bind_fails_at_the_run(social: zudb.Connection) -> None: + with social.prepare(BY_NAME) as find: + with pytest.raises(zudb.SyntaxError, match=r"\$name"): + find.execute() + # And the statement is still there to be run properly, since a + # missing binding is nothing to do with the statement. + assert find.execute({"name": "ada"}).fetchall() == [(10,)] + + +def test_closing_it_twice_does_nothing(social: zudb.Connection) -> None: + find = social.prepare(BY_NAME) + find.close() + assert find.closed is True + find.close() + assert find.closed is True + + +def test_a_closed_prepared_statement_refuses_to_run(social: zudb.Connection) -> None: + find = social.prepare(BY_NAME) + find.close() + + with pytest.raises(zudb.ProgrammingError, match="closed"): + find.execute({"name": "ada"}) + with pytest.raises(zudb.ProgrammingError, match="closed"): + find.sql({"name": "ada"}) + + +def test_the_block_closes_it_at_the_end(social: zudb.Connection) -> None: + with social.prepare(BY_NAME) as find: + assert find.closed is False + assert find.closed is True + + +def test_an_exception_inside_the_block_closes_it_and_carries_on(social: zudb.Connection) -> None: + find = social.prepare(BY_NAME) + with pytest.raises(ZeroDivisionError), find: + raise ZeroDivisionError + assert find.closed is True + + +def test_one_whose_connection_closed_says_the_connection_is_closed(tmp_path: Path) -> None: + conn = zudb.connect(tmp_path / "gone.zu1") + conn.execute("INSERT (p:person {uid: 1, name: 'ada'})") + find = conn.prepare(BY_NAME) + conn.close() + + with pytest.raises(zudb.ProgrammingError, match="closed"): + find.execute({"name": "ada"}) + # Closing it is still fine, and still does nothing: the session that + # was holding the id went when the connection did. + find.close() + assert find.closed is True + + +def test_a_read_only_connection_prepares_and_runs_a_read(loaded: zudb.Connection) -> None: + with loaded.prepare("MATCH (p:person) RETURN p.name AS name") as everyone: + assert len(everyone.execute()) == 3 + + +def test_a_connection_prepares_as_many_statements_as_it_likes(social: zudb.Connection) -> None: + first = social.prepare(BY_NAME) + second = social.prepare("MATCH (p:person) RETURN count(*) AS n") + third = social.prepare("MATCH (p:person) RETURN p.name AS name") + + assert first.execute({"name": "grace"}).fetchall() == [(20,)] + assert second.execute().fetchall() == [(3,)] + assert len(third.execute()) == 3 + + for statement in (first, second, third): + statement.close() + + +def test_preparing_on_a_closed_connection_is_refused(tmp_path: Path) -> None: + conn = zudb.connect(tmp_path / "shut.zu1") + conn.close() + + with pytest.raises(zudb.ProgrammingError, match="closed"): + conn.prepare("MATCH (p:person) RETURN p.name AS name") + + +def test_a_statement_that_is_not_a_string_is_refused(social: zudb.Connection) -> None: + with pytest.raises(TypeError): + social.prepare(42) + + +def test_it_says_what_it_is_and_whether_it_is_closed(social: zudb.Connection) -> None: + find = social.prepare(BY_NAME) + assert repr(find) == f'' + find.close() + assert repr(find) == f'' + + +def test_the_class_is_the_one_the_package_exports(social: zudb.Connection) -> None: + with social.prepare(BY_NAME) as find: + assert type(find) is zudb.Prepared + assert type(find).__module__ == "zudb"