feat: adding support for spanner request options tags (#276)

* feat: added support for request options with request tag and transaction tag in supported classes

* feat: corrected import for RequestOptions

* feat: request options added lint corrections

* feat: added system test for request tagging

* feat: added annotation to skip request tags validation test while using emulator

* feat: lint fix

* fix: remove request_option from batch

* lint: lint fixes

* refactor: undo changes

* refactor: undo changes

* refactor: remove test_system file, as it has been removed in master

* refactor: update code to latest changes

* feat: added support for request options with request tag and transaction tag in supported classes

* feat: corrected import for RequestOptions

* fix: add transaction_tag test for transaction_tag set in transaction class

* fix: lint fixes

* refactor: lint fixes

* fix: change request_options dictionary to RequestOptions object

* refactor: fix lint issues

* refactor: lint fixes

* refactor: move write txn properties to BatchBase

* fix: use transaction tag on all write methods

* feat: add support for batch commit

* feat: add support for setting a transaction tag on batch checkout

* refactor: update checks for readability

* test: use separate expectation object for readability

* test: add run_in_transaction test

* test: remove test for unsupported behaviour

* style: lint fixes

Co-authored-by: larkee <larkee@users.noreply.github.com>

google/cloud/spanner_v1/batch.py

@@ -32,6 +32,9 @@ class _BatchBase(_SessionWrapper):
     :param session: the session used to perform the commit
     """
 
+    transaction_tag = None
+    _read_only = False
+
     def __init__(self, session):
         super(_BatchBase, self).__init__(session)
         self._mutations = []
@@ -118,8 +121,7 @@ def delete(self, table, keyset):
 
 
 class Batch(_BatchBase):
-    """Accumulate mutations for transmission during :meth:`commit`.
-    """
+    """Accumulate mutations for transmission during :meth:`commit`."""
 
     committed = None
     commit_stats = None
@@ -160,8 +162,14 @@ def commit(self, return_commit_stats=False, request_options=None):
         txn_options = TransactionOptions(read_write=TransactionOptions.ReadWrite())
         trace_attributes = {"num_mutations": len(self._mutations)}
 
-        if type(request_options) == dict:
+        if request_options is None:
+            request_options = RequestOptions()
+        elif type(request_options) == dict:
             request_options = RequestOptions(request_options)
+        request_options.transaction_tag = self.transaction_tag
+
+        # Request tags are not supported for commit requests.
+        request_options.request_tag = None
 
         request = CommitRequest(
             session=self._session.name,

google/cloud/spanner_v1/database.py

@@ -494,15 +494,20 @@ def execute_partitioned_dml(
             (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`.
+            Please note, the `transactionTag` setting will be ignored as it is
+            not supported for partitioned DML.
 
         :rtype: int
         :returns: Count of rows affected by the DML statement.
         """
         query_options = _merge_query_options(
             self._instance._client._query_options, query_options
         )
-        if type(request_options) == dict:
+        if request_options is None:
+            request_options = RequestOptions()
+        elif type(request_options) == dict:
             request_options = RequestOptions(request_options)
+        request_options.transaction_tag = None
 
         if params is not None:
             from google.cloud.spanner_v1.transaction import Transaction
@@ -796,12 +801,19 @@ class BatchCheckout(object):
     def __init__(self, database, request_options=None):
         self._database = database
         self._session = self._batch = None
-        self._request_options = request_options
+        if request_options is None:
+            self._request_options = RequestOptions()
+        elif type(request_options) == dict:
+            self._request_options = RequestOptions(request_options)
+        else:
+            self._request_options = request_options
 
     def __enter__(self):
         """Begin ``with`` block."""
         session = self._session = self._database._pool.get()
         batch = self._batch = Batch(session)
+        if self._request_options.transaction_tag:
+            batch.transaction_tag = self._request_options.transaction_tag
         return batch
 
     def __exit__(self, exc_type, exc_val, exc_tb):

google/cloud/spanner_v1/session.py

@@ -340,11 +340,13 @@ def run_in_transaction(self, func, *args, **kw):
         """
         deadline = time.time() + kw.pop("timeout_secs", DEFAULT_RETRY_TIMEOUT_SECS)
         commit_request_options = kw.pop("commit_request_options", None)
+        transaction_tag = kw.pop("transaction_tag", None)
         attempts = 0
 
         while True:
             if self._transaction is None:
                 txn = self.transaction()
+                txn.transaction_tag = transaction_tag
             else:
                 txn = self._transaction
             if txn._transaction_id is None:

google/cloud/spanner_v1/snapshot.py

@@ -102,6 +102,7 @@ class _SnapshotBase(_SessionWrapper):
     """
 
     _multi_use = False
+    _read_only = True
     _transaction_id = None
     _read_request_count = 0
     _execute_sql_count = 0
@@ -160,6 +161,8 @@ def read(
                 (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`.
+                Please note, the `transactionTag` setting will be ignored for
+                snapshot as it's not supported for read-only transactions.
 
         :type retry: :class:`~google.api_core.retry.Retry`
         :param retry: (Optional) The retry settings for this request.
@@ -185,9 +188,17 @@ def read(
         metadata = _metadata_with_prefix(database.name)
         transaction = self._make_txn_selector()
 
-        if type(request_options) == dict:
+        if request_options is None:
+            request_options = RequestOptions()
+        elif type(request_options) == dict:
             request_options = RequestOptions(request_options)
 
+        if self._read_only:
+            # Transaction tags are not supported for read only transactions.
+            request_options.transaction_tag = None
+        else:
+            request_options.transaction_tag = self.transaction_tag
+
         request = ReadRequest(
             session=self._session.name,
             table=table,
@@ -312,8 +323,15 @@ def execute_sql(
         default_query_options = database._instance._client._query_options
         query_options = _merge_query_options(default_query_options, query_options)
 
-        if type(request_options) == dict:
+        if request_options is None:
+            request_options = RequestOptions()
+        elif type(request_options) == dict:
             request_options = RequestOptions(request_options)
+        if self._read_only:
+            # Transaction tags are not supported for read only transactions.
+            request_options.transaction_tag = None
+        else:
+            request_options.transaction_tag = self.transaction_tag
 
         request = ExecuteSqlRequest(
             session=self._session.name,

google/cloud/spanner_v1/transaction.py

@@ -148,8 +148,15 @@ def commit(self, return_commit_stats=False, request_options=None):
         metadata = _metadata_with_prefix(database.name)
         trace_attributes = {"num_mutations": len(self._mutations)}
 
-        if type(request_options) == dict:
+        if request_options is None:
+            request_options = RequestOptions()
+        elif type(request_options) == dict:
             request_options = RequestOptions(request_options)
+        if self.transaction_tag is not None:
+            request_options.transaction_tag = self.transaction_tag
+
+        # Request tags are not supported for commit requests.
+        request_options.request_tag = None
 
         request = CommitRequest(
             session=self._session.name,
@@ -267,8 +274,11 @@ def execute_update(
         default_query_options = database._instance._client._query_options
         query_options = _merge_query_options(default_query_options, query_options)
 
-        if type(request_options) == dict:
+        if request_options is None:
+            request_options = RequestOptions()
+        elif type(request_options) == dict:
             request_options = RequestOptions(request_options)
+        request_options.transaction_tag = self.transaction_tag
 
         trace_attributes = {"db.statement": dml}
 
@@ -343,8 +353,11 @@ def batch_update(self, statements, request_options=None):
             self._execute_sql_count + 1,
         )
 
-        if type(request_options) == dict:
+        if request_options is None:
+            request_options = RequestOptions()
+        elif type(request_options) == dict:
             request_options = RequestOptions(request_options)
+        request_options.transaction_tag = self.transaction_tag
 
         trace_attributes = {
             # Get just the queries from the DML statement batch

tests/unit/test_batch.py

@@ -15,6 +15,7 @@
 
 import unittest
 from tests._helpers import OpenTelemetryBase, StatusCode
+from google.cloud.spanner_v1 import RequestOptions
 
 TABLE_NAME = "citizens"
 COLUMNS = ["email", "first_name", "last_name", "age"]
@@ -39,6 +40,7 @@ class _BaseTest(unittest.TestCase):
     DATABASE_NAME = INSTANCE_NAME + "/databases/" + DATABASE_ID
     SESSION_ID = "session-id"
     SESSION_NAME = DATABASE_NAME + "/sessions/" + SESSION_ID
+    TRANSACTION_TAG = "transaction-tag"
 
     def _make_one(self, *args, **kwargs):
         return self._getTargetClass()(*args, **kwargs)
@@ -232,18 +234,87 @@ def test_commit_ok(self):
         self.assertEqual(committed, now)
         self.assertEqual(batch.committed, committed)
 
-        (session, mutations, single_use_txn, metadata, request_options) = 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)
         self.assertTrue(type(single_use_txn).pb(single_use_txn).HasField("read_write"))
         self.assertEqual(metadata, [("google-cloud-resource-prefix", database.name)])
-        self.assertEqual(request_options, None)
+        self.assertEqual(request_options, RequestOptions())
 
         self.assertSpanAttributes(
             "CloudSpanner.Commit", attributes=dict(BASE_ATTRIBUTES, num_mutations=1)
         )
 
+    def _test_commit_with_request_options(self, request_options=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)
+        database = _Database()
+        api = database.spanner_api = _FauxSpannerAPI(_commit_response=response)
+        session = _Session(database)
+        batch = self._make_one(session)
+        batch.transaction_tag = self.TRANSACTION_TAG
+        batch.insert(TABLE_NAME, COLUMNS, VALUES)
+        committed = batch.commit(request_options=request_options)
+
+        self.assertEqual(committed, now)
+        self.assertEqual(batch.committed, committed)
+
+        if type(request_options) == dict:
+            expected_request_options = RequestOptions(request_options)
+        else:
+            expected_request_options = request_options
+        expected_request_options.transaction_tag = self.TRANSACTION_TAG
+        expected_request_options.request_tag = None
+
+        (
+            session,
+            mutations,
+            single_use_txn,
+            actual_request_options,
+            metadata,
+        ) = api._committed
+        self.assertEqual(session, self.SESSION_NAME)
+        self.assertEqual(mutations, batch._mutations)
+        self.assertIsInstance(single_use_txn, TransactionOptions)
+        self.assertTrue(type(single_use_txn).pb(single_use_txn).HasField("read_write"))
+        self.assertEqual(metadata, [("google-cloud-resource-prefix", database.name)])
+        self.assertEqual(actual_request_options, expected_request_options)
+
+        self.assertSpanAttributes(
+            "CloudSpanner.Commit", attributes=dict(BASE_ATTRIBUTES, num_mutations=1)
+        )
+
+    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)
+
+    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)
+
+    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)
+
+    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)
+
+    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)
+
     def test_context_mgr_already_committed(self):
         import datetime
         from google.cloud._helpers import UTC
@@ -281,13 +352,13 @@ def test_context_mgr_success(self):
 
         self.assertEqual(batch.committed, now)
 
-        (session, mutations, single_use_txn, metadata, request_options) = 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)
         self.assertTrue(type(single_use_txn).pb(single_use_txn).HasField("read_write"))
         self.assertEqual(metadata, [("google-cloud-resource-prefix", database.name)])
-        self.assertEqual(request_options, None)
+        self.assertEqual(request_options, RequestOptions())
 
         self.assertSpanAttributes(
             "CloudSpanner.Commit", attributes=dict(BASE_ATTRIBUTES, num_mutations=1)
@@ -341,7 +412,7 @@ def __init__(self, **kwargs):
         self.__dict__.update(**kwargs)
 
     def commit(
-        self, request=None, metadata=None, request_options=None,
+        self, request=None, metadata=None,
     ):
         from google.api_core.exceptions import Unknown
 
@@ -350,8 +421,8 @@ def commit(
             request.session,
             request.mutations,
             request.single_use_transaction,
+            request.request_options,
             metadata,
-            request_options,
         )
         if self._rpc_error:
             raise Unknown("error")