dataclass for default transaction options

Author: surbhigarg92

google/cloud/spanner_v1/__init__.py

@@ -64,7 +64,10 @@
 from .types.type import TypeAnnotationCode
 from .types.type import TypeCode
 from .data_types import JsonObject
-from .transaction import BatchTransactionId
+from .transaction import (
+    BatchTransactionId,
+    DefaultTransactionOptions
+)
 
 from google.cloud.spanner_v1 import param_types
 from google.cloud.spanner_v1.client import Client
@@ -149,4 +152,5 @@
     "SpannerClient",
     "SpannerAsyncClient",
     "BatchTransactionId",
+    "DefaultTransactionOptions"
 )

google/cloud/spanner_v1/client.py

@@ -31,6 +31,7 @@
 from google.auth.credentials import AnonymousCredentials
 import google.api_core.client_options
 from google.cloud.client import ClientWithProject
+from typing import Optional
 
 
 from google.cloud.spanner_admin_database_v1 import DatabaseAdminClient
@@ -45,6 +46,7 @@
 from google.cloud.spanner_admin_instance_v1 import ListInstancesRequest
 from google.cloud.spanner_v1 import __version__
 from google.cloud.spanner_v1 import ExecuteSqlRequest
+from google.cloud.spanner_v1 import DefaultTransactionOptions
 from google.cloud.spanner_v1._helpers import _merge_query_options
 from google.cloud.spanner_v1._helpers import _metadata_with_prefix
 from google.cloud.spanner_v1.instance import Instance
@@ -161,11 +163,10 @@ 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`
+    :type default_transaction_options: :class:`~google.cloud.spanner_v1.DefaultTransactionOptions`
         or :class:`dict`
-    :param default_transaction_options: (Optional) Default options to use for all read/write transactions.
-        Any fields other than `isolation_level` will be ignored.
-
+    :param default_transaction_options: (Optional) Default options to use for all transactions.
+    
     :raises: :class:`ValueError <exceptions.ValueError>` if both ``read_only``
              and ``admin`` are :data:`True`
     """
@@ -187,7 +188,7 @@ def __init__(
         route_to_leader_enabled=True,
         directed_read_options=None,
         observability_options=None,
-        default_transaction_options=None,
+        default_transaction_options: Optional[DefaultTransactionOptions] =None,
     ):
         self._emulator_host = _get_spanner_emulator_host()
 
@@ -249,6 +250,11 @@ def __init__(
         self._route_to_leader_enabled = route_to_leader_enabled
         self._directed_read_options = directed_read_options
         self._observability_options = observability_options
+        if default_transaction_options is None:
+            default_transaction_options = DefaultTransactionOptions()
+        elif not isinstance(default_transaction_options, DefaultTransactionOptions):
+            raise TypeError("default_transaction_options must be an instance of DefaultTransactionOptions")
+
         self._default_transaction_options = default_transaction_options
 
     @property
@@ -345,9 +351,9 @@ def default_transaction_options(self):
         """Getter for default_transaction_options.
 
         :rtype:
-            :class:`~google.cloud.spanner_v1.TransactionOptions`
+            :class:`~google.cloud.spanner_v1.DefaultTransactionOptions`
             or :class:`dict`
-        :returns: The default transaction options that are used by this client for all read/write transactions.
+        :returns: The default transaction options that are used by this client for all transactions.
         """
         return self._default_transaction_options
 
@@ -498,11 +504,15 @@ def directed_read_options(self, directed_read_options):
         self._directed_read_options = directed_read_options
 
     @default_transaction_options.setter
-    def default_transaction_options(self, default_transaction_options):
+    def default_transaction_options(self, default_transaction_options:DefaultTransactionOptions):
         """Sets default_transaction_options for the client
-        :type default_transaction_options: :class:`~google.cloud.spanner_v1.TransactionOptions`
+        :type default_transaction_options: :class:`~google.cloud.spanner_v1.DefaultTransactionOptions`
             or :class:`dict`
-        :param default_transaction_options: Default options to use for all read/write transactions.
-            Any fields other than `isolation_level` will be ignored.
+        :param default_transaction_options: Default options to use for transactions.
         """
+        if default_transaction_options is None:
+            default_transaction_options = DefaultTransactionOptions()
+        elif not isinstance(default_transaction_options, DefaultTransactionOptions):
+            raise TypeError("default_transaction_options must be an instance of DefaultTransactionOptions")
+
         self._default_transaction_options = default_transaction_options

google/cloud/spanner_v1/database.py

@@ -46,6 +46,7 @@
 from google.cloud.spanner_v1 import TypeCode
 from google.cloud.spanner_v1 import TransactionSelector
 from google.cloud.spanner_v1 import TransactionOptions
+from google.cloud.spanner_v1 import DefaultTransactionOptions
 from google.cloud.spanner_v1 import RequestOptions
 from google.cloud.spanner_v1 import SpannerClient
 from google.cloud.spanner_v1._helpers import _merge_query_options
@@ -183,7 +184,7 @@ 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.default_transaction_options: DefaultTransactionOptions = (
             self._instance._client.default_transaction_options
         )
         self._proto_descriptors = proto_descriptors
@@ -823,7 +824,7 @@ def batch(
 
         # Set isolation level
         if isolation_level is None:
-            isolation_level = self.get_default_isolation_level()
+            isolation_level = self.default_transaction_options.isolation_level
         return BatchCheckout(
             self,
             request_options,
@@ -1144,23 +1145,6 @@ def set_iam_policy(self, policy):
         response = api.set_iam_policy(request=request, metadata=metadata)
         return response
 
-    def get_default_isolation_level(self):
-        """
-        Returns the isolation level set in default transaction options when creating
-        the SpannerClient.
-        """
-        if isinstance(self.default_transaction_options, dict):
-            return self.default_transaction_options.get(
-                "isolation_level",
-                TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED,
-            )
-
-        return getattr(
-            self.default_transaction_options,
-            "isolation_level",
-            TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED,
-        )
-
     @property
     def observability_options(self):
         """

google/cloud/spanner_v1/session.py

@@ -465,7 +465,7 @@ def run_in_transaction(self, func, *args, **kw):
         isolation_level = kw.pop("isolation_level", None)
 
         if isolation_level is None:
-            isolation_level = self._database.get_default_isolation_level()
+            isolation_level = self._database.default_transaction_options.isolation_level
         attempts = 0
 
         observability_options = getattr(self._database, "observability_options", None)

google/cloud/spanner_v1/transaction.py

@@ -655,3 +655,7 @@ class BatchTransactionId:
     transaction_id: str
     session_id: str
     read_timestamp: Any
+
+@dataclass
+class DefaultTransactionOptions:
+    isolation_level: TransactionOptions.IsolationLevel = TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED

tests/unit/test_client.py

@@ -16,7 +16,11 @@
 
 import os
 import mock
-from google.cloud.spanner_v1 import DirectedReadOptions
+from google.cloud.spanner_v1 import (
+    DirectedReadOptions,
+    DefaultTransactionOptions
+)
+
 
 
 def _make_credentials():
@@ -53,7 +57,7 @@ class TestClient(unittest.TestCase):
             "auto_failover_disabled": True,
         },
     }
-    DEFAULT_TRANSACTION_OPTIONS = {"isolation_level": "SERIALIZABLE"}
+    DEFAULT_TRANSACTION_OPTIONS = DefaultTransactionOptions(isolation_level="SERIALIZABLE")
 
     def _get_target_class(self):
         from google.cloud import spanner

tests/unit/test_database.py

@@ -25,7 +25,7 @@
 from google.api_core.retry import Retry
 from google.protobuf.field_mask_pb2 import FieldMask
 
-from google.cloud.spanner_v1 import RequestOptions, DirectedReadOptions
+from google.cloud.spanner_v1 import RequestOptions, DirectedReadOptions, DefaultTransactionOptions
 
 DML_WO_PARAM = """
 DELETE FROM citizens
@@ -3116,7 +3116,7 @@ def __init__(
         project=TestDatabase.PROJECT_ID,
         route_to_leader_enabled=True,
         directed_read_options=None,
-        default_transaction_options=None,
+        default_transaction_options=DefaultTransactionOptions(),
     ):
         from google.cloud.spanner_v1 import ExecuteSqlRequest
 
@@ -3157,7 +3157,7 @@ def __init__(self, name, instance=None):
         from logging import Logger
 
         self.logger = mock.create_autospec(Logger, instance=True)
-        self._directed_read_options = None
+        self._directed_read_options = DefaultTransactionOptions()
         self.default_transaction_options = None
 
 

tests/unit/test_instance.py

@@ -14,6 +14,7 @@
 
 import unittest
 import mock
+from google.cloud.spanner_v1 import DefaultTransactionOptions
 
 
 class TestInstance(unittest.TestCase):
@@ -1019,7 +1020,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
+        self.default_transaction_options = DefaultTransactionOptions()
 
     def copy(self):
         from copy import deepcopy

tests/unit/test_session.py

@@ -48,6 +48,7 @@
 from google.api_core.exceptions import Unknown, Aborted, NotFound, Cancelled
 from google.protobuf.struct_pb2 import Struct, Value
 from google.cloud.spanner_v1.batch import Batch
+from google.cloud.spanner_v1 import DefaultTransactionOptions
 
 
 def _make_rpc_error(error_cls, trailing_metadata=None):
@@ -84,7 +85,7 @@ def _make_one(self, *args, **kwargs):
 
     @staticmethod
     def _make_database(
-        name=DATABASE_NAME, database_role=None, default_transaction_options=None
+        name=DATABASE_NAME, database_role=None, default_transaction_options=DefaultTransactionOptions()
     ):
         database = mock.create_autospec(Database, instance=True)
         database.name = name
@@ -93,25 +94,6 @@ def _make_database(
         database._route_to_leader_enabled = True
         database.default_transaction_options = default_transaction_options
 
-        # Define side_effect function to use injected default_transaction_options
-        def get_default_isolation_level_side_effect():
-            if isinstance(database.default_transaction_options, dict):
-                return database.default_transaction_options.get(
-                    "isolation_level",
-                    TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED,
-                )
-
-            return getattr(
-                database.default_transaction_options,
-                "isolation_level",
-                TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED,
-            )
-
-        # Mock get_default_isolation_level method
-        database.get_default_isolation_level = mock.Mock(
-            side_effect=get_default_isolation_level_side_effect
-        )
-
         return database
 
     @staticmethod
@@ -1794,7 +1776,7 @@ def test_run_in_transaction_w_isolation_level_at_client(self):
         gax_api = self._make_spanner_api()
         gax_api.begin_transaction.return_value = TransactionPB(id=b"FACEDACE")
         database = self._make_database(
-            default_transaction_options={"isolation_level": "SERIALIZABLE"}
+            default_transaction_options=DefaultTransactionOptions(isolation_level="SERIALIZABLE")
         )
         database.spanner_api = gax_api
         session = self._make_one(database)

tests/unit/test_spanner.py

@@ -32,6 +32,7 @@
     ExecuteBatchDmlRequest,
     ExecuteBatchDmlResponse,
     param_types,
+    DefaultTransactionOptions,
 )
 from google.cloud.spanner_v1.types import transaction as transaction_type
 from google.cloud.spanner_v1.keyset import KeySet
@@ -1090,7 +1091,7 @@ def __init__(self):
 
         self._query_options = ExecuteSqlRequest.QueryOptions(optimizer_version="1")
         self.directed_read_options = None
-        self.default_transaction_options = None
+        self.default_transaction_options = DefaultTransactionOptions()
 
 
 class _Instance(object):