Made the transactional database functions accept an optional external transaction.

[didex.git] / didex / index.py

diff --git a/didex/index.py b/didex/index.py
index 5d206ff..2f3fdef 100644 (file)
--- a/didex/index.py
+++ b/didex/index.py
@@ -1,6 +1,6 @@
 import struct, contextlib
 from . import db, lib
-from .db import bd
+from .db import bd, txnfun
 
 deadlock = bd.DBLockDeadlockError
 notfound = bd.DBNotFoundError
@@ -177,39 +177,29 @@ class ordered(index, lib.closable):
             except deadlock:
                 continue
 
-    def put(self, key, id):
-        while True:
-            try:
-                with db.txn(self.db.env.env) as tx:
-                    obid = struct.pack(">Q", id)
-                    if not self.db.ob.has_key(obid, txn=tx.tx):
-                        raise ValueError("no such object in database: " + str(id))
-                    try:
-                        self.bk.put(self.typ.encode(key), obid, txn=tx.tx, flags=bd.DB_NODUPDATA)
-                    except bd.DBKeyExistError:
-                        return False
-                    tx.commit()
-                    return True
-            except deadlock:
-                continue
-
-    def remove(self, key, id):
-        while True:
+    @txnfun(lambda self: self.db.env.env)
+    def put(self, key, id, *, tx):
+        obid = struct.pack(">Q", id)
+        if not self.db.ob.has_key(obid, txn=tx.tx):
+            raise ValueError("no such object in database: " + str(id))
+        try:
+            self.bk.put(self.typ.encode(key), obid, txn=tx.tx, flags=bd.DB_NODUPDATA)
+        except bd.DBKeyExistError:
+            return False
+        return True
+
+    @txnfun(lambda self: self.db.env.env)
+    def remove(self, key, id, *, tx):
+        obid = struct.pack(">Q", id)
+        if not self.db.ob.has_key(obid, txn=tx.tx):
+            raise ValueError("no such object in database: " + str(id))
+        cur = self.bk.cursor(txn=tx.tx)
+        try:
             try:
-                with db.txn(self.db.env.env) as tx:
-                    obid = struct.pack(">Q", id)
-                    if not self.db.ob.has_key(obid, txn=tx.tx):
-                        raise ValueError("no such object in database: " + str(id))
-                    cur = self.bk.cursor(txn=tx.tx)
-                    try:
-                        try:
-                            cur.get_both(self.typ.encode(key), obid)
-                        except notfound:
-                            return False
-                        cur.delete()
-                    finally:
-                        cur.close()
-                    tx.commit()
-                    return True
-            except deadlock:
-                continue
+                cur.get_both(self.typ.encode(key), obid)
+            except notfound:
+                return False
+            cur.delete()
+        finally:
+            cur.close()
+        return True