Implement PreferRepeatableRead, apply TransactionOptions to config state (#609)

Currently it is enabled by setting default_transaction_isolation or isolation in TransactionOptions
to the enum value PreferRepeatableRead.

If that is set, the parse and execute paths will strip it from the
state that is sent, and then potentially inject RepeatableRead as the
value when doing execute, if it is safe.

Author: msullivan

gel/abstract.py

@@ -69,6 +69,7 @@ class QueryContext(typing.NamedTuple):
     state: typing.Optional[options.State]
     warning_handler: options.WarningHandler
     annotations: typing.Dict[str, str]
+    transaction_options: typing.Optional[options.TransactionOptions]
 
     def lower(
         self, *, allow_capabilities: enums.Capability
@@ -86,6 +87,7 @@ def lower(
             allow_capabilities=allow_capabilities,
             state=self.state.as_dict() if self.state else None,
             annotations=self.annotations,
+            transaction_options=self.transaction_options,
         )
 
 
@@ -96,6 +98,7 @@ class ExecuteContext(typing.NamedTuple):
     state: typing.Optional[options.State]
     warning_handler: options.WarningHandler
     annotations: typing.Dict[str, str]
+    transaction_options: typing.Optional[options.TransactionOptions]
 
     def lower(
         self, *, allow_capabilities: enums.Capability
@@ -111,6 +114,7 @@ def lower(
             allow_capabilities=allow_capabilities,
             state=self.state.as_dict() if self.state else None,
             annotations=self.annotations,
+            transaction_options=self.transaction_options,
         )
 
 
@@ -220,6 +224,7 @@ def query(self, query: str, *args, **kwargs) -> list:
             query_options=_query_opts,
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -233,6 +238,7 @@ def query_single(
             query_options=_query_single_opts,
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -244,6 +250,7 @@ def query_required_single(self, query: str, *args, **kwargs) -> typing.Any:
             query_options=_query_required_single_opts,
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -255,6 +262,7 @@ def query_json(self, query: str, *args, **kwargs) -> str:
             query_options=_query_json_opts,
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -266,6 +274,7 @@ def query_single_json(self, query: str, *args, **kwargs) -> str:
             query_options=_query_single_json_opts,
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -277,6 +286,7 @@ def query_required_single_json(self, query: str, *args, **kwargs) -> str:
             query_options=_query_required_single_json_opts,
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -293,6 +303,7 @@ def query_sql(self, query: str, *args, **kwargs) -> list[datatypes.Record]:
             query_options=_query_opts,
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -307,6 +318,7 @@ def execute(self, commands: str, *args, **kwargs) -> None:
             cache=self._get_query_cache(),
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -322,6 +334,7 @@ def execute_sql(self, commands: str, *args, **kwargs) -> None:
             cache=self._get_query_cache(),
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -349,6 +362,7 @@ async def query(self, query: str, *args, **kwargs) -> list:
             query_options=_query_opts,
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -360,6 +374,7 @@ async def query_single(self, query: str, *args, **kwargs) -> typing.Any:
             query_options=_query_single_opts,
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -376,6 +391,7 @@ async def query_required_single(
             query_options=_query_required_single_opts,
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -387,6 +403,7 @@ async def query_json(self, query: str, *args, **kwargs) -> str:
             query_options=_query_json_opts,
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -398,6 +415,7 @@ async def query_single_json(self, query: str, *args, **kwargs) -> str:
             query_options=_query_single_json_opts,
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -414,6 +432,7 @@ async def query_required_single_json(
             query_options=_query_required_single_json_opts,
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -430,6 +449,7 @@ async def query_sql(self, query: str, *args, **kwargs) -> typing.Any:
             query_options=_query_opts,
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -444,6 +464,7 @@ async def execute(self, commands: str, *args, **kwargs) -> None:
             cache=self._get_query_cache(),
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))
@@ -459,6 +480,7 @@ async def execute_sql(self, commands: str, *args, **kwargs) -> None:
             cache=self._get_query_cache(),
             retry_options=self._get_retry_options(),
             state=self._get_state(),
+            transaction_options=self._get_active_tx_options(),
             warning_handler=self._get_warning_handler(),
             annotations=self._get_annotations(),
         ))

gel/base_client.py

@@ -701,8 +701,18 @@ def _get_query_cache(self) -> abstract.QueryCache:
         )
 
     def _get_retry_options(self) -> typing.Optional[_options.RetryOptions]:
+        # This is overloaded in transaction.py to return None, to prevent
+        # retrying *inside* a transaction.
         return self._options.retry_options
 
+    def _get_active_tx_options(self) -> typing.Optional[
+        _options.TransactionOptions
+    ]:
+        # This is overloaded in transaction.py to return None, since
+        # the tx options are applied at the *start* of transactions,
+        # not inside them.
+        return self._options.transaction_options
+
     def _get_state(self) -> _options.State:
         return self._options.state
 

gel/options.py

@@ -54,6 +54,7 @@ class IsolationLevel:
     """Isolation level for transaction"""
     Serializable = "Serializable"
     RepeatableRead = "RepeatableRead"
+    PreferRepeatableRead = "PreferRepeatableRead"
 
     @staticmethod
     def _to_start_tx_str(v):

gel/protocol/protocol.pxd

@@ -90,13 +90,15 @@ cdef class ExecuteContext:
         uint64_t allow_capabilities
         object state
         object annotations
+        object tx_options
 
         # Contextual variables
         readonly bytes cardinality
         readonly BaseCodec in_dc
         readonly BaseCodec out_dc
         readonly uint64_t capabilities
         readonly tuple warnings
+        readonly tuple unsafe_isolation_dangers
 
     cdef inline bint has_na_cardinality(self)
     cdef bint load_from_cache(self)
@@ -168,7 +170,7 @@ cdef class SansIOProtocol:
 
     cdef ensure_connected(self)
 
-    cdef WriteBuffer encode_parse_params(self, ExecuteContext ctx)
+    cdef WriteBuffer encode_parse_params(self, ExecuteContext ctx, dict state)
 
 
 include "protocol_v0.pxd"