From 2b7aba88407da2a87f2547b794bcd57e329a39eb Mon Sep 17 00:00:00 2001 From: Tam Nguyen Duc <1218621+tamnd@users.noreply.github.com> Date: Tue, 18 Aug 2026 22:39:37 +0700 Subject: [PATCH] the same calls, awaited zudb.aio gives every connection a thread of its own and hands it every call that could wait. A statement is inside Rust with the GIL down, so one run straight from a coroutine stops the loop for as long as it takes, including the tasks answering requests that have nothing to do with the database. A thread per connection rather than a pool shared between them, because a connection is one lock and statements on it queue anyway. A pool would add no parallelism the engine can use and would let two statements written one after the other run in the other order. Two connections do run at the same time, and two statements together cost 1.09 times what one costs alone. What comes back is a zudb.Result, which is rows already in memory, so reading them is not awaited and never was. Only the ways in are, and only the ones that can wait: path, read_only, closed, rows_read and interrupt stay properties, answered from beside the lock, so a progress bar still reads while the statement it measures runs. Cancelling the task that awaits a statement interrupts the statement and waits for it to stop, so the connection is idle by the time the CancelledError arrives rather than busy with work nobody wants. A statement still queued is dropped without running, and a transaction block cancelled partway leaves through the rollback. The module is a submodule to ask for by name, not one the package imports, so a script that never awaits anything pays for neither asyncio nor the thread pool, and there is a test that says so. --- README.md | 31 ++- python/zudb/__init__.py | 4 + python/zudb/aio.py | 498 ++++++++++++++++++++++++++++++++++++++++ tests/test_aio.py | 413 +++++++++++++++++++++++++++++++++ tests/test_import.py | 11 + tests/test_readme.py | 25 +- 6 files changed, 973 insertions(+), 9 deletions(-) create mode 100644 python/zudb/aio.py create mode 100644 tests/test_aio.py diff --git a/README.md b/README.md index 2974824..eab13c2 100644 --- a/README.md +++ b/README.md @@ -140,6 +140,35 @@ conn.rows_read # how far the statement running now has got `interrupt()` is the one a program calls, from a thread that is not the one inside `execute`, and it raises `zudb.Interrupted` there. It is one of the three calls that may be made on a connection while a statement is running, with `rows_read` and `closed`, and none of the three waits for it: a progress bar drawn from `rows_read` is a poll of an atomic, not a queue behind the executor. +## On an event loop + +A statement runs inside Rust with the GIL down and comes back when it comes back. Called straight from a coroutine it stops the loop for that whole time, including the tasks answering requests that have nothing to do with the database, so `zudb.aio` gives each connection a thread of its own and hands it every call that could wait. + +```python +import asyncio + +import zudb.aio + + +async def main(): + async with zudb.aio.connect("social.zu1") as conn: + await conn.execute("INSERT (p:person {uid: 1, name: 'ada'})") + rows = await conn.execute("MATCH (p:person) RETURN p.name AS name") + for (name,) in rows: + print(name) + + +asyncio.run(main()) +``` + +A thread per connection rather than a pool shared between them, because a connection is one lock and statements on it queue anyway: a pool would add no parallelism the engine can use and would let two statements written one after the other run in the other order. Two connections do run at the same time, since the engine puts the GIL down for the work, and two statements together cost 1.09 times what one costs alone on this machine. + +What comes back is a `zudb.Result`, which is rows already in memory, so reading them is the call it was: `for row in rows` and `rows.to_arrow()` are not awaited and never were. Only the ways in are, and only the ones that can wait. `path`, `read_only`, `closed`, `rows_read` and `interrupt()` are answered from beside the lock rather than through it, so they stay properties and a progress bar drawn from `rows_read` still reads while the statement it is measuring runs. + +Cancelling the task that awaits a statement interrupts the statement. The engine is asked to stop and the coroutine does not return until it has, so the connection is idle again by the time the `CancelledError` reaches the caller rather than busy with work nobody is waiting for, and a statement still queued when the cancellation arrives is dropped without running. A transaction block cancelled partway is a block that raised, so it leaves through the rollback. + +`transaction()` and `appender()` are opened with `async with`, `in_transaction()` and `registered()` are methods here because both answers live behind the lock, and everything else is the sync call with an `await` in front of it. + ## When a program is wrong Every condition arrives as an exception class carrying its GQLSTATUS code, its position and a link to what the standard says about it, and the class is the class a Python caller would have written: a mistake the program made is a `zudb.ProgrammingError`, a value Python has and zu does not is a `TypeError`, a value of the right type and the wrong shape is a `ValueError`, and a file that is not a database is a `zudb.ConnectionError`, the same as a file that is not there. Both of those are a path that does not lead to a database, and telling a caller who mistyped one to file a bug would be the wrong answer twice. @@ -154,7 +183,7 @@ The stub is checked against the module it describes in CI: griffe reads the stub ## 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, and `Ctrl-C` and `interrupt()` stopping a statement without touching the connection under it. `zudb.aio` is next, and each one lands 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, and `zudb.aio` for the same calls awaited on an event loop. A DB-API 2.0 wrapper is next, and each one lands with the tests that say it works. ## Wheels diff --git a/python/zudb/__init__.py b/python/zudb/__init__.py index aec3fbe..bca1237 100644 --- a/python/zudb/__init__.py +++ b/python/zudb/__init__.py @@ -10,6 +10,10 @@ The engine is compiled into the wheel, so there is nothing to install, nothing to run, and no server to connect to. Statements are ISO/IEC 39075 GQL. + +On an event loop the same calls are awaited, from `zudb.aio`. It is a +submodule to ask for by name rather than one imported here, so a script +that never awaits anything pays for none of it. """ from __future__ import annotations diff --git a/python/zudb/aio.py b/python/zudb/aio.py new file mode 100644 index 0000000..c0f07d2 --- /dev/null +++ b/python/zudb/aio.py @@ -0,0 +1,498 @@ +"""zu on an event loop: the same calls, awaited. + + import zudb.aio + + async with zudb.aio.connect("social.zu1") as conn: + rows = await conn.execute("MATCH (p:person) RETURN p.name AS name") + for (name,) in rows: + print(name) + +A statement runs inside Rust with the GIL down, and a call like that +comes back when it comes back. Run one straight from a coroutine and +every other task on the loop waits for it, including the ones answering +requests that have nothing to do with the database. So each connection +here keeps a thread of its own, and every call that could wait is handed +to that thread and awaited. + +A thread per connection rather than a pool shared by all of them, +because a connection is one lock and statements on it queue anyway. A +pool would add no parallelism the engine can use and would let two +statements on one connection run in an order neither caller wrote. Two +connections do run at the same time, on their own threads, since the +engine puts the GIL down for the work. + +What comes back is a `zudb.Result`, which is rows already in memory, so +nothing about reading them waits and nothing about reading them is +awaited: `for row in rows` and `rows.to_arrow()` are the calls they were. +Only the ways in are awaited, and only the ones that can wait: `path`, +`read_only`, `closed` and `rows_read` are answered from beside the lock +rather than through it, so they stay properties here too. + +Cancelling the task that awaits a statement interrupts the statement. +The engine is asked to stop, and the coroutine does not return until it +has, so the connection is idle again by the time the `CancelledError` +reaches the caller rather than busy with a statement nobody is waiting +for. A statement still queued when the cancellation arrives is dropped +without running. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import functools +import os +import pathlib +import threading +from collections.abc import Callable, Coroutine, Generator, Iterable, Mapping +from concurrent.futures import Future, ThreadPoolExecutor +from typing import Any, Generic, TypeVar + +from . import _zudb +from ._zudb import Appender, Connection, Result, Transaction +from .types import Value + +__all__ = ["connect", "AsyncConnection", "AsyncTransaction", "AsyncAppender"] + +T = TypeVar("T") + + +def connect( + path: str | os.PathLike[str], + *, + read_only: bool = False, + memory_limit: int | None = None, + threads: int | None = None, +) -> _Opening[AsyncConnection]: + """Opens the database at `path` and connects to it. + + The arguments are `zudb.connect`'s, and so is the behaviour: a path + holding nothing becomes a new database unless the connection is + read-only. Opening reads the file, so it happens on the connection's + own thread like everything else. + + Await it for a connection to close yourself, or open it with `async + with` for one that closes at the end of the block: + + conn = await zudb.aio.connect("social.zu1") + async with zudb.aio.connect("social.zu1") as conn: + ... + """ + return _Opening( + functools.partial( + _open, + path, + read_only=read_only, + memory_limit=memory_limit, + threads=threads, + ) + ) + + +async def _open( + path: str | os.PathLike[str], + *, + read_only: bool, + memory_limit: int | None, + threads: int | None, +) -> AsyncConnection: + """Starts the thread, opens the database on it, and joins the two.""" + pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="zudb aio") + opening = pool.submit( + functools.partial( + _zudb.connect, + path, + read_only=read_only, + memory_limit=memory_limit, + threads=threads, + ) + ) + try: + conn = await asyncio.wrap_future(opening) + except BaseException: + # An open that was cancelled is still an open file once the + # thread reaches it, so what arrives is closed rather than left + # for the collector to find. + opening.add_done_callback(_discard) + pool.shutdown(wait=False) + raise + return AsyncConnection(conn, pool) + + +def _discard(opened: Future[Connection]) -> None: + """Closes a connection nobody is waiting for any more.""" + if not opened.cancelled() and opened.exception() is None: + opened.result().close() + + +class _Opening(Generic[T]): + """Something to `await`, or to open with `async with`. + + Opening a connection, starting a transaction and taking an appender + all reach the engine, so all three are coroutines, and a coroutine + is not a context manager. This is the both of them: awaiting it + opens the thing, and `async with` opens it and then enters it, so + the two spellings a caller might reach for are the same call. + """ + + __slots__ = ("_open", "_held") + + def __init__(self, open_: Callable[[], Coroutine[Any, Any, T]]) -> None: + self._open = open_ + self._held: T | None = None + + def __await__(self) -> Generator[Any, None, T]: + return self._open().__await__() + + async def __aenter__(self) -> T: + self._held = await self._open() + return await self._held.__aenter__() # type: ignore[attr-defined,no-any-return] + + async def __aexit__(self, *exception: Any) -> bool: + return await self._held.__aexit__(*exception) # type: ignore[attr-defined,no-any-return] + + +class AsyncConnection: + """One connection to one database, and the thread its work runs on. + + Take one with `zudb.aio.connect`. Every method that could wait is a + coroutine, and they queue in the order they were awaited, because + the thread underneath runs one at a time. + """ + + __slots__ = ("_conn", "_pool", "_guard", "_running", "_shut") + + def __init__(self, conn: Connection, pool: ThreadPoolExecutor) -> None: + self._conn = conn + self._pool = pool + # Held across the two lines that say which job is on the thread, + # so that a cancellation either interrupts its own statement or + # interrupts nothing. Without it a job that finished a moment + # ago could have its stop land on the one after it. + self._guard = threading.Lock() + self._running: object | None = None + self._shut = False + + @property + def path(self) -> pathlib.Path: + """The file this connection was opened on.""" + return self._conn.path + + @property + def read_only(self) -> bool: + """Whether it was opened read-only.""" + return self._conn.read_only + + @property + def closed(self) -> bool: + """Whether this connection is still open.""" + return self._conn.closed + + @property + def rows_read(self) -> int: + """How many rows the statement running on this connection has + read out of storage. + + A property and not a coroutine on purpose: this is what a + progress bar reads while a statement runs, and a call that had + to queue behind that statement could only answer once it was + over. + """ + return self._conn.rows_read + + def interrupt(self) -> None: + """Asks the statement running on this connection to stop. + + Not awaited, for the same reason: an ask that queued behind the + statement it means to stop would arrive after it. Cancelling the + task that awaits a statement does this for you. + """ + self._conn.interrupt() + + async def execute(self, statement: str, params: Mapping[str, Value] | None = None) -> Result: + """Runs one statement and gives back its rows.""" + return await self._call(functools.partial(self._conn.execute, statement, params)) + + async def sql(self, statement: str, params: Mapping[str, Value] | None = None) -> Result: + """The same call, named for the way it reads in a notebook.""" + return await self._call(functools.partial(self._conn.sql, statement, params)) + + def transaction(self, *, read_only: bool = False) -> _Opening[AsyncTransaction]: + """Starts a transaction and hands it back for an `async with` + block. + + async with conn.transaction(): + await conn.execute("INSERT (a:account {uid: 1, balance: 100})") + await conn.execute("INSERT (b:account {uid: 2, balance: 0})") + + The block commits when it ends and rolls back when it raises, + and a task cancelled inside one leaves through the rollback like + any other exception. + """ + return _Opening(functools.partial(self._transaction, read_only)) + + async def _transaction(self, read_only: bool) -> AsyncTransaction: + started = await self._call(functools.partial(self._conn.transaction, read_only=read_only)) + return AsyncTransaction(self, started) + + async def in_transaction(self) -> bool: + """Whether an explicit transaction is running on this + connection. + + A method where the sync client has a property, because the + answer is inside the lock and so has to be awaited, and a + property that gave back a coroutine would read as `if + conn.in_transaction:`, which is true whatever the answer. + """ + return await self._call(lambda: self._conn.in_transaction) + + def appender(self, table: str) -> _Opening[AsyncAppender]: + """Opens an appender on `table`, for loading rows into a + database that already exists. + """ + return _Opening(functools.partial(self._appender, table)) + + async def _appender(self, table: str) -> AsyncAppender: + opened = await self._call(functools.partial(self._conn.appender, table)) + return AsyncAppender(self, opened) + + async def register(self, name: str, data: Any) -> int: + """Puts a DataFrame under a name a statement can match on.""" + return await self._call(functools.partial(self._conn.register, name, data)) + + async def unregister(self, name: str) -> None: + """Takes a registered frame's name away.""" + await self._call(functools.partial(self._conn.unregister, name)) + + async def registered(self) -> list[str]: + """The names this connection has registered frames under, + sorted. + + A method where the sync client has a property, for the reason + `in_transaction` is one: the names live in the engine, behind + the lock, so asking has to be awaited. + """ + return await self._call(lambda: self._conn.registered) + + async def close(self) -> None: + """Closes the connection, frees what it held, and ends the + thread it ran on. + + Closing waits for whatever is still running, which is why it is + awaited. Doing it twice is not an error. + """ + if self._shut: + return + await self._call(self._conn.close) + self._shut = True + # False, because the thread has just finished the close and has + # nothing else to do, and waiting for a thread from inside a + # coroutine is the one thing this module exists to avoid. + self._pool.shutdown(wait=False) + + async def __aenter__(self) -> AsyncConnection: + return self + + async def __aexit__(self, *_exception: Any) -> bool: + await self.close() + # False, so an exception raised inside the block carries on out + # of it. + return False + + def __repr__(self) -> str: + state = ", closed" if self.closed else "" + return f"" + + async def _call(self, work: Callable[[], T]) -> T: + """Runs `work` on this connection's thread and waits for it. + + Everything here that can wait goes through this, so the order + statements run in is the order they were awaited in, and the + loop is free the whole time any of them is running. + """ + if self._shut: + # The thread is gone and there is nothing to run on, but + # every call into a closed connection is refused before it + # reaches the engine, so running it here raises exactly what + # the thread would have raised, without one. + return work() + token = object() + queued = self._pool.submit(self._job, token, work) + try: + return await asyncio.shield(asyncio.wrap_future(queued)) + except asyncio.CancelledError: + if not queued.cancel(): + self._stop(token) + # Waited for rather than abandoned, so that the next + # statement finds the connection free rather than locked + # by one nobody is listening to. Whatever it ends with, + # including the rows it managed to finish, is dropped: + # the caller asked to stop, and the answer to that is + # the cancellation. + with contextlib.suppress(BaseException): + await asyncio.shield(asyncio.wrap_future(queued)) + raise + + def _job(self, token: object, work: Callable[[], T]) -> T: + """The work, with a note of whose it is while it runs.""" + with self._guard: + self._running = token + try: + return work() + finally: + with self._guard: + self._running = None + + def _stop(self, token: object) -> None: + """Interrupts the statement `token` names, if it is the one + running. + """ + with self._guard: + if self._running is token: + # Refused on a connection that was closed underneath + # us, which is not something a cancellation should + # raise about: the statement is over either way. + with contextlib.suppress(Exception): + self._conn.interrupt() + + +class AsyncTransaction: + """A transaction that has been started and not yet ended. + + The transaction underneath is `zudb.Transaction` and the rules are + its rules: it starts when it is taken, ending it twice is refused, + and the statements inside it are the ones written on the connection + it came from. + """ + + __slots__ = ("_conn", "_txn") + + def __init__(self, conn: AsyncConnection, txn: Transaction) -> None: + self._conn = conn + self._txn = txn + + @property + def read_only(self) -> bool: + """Whether it was started `READ ONLY`.""" + return self._txn.read_only + + @property + def done(self) -> bool: + """Whether this transaction has already been committed or + rolled back. + """ + return self._txn.done + + async def commit(self) -> None: + """Ends the transaction and keeps what it wrote.""" + await self._conn._call(self._txn.commit) + + async def rollback(self) -> None: + """Ends the transaction and throws away what it wrote.""" + await self._conn._call(self._txn.rollback) + + async def __aenter__(self) -> AsyncTransaction: + return self + + async def __aexit__(self, *exception: Any) -> bool: + """Commits at the end of the block, and rolls back when the + block raised. + + A block cancelled partway is a block that raised, so it unwinds: + the rollback is the reason to write the block this way, and a + task that stopped halfway through two writes is exactly the case + it is there for. + """ + raised = bool(exception) and exception[0] is not None + if not self.done: + await (self.rollback() if raised else self.commit()) + return False + + def __repr__(self) -> str: + read_only = " read only" if self.read_only else "" + done = ", done" if self.done else "" + return f"" + + +class AsyncAppender: + """Rows on their way into a table, buffered until they are flushed. + + Every call here is awaited, the ones that only touch the buffer + included. The buffer's lock is held for the whole of a flush, and a + flush is inside the engine with the GIL down: an append from the + loop's own thread would wait on that lock while holding the GIL, + which is the one wait the flush cannot get out from under. Handing + them all to the connection's thread keeps them in one queue where + they cannot meet. + """ + + __slots__ = ("_conn", "_appender") + + def __init__(self, conn: AsyncConnection, appender: Appender) -> None: + self._conn = conn + self._appender = appender + + @property + def table(self) -> str: + """The table this appender writes to.""" + return self._appender.table + + async def buffered(self) -> int: + """Rows buffered and not yet written.""" + return await self._conn._call(lambda: self._appender.buffered) + + async def committed(self) -> int: + """Rows this appender has committed, across every flush.""" + return await self._conn._call(lambda: self._appender.committed) + + async def closed(self) -> bool: + """Whether this appender has been closed.""" + return await self._conn._call(lambda: self._appender.closed) + + async def append_row(self, row: Iterable[Value]) -> None: + """Appends one row, which is a sequence of one value per column + of the table. + + A row is a conversion and a push per column, so a loader with + rows to hand should reach for `append_rows` and pay for one + crossing rather than one per row. + """ + await self._conn._call(functools.partial(self._appender.append_row, row)) + + async def append_rows(self, rows: Iterable[Iterable[Value]]) -> int: + """Appends every row of an iterable of rows. + + The iterable is read on the connection's thread, so a generator + passed here runs there: keep it to arranging values, and do the + awaiting that produced them before the call. + """ + return await self._conn._call(functools.partial(self._appender.append_rows, rows)) + + async def flush(self) -> int: + """Writes every buffered row and makes it readable.""" + return await self._conn._call(self._appender.flush) + + async def discard(self) -> int: + """Throws away what is buffered and answers how many rows that + was. + """ + return await self._conn._call(self._appender.discard) + + async def close(self) -> int: + """Flushes what is left and answers how many rows this appender + committed in all. + """ + return await self._conn._call(self._appender.close) + + async def __aenter__(self) -> AsyncAppender: + return self + + async def __aexit__(self, *_exception: Any) -> bool: + """Closes on the way out, which flushes, and does it whether the + block ended well or badly, like the appender underneath. + """ + await self.close() + return False + + def __repr__(self) -> str: + return f"" diff --git a/tests/test_aio.py b/tests/test_aio.py new file mode 100644 index 0000000..aab5395 --- /dev/null +++ b/tests/test_aio.py @@ -0,0 +1,413 @@ +"""The same calls, awaited. + +There are three claims in `zudb.aio` and everything here is one of +them. The loop keeps running while a statement does, which is the whole +reason the module exists. Statements on one connection arrive in the +order they were awaited, because the thread underneath runs one at a +time. Cancelling the task that awaits a statement stops the statement +and leaves the connection free rather than busy with work nobody is +waiting for. + +The rest is the sync client's behaviour, checked once through the async +spelling to show it survived the crossing. +""" + +from __future__ import annotations + +import asyncio +import functools +import time +from collections.abc import Callable, Coroutine +from pathlib import Path +from typing import Any, TypeVar + +import pytest +import zudb +import zudb.aio + +T = TypeVar("T") + +# Every pair of people, filtered, which is a statement that runs for +# long enough to watch rather than one that is over before the loop is +# scheduled. +WORK = "MATCH (a:person), (b:person) WHERE a.uid < b.uid RETURN count(a) AS n" + +# Three thousand people is about a second of that work on the machine +# this was written on, and six thousand about three seconds. The first +# is for the tests that wait for the statement to end and the second +# for the ones that stop it partway, which need it still running a +# fifth of a second in on a machine some multiple faster than this one. +WATCHED = 3_000 +LONG = 6_000 + +# Two connections are timed against each other rather than against a +# budget, so what matters is that the statement is long enough to time +# and short enough to run three times. +TIMED = 1_500 + + +def pairs(people: int) -> int: + return people * (people - 1) // 2 + + +def run(test: Callable[..., Coroutine[Any, Any, None]]) -> Callable[..., None]: + """Runs a coroutine test on a loop of its own. + + pytest does not await, and this suite has no plugin that teaches it + to, so each test is a coroutine and this is the call that runs one. + A loop per test rather than one shared, because a loop that outlived + a failure would carry whatever that failure left running into the + test after it. `functools.wraps` keeps the signature, which is what + pytest reads the fixtures out of. + """ + + @functools.wraps(test) + def wrapper(*args: Any, **kwargs: Any) -> None: + asyncio.run(test(*args, **kwargs)) + + return wrapper + + +async def crowded(path: Path, people: int) -> zudb.aio.AsyncConnection: + """A connection to a database with `people` people in it. + + Loaded rather than inserted because a row at a time is a commit at + a time, and this is scaffolding rather than the thing under test. + The load is the sync call, since there is nothing to overlap it + with and the connection comes after it either way. + """ + zudb.load( + path, + nodes="person", + rels="knows", + columns={"uid": list(range(people))}, + edges=[(0, 1)], + ) + return await zudb.aio.connect(path) + + +@run +async def test_a_statement_runs_and_gives_back_its_rows(tmp_path: Path) -> None: + async with zudb.aio.connect(tmp_path / "one.zu1") as conn: + await conn.execute("INSERT (p:person {uid: 10, name: 'ada'})") + rows = await conn.execute("MATCH (p:person) RETURN p.name AS name") + assert rows.fetchall() == [("ada",)] + + +@run +async def test_what_comes_back_is_read_without_awaiting(tmp_path: Path) -> None: + """Rows are already in memory, so nothing about reading them waits.""" + async with zudb.aio.connect(tmp_path / "read.zu1") as conn: + await conn.execute("INSERT (p:person {uid: 10, name: 'ada'})") + rows = await conn.execute("MATCH (p:person) RETURN p.uid AS uid, p.name AS name") + assert rows.columns == ["uid", "name"] + assert len(rows) == 1 + assert [row for row in rows] == [(10, "ada")] + + +@run +async def test_the_loop_runs_while_a_statement_does(tmp_path: Path) -> None: + """The claim the module is for. + + A statement is inside Rust with the GIL down, and while it is there + a task that only wants the loop keeps getting it. The count is + asserted low enough to be about whether the loop moved at all + rather than about how fast this machine is. + """ + async with await crowded(tmp_path / "ticking.zu1", WATCHED) as conn: + ticks = 0 + statement = asyncio.ensure_future(conn.execute(WORK)) + while not statement.done(): + await asyncio.sleep(0.001) + ticks += 1 + assert (await statement).fetchone() == (pairs(WATCHED),) + assert ticks > 20, f"the loop only got round {ticks} times" + + +@run +async def test_two_connections_run_at_the_same_time(tmp_path: Path) -> None: + """A thread each, and the engine puts the GIL down for the work, so + two statements together cost about what one costs rather than two. + """ + first = await crowded(tmp_path / "first.zu1", TIMED) + second = await crowded(tmp_path / "second.zu1", TIMED) + async with first, second: + # Once each before the clock starts, because the first + # statement on a connection reads the file in and the number + # wanted here is about two threads and not about a cold cache. + await asyncio.gather(first.execute(WORK), second.execute(WORK)) + + started = time.perf_counter() + await first.execute(WORK) + alone = time.perf_counter() - started + + started = time.perf_counter() + await asyncio.gather(first.execute(WORK), second.execute(WORK)) + together = time.perf_counter() - started + assert together < alone * 1.8, f"{together:.3f}s together against {alone:.3f}s alone" + + +@run +async def test_statements_arrive_in_the_order_they_were_awaited(tmp_path: Path) -> None: + async with zudb.aio.connect(tmp_path / "ordered.zu1") as conn: + await conn.execute("INSERT (p:person {uid: 0, name: 'seed'})") + await asyncio.gather( + *( + conn.execute("INSERT (p:person {uid: $uid, name: 'p'})", {"uid": uid}) + for uid in range(1, 6) + ) + ) + rows = await conn.execute("MATCH (p:person) RETURN p.uid AS uid") + assert [uid for (uid,) in rows] == [0, 1, 2, 3, 4, 5] + + +@run +async def test_cancelling_the_task_stops_the_statement(tmp_path: Path) -> None: + async with await crowded(tmp_path / "cancelled.zu1", LONG) as conn: + statement = asyncio.ensure_future(conn.execute(WORK)) + # Long enough that it is inside the executor rather than still + # being parsed when the cancellation arrives. + await asyncio.sleep(0.2) + statement.cancel() + with pytest.raises(asyncio.CancelledError): + await statement + + +@run +async def test_a_cancelled_statement_leaves_the_connection_free(tmp_path: Path) -> None: + """The reason the cancellation waits for the statement it stopped. + + A connection still running work nobody is listening to would make + the next statement queue behind all of it, so the one after the + cancellation is timed rather than merely run. + """ + async with await crowded(tmp_path / "free.zu1", LONG) as conn: + statement = asyncio.ensure_future(conn.execute(WORK)) + await asyncio.sleep(0.2) + statement.cancel() + with pytest.raises(asyncio.CancelledError): + await statement + + started = time.perf_counter() + rows = await conn.execute("MATCH (p:person) RETURN count(p) AS n") + took = time.perf_counter() - started + assert rows.fetchone() == (LONG,) + assert took < 1.0, f"the next statement waited {took:.3f}s" + + +@run +async def test_a_statement_still_queued_is_dropped_without_running(tmp_path: Path) -> None: + async with await crowded(tmp_path / "queued.zu1", LONG) as conn: + running = asyncio.ensure_future(conn.execute(WORK)) + waiting = asyncio.ensure_future(conn.execute("INSERT (p:person {uid: 999999})")) + await asyncio.sleep(0.2) + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + running.cancel() + with pytest.raises(asyncio.CancelledError): + await running + rows = await conn.execute("MATCH (p:person) WHERE p.uid = 999999 RETURN p.uid AS uid") + assert rows.fetchall() == [] + + +@run +async def test_interrupt_is_not_awaited_and_stops_what_is_running(tmp_path: Path) -> None: + async with await crowded(tmp_path / "interrupted.zu1", LONG) as conn: + statement = asyncio.ensure_future(conn.execute(WORK)) + await asyncio.sleep(0.2) + conn.interrupt() + with pytest.raises(zudb.Interrupted): + await statement + + +@run +async def test_a_transaction_block_commits_what_it_wrote(tmp_path: Path) -> None: + async with zudb.aio.connect(tmp_path / "committed.zu1") as conn: + await conn.execute("INSERT (p:person {uid: 10, name: 'ada'})") + async with conn.transaction(): + await conn.execute("INSERT (p:person {uid: 20, name: 'grace'})") + await conn.execute("INSERT (p:person {uid: 30, name: 'kay'})") + rows = await conn.execute("MATCH (p:person) RETURN p.uid AS uid") + assert [uid for (uid,) in rows] == [10, 20, 30] + + +@run +async def test_a_transaction_block_that_raises_rolls_back(tmp_path: Path) -> None: + async with zudb.aio.connect(tmp_path / "rolled.zu1") as conn: + await conn.execute("INSERT (p:person {uid: 10, name: 'ada'})") + with pytest.raises(RuntimeError, match="halfway"): + async with conn.transaction(): + await conn.execute("INSERT (p:person {uid: 20, name: 'grace'})") + raise RuntimeError("halfway through") + rows = await conn.execute("MATCH (p:person) RETURN p.uid AS uid") + assert [uid for (uid,) in rows] == [10] + + +@run +async def test_a_transaction_cancelled_partway_rolls_back(tmp_path: Path) -> None: + """The case the block is written this way for. + + A task stopped between two writes is a task that raised, so the + block unwinds through the rollback and neither write is kept. + """ + conn = await zudb.aio.connect(tmp_path / "half.zu1") + await conn.execute("INSERT (p:person {uid: 10, name: 'ada'})") + reached = asyncio.Event() + + async def both() -> None: + async with conn.transaction(): + await conn.execute("INSERT (p:person {uid: 20, name: 'grace'})") + reached.set() + await asyncio.sleep(30) + await conn.execute("INSERT (p:person {uid: 30, name: 'kay'})") + + task = asyncio.ensure_future(both()) + await reached.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + rows = await conn.execute("MATCH (p:person) RETURN p.uid AS uid") + assert [uid for (uid,) in rows] == [10] + await conn.close() + + +@run +async def test_in_transaction_is_awaited(tmp_path: Path) -> None: + async with zudb.aio.connect(tmp_path / "asking.zu1") as conn: + assert await conn.in_transaction() is False + async with conn.transaction() as txn: + assert await conn.in_transaction() is True + assert txn.read_only is False + assert txn.done is False + assert await conn.in_transaction() is False + + +@run +async def test_a_transaction_ended_by_hand_is_done(tmp_path: Path) -> None: + async with zudb.aio.connect(tmp_path / "byhand.zu1") as conn: + txn = await conn.transaction() + await conn.execute("INSERT (p:person {uid: 10, name: 'ada'})") + await txn.rollback() + assert txn.done is True + rows = await conn.execute("MATCH (p:person) RETURN p.uid AS uid") + assert rows.fetchall() == [] + + +@run +async def test_an_appender_loads_rows(tmp_path: Path) -> None: + async with zudb.aio.connect(tmp_path / "appended.zu1") as conn: + await conn.execute("INSERT (p:person {uid: 0, name: 'seed'})") + async with conn.appender("person") as appender: + assert appender.table == "person" + await appender.append_row([1, "ada"]) + await appender.append_rows([[2, "grace"], [3, "kay"]]) + assert await appender.buffered() == 3 + assert await appender.flush() == 3 + assert await appender.committed() == 3 + rows = await conn.execute("MATCH (p:person) RETURN p.name AS name") + assert [name for (name,) in rows] == ["seed", "ada", "grace", "kay"] + + +@run +async def test_a_frame_registers_and_a_statement_matches_it(tmp_path: Path) -> None: + pa = pytest.importorskip("pyarrow") + async with zudb.aio.connect(tmp_path / "framed.zu1") as conn: + frame = pa.table({"uid": [1, 2, 3], "score": [10.0, 20.0, 30.0]}) + assert await conn.register("people", frame) == 3 + assert await conn.registered() == ["people"] + rows = await conn.execute("MATCH (p:people) WHERE p.uid > 1 RETURN sum(p.score) AS total") + assert rows.fetchone() == (50.0,) + await conn.unregister("people") + assert await conn.registered() == [] + + +@run +async def test_the_ways_that_cannot_wait_are_properties(tmp_path: Path) -> None: + """`rows_read` is what a progress bar reads while a statement runs, + so it is answered from beside the lock rather than through it. + """ + path = tmp_path / "beside.zu1" + async with await crowded(path, LONG) as conn: + assert conn.path == path + assert conn.read_only is False + assert conn.closed is False + statement = asyncio.ensure_future(conn.execute(WORK)) + await asyncio.sleep(0.2) + assert conn.rows_read > 0 + statement.cancel() + with pytest.raises(asyncio.CancelledError): + await statement + + +@run +async def test_a_connection_closes_at_the_end_of_the_block(tmp_path: Path) -> None: + async with zudb.aio.connect(tmp_path / "block.zu1") as conn: + assert conn.closed is False + assert conn.closed is True + + +@run +async def test_awaiting_the_open_gives_a_connection_to_close_yourself(tmp_path: Path) -> None: + conn = await zudb.aio.connect(tmp_path / "byhand.zu1") + assert conn.closed is False + await conn.close() + assert conn.closed is True + # Twice is not an error. + await conn.close() + assert conn.closed is True + + +@run +async def test_a_call_on_a_closed_connection_is_refused(tmp_path: Path) -> None: + """The thread is gone, so the refusal has to come from here, and it + is the refusal the engine would have given. + """ + conn = await zudb.aio.connect(tmp_path / "gone.zu1") + await conn.close() + with pytest.raises(zudb.ProgrammingError, match="closed"): + await conn.execute("MATCH (p:person) RETURN p.uid AS uid") + + +@run +async def test_a_read_only_connection_says_so(tmp_path: Path) -> None: + path = tmp_path / "readonly.zu1" + async with zudb.aio.connect(path) as writer: + await writer.execute("INSERT (p:person {uid: 10, name: 'ada'})") + async with zudb.aio.connect(path, read_only=True) as reader: + assert reader.read_only is True + rows = await reader.execute("MATCH (p:person) RETURN p.name AS name") + assert rows.fetchall() == [("ada",)] + + +@run +async def test_opening_a_database_that_is_not_one_raises(tmp_path: Path) -> None: + """The open is on the thread too, so its failure arrives awaited.""" + path = tmp_path / "junk.zu1" + path.write_bytes(b"not a database") + with pytest.raises(zudb.Error): + await zudb.aio.connect(path) + + +@run +async def test_a_statement_that_is_wrong_raises_what_it_would_have(tmp_path: Path) -> None: + async with zudb.aio.connect(tmp_path / "wrong.zu1") as conn: + with pytest.raises(zudb.SyntaxError, match="expected"): + await conn.execute("MATCH (p:person RETURN p") + + +@run +async def test_repr_names_the_file_and_says_when_it_is_closed(tmp_path: Path) -> None: + conn = await zudb.aio.connect(tmp_path / "shown.zu1") + assert "shown.zu1" in repr(conn) + assert "closed" not in repr(conn) + await conn.close() + assert "closed" in repr(conn) + + +@run +async def test_sql_is_execute_under_another_name(tmp_path: Path) -> None: + async with zudb.aio.connect(tmp_path / "notebook.zu1") as conn: + await conn.sql("INSERT (p:person {uid: 10, name: 'ada'})") + rows = await conn.sql("MATCH (p:person) RETURN p.name AS name") + assert rows.fetchall() == [("ada",)] diff --git a/tests/test_import.py b/tests/test_import.py index 553462a..c6c5c24 100644 --- a/tests/test_import.py +++ b/tests/test_import.py @@ -111,6 +111,17 @@ def test_no_dataframe_library_is_imported_by_importing_this_one() -> None: assert done.stdout.strip() == "", f"imported without being asked: {done.stdout.strip()}" +def test_the_event_loop_module_is_not_imported_by_importing_this_one() -> None: + """`zudb.aio` is a submodule a caller asks for by name. + + It pulls in asyncio and a thread pool, and a script that never + awaits anything should not pay for either, so it is left out of the + package's own imports. + """ + done = run("import sys, zudb; print('asyncio' in sys.modules, 'zudb.aio' in sys.modules)") + assert done.stdout.strip() == "False False" + + def test_pyarrow_arrives_when_a_result_is_asked_for_its_columns(tmp_path: Path) -> None: pytest.importorskip("pyarrow") path = tmp_path / "columns.zu1" diff --git a/tests/test_readme.py b/tests/test_readme.py index bc8c515..fb7d295 100644 --- a/tests/test_readme.py +++ b/tests/test_readme.py @@ -6,12 +6,12 @@ spent on a traceback. So the blocks that are whole programs are run here, as printed, character for character. -A block is a whole program when it starts with `import zudb`, which is -the rule the README follows: a block that stands on its own carries its -import, and a block that shows one call in the middle of a session does -not. Each program runs in an interpreter of its own with a temporary -directory as its working directory, because the file it writes is the -one a reader would find beside them afterwards. +A block is a whole program when it starts with an import, which is the +rule the README follows: a block that stands on its own opens with what +it imports, and a block that shows one call in the middle of a session +opens with the call. Each program runs in an interpreter of its own +with a temporary directory as its working directory, because the file +it writes is the one a reader would find beside them afterwards. """ from __future__ import annotations @@ -44,7 +44,7 @@ def blocks(language: str) -> list[str]: def programs() -> list[str]: """The blocks that are whole programs.""" - return [block for block in blocks("python") if block.startswith("import zudb")] + return [block for block in blocks("python") if block.startswith("import ")] def run(program: str, where: Path) -> subprocess.CompletedProcess[str]: @@ -60,7 +60,7 @@ def run(program: str, where: Path) -> subprocess.CompletedProcess[str]: def test_the_readme_prints_programs_and_not_fragments() -> None: """The rule above holds: the page has both kinds and knows which.""" - assert len(programs()) == 2, "the README's whole programs" + assert len(programs()) == 3, "the README's whole programs" assert len(blocks("python")) > len(programs()), "and its fragments" @@ -79,6 +79,15 @@ def test_the_sixty_second_snippet_runs_as_printed(tmp_path: Path) -> None: assert (tmp_path / "social.zu1").is_file() +def test_the_event_loop_snippet_runs_as_printed(tmp_path: Path) -> None: + """The third block: the same two calls, awaited.""" + snippet = programs()[2] + assert "zudb.aio.connect" in snippet and "asyncio.run" in snippet + done = run(snippet, tmp_path) + assert done.returncode == 0, done.stderr + assert done.stdout.split() == ["ada"] + + @pytest.mark.parametrize("index", range(len(programs()))) def test_every_whole_program_in_the_readme_runs(index: int, tmp_path: Path) -> None: """Including the ones no other test looks at the output of."""