feat: add requestID info in error exceptions

Author: rahul2393

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 SpannerError, wrap_with_request_id
 
 from google.cloud.spanner_v1 import param_types
 from google.cloud.spanner_v1.client import Client
@@ -88,6 +89,9 @@
     # google.cloud.spanner_v1
     "__version__",
     "param_types",
+    # google.cloud.spanner_v1.exceptions
+    "SpannerError",
+    "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
 
@@ -568,7 +573,10 @@ def _retry_on_aborted_exception(
 ):
     """
     Handles retry logic for Aborted exceptions, considering the deadline.
+    Also handles SpannerError that wraps Aborted exceptions.
     """
+    from google.cloud.spanner_v1.exceptions import SpannerError
+
     attempts = 0
     while True:
         try:
@@ -582,6 +590,17 @@ def _retry_on_aborted_exception(
                 default_retry_delay=default_retry_delay,
             )
             continue
+        except SpannerError as exc:
+            # Check if the wrapped error is Aborted
+            if isinstance(exc._error, Aborted):
+                _delay_until_retry(
+                    exc._error,
+                    deadline=deadline,
+                    attempts=attempts,
+                    default_retry_delay=default_retry_delay,
+                )
+                continue
+            raise
 
 
 def _retry(
@@ -600,10 +619,13 @@ def _retry(
         delay: The delay in seconds between retries.
         allowed_exceptions: A tuple of exceptions that are allowed to occur without triggering a retry.
                             Passing allowed_exceptions as None will lead to retrying for all exceptions.
+                            Also handles SpannerError wrapping allowed exceptions.
 
     Returns:
         The result of the function if it is successful, or raises the last exception if all retries fail.
     """
+    from google.cloud.spanner_v1.exceptions import SpannerError
+
     retries = 0
     while retries <= retry_count:
         if retries > 0 and before_next_retry:
@@ -612,14 +634,23 @@ def _retry(
         try:
             return func()
         except Exception as exc:
-            if (
-                allowed_exceptions is None or exc.__class__ in allowed_exceptions
-            ) and retries < retry_count:
+            # Check if exception is allowed directly or wrapped in SpannerError
+            exc_to_check = exc
+            if isinstance(exc, SpannerError):
+                exc_to_check = exc._error
+
+            is_allowed = (
+                allowed_exceptions is None
+                or exc_to_check.__class__ in allowed_exceptions
+            )
+
+            if is_allowed and retries < retry_count:
                 if (
                     allowed_exceptions is not None
-                    and allowed_exceptions[exc.__class__] is not None
+                    and exc_to_check.__class__ in allowed_exceptions
+                    and allowed_exceptions[exc_to_check.__class__] is not None
                 ):
-                    allowed_exceptions[exc.__class__](exc)
+                    allowed_exceptions[exc_to_check.__class__](exc_to_check)
                 time.sleep(delay)
                 delay = delay * 2
                 retries = retries + 1
@@ -767,9 +798,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)
 
@@ -2052,13 +2115,24 @@ def _retry_on_aborted(func, retry_config):
     """Helper for :meth:`Database.execute_partitioned_dml`.
 
     Wrap function in a Retry that will retry on Aborted exceptions
-    with the retry config specified.
+    with the retry config specified. Also handles SpannerError that
+    wraps Aborted exceptions.
 
     :type func: callable
     :param func: the function to be retried on Aborted exceptions
 
     :type retry_config: Retry
     :param retry_config: retry object with the settings to be used
     """
-    retry = retry_config.with_predicate(if_exception_type(Aborted))
+    from google.cloud.spanner_v1.exceptions import SpannerError
+
+    def _is_aborted_or_wrapped_aborted(exc):
+        """Check if exception is Aborted or SpannerError wrapping Aborted."""
+        if isinstance(exc, Aborted):
+            return True
+        if isinstance(exc, SpannerError) and isinstance(exc._error, Aborted):
+            return True
+        return False
+
+    retry = retry_config.with_predicate(_is_aborted_or_wrapped_aborted)
     return retry(func)