feat: Add support for max commit delay

google/cloud/spanner_v1/batch.py

@@ -31,6 +31,7 @@
 from google.cloud.spanner_v1._helpers import _retry
 from google.cloud.spanner_v1._helpers import _check_rst_stream_error
 from google.api_core.exceptions import InternalServerError
+from google.protobuf.duration_pb2 import Duration
 
 
 class _BatchBase(_SessionWrapper):
@@ -146,7 +147,9 @@ def _check_state(self):
         if self.committed is not None:
             raise ValueError("Batch already committed")
 
-    def commit(self, return_commit_stats=False, request_options=None):
+    def commit(
+        self, return_commit_stats=False, request_options=None, max_commit_delay_ms=None
+    ):
         """Commit mutations to the database.
 
         :type return_commit_stats: bool
@@ -159,6 +162,9 @@ def commit(self, return_commit_stats=False, request_options=None):
                 (Optional) Common options for this request.
                 If a dict is provided, it must be of the same form as the protobuf
                 message :class:`~google.cloud.spanner_v1.types.RequestOptions`.
+        :param max_commit_delay_ms:
+                (Optional) The amount of latency this request is willing to incur
+                in order to improve throughput.
 
         :rtype: datetime
         :returns: timestamp of the committed changes.
@@ -183,11 +189,17 @@ def commit(self, return_commit_stats=False, request_options=None):
         # Request tags are not supported for commit requests.
         request_options.request_tag = None
 
+        max_commit_delay = None
+        if max_commit_delay_ms:
+            max_commit_delay = Duration()
+            max_commit_delay.FromMilliseconds(max_commit_delay_ms)
+
         request = CommitRequest(
             session=self._session.name,
             mutations=self._mutations,
             single_use_transaction=txn_options,
             return_commit_stats=return_commit_stats,
+            max_commit_delay=max_commit_delay,
             request_options=request_options,
         )
         with trace_call("CloudSpanner.Commit", self._session, trace_attributes):

google/cloud/spanner_v1/transaction.py

@@ -36,6 +36,7 @@
 from google.cloud.spanner_v1 import RequestOptions
 from google.api_core import gapic_v1
 from google.api_core.exceptions import InternalServerError
+from google.protobuf.duration_pb2 import Duration
 
 
 class Transaction(_SnapshotBase, _BatchBase):
@@ -180,7 +181,9 @@ def rollback(self):
         self.rolled_back = True
         del self._session._transaction
 
-    def commit(self, return_commit_stats=False, request_options=None):
+    def commit(
+        self, return_commit_stats=False, request_options=None, max_commit_delay_ms=None
+    ):
         """Commit mutations to the database.
 
         :type return_commit_stats: bool
@@ -193,6 +196,10 @@ def commit(self, return_commit_stats=False, request_options=None):
                 (Optional) Common options for this request.
                 If a dict is provided, it must be of the same form as the protobuf
                 message :class:`~google.cloud.spanner_v1.types.RequestOptions`.
+        :param max_commit_delay_ms:
+                (Optional) The amount of latency this request is willing to incur
+                in order to improve throughput.
+                :class:`~google.cloud.spanner_v1.types.MaxCommitDelay`.
 
         :rtype: datetime
         :returns: timestamp of the committed changes.
@@ -223,11 +230,17 @@ def commit(self, return_commit_stats=False, request_options=None):
         # Request tags are not supported for commit requests.
         request_options.request_tag = None
 
+        max_commit_delay = None
+        if max_commit_delay_ms:
+            max_commit_delay = Duration()
+            max_commit_delay.FromMilliseconds(millis=max_commit_delay_ms)
+
         request = CommitRequest(
             session=self._session.name,
             mutations=self._mutations,
             transaction_id=self._transaction_id,
             return_commit_stats=return_commit_stats,
+            max_commit_delay=max_commit_delay,
             request_options=request_options,
         )
         with trace_call("CloudSpanner.Commit", self._session, trace_attributes):

google/cloud/spanner_v1/types/spanner.py

@@ -24,6 +24,7 @@
 from google.cloud.spanner_v1.types import result_set
 from google.cloud.spanner_v1.types import transaction as gs_transaction
 from google.cloud.spanner_v1.types import type as gs_type
+from google.protobuf import duration_pb2  # type: ignore
 from google.protobuf import struct_pb2  # type: ignore
 from google.protobuf import timestamp_pb2  # type: ignore
 from google.rpc import status_pb2  # type: ignore
@@ -1434,6 +1435,9 @@ class CommitRequest(proto.Message):
             be included in the
             [CommitResponse][google.spanner.v1.CommitResponse.commit_stats].
             Default value is ``false``.
+        max_commit_delay (google.protobuf.duration_pb2.Duration):
+
+            This field is a member of `oneof`_ ``_max_commit_delay``.
         request_options (google.cloud.spanner_v1.types.RequestOptions):
             Common options for this request.
     """
@@ -1462,6 +1466,12 @@ class CommitRequest(proto.Message):
         proto.BOOL,
         number=5,
     )
+    max_commit_delay: duration_pb2.Duration = proto.Field(
+        proto.MESSAGE,
+        number=8,
+        optional=True,
+        message=duration_pb2.Duration,
+    )
     request_options: "RequestOptions" = proto.Field(
         proto.MESSAGE,
         number=6,

scripts/fixup_spanner_v1_keywords.py

@@ -42,7 +42,7 @@ class spannerCallTransformer(cst.CSTTransformer):
         'batch_create_sessions': ('database', 'session_count', 'session_template', ),
         'batch_write': ('session', 'mutation_groups', 'request_options', ),
         'begin_transaction': ('session', 'options', 'request_options', ),
-        'commit': ('session', 'transaction_id', 'single_use_transaction', 'mutations', 'return_commit_stats', 'request_options', ),
+        'commit': ('session', 'transaction_id', 'single_use_transaction', 'mutations', 'return_commit_stats', 'max_commit_delay', 'request_options', ),
         'create_session': ('database', 'session', ),
         'delete_session': ('name', ),
         'execute_batch_dml': ('session', 'transaction', 'statements', 'seqno', 'request_options', ),

tests/unit/test_batch.py

@@ -233,7 +233,14 @@ def test_commit_ok(self):
         self.assertEqual(committed, now)
         self.assertEqual(batch.committed, committed)
 
-        (session, mutations, single_use_txn, request_options, metadata) = api._committed
+        (
+            session,
+            mutations,
+            single_use_txn,
+            request_options,
+            max_commit_delay,
+            metadata,
+        ) = api._committed
         self.assertEqual(session, self.SESSION_NAME)
         self.assertEqual(mutations, batch._mutations)
         self.assertIsInstance(single_use_txn, TransactionOptions)
@@ -246,12 +253,13 @@ def test_commit_ok(self):
             ],
         )
         self.assertEqual(request_options, RequestOptions())
+        self.assertEqual(max_commit_delay, None)
 
         self.assertSpanAttributes(
             "CloudSpanner.Commit", attributes=dict(BASE_ATTRIBUTES, num_mutations=1)
         )
 
-    def _test_commit_with_request_options(self, request_options=None):
+    def _test_commit_with_options(self, request_options=None, max_commit_delay_ms=None):
         import datetime
         from google.cloud.spanner_v1 import CommitResponse
         from google.cloud.spanner_v1 import TransactionOptions
@@ -267,7 +275,9 @@ def _test_commit_with_request_options(self, request_options=None):
         batch = self._make_one(session)
         batch.transaction_tag = self.TRANSACTION_TAG
         batch.insert(TABLE_NAME, COLUMNS, VALUES)
-        committed = batch.commit(request_options=request_options)
+        committed = batch.commit(
+            request_options=request_options, max_commit_delay_ms=max_commit_delay_ms
+        )
 
         self.assertEqual(committed, now)
         self.assertEqual(batch.committed, committed)
@@ -284,6 +294,7 @@ def _test_commit_with_request_options(self, request_options=None):
             mutations,
             single_use_txn,
             actual_request_options,
+            max_commit_delay,
             metadata,
         ) = api._committed
         self.assertEqual(session, self.SESSION_NAME)
@@ -303,33 +314,48 @@ def _test_commit_with_request_options(self, request_options=None):
             "CloudSpanner.Commit", attributes=dict(BASE_ATTRIBUTES, num_mutations=1)
         )
 
+        expected_max_commit_delay = None
+        if max_commit_delay_ms:
+            expected_max_commit_delay = datetime.timedelta(
+                milliseconds=max_commit_delay_ms
+            )
+        self.assertEqual(expected_max_commit_delay, max_commit_delay)
+
     def test_commit_w_request_tag_success(self):
         request_options = RequestOptions(
             request_tag="tag-1",
         )
-        self._test_commit_with_request_options(request_options=request_options)
+        self._test_commit_with_options(request_options=request_options)
 
     def test_commit_w_transaction_tag_success(self):
         request_options = RequestOptions(
             transaction_tag="tag-1-1",
         )
-        self._test_commit_with_request_options(request_options=request_options)
+        self._test_commit_with_options(request_options=request_options)
 
     def test_commit_w_request_and_transaction_tag_success(self):
         request_options = RequestOptions(
             request_tag="tag-1",
             transaction_tag="tag-1-1",
         )
-        self._test_commit_with_request_options(request_options=request_options)
+        self._test_commit_with_options(request_options=request_options)
 
     def test_commit_w_request_and_transaction_tag_dictionary_success(self):
         request_options = {"request_tag": "tag-1", "transaction_tag": "tag-1-1"}
-        self._test_commit_with_request_options(request_options=request_options)
+        self._test_commit_with_options(request_options=request_options)
 
     def test_commit_w_incorrect_tag_dictionary_error(self):
         request_options = {"incorrect_tag": "tag-1-1"}
         with self.assertRaises(ValueError):
-            self._test_commit_with_request_options(request_options=request_options)
+            self._test_commit_with_options(request_options=request_options)
+
+    def test_commit_w_max_commit_delay(self):
+        request_options = RequestOptions(
+            request_tag="tag-1",
+        )
+        self._test_commit_with_options(
+            request_options=request_options, max_commit_delay_ms=100
+        )
 
     def test_context_mgr_already_committed(self):
         import datetime
@@ -368,7 +394,14 @@ def test_context_mgr_success(self):
 
         self.assertEqual(batch.committed, now)
 
-        (session, mutations, single_use_txn, request_options, metadata) = api._committed
+        (
+            session,
+            mutations,
+            single_use_txn,
+            request_options,
+            _,
+            metadata,
+        ) = api._committed
         self.assertEqual(session, self.SESSION_NAME)
         self.assertEqual(mutations, batch._mutations)
         self.assertIsInstance(single_use_txn, TransactionOptions)
@@ -564,13 +597,19 @@ def commit(
         metadata=None,
     ):
         from google.api_core.exceptions import Unknown
+        from google.cloud.spanner_v1 import CommitRequest
+
+        max_commit_delay = None
+        if CommitRequest.max_commit_delay in request:
+            max_commit_delay = request.max_commit_delay
 
         assert request.transaction_id == b""
         self._committed = (
             request.session,
             request.mutations,
             request.single_use_transaction,
             request.request_options,
+            max_commit_delay,
             metadata,
         )
         if self._rpc_error:

tests/unit/test_transaction.py

@@ -346,12 +346,17 @@ def test_commit_w_other_error(self):
         )
 
     def _commit_helper(
-        self, mutate=True, return_commit_stats=False, request_options=None
+        self,
+        mutate=True,
+        return_commit_stats=False,
+        request_options=None,
+        max_commit_delay_ms=None,
     ):
         import datetime
         from google.cloud.spanner_v1 import CommitResponse
         from google.cloud.spanner_v1.keyset import KeySet
         from google.cloud._helpers import UTC
+        from google.protobuf.duration_pb2 import Duration
 
         now = datetime.datetime.utcnow().replace(tzinfo=UTC)
         keys = [[0], [1], [2]]
@@ -370,13 +375,22 @@ def _commit_helper(
             transaction.delete(TABLE_NAME, keyset)
 
         transaction.commit(
-            return_commit_stats=return_commit_stats, request_options=request_options
+            return_commit_stats=return_commit_stats,
+            request_options=request_options,
+            max_commit_delay_ms=max_commit_delay_ms,
         )
 
         self.assertEqual(transaction.committed, now)
         self.assertIsNone(session._transaction)
 
-        session_id, mutations, txn_id, actual_request_options, metadata = api._committed
+        (
+            session_id,
+            mutations,
+            txn_id,
+            actual_request_options,
+            max_commit_delay,
+            metadata,
+        ) = api._committed
 
         if request_options is None:
             expected_request_options = RequestOptions(
@@ -391,6 +405,13 @@ def _commit_helper(
             expected_request_options.transaction_tag = self.TRANSACTION_TAG
             expected_request_options.request_tag = None
 
+        expected_max_commit_delay = None
+        if max_commit_delay_ms:
+            expected_max_commit_delay = datetime.timedelta(
+                milliseconds=max_commit_delay_ms
+            )
+
+        self.assertEqual(expected_max_commit_delay, max_commit_delay)
         self.assertEqual(session_id, session.name)
         self.assertEqual(txn_id, self.TRANSACTION_ID)
         self.assertEqual(mutations, transaction._mutations)
@@ -423,6 +444,9 @@ def test_commit_w_mutations(self):
     def test_commit_w_return_commit_stats(self):
         self._commit_helper(return_commit_stats=True)
 
+    def test_commit_w_max_commit_delay(self):
+        self._commit_helper(max_commit_delay_ms=100)
+
     def test_commit_w_request_tag_success(self):
         request_options = RequestOptions(
             request_tag="tag-1",
@@ -851,7 +875,7 @@ def test_context_mgr_success(self):
 
         self.assertEqual(transaction.committed, now)
 
-        session_id, mutations, txn_id, _, metadata = api._committed
+        session_id, mutations, txn_id, _, _, metadata = api._committed
         self.assertEqual(session_id, self.SESSION_NAME)
         self.assertEqual(txn_id, self.TRANSACTION_ID)
         self.assertEqual(mutations, transaction._mutations)
@@ -935,12 +959,20 @@ def commit(
         request=None,
         metadata=None,
     ):
+        from google.cloud.spanner_v1 import CommitRequest
+
         assert not request.single_use_transaction
+
+        max_commit_delay = None
+        if CommitRequest.max_commit_delay in request:
+            max_commit_delay = request.max_commit_delay
+
         self._committed = (
             request.session,
             request.mutations,
             request.transaction_id,
             request.request_options,
+            max_commit_delay,
             metadata,
         )
         return self._commit_response