proto merge for transaction options

Author: surbhigarg92

google/cloud/spanner_v1/_helpers.py

@@ -32,6 +32,7 @@
 from google.cloud.spanner_v1 import TypeCode
 from google.cloud.spanner_v1 import ExecuteSqlRequest
 from google.cloud.spanner_v1 import JsonObject
+from google.cloud.spanner_v1 import TransactionOptions
 from google.cloud.spanner_v1.request_id_header import with_request_id
 from google.rpc.error_details_pb2 import RetryInfo
 
@@ -690,3 +691,37 @@ def __radd__(self, n):
 
 def _metadata_with_request_id(*args, **kwargs):
     return with_request_id(*args, **kwargs)
+
+
+def _merge_Transaction_Options(
+    defaultTransactionOptions: TransactionOptions,
+    mergeTransactionOptions: TransactionOptions,
+) -> TransactionOptions:
+    """Merges two TransactionOptions objects.
+
+    - Values from `mergeTransactionOptions` take precedence if set.
+    - Values from `defaultTransactionOptions` are used only if missing.
+
+    Args:
+        defaultTransactionOptions (TransactionOptions): The default transaction options (fallback values).
+        mergeTransactionOptions (TransactionOptions): The main transaction options (overrides when set).
+
+    Returns:
+        TransactionOptions: A merged TransactionOptions object.
+    """
+
+    if defaultTransactionOptions is None:
+        return mergeTransactionOptions
+
+    if mergeTransactionOptions is None:
+        return defaultTransactionOptions
+
+    merged_pb = TransactionOptions()._pb  # Create a new protobuf object
+
+    # Merge defaultTransactionOptions first
+    merged_pb.MergeFrom(defaultTransactionOptions._pb)
+
+    # Merge transactionOptions, ensuring it overrides default values
+    merged_pb.MergeFrom(mergeTransactionOptions._pb)
+
+    return merged_pb

google/cloud/spanner_v1/batch.py

@@ -25,6 +25,7 @@
 from google.cloud.spanner_v1._helpers import (
     _metadata_with_prefix,
     _metadata_with_leader_aware_routing,
+    _merge_Transaction_Options,
 )
 from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
 from google.cloud.spanner_v1 import RequestOptions
@@ -216,6 +217,11 @@ def commit(
             exclude_txn_from_change_streams=exclude_txn_from_change_streams,
             isolation_level=isolation_level,
         )
+
+        txn_options = _merge_Transaction_Options(
+            database.default_transaction_options.default_read_write_transaction_options,
+            txn_options,
+        )
         trace_attributes = {"num_mutations": len(self._mutations)}
 
         if request_options is None:

google/cloud/spanner_v1/database.py

@@ -786,7 +786,7 @@ def batch(
         request_options=None,
         max_commit_delay=None,
         exclude_txn_from_change_streams=False,
-        isolation_level=None,
+        isolation_level=TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED,
         **kw,
     ):
         """Return an object which wraps a batch.
@@ -823,9 +823,6 @@ def batch(
         :returns: new wrapper
         """
 
-        # Set isolation level
-        if isolation_level is None:
-            isolation_level = self.default_transaction_options.isolation_level
         return BatchCheckout(
             self,
             request_options,

google/cloud/spanner_v1/session.py

@@ -469,8 +469,6 @@ def run_in_transaction(self, func, *args, **kw):
         )
         isolation_level = kw.pop("isolation_level", None)
 
-        if isolation_level is None:
-            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

@@ -16,6 +16,7 @@
 import functools
 import threading
 from google.protobuf.struct_pb2 import Struct
+from typing import Optional
 
 from google.cloud.spanner_v1._helpers import (
     _make_value_pb,
@@ -24,6 +25,7 @@
     _metadata_with_leader_aware_routing,
     _retry,
     _check_rst_stream_error,
+    _merge_Transaction_Options,
 )
 from google.cloud.spanner_v1 import CommitRequest
 from google.cloud.spanner_v1 import ExecuteBatchDmlRequest
@@ -37,7 +39,7 @@
 from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
 from google.api_core import gapic_v1
 from google.api_core.exceptions import InternalServerError
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 
@@ -90,13 +92,17 @@ def _make_txn_selector(self):
         self._check_state()
 
         if self._transaction_id is None:
-            return TransactionSelector(
-                begin=TransactionOptions(
-                    read_write=TransactionOptions.ReadWrite(),
-                    exclude_txn_from_change_streams=self.exclude_txn_from_change_streams,
-                    isolation_level=self.isolation_level,
-                )
+            txn_options = TransactionOptions(
+                read_write=TransactionOptions.ReadWrite(),
+                exclude_txn_from_change_streams=self.exclude_txn_from_change_streams,
+                isolation_level=self.isolation_level,
+            )
+
+            txn_options = _merge_Transaction_Options(
+                self._session._database.default_transaction_options.default_read_write_transaction_options,
+                txn_options,
             )
+            return TransactionSelector(begin=txn_options)
         else:
             return TransactionSelector(id=self._transaction_id)
 
@@ -164,6 +170,10 @@ def begin(self):
             exclude_txn_from_change_streams=self.exclude_txn_from_change_streams,
             isolation_level=self.isolation_level,
         )
+        txn_options = _merge_Transaction_Options(
+            database.default_transaction_options.default_read_write_transaction_options,
+            txn_options,
+        )
         observability_options = getattr(database, "observability_options", None)
         with trace_call(
             f"CloudSpanner.{type(self).__name__}.begin",
@@ -668,6 +678,18 @@ class BatchTransactionId:
 
 @dataclass
 class DefaultTransactionOptions:
-    isolation_level: TransactionOptions.IsolationLevel = (
-        TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED
+    isolation_level: str = TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED
+    _defaultReadWriteTransactionOptions: Optional[TransactionOptions] = field(
+        init=False, repr=False
     )
+
+    def __post_init__(self):
+        """Initialize _defaultReadWriteTransactionOptions automatically"""
+        self._defaultReadWriteTransactionOptions = TransactionOptions(
+            isolation_level=self.isolation_level
+        )
+
+    @property
+    def default_read_write_transaction_options(self) -> TransactionOptions:
+        """Public accessor for _defaultReadWriteTransactionOptions"""
+        return self._defaultReadWriteTransactionOptions

tests/unit/test__helpers.py

@@ -15,6 +15,7 @@
 
 import unittest
 import mock
+from google.cloud.spanner_v1 import TransactionOptions
 
 
 class Test_merge_query_options(unittest.TestCase):
@@ -955,3 +956,83 @@ def test(self):
         self.assertEqual(
             metadata, ("x-goog-spanner-route-to-leader", str(value).lower())
         )
+
+
+class Test_merge_transaction_options(unittest.TestCase):
+    def _callFUT(self, *args, **kw):
+        from google.cloud.spanner_v1._helpers import _merge_Transaction_Options
+
+        return _merge_Transaction_Options(*args, **kw)
+
+    def test_default_none_and_merge_none(self):
+        default = merge = None
+        result = self._callFUT(default, merge)
+        self.assertIsNone(result)
+
+    def test_default_options_and_merge_none(self):
+        default = TransactionOptions(
+            isolation_level=TransactionOptions.IsolationLevel.REPEATABLE_READ
+        )
+        merge = None
+        result = self._callFUT(default, merge)
+        expected = default
+        self.assertEqual(result, expected)
+
+    def test_default_none_and_merge_options(self):
+        default = None
+        merge = TransactionOptions(
+            isolation_level=TransactionOptions.IsolationLevel.SERIALIZABLE
+        )
+        expected = merge
+        result = self._callFUT(default, merge)
+        self.assertEqual(result, expected)
+
+    def test_default_and_merge_isolation_options(self):
+        default = TransactionOptions(
+            isolation_level=TransactionOptions.IsolationLevel.SERIALIZABLE,
+            read_write=TransactionOptions.ReadWrite(),
+        )
+        merge = TransactionOptions(
+            isolation_level=TransactionOptions.IsolationLevel.REPEATABLE_READ,
+            exclude_txn_from_change_streams=True,
+        )
+        expected = TransactionOptions(
+            isolation_level=TransactionOptions.IsolationLevel.REPEATABLE_READ,
+            read_write=TransactionOptions.ReadWrite(),
+            exclude_txn_from_change_streams=True,
+        )
+        result = self._callFUT(default, merge)
+        self.assertEqual(result, expected)
+
+    def test_default_isolation_and_merge_options(self):
+        default = TransactionOptions(
+            isolation_level=TransactionOptions.IsolationLevel.SERIALIZABLE
+        )
+        merge = TransactionOptions(
+            read_write=TransactionOptions.ReadWrite(),
+            exclude_txn_from_change_streams=True,
+        )
+        expected = TransactionOptions(
+            isolation_level=TransactionOptions.IsolationLevel.SERIALIZABLE,
+            read_write=TransactionOptions.ReadWrite(),
+            exclude_txn_from_change_streams=True,
+        )
+        result = self._callFUT(default, merge)
+        self.assertEqual(result, expected)
+
+    def test_default_isolation_and_merge_options_isolation_unspecified(self):
+        default = TransactionOptions(
+            isolation_level=TransactionOptions.IsolationLevel.SERIALIZABLE
+        )
+        merge = TransactionOptions(
+            read_write=TransactionOptions.ReadWrite(),
+            exclude_txn_from_change_streams=True,
+            isolation_level=TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED,
+        )
+        expected = TransactionOptions(
+            isolation_level=TransactionOptions.IsolationLevel.SERIALIZABLE,
+            read_write=TransactionOptions.ReadWrite(),
+            exclude_txn_from_change_streams=True,
+        )
+        result = self._callFUT(default, merge)
+        self.assertEqual(result, expected)

tests/unit/test_batch.py

@@ -28,6 +28,7 @@
     TransactionOptions,
     Mutation,
     BatchWriteResponse,
+    DefaultTransactionOptions,
 )
 from google.cloud._helpers import UTC, _datetime_to_pb_timestamp
 import datetime
@@ -643,6 +644,7 @@ def __init__(self, enable_end_to_end_tracing=False):
         self._route_to_leader_enabled = True
         if enable_end_to_end_tracing:
             self.observability_options = dict(enable_end_to_end_tracing=True)
+        default_transaction_options = DefaultTransactionOptions()
 
 
 class _FauxSpannerAPI:

tests/unit/test_database.py

@@ -3161,8 +3161,8 @@ def __init__(self, name, instance=None):
         from logging import Logger
 
         self.logger = mock.create_autospec(Logger, instance=True)
-        self._directed_read_options = DefaultTransactionOptions()
-        self.default_transaction_options = None
+        self._directed_read_options = None
+        self.default_transaction_options = DefaultTransactionOptions()
 
 
 class _Pool(object):

tests/unit/test_spanner.py

@@ -1105,6 +1105,7 @@ def __init__(self):
         self._instance = _Instance()
         self._route_to_leader_enabled = True
         self._directed_read_options = None
+        self.default_transaction_options = DefaultTransactionOptions()
 
 
 class _Session(object):

tests/unit/test_transaction.py

@@ -16,6 +16,7 @@
 import mock
 
 from google.cloud.spanner_v1 import RequestOptions
+from google.cloud.spanner_v1 import DefaultTransactionOptions
 from google.cloud.spanner_v1 import Type
 from google.cloud.spanner_v1 import TypeCode
 from google.api_core.retry import Retry
@@ -1021,6 +1022,7 @@ def __init__(self):
         self._instance = _Instance()
         self._route_to_leader_enabled = True
         self._directed_read_options = None
+        self.default_transaction_options = DefaultTransactionOptions()
 
 
 class _Session(object):