feat: snapshot isolation

Author: surbhigarg92

google/cloud/spanner_v1/batch.py

@@ -166,6 +166,7 @@ def commit(
         request_options=None,
         max_commit_delay=None,
         exclude_txn_from_change_streams=False,
+        isolation_level=TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED,
         **kwargs,
     ):
         """Commit mutations to the database.
@@ -186,6 +187,18 @@ def commit(
                 (Optional) The amount of latency this request is willing to incur
                 in order to improve throughput.
 
+        :type exclude_txn_from_change_streams: bool
+        :param exclude_txn_from_change_streams:
+          (Optional) If true, instructs the transaction to be excluded from being recorded in change streams
+          with the DDL option `allow_txn_exclusion=true`. This does not exclude the transaction from
+          being recorded in the change streams with the DDL option `allow_txn_exclusion` being false or
+          unset.
+
+        :type isolation_level:
+            :class:`google.cloud.spanner_v1.types.TransactionOptions.IsolationLevel`
+        :param isolation_level:
+                (Optional) Sets isolation level for the transaction.
+
         :rtype: datetime
         :returns: timestamp of the committed changes.
         """
@@ -200,6 +213,7 @@ def commit(
         txn_options = TransactionOptions(
             read_write=TransactionOptions.ReadWrite(),
             exclude_txn_from_change_streams=exclude_txn_from_change_streams,
+            isolation_level=isolation_level,
         )
         trace_attributes = {"num_mutations": len(self._mutations)}
 

google/cloud/spanner_v1/client.py

@@ -41,6 +41,7 @@
 from google.cloud.spanner_admin_instance_v1.services.instance_admin.transports.grpc import (
     InstanceAdminGrpcTransport,
 )
+from google.cloud.spanner_v1 import TransactionOptions
 from google.cloud.spanner_admin_instance_v1 import ListInstanceConfigsRequest
 from google.cloud.spanner_admin_instance_v1 import ListInstancesRequest
 from google.cloud.spanner_v1 import __version__
@@ -136,6 +137,11 @@ class Client(ClientWithProject):
            or you can use the environment variable `SPANNER_ENABLE_EXTENDED_TRACING=<boolean>`
            to control it.
 
+    :type default_transaction_options: :class:`~google.cloud.spanner_v1.TransactionOptions`
+        or :class:`dict`
+    :param default_transaction_options: (Optional) Client options used to set the default_transaction_options
+            used for all Read Write Transactions.
+
     :raises: :class:`ValueError <exceptions.ValueError>` if both ``read_only``
              and ``admin`` are :data:`True`
     """
@@ -157,6 +163,7 @@ def __init__(
         route_to_leader_enabled=True,
         directed_read_options=None,
         observability_options=None,
+        default_transaction_options=None,
     ):
         self._emulator_host = _get_spanner_emulator_host()
 
@@ -199,6 +206,7 @@ def __init__(
         self._route_to_leader_enabled = route_to_leader_enabled
         self._directed_read_options = directed_read_options
         self._observability_options = observability_options
+        self._default_transaction_options = default_transaction_options
 
     @property
     def credentials(self):
@@ -289,6 +297,17 @@ def observability_options(self):
         """
         return self._observability_options
 
+    @property
+    def default_transaction_options(self):
+        """Getter for default_transaction_options.
+
+        :rtype:
+            :class:`~google.cloud.spanner_v1.TransactionOptions`
+            or :class:`dict`
+        :returns: The default_transaction_options for the client.
+        """
+        return self._default_transaction_options
+
     @property
     def directed_read_options(self):
         """Getter for directed_read_options.
@@ -434,3 +453,19 @@ def directed_read_options(self, directed_read_options):
             or regions should be used for non-transactional reads or queries.
         """
         self._directed_read_options = directed_read_options
+
+    @default_transaction_options.setter
+    def default_transaction_options(self, default_transaction_options):
+        """Sets default_transaction_options for the client
+        :type default_transaction_options: :class:`~google.cloud.spanner_v1.TransactionOptions`
+            or :class:`dict`
+        :param default_transaction_options: Client options used to set the default_transaction_options
+            used for all Read Write Transactions.
+        """
+        if default_transaction_options is not None and not isinstance(
+            default_transaction_options, TransactionOptions
+        ):
+            raise TypeError(
+                "default_transaction_option must be a TransactionOptions instance"
+            )
+        self._default_transaction_options = default_transaction_options

google/cloud/spanner_v1/database.py

@@ -182,6 +182,9 @@ def __init__(
         self._enable_drop_protection = enable_drop_protection
         self._reconciling = False
         self._directed_read_options = self._instance._client.directed_read_options
+        self.default_transaction_options = (
+            self._instance._client.default_transaction_options
+        )
         self._proto_descriptors = proto_descriptors
 
         if pool is None:
@@ -780,6 +783,7 @@ def batch(
         request_options=None,
         max_commit_delay=None,
         exclude_txn_from_change_streams=False,
+        isolation_level=None,
         **kw,
     ):
         """Return an object which wraps a batch.
@@ -807,14 +811,28 @@ def batch(
           being recorded in the change streams with the DDL option `allow_txn_exclusion` being false or
           unset.
 
+        :type isolation_level:
+            :class:`google.cloud.spanner_v1.types.TransactionOptions.IsolationLevel`
+        :param isolation_level:
+                (Optional) Sets isolation level for the transaction and overrides the isolation level set at the client.
+
         :rtype: :class:`~google.cloud.spanner_v1.database.BatchCheckout`
         :returns: new wrapper
         """
+
+        # Set isolation level
+        if isolation_level is None:
+            isolation_level = (
+                self.default_transaction_options.isolation_level
+                if isinstance(self.default_transaction_options, TransactionOptions)
+                else TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED
+            )
         return BatchCheckout(
             self,
             request_options,
             max_commit_delay,
             exclude_txn_from_change_streams,
+            isolation_level,
             **kw,
         )
 
@@ -886,6 +904,7 @@ def run_in_transaction(self, func, *args, **kw):
                    from being recorded in change streams with the DDL option `allow_txn_exclusion=true`.
                    This does not exclude the transaction from being recorded in the change streams with
                    the DDL option `allow_txn_exclusion` being false or unset.
+                   "isolation_level" sets the isolation level for transaction.
 
         :rtype: Any
         :returns: The return value of ``func``.
@@ -1176,6 +1195,7 @@ def __init__(
         request_options=None,
         max_commit_delay=None,
         exclude_txn_from_change_streams=False,
+        isolation_level=TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED,
         **kw,
     ):
         self._database = database
@@ -1188,6 +1208,7 @@ def __init__(
             self._request_options = request_options
         self._max_commit_delay = max_commit_delay
         self._exclude_txn_from_change_streams = exclude_txn_from_change_streams
+        self._isolation_level = isolation_level
         self._kw = kw
 
     def __enter__(self):
@@ -1209,7 +1230,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                     request_options=self._request_options,
                     max_commit_delay=self._max_commit_delay,
                     exclude_txn_from_change_streams=self._exclude_txn_from_change_streams,
-                    **self._kw,
+                    isolation_level=self._isolation_level,
                 )
         finally:
             if self._database.log_commit_stats and self._batch.commit_stats:

google/cloud/spanner_v1/session.py

@@ -39,6 +39,7 @@
 from google.cloud.spanner_v1.batch import Batch
 from google.cloud.spanner_v1.snapshot import Snapshot
 from google.cloud.spanner_v1.transaction import Transaction
+from google.cloud.spanner_v1 import TransactionOptions
 
 
 DEFAULT_RETRY_TIMEOUT_SECS = 30
@@ -446,6 +447,7 @@ def run_in_transaction(self, func, *args, **kw):
                    from being recorded in change streams with the DDL option `allow_txn_exclusion=true`.
                    This does not exclude the transaction from being recorded in the change streams with
                    the DDL option `allow_txn_exclusion` being false or unset.
+                   "isolation_level" sets the isolation level for transaction.
 
         :rtype: Any
         :returns: The return value of ``func``.
@@ -460,6 +462,16 @@ def run_in_transaction(self, func, *args, **kw):
         exclude_txn_from_change_streams = kw.pop(
             "exclude_txn_from_change_streams", None
         )
+        isolation_level = kw.pop("isolation_level", None)
+
+        if isolation_level is None:
+            isolation_level = (
+                self._database.default_transaction_options.isolation_level
+                if isinstance(
+                    self._database.default_transaction_options, TransactionOptions
+                )
+                else TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED
+            )
         attempts = 0
 
         observability_options = getattr(self._database, "observability_options", None)
@@ -475,6 +487,7 @@ def run_in_transaction(self, func, *args, **kw):
                     txn.exclude_txn_from_change_streams = (
                         exclude_txn_from_change_streams
                     )
+                    txn.isolation_level = isolation_level
                 else:
                     txn = self._transaction
 

google/cloud/spanner_v1/transaction.py

@@ -58,6 +58,7 @@ class Transaction(_SnapshotBase, _BatchBase):
     _lock = threading.Lock()
     _read_only = False
     exclude_txn_from_change_streams = False
+    isolation_level = TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED
 
     def __init__(self, session):
         if session._transaction is not None:
@@ -92,6 +93,7 @@ def _make_txn_selector(self):
                 begin=TransactionOptions(
                     read_write=TransactionOptions.ReadWrite(),
                     exclude_txn_from_change_streams=self.exclude_txn_from_change_streams,
+                    isolation_level=self.isolation_level,
                 )
             )
         else:
@@ -154,6 +156,7 @@ def begin(self):
         txn_options = TransactionOptions(
             read_write=TransactionOptions.ReadWrite(),
             exclude_txn_from_change_streams=self.exclude_txn_from_change_streams,
+            isolation_level=self.isolation_level,
         )
         observability_options = getattr(database, "observability_options", None)
         with trace_call(

tests/unit/test_database.py

@@ -3116,6 +3116,7 @@ def __init__(
         project=TestDatabase.PROJECT_ID,
         route_to_leader_enabled=True,
         directed_read_options=None,
+        default_transaction_options=None,
     ):
         from google.cloud.spanner_v1 import ExecuteSqlRequest
 
@@ -3129,6 +3130,7 @@ def __init__(
         self._query_options = ExecuteSqlRequest.QueryOptions(optimizer_version="1")
         self.route_to_leader_enabled = route_to_leader_enabled
         self.directed_read_options = directed_read_options
+        self.default_transaction_options = default_transaction_options
 
 
 class _Instance(object):
@@ -3156,6 +3158,7 @@ def __init__(self, name, instance=None):
 
         self.logger = mock.create_autospec(Logger, instance=True)
         self._directed_read_options = None
+        self.default_transaction_options = None
 
 
 class _Pool(object):

tests/unit/test_instance.py

@@ -1019,6 +1019,7 @@ def __init__(self, project, timeout_seconds=None):
         self.timeout_seconds = timeout_seconds
         self.route_to_leader_enabled = True
         self.directed_read_options = None
+        self.default_transaction_options = None
 
     def copy(self):
         from copy import deepcopy

tests/unit/test_session.py

@@ -70,6 +70,7 @@ def _make_database(name=DATABASE_NAME, database_role=None):
         database.log_commit_stats = False
         database.database_role = database_role
         database._route_to_leader_enabled = True
+        database.default_transaction_options = None
         return database
 
     @staticmethod