From 4da17cb71e6231c4599435030609ca5698f1c08f Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Wed, 5 Jun 2024 15:45:59 -0500
Subject: [PATCH] add _exec_sql_fn

---
 pytools/persistent_dict.py | 51 ++++++++++++++++++--------------------
 1 file changed, 24 insertions(+), 27 deletions(-)

diff --git a/pytools/persistent_dict.py b/pytools/persistent_dict.py
index eaf873d0..56dac063 100644
--- a/pytools/persistent_dict.py
+++ b/pytools/persistent_dict.py
@@ -522,13 +522,16 @@ def _collision_check(self, key: K, stored_key: K) -> None:
             stored_key == key  # pylint:disable=pointless-statement  # noqa: B015
             raise NoSuchEntryCollisionError(key)
 
-    def _exec_sql(self, *args: Any) -> Any:
+    def _exec_sql(self, *args: Any) -> sqlite3.Cursor:
+        return self._exec_sql_fn(lambda: self.conn.execute(*args))
+
+    def _exec_sql_fn(self, fn: Any) -> Any:
         n = 0
 
         while True:
             n += 1
             try:
-                return self.conn.execute(*args)
+                return fn()
             except sqlite3.OperationalError as e:
                 # If the database is busy, retry
                 if (hasattr(e, "sqlite_errorcode")
@@ -771,32 +774,26 @@ def remove(self, key: K) -> None:
         """Remove the entry associated with *key* from the dictionary."""
         keyhash = self.key_builder(key)
 
-        while True:
+        def remove_inner() -> None:
+            self.conn.execute("BEGIN EXCLUSIVE TRANSACTION")
             try:
-                self.conn.execute("BEGIN EXCLUSIVE TRANSACTION")
-                try:
-                    # This is split into SELECT/DELETE to allow for a collision check
-                    c = self.conn.execute("SELECT key_value FROM dict WHERE "
-                                          "keyhash=?", (keyhash,))
-                    row = c.fetchone()
-                    if row is None:
-                        raise NoSuchEntryError(key)
-
-                    stored_key, _value = pickle.loads(row[0])
-                    self._collision_check(key, stored_key)
-
-                    self.conn.execute("DELETE FROM dict WHERE keyhash=?", (keyhash,))
-                    self.conn.execute("COMMIT")
-                except Exception as e:
-                    self.conn.execute("ROLLBACK")
-                    raise e
-            except sqlite3.OperationalError as e:
-                # If the database is busy, retry
-                if (hasattr(e, "sqlite_errorcode")
-                        and not e.sqlite_errorcode == sqlite3.SQLITE_BUSY):
-                    raise
-            else:
-                break
+                # This is split into SELECT/DELETE to allow for a collision check
+                c = self.conn.execute("SELECT key_value FROM dict WHERE "
+                                        "keyhash=?", (keyhash,))
+                row = c.fetchone()
+                if row is None:
+                    raise NoSuchEntryError(key)
+
+                stored_key, _value = pickle.loads(row[0])
+                self._collision_check(key, stored_key)
+
+                self.conn.execute("DELETE FROM dict WHERE keyhash=?", (keyhash,))
+                self.conn.execute("COMMIT")
+            except Exception as e:
+                self.conn.execute("ROLLBACK")
+                raise e
+
+        self._exec_sql_fn(remove_inner)
 
     def __delitem__(self, key: K) -> None:
         """Remove the entry associated with *key* from the dictionary."""