Several statements as one unit of work

README.md

@@ -75,6 +75,24 @@ An appender nobody closed is the one mistake that cannot be reported where it ha
 
 On this machine 200,000 rows of an integer and a string take 11 ms to buffer through `append_rows` and 93 ms to flush, which is 1.9 million rows a second including the commit. The same rows one call at a time cost 32 ms of buffering, since a call is a call. Against `INSERT`, 2,000 rows take 32 seconds a row at a time and 28 ms through an appender, and the gap widens with the table because every `INSERT` is a commit and a fold.
 
+## Several statements as one unit of work
+
+Every statement already runs in a transaction of its own, so this is not what makes a write atomic. What it holds is the span: two statements are one unit, and the block rolls back when it raises, which is the failure worth writing it for since it is the one nobody wrote a handler for.
+
+```python
+with conn.transaction():
+    conn.execute("INSERT (a:account {uid: 1, balance: 100})")
+    conn.execute("INSERT (b:account {uid: 2, balance: 0})")
+```
+
+It starts at the call rather than at the `with`, so a transaction that cannot start says so at the line that asked for one, and `commit()` and `rollback()` are there for a caller who would rather say when. A block whose transaction has already ended is left alone on the way out, which is what lets a program commit early and carry on, and ending one twice is refused rather than ignored, because the statements between the two are in neither of them. `conn.in_transaction` answers which side of the block a program is on.
+
+`transaction(read_only=True)` starts one that refuses to write, at the statement that writes rather than at the block that would have written. One transaction runs at a time and a second is refused rather than nested, since a rollback of an inner one would have to invent an answer for what it undoes. The three words underneath are `START TRANSACTION`, `COMMIT` and `ROLLBACK`, and they all still work written out.
+
+An appender is refused inside a transaction. Its batches are commits of their own and a rollback does not take them back, so an appender opened in a block would promise a span it is not in: load first, then transact. A statement that failed leaves the transaction running, because the engine does not end one on a failed statement and the block that opened it is still the thing that closes it. A connection closed with work uncommitted drops that work, which is the same answer the block would have given.
+
+The wrapper costs 5 microseconds for an empty transaction, so what it costs is what the engine charges. On this machine that is more rather than less: 200 `INSERT`s cost 2.2 seconds each committing on its own and 3.3 seconds inside one transaction, and reads cost the same either way. A transaction here is worth taking for the span it holds and not for the time it saves, and the v0 write path is where that number has to change.
+
 ## Reading a result as columns
 
 A result is rows to iterate and columns to hand to something else. The columns go out over the Arrow C Data Interface, so pyarrow, pandas and polars each read the same buffers and none of them gets a Python object per cell.
@@ -117,7 +135,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, 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, 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. `register` 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, 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. `register` is next, and each one lands with the tests that say it works.
 
 ## Wheels
 

python/zudb/__init__.py

@@ -22,6 +22,7 @@
     Path,
     Rel,
     Result,
+    Transaction,
     __abi_version__,
     connect,
     load,
@@ -44,6 +45,7 @@
     "connect",
     "load",
     "Connection",
+    "Transaction",
     "Appender",
     "Result",
     "Node",

python/zudb/_zudb.pyi

@@ -76,6 +76,13 @@ 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 transaction(self, *, read_only: bool = False) -> Transaction:
+        """Starts a transaction and hands it back for a `with` block."""
+
+    @property
+    def in_transaction(self) -> bool:
+        """Whether an explicit transaction is running on this connection."""
+
     def appender(self, table: str) -> Appender:
         """Opens an appender on `table`, for loading rows into a database that already exists."""
 
@@ -86,6 +93,27 @@ class Connection:
     def __exit__(self, *_exception: object) -> bool: ...
     def __repr__(self) -> str: ...
 
+class Transaction:
+    """A transaction that has been started and not yet ended."""
+
+    @property
+    def read_only(self) -> bool:
+        """Whether it was started `READ ONLY`."""
+
+    @property
+    def done(self) -> bool:
+        """Whether this transaction has already been committed or rolled back."""
+
+    def commit(self) -> None:
+        """Ends the transaction and keeps what it wrote."""
+
+    def rollback(self) -> None:
+        """Ends the transaction and throws away what it wrote."""
+
+    def __enter__(self) -> Transaction: ...
+    def __exit__(self, *_exception: object) -> bool: ...
+    def __repr__(self) -> str: ...
+
 class Appender:
     """Rows on their way into a table, buffered until they are flushed."""
 

src/conn.rs

@@ -19,8 +19,9 @@ use zudb::{Config, Database, Interrupt};
 
 use crate::appender::Appender;
 use crate::columns;
-use crate::error::{closed, to_py_err};
+use crate::error::{closed, programming, to_py_err};
 use crate::interrupt;
+use crate::txn::Transaction;
 use crate::value::{Names, from_py, to_py};
 
 /// What a capsule holding an Arrow stream is called. The name is part
@@ -130,6 +131,52 @@ impl Connection {
         self.execute(py, statement, params)
     }
 
+    /// Starts a transaction and hands it back for a `with` block.
+    ///
+    /// Several statements as one unit of work: the block commits when
+    /// it ends and rolls back when it raises, which is the failure
+    /// worth writing this for, since it is the one nobody wrote a
+    /// handler for.
+    ///
+    /// It starts here rather than at the `with`, so a transaction that
+    /// cannot start says so at the line that asked for one. A
+    /// connection is inside one transaction at a time, and asking for a
+    /// second while the first is running is refused by the engine
+    /// rather than nested, because a transaction inside a transaction
+    /// would have to invent an answer for what a rollback of the inner
+    /// one undoes.
+    ///
+    /// `read_only=True` starts one that refuses to write, at the
+    /// statement that writes rather than at the block that would have
+    /// written.
+    #[pyo3(signature = (*, read_only = false))]
+    fn transaction(slf: Py<Self>, py: Python<'_>, read_only: bool) -> PyResult<Transaction> {
+        Transaction::start(py, slf, read_only)
+    }
+
+    /// Whether an explicit transaction is running on this connection.
+    ///
+    /// The statements a program writes outside one are each a
+    /// transaction of their own, so this is false on a connection that
+    /// is writing perfectly well. It answers what is true between
+    /// statements, which is the only moment a caller can ask in: asking
+    /// while another thread is inside a statement waits for that
+    /// statement, since the answer belongs to the session and the
+    /// session is what the statement is holding.
+    #[getter]
+    fn in_transaction(&self, py: Python<'_>) -> PyResult<bool> {
+        if !self.alive.load(Ordering::Acquire) {
+            return Err(closed(py, "this connection"));
+        }
+        py.detach(|| {
+            self.inner.lock().ok().and_then(|mut held| {
+                held.as_mut()
+                    .map(|conn| conn.session_mut().in_transaction())
+            })
+        })
+        .ok_or_else(|| closed(py, "this connection"))
+    }
+
     /// Opens an appender on `table`, for loading rows into a database
     /// that already exists.
     ///
@@ -138,7 +185,20 @@ impl Connection {
     /// in columns and writes each batch as one commit. The table has to
     /// be there already, since an appender adds rows to a table and
     /// does not make one.
+    ///
+    /// It is refused inside a transaction. An appender's batches are
+    /// commits of their own and a rollback does not take them back, so
+    /// an appender opened inside a `with conn.transaction()` would
+    /// promise a span it is not in. Load first, then transact, or write
+    /// the rows as statements.
     fn appender(slf: Py<Self>, py: Python<'_>, table: &str) -> PyResult<Appender> {
+        if slf.borrow(py).in_transaction(py)? {
+            return Err(programming(
+                py,
+                "an appender writes its own commits, which a rollback does not take \
+                 back, so one cannot be opened inside a transaction",
+            ));
+        }
         Appender::open(py, slf, table)
     }
 
@@ -267,6 +327,17 @@ impl Connection {
             read_only,
         })
     }
+
+    /// Runs a statement that takes no parameters and gives back
+    /// nothing, which is what the three transaction words are.
+    ///
+    /// It goes through `execute` rather than around it, so a `COMMIT`
+    /// waits for the connection's lock, releases the GIL and feels a
+    /// `Ctrl-C` exactly as every other statement does. The result is
+    /// dropped, since the words return no columns.
+    pub(crate) fn run(&self, py: Python<'_>, statement: &str) -> PyResult<()> {
+        self.execute(py, statement, None).map(|_| ())
+    }
 }
 
 /// The rows a statement gave back.

src/lib.rs

@@ -19,6 +19,7 @@ mod conn;
 mod error;
 mod interrupt;
 mod load;
+mod txn;
 mod value;
 
 use std::path::PathBuf;
@@ -55,6 +56,7 @@ fn _zudb(module: &Bound<'_, PyModule>) -> PyResult<()> {
     module.add_class::<conn::Connection>()?;
     module.add_class::<conn::Result>()?;
     module.add_class::<appender::Appender>()?;
+    module.add_class::<txn::Transaction>()?;
     module.add_class::<value::Node>()?;
     module.add_class::<value::Rel>()?;
     module.add_class::<value::Path>()?;

src/txn.rs

@@ -0,0 +1,143 @@
+//! Several statements as one unit of work.
+//!
+//! A statement written on its own already runs in a transaction of its
+//! own, so this is not what makes a write atomic. What it holds is the
+//! span: two statements are one unit, and the file keeps the state they
+//! started from until one word or the other ends them.
+//!
+//! ```python
+//! with conn.transaction():
+//!     conn.execute("INSERT (a:account {uid: 1, balance: 100})")
+//!     conn.execute("INSERT (b:account {uid: 2, balance: 0})")
+//! ```
+//!
+//! The block commits when it ends and rolls back when it raises, which
+//! is the whole reason this is a context manager rather than a pair of
+//! methods: the failure that has to roll back is the one nobody wrote a
+//! handler for.
+//!
+//! The three words underneath are `START TRANSACTION`, `COMMIT` and
+//! `ROLLBACK`, and a caller who would rather write them can, on this
+//! connection, today. What this adds is the rollback nobody remembers
+//! to write, and a name for the state, so a program can ask whether it
+//! is inside one rather than remembering.
+
+use pyo3::prelude::*;
+use pyo3::types::PyTuple;
+
+use crate::conn::Connection;
+use crate::error::programming;
+
+/// A transaction that has been started and not yet ended.
+///
+/// Take one with `Connection.transaction`. It starts when it is taken,
+/// so a transaction that cannot start says so at the line that asked
+/// rather than at the `with` below it, and the statements that run
+/// inside it are the ones written on the connection it came from.
+#[pyclass(module = "zudb")]
+pub struct Transaction {
+    /// The connection this runs on, held rather than borrowed, because
+    /// a transaction whose connection was collected would be a span
+    /// with nothing left to commit.
+    conn: Py<Connection>,
+    /// Whether it was started `READ ONLY`, which the engine refuses a
+    /// write inside of at the statement that writes.
+    #[pyo3(get)]
+    read_only: bool,
+    /// Whether a `COMMIT` or a `ROLLBACK` has already run here. A
+    /// transaction ended inside its own block is ended, and the block
+    /// closing over it does nothing, which is what lets a caller commit
+    /// early and carry on.
+    done: bool,
+}
+
+#[pymethods]
+impl Transaction {
+    /// Ends the transaction and keeps what it wrote.
+    ///
+    /// Doing it twice is refused rather than ignored. A second commit
+    /// is a program that has lost track of where its transaction ends,
+    /// and the statements between the two are in neither of them.
+    fn commit(&mut self, py: Python<'_>) -> PyResult<()> {
+        self.end(py, "COMMIT")
+    }
+
+    /// Ends the transaction and throws away what it wrote.
+    fn rollback(&mut self, py: Python<'_>) -> PyResult<()> {
+        self.end(py, "ROLLBACK")
+    }
+
+    /// Whether this transaction has already been committed or rolled
+    /// back.
+    #[getter]
+    fn done(&self) -> bool {
+        self.done
+    }
+
+    fn __enter__(slf: Py<Self>) -> Py<Self> {
+        slf
+    }
+
+    /// Commits at the end of the block, and rolls back when the block
+    /// raised.
+    ///
+    /// A rollback that fails raises out of here in place of what was
+    /// being raised, which reads badly and is still right: the first
+    /// exception is the `__context__` of the second, and a program told
+    /// only about the failed statement would go on believing the
+    /// transaction unwound.
+    #[pyo3(signature = (*exception))]
+    fn __exit__(&mut self, py: Python<'_>, exception: &Bound<'_, PyTuple>) -> PyResult<bool> {
+        // The first of the three is the exception's type, and it is
+        // `None` when the block ended by falling off its own end. That
+        // is the whole question here, so the other two go unread.
+        let raised = exception.get_item(0).is_ok_and(|kind| !kind.is_none());
+        if !self.done {
+            self.end(py, if raised { "ROLLBACK" } else { "COMMIT" })?;
+        }
+        // False, so an exception raised inside the block carries on out
+        // of it. A context manager that swallowed one would be a very
+        // quiet way to lose an error, and a quieter way to lose the
+        // rollback that answered it.
+        Ok(false)
+    }
+
+    fn __repr__(&self) -> String {
+        let read_only = if self.read_only { " read only" } else { "" };
+        let done = if self.done { ", done" } else { "" };
+        format!("<zudb.Transaction{read_only}{done}>")
+    }
+}
+
+impl Transaction {
+    /// Starts one, which is the statement that starts one.
+    pub fn start(py: Python<'_>, conn: Py<Connection>, read_only: bool) -> PyResult<Transaction> {
+        let statement = if read_only {
+            "START TRANSACTION READ ONLY"
+        } else {
+            "START TRANSACTION"
+        };
+        conn.borrow(py).run(py, statement)?;
+        Ok(Transaction {
+            conn,
+            read_only,
+            done: false,
+        })
+    }
+
+    fn end(&mut self, py: Python<'_>, statement: &str) -> PyResult<()> {
+        if self.done {
+            return Err(programming(
+                py,
+                "this transaction has already ended, and the statements after it \
+                 belong to no transaction of yours",
+            ));
+        }
+        self.conn.borrow(py).run(py, statement)?;
+        // Written after the statement rather than before it, so a
+        // commit the engine refused leaves a transaction a caller can
+        // still roll back.
+        self.done = true;
+        Ok(())
+    }
+}