Support PreferRepeatableRead on explicit transactions by retrying

Author: msullivan

gel/options.py

@@ -57,17 +57,18 @@ class IsolationLevel:
     PreferRepeatableRead = "PreferRepeatableRead"
 
     @staticmethod
-    def _to_start_tx_str(v):
+    def _to_start_tx_str(v, optimistic_isolation):
         if (
             v == IsolationLevel.Serializable
-            # We may *prefer* repeatable read, but we have not yet
-            # implemented it for explicit transactions. (Which will
-            # require some gnarly retry-and-caching logic.)
-            or v == IsolationLevel.PreferRepeatableRead
         ):
             return 'SERIALIZABLE'
         elif v == IsolationLevel.RepeatableRead:
             return 'REPEATABLE READ'
+        elif v == IsolationLevel.PreferRepeatableRead:
+            if optimistic_isolation:
+                return 'REPEATABLE READ'
+            else:
+                return 'SERIALIZABLE'
         else:
             raise ValueError(
                 f"Invalid isolation_level value for transaction(): {self}"
@@ -134,10 +135,11 @@ def __init__(
     def defaults(cls):
         return cls()
 
-    def start_transaction_query(self):
+    def start_transaction_query(self, optimistic_isolation: bool):
         options = []
         if self._isolation is not None:
-            level = IsolationLevel._to_start_tx_str(self._isolation)
+            level = IsolationLevel._to_start_tx_str(
+                self._isolation, optimistic_isolation)
             options.append(f'ISOLATION {level}')
 
         if self._readonly is not None:

gel/transaction.py

@@ -86,7 +86,9 @@ def _make_start_query(self):
             raise errors.InterfaceError(
                 'cannot start; the transaction is already started')
 
-        return self._options.start_transaction_query()
+        return self._options.start_transaction_query(
+            optimistic_isolation=self.__retry._optimistic_rr
+        )
 
     def _make_commit_query(self):
         self.__check_state('commit')
@@ -175,9 +177,17 @@ async def _exit(self, extype, ex):
             await self._client._impl.release(self._connection)
 
         if (
-            extype is not None and
-            issubclass(extype, errors.EdgeDBError) and
-            ex.has_tag(errors.SHOULD_RETRY)
+            extype is not None
+            and issubclass(extype, errors.CapabilityError)
+            # XXX: This is not the best way to check this
+            and 'REPEATABLE READ' in str(ex)
+        ):
+            return self.__retry._retry_rr_failure(ex)
+
+        if (
+            extype is not None
+            and issubclass(extype, errors.EdgeDBError)
+            and ex.has_tag(errors.SHOULD_RETRY)
         ):
             return self.__retry._retry(ex)
 
@@ -233,6 +243,12 @@ def __init__(self, owner):
         self._next_backoff = 0
         self._options = owner._options
 
+        prefer_rr = (
+            self._options.transaction_options._isolation
+            == options.IsolationLevel.PreferRepeatableRead
+        )
+        self._optimistic_rr = prefer_rr
+
     def _retry(self, exc):
         self._last_exception = exc
         rule = self._options.retry_options.get_rule_for_exception(exc)
@@ -241,3 +257,13 @@ def _retry(self, exc):
         self._done = False
         self._next_backoff = rule.backoff(self._iteration)
         return True
+
+    def _retry_rr_failure(self, exc):
+        # Retry a failure due to REPEATABLE READ not working
+        if not self._optimistic_rr:
+            return False
+        self._optimistic_rr = False
+        # Decrement _iteration count, since this one doesn't really count.
+        self._iteration -= 1
+        self._done = False
+        return True

tests/test_sync_tx.py

@@ -31,10 +31,17 @@ class TestSyncTx(tb.SyncQueryTestCase):
         CREATE TYPE test::TransactionTest EXTENDING std::Object {
             CREATE PROPERTY name -> std::str;
         };
-    '''
 
-    TEARDOWN = '''
-        DROP TYPE test::TransactionTest;
+        CREATE TYPE test::Tmp {
+            CREATE REQUIRED PROPERTY tmp -> std::str;
+        };
+        CREATE TYPE test::TmpConflict {
+            CREATE REQUIRED PROPERTY tmp -> std::str {
+                CREATE CONSTRAINT exclusive;
+            }
+        };
+
+        CREATE TYPE test::TmpConflictChild extending test::TmpConflict;
     '''
 
     def test_sync_transaction_regular_01(self):
@@ -101,6 +108,51 @@ async def test_sync_transaction_kinds(self):
                 with tx:
                     pass
 
+    async def test_sync_transaction_prefer_rr(self):
+        if (
+            str(self.server_version.stage) != 'dev'
+            and (self.server_version.major, self.server_version.minor) < (6, 5)
+        ):
+            self.skipTest("DML in RepeatableRead not supported yet")
+        con = self.client.with_transaction_options(
+            edgedb.TransactionOptions(
+                isolation=edgedb.IsolationLevel.PreferRepeatableRead
+            )
+        )
+        # A transaction that needs to be serializable
+        for tx in con.transaction():
+            with tx:
+                res1 = tx.query_single('''
+                    select {
+                        ins := (insert test::TmpConflict { tmp := "test1" }),
+                        level := sys::get_transaction_isolation(),
+                    }
+                ''')
+
+                res2 = tx.query_single('''
+                    select {
+                        ins := (insert test::TmpConflict { tmp := "test2" }),
+                        level := sys::get_transaction_isolation(),
+                    }
+                ''')
+
+                # N.B: res1 will be RepeatableRead on the first
+                # iteration, maybe, but contingent on the second query
+                # succeeding it will be Serializable!
+                self.assertEqual(str(res1.level), 'Serializable')
+                self.assertEqual(str(res2.level), 'Serializable')
+
+        # And one that doesn't
+        for tx in con.transaction():
+            with tx:
+                res = tx.query_single('''
+                    select {
+                        ins := (insert test::Tmp { tmp := "test" }),
+                        level := sys::get_transaction_isolation(),
+                    }
+                ''')
+            self.assertEqual(str(res.level), 'RepeatableRead')
+
     def test_sync_transaction_commit_failure(self):
         with self.assertRaises(edgedb.errors.QueryError):
             for tx in self.client.transaction():