Address review feedback

Author: odeke-em

google/cloud/spanner_v1/database.py

@@ -806,7 +806,9 @@ def _next_nth_request(self):
 
     @property
     def _nth_client_id(self):
-        return self._instance._client._nth_client_id
+        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.
@@ -1028,7 +1030,6 @@ def restore(self, source):
         )
         future = api.restore_database(
             request=request,
-            # TODO: Infer the channel_id being used.
             metadata=self.metadata_with_request_id(self._next_nth_request, 1, metadata),
         )
         return future

google/cloud/spanner_v1/session.py

@@ -252,9 +252,7 @@ def delete(self):
 
         database = self._database
         api = database.spanner_api
-        metadata = database.metadata_with_request_id(
-            database._next_nth_request, 1, _metadata_with_prefix(database.name)
-        )
+        metadata = _metadata_with_prefix(database.name)
         observability_options = getattr(self._database, "observability_options", None)
         with trace_call(
             "CloudSpanner.DeleteSession",
@@ -268,7 +266,9 @@ def delete(self):
         ), MetricsCapture():
             api.delete_session(
                 name=self.name,
-                metadata=metadata,
+                metadata=database.metadata_with_request_id(
+                    database._next_nth_request, 1, metadata
+                ),
             )
 
     def ping(self):

google/cloud/spanner_v1/snapshot.py

@@ -38,6 +38,7 @@
     _retry,
     _check_rst_stream_error,
     _SessionWrapper,
+    AtomicCounter,
 )
 from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
 from google.cloud.spanner_v1.streamed import StreamedResultSet
@@ -82,11 +83,6 @@ def _restart_on_unavailable(
     resume_token = b""
     item_buffer = []
 
-    def next_nth_request():
-        return getattr(request_id_manager, "_next_nth_request", 0)
-
-    nth_request = next_nth_request()
-
     if transaction is not None:
         transaction_selector = transaction._make_txn_selector()
     elif transaction_selector is None:
@@ -96,7 +92,8 @@ def next_nth_request():
 
     request.transaction = transaction_selector
     iterator = None
-    attempt = 0
+    attempt = 1
+    nth_request = getattr(request_id_manager, "_next_nth_request", 0)
 
     while True:
         try:
@@ -108,7 +105,6 @@ def next_nth_request():
                     observability_options=observability_options,
                     metadata=metadata,
                 ), MetricsCapture():
-                    attempt += 1
                     iterator = method(
                         request=request,
                         metadata=request_id_manager.metadata_with_request_id(
@@ -142,8 +138,7 @@ def next_nth_request():
                 if transaction is not None:
                     transaction_selector = transaction._make_txn_selector()
                 request.transaction = transaction_selector
-                nth_request = next_nth_request()
-                attempt = 1
+                attempt += 1
                 iterator = method(
                     request=request,
                     metadata=request_id_manager.metadata_with_request_id(
@@ -169,8 +164,7 @@ def next_nth_request():
                 request.resume_token = resume_token
                 if transaction is not None:
                     transaction_selector = transaction._make_txn_selector()
-                nth_request = next_nth_request()
-                attempt = 1
+                attempt += 1
                 request.transaction = transaction_selector
                 iterator = method(
                     request=request,
@@ -753,12 +747,11 @@ def partition_read(
             metadata=metadata,
         ), MetricsCapture():
             nth_request = getattr(database, "_next_nth_request", 0)
-            counters = dict(attempt=0)
+            attempt = AtomicCounter()
 
             def attempt_tracking_method():
-                counters["attempt"] += 1
                 all_metadata = database.metadata_with_request_id(
-                    nth_request, counters["attempt"], metadata
+                    nth_request, attempt.increment(), metadata
                 )
                 method = functools.partial(
                     api.partition_read,
@@ -867,12 +860,11 @@ def partition_query(
             metadata=metadata,
         ), MetricsCapture():
             nth_request = getattr(database, "_next_nth_request", 0)
-            counters = dict(attempt=0)
+            attempt = AtomicCounter()
 
             def attempt_tracking_method():
-                counters["attempt"] += 1
                 all_metadata = database.metadata_with_request_id(
-                    nth_request, counters["attempt"], metadata
+                    nth_request, attempt.increment(), metadata
                 )
                 method = functools.partial(
                     api.partition_query,
@@ -1024,12 +1016,11 @@ def begin(self):
             metadata=metadata,
         ), MetricsCapture():
             nth_request = getattr(database, "_next_nth_request", 0)
-            counters = dict(attempt=0)
+            attempt = AtomicCounter()
 
             def attempt_tracking_method():
-                counters["attempt"] += 1
                 all_metadata = database.metadata_with_request_id(
-                    nth_request, counters["attempt"], metadata
+                    nth_request, attempt.increment(), metadata
                 )
                 method = functools.partial(
                     api.begin_transaction,

google/cloud/spanner_v1/transaction.py

@@ -186,13 +186,12 @@ def begin(self):
             nth_request = database._next_nth_request
 
             def wrapped_method(*args, **kwargs):
-                attempt.increment()
                 method = functools.partial(
                     api.begin_transaction,
                     session=self._session.name,
                     options=txn_options,
                     metadata=database.metadata_with_request_id(
-                        nth_request, attempt.value, metadata
+                        nth_request, attempt.increment(), metadata
                     ),
                 )
                 return method(*args, **kwargs)

tests/unit/test_spanner.py

@@ -1008,7 +1008,6 @@ def test_transaction_for_concurrent_statement_should_begin_one_transaction_with_
         )
 
         self.assertEqual(api.execute_batch_dml.call_count, 2)
-        print("\033[34marg_list", api.execute_batch_dml.call_args_list, "\033[00m")
         self.assertEqual(
             api.execute_batch_dml.call_args_list,
             [