Officially inject the first set of x-goog-spanner-request-id values into header metadata

Author: odeke-em

google/cloud/spanner_v1/_helpers.py

@@ -31,6 +31,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.request_id_header import with_request_id
 
 # Validation error messages
 NUMERIC_MAX_SCALE_ERR_MSG = (
@@ -467,3 +468,7 @@ def value(self):
 
     def increment(self, value=1):
         return self.__iadd__(value)
+
+
+def _metadata_with_request_id(*args, **kwargs):
+    return with_request_id(*args, **kwargs)

google/cloud/spanner_v1/client.py

@@ -26,7 +26,6 @@
 import grpc
 import os
 import warnings
-import threading
 
 from google.api_core.gapic_v1 import client_info
 from google.auth.credentials import AnonymousCredentials

google/cloud/spanner_v1/database.py

@@ -53,6 +53,7 @@
     AtomicInt,
     _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
@@ -694,17 +695,18 @@ def execute_partitioned_dml(
                 _metadata_with_leader_aware_routing(self._route_to_leader_enabled)
             )
 
-        nth_request = self._next_nth_request()
+        nth_request = getattr(self, "_next_nth_request", 0)
         attempt = AtomicInt(1)  # It'll be incremented inside _restart_on_unavailable
 
         def execute_pdml():
             with SessionCheckout(self._pool) as session:
-                channel_id = session._channel_id
-                metadata = with_request_id(
-                    self._client._nth_client_id, nth_request, attempt.value, metadata
+                channel_id = getattr(session, "_channel_id", 0)
+                client_id = getattr(self, "_nth_client_id", 0)
+                all_metadata = _metadata_with_request_id(
+                    client_id, channel_id, nth_request, attempt.value, metadata
                 )
                 txn = api.begin_transaction(
-                    session=session.name, options=txn_options, metadata=metadata
+                    session=session.name, options=txn_options, metadata=all_metadata
                 )
 
                 txn_selector = TransactionSelector(id=txn.id)
@@ -737,9 +739,14 @@ def execute_pdml():
 
         return _retry_on_aborted(execute_pdml, DEFAULT_RETRY_BACKOFF)()
 
+    @property
     def _next_nth_request(self):
         return self._instance._client._next_nth_request
 
+    @property
+    def _nth_client_id(self):
+        return self._instance._client._nth_client_id
+
     def session(self, labels=None, database_role=None):
         """Factory to create a session for this database.
 

google/cloud/spanner_v1/pool.py

@@ -48,7 +48,7 @@ def __init__(self, labels=None, database_role=None):
             labels = {}
         self._labels = labels
         self._database_role = database_role
-        self.__lock = threading.lock()
+        self.__lock = threading.Lock()
         self._session_id_to_channel_id = dict()
 
     @property
@@ -129,10 +129,12 @@ def _new_session(self):
             labels=self.labels, database_role=self.database_role
         )
 
-        with self.__lock:
-            channel_id = len(self._session_id_to_channel_id) + 1
-            self._session_id_to_channel_id[session._session.id] = channel_id
-            session._channel_id = channel_id
+        session_id = getattr(session, "_session_id", None)
+        if session_id:
+            with self.__lock:
+                channel_id = len(self._session_id_to_channel_id) + 1
+                self._session_id_to_channel_id[session._session_id] = channel_id
+                session._channel_id = channel_id
 
         return session
 

google/cloud/spanner_v1/request_id_header.py

@@ -36,7 +36,7 @@ def generate_rand_uint64():
 REQ_RAND_PROCESS_ID = generate_rand_uint64()
 
 
-def with_request_id(client_id, nth_request, attempt, other_metadata=[]):
+def with_request_id(client_id, channel_id, nth_request, attempt, other_metadata=[]):
     req_id = f"{REQ_ID_VERSION}.{REQ_RAND_PROCESS_ID}.{client_id}.{channel_id}.{nth_request}.{attempt}"
     other_metadata.append((REQ_ID_HEADER_KEY, req_id))
     return other_metadata