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.

1 files changed, 15 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: