test and refactoring

Author: surbhigarg92

google/cloud/spanner_v1/client.py

@@ -462,10 +462,4 @@ def default_transaction_options(self, default_transaction_options):
         :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

@@ -823,7 +823,12 @@ def batch(
         # Set isolation level
         if isolation_level is None:
             isolation_level = (
-                self.default_transaction_options.isolation_level
+                self.default_transaction_options.get(
+                    "isolation_level",
+                    TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED,
+                )
+                if isinstance(self.default_transaction_options, dict)
+                else self.default_transaction_options.isolation_level
                 if isinstance(self.default_transaction_options, TransactionOptions)
                 else TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED
             )

google/cloud/spanner_v1/session.py

@@ -466,7 +466,12 @@ def run_in_transaction(self, func, *args, **kw):
 
         if isolation_level is None:
             isolation_level = (
-                self._database.default_transaction_options.isolation_level
+                self._database.default_transaction_options.get(
+                    "isolation_level",
+                    TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED,
+                )
+                if isinstance(self._database.default_transaction_options, dict)
+                else self._database.default_transaction_options.isolation_level
                 if isinstance(
                     self._database.default_transaction_options, TransactionOptions
                 )

tests/unit/test_batch.py

@@ -21,7 +21,20 @@
     StatusCode,
     enrich_with_otel_scope,
 )
-from google.cloud.spanner_v1 import RequestOptions
+from google.cloud.spanner_v1 import (
+    RequestOptions,
+    CommitResponse,
+    TransactionOptions,
+    Mutation,
+    BatchWriteResponse,
+)
+from google.cloud._helpers import UTC, _datetime_to_pb_timestamp
+import datetime
+from google.api_core.exceptions import Aborted, Unknown
+from google.cloud.spanner_v1.batch import MutationGroups, _BatchBase, Batch
+from google.cloud.spanner_v1.keyset import KeySet
+from google.rpc.status_pb2 import Status
+
 
 TABLE_NAME = "citizens"
 COLUMNS = ["email", "first_name", "last_name", "age"]
@@ -57,8 +70,6 @@ def _make_one(self, *args, **kwargs):
 
 class Test_BatchBase(_BaseTest):
     def _getTargetClass(self):
-        from google.cloud.spanner_v1.batch import _BatchBase
-
         return _BatchBase
 
     def _compare_values(self, result, source):
@@ -83,8 +94,6 @@ def test__check_state_virtual(self):
             base._check_state()
 
     def test_insert(self):
-        from google.cloud.spanner_v1 import Mutation
-
         session = _Session()
         base = self._make_one(session)
 
@@ -100,8 +109,6 @@ def test_insert(self):
         self._compare_values(write.values, VALUES)
 
     def test_update(self):
-        from google.cloud.spanner_v1 import Mutation
-
         session = _Session()
         base = self._make_one(session)
 
@@ -117,8 +124,6 @@ def test_update(self):
         self._compare_values(write.values, VALUES)
 
     def test_insert_or_update(self):
-        from google.cloud.spanner_v1 import Mutation
-
         session = _Session()
         base = self._make_one(session)
 
@@ -134,8 +139,6 @@ def test_insert_or_update(self):
         self._compare_values(write.values, VALUES)
 
     def test_replace(self):
-        from google.cloud.spanner_v1 import Mutation
-
         session = _Session()
         base = self._make_one(session)
 
@@ -151,9 +154,6 @@ def test_replace(self):
         self._compare_values(write.values, VALUES)
 
     def test_delete(self):
-        from google.cloud.spanner_v1 import Mutation
-        from google.cloud.spanner_v1.keyset import KeySet
-
         keys = [[0], [1], [2]]
         keyset = KeySet(keys=keys)
         session = _Session()
@@ -176,8 +176,6 @@ def test_delete(self):
 
 class TestBatch(_BaseTest, OpenTelemetryBase):
     def _getTargetClass(self):
-        from google.cloud.spanner_v1.batch import Batch
-
         return Batch
 
     def test_ctor(self):
@@ -186,8 +184,6 @@ def test_ctor(self):
         self.assertIs(batch._session, session)
 
     def test_commit_already_committed(self):
-        from google.cloud.spanner_v1.keyset import KeySet
-
         keys = [[0], [1], [2]]
         keyset = KeySet(keys=keys)
         database = _Database()
@@ -202,9 +198,6 @@ def test_commit_already_committed(self):
         self.assertNoSpans()
 
     def test_commit_grpc_error(self):
-        from google.api_core.exceptions import Unknown
-        from google.cloud.spanner_v1.keyset import KeySet
-
         keys = [[0], [1], [2]]
         keyset = KeySet(keys=keys)
         database = _Database()
@@ -223,12 +216,6 @@ def test_commit_grpc_error(self):
         )
 
     def test_commit_ok(self):
-        import datetime
-        from google.cloud.spanner_v1 import CommitResponse
-        from google.cloud.spanner_v1 import TransactionOptions
-        from google.cloud._helpers import UTC
-        from google.cloud._helpers import _datetime_to_pb_timestamp
-
         now = datetime.datetime.utcnow().replace(tzinfo=UTC)
         now_pb = _datetime_to_pb_timestamp(now)
         response = CommitResponse(commit_timestamp=now_pb)
@@ -273,7 +260,6 @@ def test_commit_ok(self):
     def test_aborted_exception_on_commit_with_retries(self):
         # Test case to verify that an Aborted exception is raised when
         # batch.commit() is called and the transaction is aborted internally.
-        from google.api_core.exceptions import Aborted
 
         database = _Database()
         # Setup the spanner API which throws Aborted exception when calling commit API.
@@ -306,13 +292,8 @@ def _test_commit_with_options(
         request_options=None,
         max_commit_delay_in=None,
         exclude_txn_from_change_streams=False,
+        isolation_level=None,
     ):
-        import datetime
-        from google.cloud.spanner_v1 import CommitResponse
-        from google.cloud.spanner_v1 import TransactionOptions
-        from google.cloud._helpers import UTC
-        from google.cloud._helpers import _datetime_to_pb_timestamp
-
         now = datetime.datetime.utcnow().replace(tzinfo=UTC)
         now_pb = _datetime_to_pb_timestamp(now)
         response = CommitResponse(commit_timestamp=now_pb)
@@ -399,8 +380,6 @@ def test_commit_w_incorrect_tag_dictionary_error(self):
             self._test_commit_with_options(request_options=request_options)
 
     def test_commit_w_max_commit_delay(self):
-        import datetime
-
         request_options = RequestOptions(
             request_tag="tag-1",
         )
@@ -418,9 +397,6 @@ def test_commit_w_exclude_txn_from_change_streams(self):
         )
 
     def test_context_mgr_already_committed(self):
-        import datetime
-        from google.cloud._helpers import UTC
-
         now = datetime.datetime.utcnow().replace(tzinfo=UTC)
         database = _Database()
         api = database.spanner_api = _FauxSpannerAPI()
@@ -435,12 +411,6 @@ def test_context_mgr_already_committed(self):
         self.assertEqual(api._committed, None)
 
     def test_context_mgr_success(self):
-        import datetime
-        from google.cloud.spanner_v1 import CommitResponse
-        from google.cloud.spanner_v1 import TransactionOptions
-        from google.cloud._helpers import UTC
-        from google.cloud._helpers import _datetime_to_pb_timestamp
-
         now = datetime.datetime.utcnow().replace(tzinfo=UTC)
         now_pb = _datetime_to_pb_timestamp(now)
         response = CommitResponse(commit_timestamp=now_pb)
@@ -481,11 +451,6 @@ def test_context_mgr_success(self):
         )
 
     def test_context_mgr_failure(self):
-        import datetime
-        from google.cloud.spanner_v1 import CommitResponse
-        from google.cloud._helpers import UTC
-        from google.cloud._helpers import _datetime_to_pb_timestamp
-
         now = datetime.datetime.utcnow().replace(tzinfo=UTC)
         now_pb = _datetime_to_pb_timestamp(now)
         response = CommitResponse(commit_timestamp=now_pb)
@@ -509,8 +474,6 @@ class _BailOut(Exception):
 
 class TestMutationGroups(_BaseTest, OpenTelemetryBase):
     def _getTargetClass(self):
-        from google.cloud.spanner_v1.batch import MutationGroups
-
         return MutationGroups
 
     def test_ctor(self):
@@ -519,8 +482,6 @@ def test_ctor(self):
         self.assertIs(groups._session, session)
 
     def test_batch_write_already_committed(self):
-        from google.cloud.spanner_v1.keyset import KeySet
-
         keys = [[0], [1], [2]]
         keyset = KeySet(keys=keys)
         database = _Database()
@@ -541,9 +502,6 @@ def test_batch_write_already_committed(self):
             groups.batch_write()
 
     def test_batch_write_grpc_error(self):
-        from google.api_core.exceptions import Unknown
-        from google.cloud.spanner_v1.keyset import KeySet
-
         keys = [[0], [1], [2]]
         keyset = KeySet(keys=keys)
         database = _Database()
@@ -565,12 +523,6 @@ def test_batch_write_grpc_error(self):
     def _test_batch_write_with_request_options(
         self, request_options=None, exclude_txn_from_change_streams=False
     ):
-        import datetime
-        from google.cloud.spanner_v1 import BatchWriteResponse
-        from google.cloud._helpers import UTC
-        from google.cloud._helpers import _datetime_to_pb_timestamp
-        from google.rpc.status_pb2 import Status
-
         now = datetime.datetime.utcnow().replace(tzinfo=UTC)
         now_pb = _datetime_to_pb_timestamp(now)
         status_pb = Status(code=200)
@@ -676,9 +628,6 @@ def commit(
         request=None,
         metadata=None,
     ):
-        from google.api_core.exceptions import Unknown
-        from google.api_core.exceptions import Aborted
-
         max_commit_delay = None
         if type(request).pb(request).HasField("max_commit_delay"):
             max_commit_delay = request.max_commit_delay
@@ -703,8 +652,6 @@ def batch_write(
         request=None,
         metadata=None,
     ):
-        from google.api_core.exceptions import Unknown
-
         self._batch_request = (
             request.session,
             request.mutation_groups,

tests/unit/test_client.py

@@ -52,6 +52,7 @@ class TestClient(unittest.TestCase):
             "auto_failover_disabled": True,
         },
     }
+    DEFAULT_TRANSACTION_OPTIONS = {"isolation_level": "SERIALIZABLE"}
 
     def _get_target_class(self):
         from google.cloud import spanner
@@ -72,6 +73,7 @@ def _constructor_test_helper(
         expected_query_options=None,
         route_to_leader_enabled=True,
         directed_read_options=None,
+        default_transaction_options=None,
     ):
         import google.api_core.client_options
         from google.cloud.spanner_v1 import client as MUT
@@ -98,6 +100,7 @@ def _constructor_test_helper(
             credentials=creds,
             query_options=query_options,
             directed_read_options=directed_read_options,
+            default_transaction_options=default_transaction_options,
             **kwargs
         )
 
@@ -128,6 +131,10 @@ def _constructor_test_helper(
             self.assertFalse(client.route_to_leader_enabled)
         if directed_read_options is not None:
             self.assertEqual(client.directed_read_options, directed_read_options)
+        if default_transaction_options is not None:
+            self.assertEqual(
+                client.default_transaction_options, default_transaction_options
+            )
 
     @mock.patch("google.cloud.spanner_v1.client._get_spanner_emulator_host")
     @mock.patch("warnings.warn")
@@ -259,6 +266,17 @@ def test_constructor_route_to_leader_disbled(self):
             expected_scopes, creds, route_to_leader_enabled=False
         )
 
+    def test_constructor_w_default_transaction_options(self):
+        from google.cloud.spanner_v1 import client as MUT
+
+        expected_scopes = (MUT.SPANNER_ADMIN_SCOPE,)
+        creds = _make_credentials()
+        self._constructor_test_helper(
+            expected_scopes,
+            creds,
+            default_transaction_options=self.DEFAULT_TRANSACTION_OPTIONS,
+        )
+
     @mock.patch("google.cloud.spanner_v1.client._get_spanner_emulator_host")
     def test_instance_admin_api(self, mock_em):
         from google.cloud.spanner_v1.client import SPANNER_ADMIN_SCOPE

tests/unit/test_session.py

@@ -1915,6 +1915,44 @@ def unit_of_work(txn, *args, **kw):
             * 2,
         )
 
+    def test_run_in_transaction_w_isolation_level(self):
+        from google.cloud.spanner_v1 import (
+            Transaction as TransactionPB,
+            TransactionOptions,
+        )
+        from google.cloud.spanner_v1.transaction import Transaction
+
+        gax_api = self._make_spanner_api()
+        gax_api.begin_transaction.return_value = TransactionPB(id=b"FACEDACE")
+        database = self._make_database()
+        database.spanner_api = gax_api
+        session = self._make_one(database)
+        session._session_id = self.SESSION_ID
+
+        def unit_of_work(txn, *args, **kw):
+            txn.insert("test", [], [])
+            return 42
+
+        return_value = session.run_in_transaction(
+            unit_of_work, "abc", isolation_level="SERIALIZABLE"
+        )
+
+        self.assertIsNone(session._transaction)
+        self.assertEqual(return_value, 42)
+
+        expected_options = TransactionOptions(
+            read_write=TransactionOptions.ReadWrite(),
+            isolation_level="SERIALIZABLE",
+        )
+        gax_api.begin_transaction.assert_called_once_with(
+            session=self.SESSION_NAME,
+            options=expected_options,
+            metadata=[
+                ("google-cloud-resource-prefix", database.name),
+                ("x-goog-spanner-route-to-leader", "true"),
+            ],
+        )
+
     def test_delay_helper_w_no_delay(self):
         from google.cloud.spanner_v1._helpers import _delay_until_retry