feat: add requestID info in error exceptions

google/cloud/spanner_v1/__init__.py

@@ -65,6 +65,7 @@
 from .types.type import TypeCode
 from .data_types import JsonObject, Interval
 from .transaction import BatchTransactionId, DefaultTransactionOptions
+from .exceptions import wrap_with_request_id
 
 from google.cloud.spanner_v1 import param_types
 from google.cloud.spanner_v1.client import Client
@@ -88,6 +89,8 @@
     # google.cloud.spanner_v1
     "__version__",
     "param_types",
+    # google.cloud.spanner_v1.exceptions
+    "wrap_with_request_id",
     # google.cloud.spanner_v1.client
     "Client",
     # google.cloud.spanner_v1.keyset

google/cloud/spanner_v1/_helpers.py

@@ -22,6 +22,7 @@
 import threading
 import logging
 import uuid
+from contextlib import contextmanager
 
 from google.protobuf.struct_pb2 import ListValue
 from google.protobuf.struct_pb2 import Value
@@ -34,8 +35,12 @@
 from google.cloud.spanner_v1.types import ExecuteSqlRequest
 from google.cloud.spanner_v1.types import TransactionOptions
 from google.cloud.spanner_v1.data_types import JsonObject, Interval
-from google.cloud.spanner_v1.request_id_header import with_request_id
+from google.cloud.spanner_v1.request_id_header import (
+    with_request_id,
+    with_request_id_metadata_only,
+)
 from google.cloud.spanner_v1.types import TypeCode
+from google.cloud.spanner_v1.exceptions import wrap_with_request_id
 
 from google.rpc.error_details_pb2 import RetryInfo
 
@@ -612,11 +617,14 @@ def _retry(
         try:
             return func()
         except Exception as exc:
-            if (
+            is_allowed = (
                 allowed_exceptions is None or exc.__class__ in allowed_exceptions
-            ) and retries < retry_count:
+            )
+
+            if is_allowed and retries < retry_count:
                 if (
                     allowed_exceptions is not None
+                    and exc.__class__ in allowed_exceptions
                     and allowed_exceptions[exc.__class__] is not None
                 ):
                     allowed_exceptions[exc.__class__](exc)
@@ -767,9 +775,67 @@ def reset(self):
 
 
 def _metadata_with_request_id(*args, **kwargs):
+    """Return metadata with request ID header.
+
+    This function returns only the metadata list (not a tuple),
+    maintaining backward compatibility with existing code.
+
+    Args:
+        *args: Arguments to pass to with_request_id
+        **kwargs: Keyword arguments to pass to with_request_id
+
+    Returns:
+        list: gRPC metadata with request ID header
+    """
+    return with_request_id_metadata_only(*args, **kwargs)
+
+
+def _metadata_with_request_id_and_req_id(*args, **kwargs):
+    """Return both metadata and request ID string.
+
+    This is used when we need to augment errors with the request ID.
+
+    Args:
+        *args: Arguments to pass to with_request_id
+        **kwargs: Keyword arguments to pass to with_request_id
+
+    Returns:
+        tuple: (metadata, request_id)
+    """
     return with_request_id(*args, **kwargs)
 
 
+def _augment_error_with_request_id(error, request_id=None):
+    """Augment an error with request ID information.
+
+    Args:
+        error: The error to augment (typically GoogleAPICallError)
+        request_id (str): The request ID to include
+
+    Returns:
+        The augmented error with request ID information
+    """
+    return wrap_with_request_id(error, request_id)
+
+
+@contextmanager
+def _augment_errors_with_request_id(request_id):
+    """Context manager to augment exceptions with request ID.
+
+    Args:
+        request_id (str): The request ID to include in exceptions
+
+    Yields:
+        None
+    """
+    try:
+        yield
+    except Exception as exc:
+        augmented = _augment_error_with_request_id(exc, request_id)
+        # Use exception chaining to preserve the original exception
+        raise augmented from exc
+
+
 def _merge_Transaction_Options(
     defaultTransactionOptions: TransactionOptions,
     mergeTransactionOptions: TransactionOptions,

google/cloud/spanner_v1/batch.py

@@ -252,20 +252,22 @@ def wrapped_method():
                     max_commit_delay=max_commit_delay,
                     request_options=request_options,
                 )
+                # This code is retried due to ABORTED, hence nth_request
+                # should be increased. attempt can only be increased if
+                # we encounter UNAVAILABLE or INTERNAL.
+                call_metadata, error_augmenter = database.with_error_augmentation(
+                    getattr(database, "_next_nth_request", 0),
+                    1,
+                    metadata,
+                    span,
+                )
                 commit_method = functools.partial(
                     api.commit,
                     request=commit_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,
-                        span,
-                    ),
+                    metadata=call_metadata,
                 )
-                return commit_method()
+                with error_augmenter:
+                    return commit_method()
 
             response = _retry_on_aborted_exception(
                 wrapped_method,

google/cloud/spanner_v1/database.py

@@ -25,7 +25,6 @@
 
 import google.auth.credentials
 from google.api_core.retry import Retry
-from google.api_core.retry import if_exception_type
 from google.cloud.exceptions import NotFound
 from google.api_core.exceptions import Aborted
 from google.api_core import gapic_v1
@@ -55,6 +54,8 @@
     _metadata_with_prefix,
     _metadata_with_leader_aware_routing,
     _metadata_with_request_id,
+    _augment_errors_with_request_id,
+    _metadata_with_request_id_and_req_id,
 )
 from google.cloud.spanner_v1.batch import Batch
 from google.cloud.spanner_v1.batch import MutationGroups
@@ -496,6 +497,66 @@ def metadata_with_request_id(
             span,
         )
 
+    def metadata_and_request_id(
+        self, nth_request, nth_attempt, prior_metadata=[], span=None
+    ):
+        """Return metadata and request ID string.
+
+        This method returns both the gRPC metadata with request ID header
+        and the request ID string itself, which can be used to augment errors.
+
+        Args:
+            nth_request: The request sequence number
+            nth_attempt: The attempt number (for retries)
+            prior_metadata: Prior metadata to include
+            span: Optional span for tracing
+
+        Returns:
+            tuple: (metadata_list, request_id_string)
+        """
+        if span is None:
+            span = get_current_span()
+
+        return _metadata_with_request_id_and_req_id(
+            self._nth_client_id,
+            self._channel_id,
+            nth_request,
+            nth_attempt,
+            prior_metadata,
+            span,
+        )
+
+    def with_error_augmentation(
+        self, nth_request, nth_attempt, prior_metadata=[], span=None
+    ):
+        """Context manager for gRPC calls with error augmentation.
+
+        This context manager provides both metadata with request ID and
+        automatically augments any exceptions with the request ID.
+
+        Args:
+            nth_request: The request sequence number
+            nth_attempt: The attempt number (for retries)
+            prior_metadata: Prior metadata to include
+            span: Optional span for tracing
+
+        Yields:
+            tuple: (metadata_list, context_manager)
+        """
+        if span is None:
+            span = get_current_span()
+
+        metadata, request_id = _metadata_with_request_id_and_req_id(
+            self._nth_client_id,
+            self._channel_id,
+            nth_request,
+            nth_attempt,
+            prior_metadata,
+            span,
+        )
+
+        return metadata, _augment_errors_with_request_id(request_id)
+
     def __eq__(self, other):
         if not isinstance(other, self.__class__):
             return NotImplemented
@@ -783,16 +844,18 @@ def execute_pdml():
 
                 try:
                     add_span_event(span, "Starting BeginTransaction")
-                    txn = api.begin_transaction(
-                        session=session.name,
-                        options=txn_options,
-                        metadata=self.metadata_with_request_id(
-                            self._next_nth_request,
-                            1,
-                            metadata,
-                            span,
-                        ),
+                    call_metadata, error_augmenter = self.with_error_augmentation(
+                        self._next_nth_request,
+                        1,
+                        metadata,
+                        span,
                     )
+                    with error_augmenter:
+                        txn = api.begin_transaction(
+                            session=session.name,
+                            options=txn_options,
+                            metadata=call_metadata,
+                        )
 
                     txn_selector = TransactionSelector(id=txn.id)
 
@@ -2060,5 +2123,10 @@ def _retry_on_aborted(func, retry_config):
     :type retry_config: Retry
     :param retry_config: retry object with the settings to be used
     """
-    retry = retry_config.with_predicate(if_exception_type(Aborted))
+
+    def _is_aborted(exc):
+        """Check if exception is Aborted."""
+        return isinstance(exc, Aborted)
+
+    retry = retry_config.with_predicate(_is_aborted)
     return retry(func)

google/cloud/spanner_v1/exceptions.py

@@ -0,0 +1,42 @@
+# Copyright 2026 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.
+
+"""Cloud Spanner exception utilities with request ID support."""
+
+from google.api_core.exceptions import GoogleAPICallError
+
+
+def wrap_with_request_id(error, request_id=None):
+    """Add request ID information to a GoogleAPICallError.
+
+    This function adds request_id as an attribute to the exception,
+    preserving the original exception type for exception handling compatibility.
+    The request_id is also appended to the error message so it appears in logs.
+
+    Args:
+        error: The error to augment. If not a GoogleAPICallError, returns as-is
+        request_id (str): The request ID to include
+
+    Returns:
+        The original error with request_id attribute added and message updated
+        (if GoogleAPICallError and request_id is provided), otherwise returns
+        the original error unchanged.
+    """
+    if isinstance(error, GoogleAPICallError) and request_id:
+        # Add request_id as an attribute for programmatic access
+        error.request_id = request_id
+        # Modify the message to include request_id so it appears in logs
+        if hasattr(error, "message") and error.message:
+            error.message = f"{error.message}, request_id = {request_id}"
+    return error

google/cloud/spanner_v1/pool.py

@@ -259,15 +259,17 @@ def bind(self, database):
                     f"Creating {request.session_count} sessions",
                     span_event_attributes,
                 )
-                resp = api.batch_create_sessions(
-                    request=request,
-                    metadata=database.metadata_with_request_id(
-                        database._next_nth_request,
-                        1,
-                        metadata,
-                        span,
-                    ),
+                call_metadata, error_augmenter = database.with_error_augmentation(
+                    database._next_nth_request,
+                    1,
+                    metadata,
+                    span,
                 )
+                with error_augmenter:
+                    resp = api.batch_create_sessions(
+                        request=request,
+                        metadata=call_metadata,
+                    )
 
                 add_span_event(
                     span,
@@ -570,15 +572,17 @@ def bind(self, database):
         ) as span, MetricsCapture():
             returned_session_count = 0
             while returned_session_count < self.size:
-                resp = api.batch_create_sessions(
-                    request=request,
-                    metadata=database.metadata_with_request_id(
-                        database._next_nth_request,
-                        1,
-                        metadata,
-                        span,
-                    ),
+                call_metadata, error_augmenter = database.with_error_augmentation(
+                    database._next_nth_request,
+                    1,
+                    metadata,
+                    span,
                 )
+                with error_augmenter:
+                    resp = api.batch_create_sessions(
+                        request=request,
+                        metadata=call_metadata,
+                    )
 
                 add_span_event(
                     span,

google/cloud/spanner_v1/request_id_header.py

@@ -43,6 +43,19 @@ def with_request_id(
     all_metadata = (other_metadata or []).copy()
     all_metadata.append((REQ_ID_HEADER_KEY, req_id))
 
+    if span:
+        span.set_attribute(X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR, req_id)
+
+    return all_metadata, req_id
+
+
+def with_request_id_metadata_only(
+    client_id, channel_id, nth_request, attempt, other_metadata=[], span=None
+):
+    req_id = build_request_id(client_id, channel_id, nth_request, attempt)
+    all_metadata = (other_metadata or []).copy()
+    all_metadata.append((REQ_ID_HEADER_KEY, req_id))
+
     if span:
         span.set_attribute(X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR, req_id)