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
 
@@ -767,9 +772,65 @@ 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:
+        raise _augment_error_with_request_id(exc, request_id)
+
+
 def _merge_Transaction_Options(
     defaultTransactionOptions: TransactionOptions,
     mergeTransactionOptions: TransactionOptions,

google/cloud/spanner_v1/database.py

@@ -55,6 +55,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 +498,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

google/cloud/spanner_v1/exceptions.py

@@ -0,0 +1,67 @@
+# 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 classes with request ID support."""
+
+from google.api_core.exceptions import GoogleAPICallError
+
+
+class SpannerError(Exception):
+    """Base Spanner exception that includes request ID.
+
+    This class wraps an error (typically GoogleAPICallError) to add request ID information
+    for better debugging and error correlation.
+
+    Args:
+        error: The original error to wrap (typically GoogleAPICallError)
+        request_id (str): The request ID associated with this error
+    """
+
+    def __init__(self, error, request_id=None):
+        self._error = error
+        self._request_id = request_id or ""
+
+    @property
+    def request_id(self):
+        return self._request_id
+
+    def __str__(self):
+        s = str(self._error)
+        if self._request_id:
+            s = f"{s}, request_id = {self._request_id!r}"
+        return s
+
+    def __repr__(self):
+        return f"SpannerError(error={self._error!r}, request_id={self._request_id!r})"
+
+    def __getattr__(self, name):
+        if name == "request_id":
+            return self._request_id
+        return getattr(self._error, name)
+
+
+def wrap_with_request_id(error, request_id=None):
+    """Wrap a GoogleAPICallError with request ID information.
+
+    Args:
+        error: The error to wrap. If not a GoogleAPICallError, returns as-is
+        request_id (str): The request ID to include
+
+    Returns:
+        SpannerError if error is GoogleAPICallError and request_id is provided,
+        otherwise returns the original error
+    """
+    if isinstance(error, GoogleAPICallError) and request_id:
+        return SpannerError(error, request_id)
+    return error

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)