Support aborting the atomic context

Since it is possible to use an atomic context to abort a number of
changes support this usage.  Because the only way to actually abort
the transaction is to close the database this must also do so.

Amended by db: Note the limitation requiring close is a limitation of
the underlying notmuch API, which should be fixed in a future notmuch
release.

2 files changed, 20 insertions, 1 deletions

diff --git a/bindings/python-cffi/notmuch2/_database.py b/bindings/python-cffi/notmuch2/_database.py
index 3c06402d..7db5a7f8 100644
--- a/bindings/python-cffi/notmuch2/_database.py
+++ b/bindings/python-cffi/notmuch2/_database.py
@@ -664,6 +664,7 @@ class AtomicContext:
     def __init__(self, db, ptr_name):
         self._db = db
         self._ptr = lambda: getattr(db, ptr_name)
+        self._exit_fn = lambda: None
 
     def __del__(self):
         self._destroy()
@@ -679,13 +680,17 @@ class AtomicContext:
         ret = capi.lib.notmuch_database_begin_atomic(self._ptr())
         if ret != capi.lib.NOTMUCH_STATUS_SUCCESS:
             raise errors.NotmuchError(ret)
+        self._exit_fn = self._end_atomic
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def _end_atomic(self):
         ret = capi.lib.notmuch_database_end_atomic(self._ptr())
         if ret != capi.lib.NOTMUCH_STATUS_SUCCESS:
             raise errors.NotmuchError(ret)
 
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._exit_fn()
+
     def force_end(self):
         """Force ending the atomic section.
 
@@ -704,6 +709,15 @@ class AtomicContext:
         if ret != capi.lib.NOTMUCH_STATUS_SUCCESS:
             raise errors.NotmuchError(ret)
 
+    def abort(self):
+        """Abort the transaction.
+
+        Aborting a transaction will not commit any of the changes, but
+        will also implicitly close the database.
+        """
+        self._exit_fn = lambda: None
+        self._db.close()
+
 
 @functools.total_ordering
 class DbRevision:
diff --git a/bindings/python-cffi/tests/test_database.py b/bindings/python-cffi/tests/test_database.py
index 55069b5e..a2c69de6 100644
--- a/bindings/python-cffi/tests/test_database.py
+++ b/bindings/python-cffi/tests/test_database.py
@@ -127,6 +127,11 @@ class TestAtomic:
         with pytest.raises(errors.UnbalancedAtomicError):
             ctx.force_end()
 
+    def test_abort(self, db):
+        with db.atomic() as txn:
+            txn.abort()
+        assert db.closed
+
 
 class TestRevision: