feat(x-goog-spanner-request-id)

request-id: generate rand_process_id, increment nth_request per client

Author: odeke-em

google/cloud/spanner_v1/client.py

@@ -26,6 +26,7 @@
 import grpc
 import os
 import warnings
+import threading
 
 from google.api_core.gapic_v1 import client_info
 from google.auth.credentials import AnonymousCredentials
@@ -147,6 +148,8 @@ class Client(ClientWithProject):
     SCOPE = (SPANNER_ADMIN_SCOPE,)
     """The scopes required for Google Cloud Spanner."""
 
+    NTH_CLIENT = 0
+
     def __init__(
         self,
         project=None,
@@ -199,6 +202,16 @@ def __init__(
         self._route_to_leader_enabled = route_to_leader_enabled
         self._directed_read_options = directed_read_options
         self._observability_options = observability_options
+        Client.NTH_CLIENT += 1
+        self._nth_client_id = Client.NTH_CLIENT
+        self._nth_request = 0
+        self.__lock = threading.lock()
+
+    @property
+    def _next_nth_request(self):
+        with self.__lock:
+            self._nth_request += 1
+            return self._nth_request
 
     @property
     def credentials(self):

google/cloud/spanner_v1/database.py

@@ -693,6 +693,9 @@ def execute_partitioned_dml(
                 _metadata_with_leader_aware_routing(self._route_to_leader_enabled)
             )
 
+        nth_request = self._increment_request_count()
+        attempt = 1
+
         def execute_pdml():
             with SessionCheckout(self._pool) as session:
                 txn = api.begin_transaction(
@@ -728,6 +731,9 @@ def execute_pdml():
 
         return _retry_on_aborted(execute_pdml, DEFAULT_RETRY_BACKOFF)()
 
+    def _next_nth_request(self):
+        return self._instance._client._next_nth_request()
+
     def session(self, labels=None, database_role=None):
         """Factory to create a session for this database.
 

google/cloud/spanner_v1/request_id_header.py

@@ -0,0 +1,39 @@
+# Copyright 2024 Google LLC All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+REQ_ID_VERSION = 1  # The version of the x-goog-spanner-request-id spec.
+REQ_ID_HEADER_KEY = "x-goog-spanner-request-id"
+
+
+def generate_rand_uint64():
+    b = os.urandom(8)
+    return (
+        b[7] & 0xFF
+        | (b[6] & 0xFF) << 8
+        | (b[5] & 0xFF) << 16
+        | (b[4] & 0xFF) << 24
+        | (b[3] & 0xFF) << 32
+        | (b[2] & 0xFF) << 36
+        | (b[1] & 0xFF) << 48
+        | (b[0] & 0xFF) << 56
+    )
+
+
+REQ_RAND_PROCESS_ID = generate_rand_uint64()
+
+
+def metadata_with_request_id(req_id):
+    return [(REQ_ID_HEADER_KEY, req_id)]