Update _execute_partitioned_dml_helper

Author: odeke-em

google/cloud/spanner_v1/database.py

@@ -766,14 +766,12 @@ def execute_pdml():
             ) as span, MetricsCapture():
                 with SessionCheckout(self._pool) as session:
                     add_span_event(span, "Starting BeginTransaction")
-                    begin_txn_nth_request = self._next_nth_request
-                    begin_txn_attempt = AtomicCounter(0)
                     txn = api.begin_transaction(
                         session=session.name,
                         options=txn_options,
                         metadata=self.metadata_with_request_id(
-                            begin_txn_nth_request,
-                            begin_txn_attempt.increment(),
+                            self._next_nth_request,
+                            1,
                             metadata,
                         ),
                     )

tests/unit/test_database.py

@@ -1298,7 +1298,8 @@ def _execute_partitioned_dml_helper(
                     ("x-goog-spanner-route-to-leader", "true"),
                     (
                         "x-goog-spanner-request-id",
-                        f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.2",
+                        # Please note that this try was by an abort and not from service unavailable.
+                        f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.3.1",
                     ),
                 ],
             )
@@ -1370,14 +1371,32 @@ def _execute_partitioned_dml_helper(
                 query_options=expected_query_options,
                 request_options=expected_request_options,
             )
+
+            api.begin_transaction.assert_called_with(
+                session=self.SESSION_NAME,
+                options=txn_options,
+                metadata=[
+                    ("google-cloud-resource-prefix", database.name),
+                    ("x-goog-spanner-route-to-leader", "true"),
+                    (
+                        "x-goog-spanner-request-id",
+                        # Retrying on an aborted response involves creating the transaction afresh
+                        # and also re-invoking execute_streaming_sql, hence the fresh request 4.1.
+                        f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.3.1",
+                    ),
+                ],
+            )
+
             api.execute_streaming_sql.assert_called_with(
                 request=expected_request,
                 metadata=[
                     ("google-cloud-resource-prefix", database.name),
                     ("x-goog-spanner-route-to-leader", "true"),
                     (
                         "x-goog-spanner-request-id",
-                        f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.2.2",
+                        # Retrying on an aborted response involves creating the transaction afresh
+                        # and also re-invoking execute_streaming_sql, hence the fresh request 4.1.
+                        f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.4.1",
                     ),
                 ],
             )