Merge branch 'main' into owl-bot-copy

Author: parthea

docs/transaction-usage.rst

@@ -5,7 +5,8 @@ A :class:`~google.cloud.spanner_v1.transaction.Transaction` represents a
 transaction:  when the transaction commits, it will send any accumulated
 mutations to the server.
 
-To understand more about how transactions work, visit [Transaction](https://cloud.google.com/spanner/docs/reference/rest/v1/Transaction).
+To understand more about how transactions work, visit
+`Transaction <https://cloud.google.com/spanner/docs/reference/rest/v1/Transaction>`_.
 To learn more about how to use them in the Python client, continue reading.
 
 
@@ -90,8 +91,8 @@ any of the records already exists.
 Update records using a Transaction
 ----------------------------------
 
-:meth:`Transaction.update` updates one or more existing records in a table.  Fails
-if any of the records does not already exist.
+:meth:`Transaction.update` updates one or more existing records in a table.
+Fails if any of the records does not already exist.
 
 .. code:: python
 
@@ -178,35 +179,40 @@ Using :meth:`~Database.run_in_transaction`
 
 Rather than calling :meth:`~Transaction.commit` or :meth:`~Transaction.rollback`
 manually, you should use :meth:`~Database.run_in_transaction` to run the
-function that you need.  The transaction's :meth:`~Transaction.commit` method
+function that you need. The transaction's :meth:`~Transaction.commit` method
 will be called automatically if the ``with`` block exits without raising an
-exception.  The function will automatically be retried for
+exception. The function will automatically be retried for
 :class:`~google.api_core.exceptions.Aborted` errors, but will raise on
 :class:`~google.api_core.exceptions.GoogleAPICallError` and
 :meth:`~Transaction.rollback` will be called on all others.
 
 .. code:: python
 
     def _unit_of_work(transaction):
-
         transaction.insert(
-            'citizens', columns=['email', 'first_name', 'last_name', 'age'],
+            'citizens',
+            columns=['email', 'first_name', 'last_name', 'age'],
             values=[
                 ['phred@exammple.com', 'Phred', 'Phlyntstone', 32],
                 ['bharney@example.com', 'Bharney', 'Rhubble', 31],
-            ])
+            ]
+        )
 
         transaction.update(
-            'citizens', columns=['email', 'age'],
+            'citizens',
+            columns=['email', 'age'],
             values=[
                 ['phred@exammple.com', 33],
                 ['bharney@example.com', 32],
-            ])
+            ]
+        )
 
         ...
 
-        transaction.delete('citizens',
-            keyset['bharney@example.com', 'nonesuch@example.com'])
+        transaction.delete(
+            'citizens',
+            keyset=['bharney@example.com', 'nonesuch@example.com']
+        )
 
     db.run_in_transaction(_unit_of_work)
 
@@ -242,7 +248,7 @@ If an exception is raised inside the ``with`` block, the transaction's
         ...
 
         transaction.delete('citizens',
-            keyset['bharney@example.com', 'nonesuch@example.com'])
+            keyset=['bharney@example.com', 'nonesuch@example.com'])
 
 
 Begin a Transaction

google/cloud/spanner_v1/batch.py

@@ -249,17 +249,28 @@ def commit(
             observability_options=observability_options,
             metadata=metadata,
         ), MetricsCapture():
-            method = functools.partial(
-                api.commit,
-                request=request,
-                metadata=metadata,
-            )
+
+            def wrapped_method(*args, **kwargs):
+                method = functools.partial(
+                    api.commit,
+                    request=request,
+                    metadata=database.metadata_with_request_id(
+                        # This code is retried due to ABORTED, hence nth_request
+                        # should be increased. attempt can only be increased if
+                        # we encounter UNAVAILABLE or INTERNAL.
+                        getattr(database, "_next_nth_request", 0),
+                        1,
+                        metadata,
+                    ),
+                )
+                return method(*args, **kwargs)
+
             deadline = time.time() + kwargs.get(
                 "timeout_secs", DEFAULT_RETRY_TIMEOUT_SECS
             )
             default_retry_delay = kwargs.get("default_retry_delay", None)
             response = _retry_on_aborted_exception(
-                method,
+                wrapped_method,
                 deadline=deadline,
                 default_retry_delay=default_retry_delay,
             )

google/cloud/spanner_v1/database.py

@@ -53,6 +53,7 @@
 from google.cloud.spanner_v1._helpers import (
     _metadata_with_prefix,
     _metadata_with_leader_aware_routing,
+    _metadata_with_request_id,
 )
 from google.cloud.spanner_v1.batch import Batch
 from google.cloud.spanner_v1.batch import MutationGroups
@@ -151,6 +152,9 @@ class Database(object):
 
     _spanner_api: SpannerClient = None
 
+    __transport_lock = threading.Lock()
+    __transports_to_channel_id = dict()
+
     def __init__(
         self,
         database_id,
@@ -188,6 +192,7 @@ def __init__(
             self._instance._client.default_transaction_options
         )
         self._proto_descriptors = proto_descriptors
+        self._channel_id = 0  # It'll be created when _spanner_api is created.
 
         if pool is None:
             pool = BurstyPool(database_role=database_role)
@@ -446,8 +451,26 @@ def spanner_api(self):
                 client_info=client_info,
                 client_options=client_options,
             )
+
+            with self.__transport_lock:
+                transport = self._spanner_api._transport
+                channel_id = self.__transports_to_channel_id.get(transport, None)
+                if channel_id is None:
+                    channel_id = len(self.__transports_to_channel_id) + 1
+                    self.__transports_to_channel_id[transport] = channel_id
+                self._channel_id = channel_id
+
         return self._spanner_api
 
+    def metadata_with_request_id(self, nth_request, nth_attempt, prior_metadata=[]):
+        return _metadata_with_request_id(
+            self._nth_client_id,
+            self._channel_id,
+            nth_request,
+            nth_attempt,
+            prior_metadata,
+        )
+
     def __eq__(self, other):
         if not isinstance(other, self.__class__):
             return NotImplemented
@@ -490,7 +513,10 @@ def create(self):
             database_dialect=self._database_dialect,
             proto_descriptors=self._proto_descriptors,
         )
-        future = api.create_database(request=request, metadata=metadata)
+        future = api.create_database(
+            request=request,
+            metadata=self.metadata_with_request_id(self._next_nth_request, 1, metadata),
+        )
         return future
 
     def exists(self):
@@ -506,7 +532,12 @@ def exists(self):
         metadata = _metadata_with_prefix(self.name)
 
         try:
-            api.get_database_ddl(database=self.name, metadata=metadata)
+            api.get_database_ddl(
+                database=self.name,
+                metadata=self.metadata_with_request_id(
+                    self._next_nth_request, 1, metadata
+                ),
+            )
         except NotFound:
             return False
         return True
@@ -523,10 +554,16 @@ def reload(self):
         """
         api = self._instance._client.database_admin_api
         metadata = _metadata_with_prefix(self.name)
-        response = api.get_database_ddl(database=self.name, metadata=metadata)
+        response = api.get_database_ddl(
+            database=self.name,
+            metadata=self.metadata_with_request_id(self._next_nth_request, 1, metadata),
+        )
         self._ddl_statements = tuple(response.statements)
         self._proto_descriptors = response.proto_descriptors
-        response = api.get_database(name=self.name, metadata=metadata)
+        response = api.get_database(
+            name=self.name,
+            metadata=self.metadata_with_request_id(self._next_nth_request, 1, metadata),
+        )
         self._state = DatabasePB.State(response.state)
         self._create_time = response.create_time
         self._restore_info = response.restore_info
@@ -571,7 +608,10 @@ def update_ddl(self, ddl_statements, operation_id="", proto_descriptors=None):
             proto_descriptors=proto_descriptors,
         )
 
-        future = api.update_database_ddl(request=request, metadata=metadata)
+        future = api.update_database_ddl(
+            request=request,
+            metadata=self.metadata_with_request_id(self._next_nth_request, 1, metadata),
+        )
         return future
 
     def update(self, fields):
@@ -609,7 +649,9 @@ def update(self, fields):
         metadata = _metadata_with_prefix(self.name)
 
         future = api.update_database(
-            database=database_pb, update_mask=field_mask, metadata=metadata
+            database=database_pb,
+            update_mask=field_mask,
+            metadata=self.metadata_with_request_id(self._next_nth_request, 1, metadata),
         )
 
         return future
@@ -622,7 +664,10 @@ def drop(self):
         """
         api = self._instance._client.database_admin_api
         metadata = _metadata_with_prefix(self.name)
-        api.drop_database(database=self.name, metadata=metadata)
+        api.drop_database(
+            database=self.name,
+            metadata=self.metadata_with_request_id(self._next_nth_request, 1, metadata),
+        )
 
     def execute_partitioned_dml(
         self,
@@ -711,7 +756,13 @@ def execute_pdml():
                 with SessionCheckout(self._pool) as session:
                     add_span_event(span, "Starting BeginTransaction")
                     txn = api.begin_transaction(
-                        session=session.name, options=txn_options, metadata=metadata
+                        session=session.name,
+                        options=txn_options,
+                        metadata=self.metadata_with_request_id(
+                            self._next_nth_request,
+                            1,
+                            metadata,
+                        ),
                     )
 
                     txn_selector = TransactionSelector(id=txn.id)
@@ -724,6 +775,7 @@ def execute_pdml():
                         query_options=query_options,
                         request_options=request_options,
                     )
+
                     method = functools.partial(
                         api.execute_streaming_sql,
                         metadata=metadata,
@@ -736,6 +788,7 @@ def execute_pdml():
                         metadata=metadata,
                         transaction_selector=txn_selector,
                         observability_options=self.observability_options,
+                        request_id_manager=self,
                     )
 
                     result_set = StreamedResultSet(iterator)
@@ -745,6 +798,18 @@ def execute_pdml():
 
         return _retry_on_aborted(execute_pdml, DEFAULT_RETRY_BACKOFF)()
 
+    @property
+    def _next_nth_request(self):
+        if self._instance and self._instance._client:
+            return self._instance._client._next_nth_request
+        return 1
+
+    @property
+    def _nth_client_id(self):
+        if self._instance and self._instance._client:
+            return self._instance._client._nth_client_id
+        return 0
+
     def session(self, labels=None, database_role=None):
         """Factory to create a session for this database.
 
@@ -965,7 +1030,7 @@ def restore(self, source):
         )
         future = api.restore_database(
             request=request,
-            metadata=metadata,
+            metadata=self.metadata_with_request_id(self._next_nth_request, 1, metadata),
         )
         return future
 
@@ -1034,7 +1099,10 @@ def list_database_roles(self, page_size=None):
             parent=self.name,
             page_size=page_size,
         )
-        return api.list_database_roles(request=request, metadata=metadata)
+        return api.list_database_roles(
+            request=request,
+            metadata=self.metadata_with_request_id(self._next_nth_request, 1, metadata),
+        )
 
     def table(self, table_id):
         """Factory to create a table object within this database.
@@ -1118,7 +1186,10 @@ def get_iam_policy(self, policy_version=None):
                 requested_policy_version=policy_version
             ),
         )
-        response = api.get_iam_policy(request=request, metadata=metadata)
+        response = api.get_iam_policy(
+            request=request,
+            metadata=self.metadata_with_request_id(self._next_nth_request, 1, metadata),
+        )
         return response
 
     def set_iam_policy(self, policy):
@@ -1140,7 +1211,10 @@ def set_iam_policy(self, policy):
             resource=self.name,
             policy=policy,
         )
-        response = api.set_iam_policy(request=request, metadata=metadata)
+        response = api.set_iam_policy(
+            request=request,
+            metadata=self.metadata_with_request_id(self._next_nth_request, 1, metadata),
+        )
         return response
 
     @property

google/cloud/spanner_v1/pool.py

@@ -256,7 +256,9 @@ def bind(self, database):
                 )
                 resp = api.batch_create_sessions(
                     request=request,
-                    metadata=metadata,
+                    metadata=database.metadata_with_request_id(
+                        database._next_nth_request, 1, metadata
+                    ),
                 )
 
                 add_span_event(
@@ -561,7 +563,9 @@ def bind(self, database):
             while returned_session_count < self.size:
                 resp = api.batch_create_sessions(
                     request=request,
-                    metadata=metadata,
+                    metadata=database.metadata_with_request_id(
+                        database._next_nth_request, 1, metadata
+                    ),
                 )
 
                 add_span_event(