Skip to content
This repository was archived by the owner on Mar 31, 2026. It is now read-only.

Commit 2c5eb96

Browse files
rahul2393rahul2393
authored
feat: add requestID info in error exceptions (#1415)
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-spanner/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes #<issue_number_goes_here> 🦕
1 parent ed4735b commit 2c5eb96

23 files changed

+673
-180
lines changed

‎google/cloud/spanner_v1/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@
6565
from .types.type import TypeCode
6666
from .data_types import JsonObject, Interval
6767
from .transaction import BatchTransactionId, DefaultTransactionOptions
68+
from .exceptions import wrap_with_request_id
6869

6970
from google.cloud.spanner_v1 import param_types
7071
from google.cloud.spanner_v1.client import Client
@@ -88,6 +89,8 @@
8889
# google.cloud.spanner_v1
8990
"__version__",
9091
"param_types",
92+
# google.cloud.spanner_v1.exceptions
93+
"wrap_with_request_id",
9194
# google.cloud.spanner_v1.client
9295
"Client",
9396
# google.cloud.spanner_v1.keyset

‎google/cloud/spanner_v1/_helpers.py‎

Lines changed: 68 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import threading
2323
import logging
2424
import uuid
25+
from contextlib import contextmanager
2526

2627
from google.protobuf.struct_pb2 import ListValue
2728
from google.protobuf.struct_pb2 import Value
@@ -34,8 +35,12 @@
3435
from google.cloud.spanner_v1.types import ExecuteSqlRequest
3536
from google.cloud.spanner_v1.types import TransactionOptions
3637
from google.cloud.spanner_v1.data_types import JsonObject, Interval
37-
from google.cloud.spanner_v1.request_id_header import with_request_id
38+
from google.cloud.spanner_v1.request_id_header import (
39+
with_request_id,
40+
with_request_id_metadata_only,
41+
)
3842
from google.cloud.spanner_v1.types import TypeCode
43+
from google.cloud.spanner_v1.exceptions import wrap_with_request_id
3944

4045
from google.rpc.error_details_pb2 import RetryInfo
4146

@@ -612,9 +617,11 @@ def _retry(
612617
try:
613618
return func()
614619
except Exception as exc:
615-
if (
620+
is_allowed = (
616621
allowed_exceptions is None or exc.__class__ in allowed_exceptions
617-
) and retries < retry_count:
622+
)
623+
624+
if is_allowed and retries < retry_count:
618625
if (
619626
allowed_exceptions is not None
620627
and allowed_exceptions[exc.__class__] is not None
@@ -767,9 +774,67 @@ def reset(self):
767774

768775

769776
def _metadata_with_request_id(*args, **kwargs):
777+
"""Return metadata with request ID header.
778+
779+
This function returns only the metadata list (not a tuple),
780+
maintaining backward compatibility with existing code.
781+
782+
Args:
783+
*args: Arguments to pass to with_request_id
784+
**kwargs: Keyword arguments to pass to with_request_id
785+
786+
Returns:
787+
list: gRPC metadata with request ID header
788+
"""
789+
return with_request_id_metadata_only(*args, **kwargs)
790+
791+
792+
def _metadata_with_request_id_and_req_id(*args, **kwargs):
793+
"""Return both metadata and request ID string.
794+
795+
This is used when we need to augment errors with the request ID.
796+
797+
Args:
798+
*args: Arguments to pass to with_request_id
799+
**kwargs: Keyword arguments to pass to with_request_id
800+
801+
Returns:
802+
tuple: (metadata, request_id)
803+
"""
770804
return with_request_id(*args, **kwargs)
771805

772806

807+
def _augment_error_with_request_id(error, request_id=None):
808+
"""Augment an error with request ID information.
809+
810+
Args:
811+
error: The error to augment (typically GoogleAPICallError)
812+
request_id (str): The request ID to include
813+
814+
Returns:
815+
The augmented error with request ID information
816+
"""
817+
return wrap_with_request_id(error, request_id)
818+
819+
820+
@contextmanager
821+
def _augment_errors_with_request_id(request_id):
822+
"""Context manager to augment exceptions with request ID.
823+
824+
Args:
825+
request_id (str): The request ID to include in exceptions
826+
827+
Yields:
828+
None
829+
"""
830+
try:
831+
yield
832+
except Exception as exc:
833+
augmented = _augment_error_with_request_id(exc, request_id)
834+
# Use exception chaining to preserve the original exception
835+
raise augmented from exc
836+
837+
773838
def _merge_Transaction_Options(
774839
defaultTransactionOptions: TransactionOptions,
775840
mergeTransactionOptions: TransactionOptions,

‎google/cloud/spanner_v1/batch.py‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -252,20 +252,22 @@ def wrapped_method():
252252
max_commit_delay=max_commit_delay,
253253
request_options=request_options,
254254
)
255+
# This code is retried due to ABORTED, hence nth_request
256+
# should be increased. attempt can only be increased if
257+
# we encounter UNAVAILABLE or INTERNAL.
258+
call_metadata, error_augmenter = database.with_error_augmentation(
259+
getattr(database, "_next_nth_request", 0),
260+
1,
261+
metadata,
262+
span,
263+
)
255264
commit_method = functools.partial(
256265
api.commit,
257266
request=commit_request,
258-
metadata=database.metadata_with_request_id(
259-
# This code is retried due to ABORTED, hence nth_request
260-
# should be increased. attempt can only be increased if
261-
# we encounter UNAVAILABLE or INTERNAL.
262-
getattr(database, "_next_nth_request", 0),
263-
1,
264-
metadata,
265-
span,
266-
),
267+
metadata=call_metadata,
267268
)
268-
return commit_method()
269+
with error_augmenter:
270+
return commit_method()
269271

270272
response = _retry_on_aborted_exception(
271273
wrapped_method,

‎google/cloud/spanner_v1/database.py‎

Lines changed: 79 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525

2626
import google.auth.credentials
2727
from google.api_core.retry import Retry
28-
from google.api_core.retry import if_exception_type
2928
from google.cloud.exceptions import NotFound
3029
from google.api_core.exceptions import Aborted
3130
from google.api_core import gapic_v1
@@ -55,6 +54,8 @@
5554
_metadata_with_prefix,
5655
_metadata_with_leader_aware_routing,
5756
_metadata_with_request_id,
57+
_augment_errors_with_request_id,
58+
_metadata_with_request_id_and_req_id,
5859
)
5960
from google.cloud.spanner_v1.batch import Batch
6061
from google.cloud.spanner_v1.batch import MutationGroups
@@ -496,6 +497,66 @@ def metadata_with_request_id(
496497
span,
497498
)
498499

500+
def metadata_and_request_id(
501+
self, nth_request, nth_attempt, prior_metadata=[], span=None
502+
):
503+
"""Return metadata and request ID string.
504+
505+
This method returns both the gRPC metadata with request ID header
506+
and the request ID string itself, which can be used to augment errors.
507+
508+
Args:
509+
nth_request: The request sequence number
510+
nth_attempt: The attempt number (for retries)
511+
prior_metadata: Prior metadata to include
512+
span: Optional span for tracing
513+
514+
Returns:
515+
tuple: (metadata_list, request_id_string)
516+
"""
517+
if span is None:
518+
span = get_current_span()
519+
520+
return _metadata_with_request_id_and_req_id(
521+
self._nth_client_id,
522+
self._channel_id,
523+
nth_request,
524+
nth_attempt,
525+
prior_metadata,
526+
span,
527+
)
528+
529+
def with_error_augmentation(
530+
self, nth_request, nth_attempt, prior_metadata=[], span=None
531+
):
532+
"""Context manager for gRPC calls with error augmentation.
533+
534+
This context manager provides both metadata with request ID and
535+
automatically augments any exceptions with the request ID.
536+
537+
Args:
538+
nth_request: The request sequence number
539+
nth_attempt: The attempt number (for retries)
540+
prior_metadata: Prior metadata to include
541+
span: Optional span for tracing
542+
543+
Yields:
544+
tuple: (metadata_list, context_manager)
545+
"""
546+
if span is None:
547+
span = get_current_span()
548+
549+
metadata, request_id = _metadata_with_request_id_and_req_id(
550+
self._nth_client_id,
551+
self._channel_id,
552+
nth_request,
553+
nth_attempt,
554+
prior_metadata,
555+
span,
556+
)
557+
558+
return metadata, _augment_errors_with_request_id(request_id)
559+
499560
def __eq__(self, other):
500561
if not isinstance(other, self.__class__):
501562
return NotImplemented
@@ -783,16 +844,18 @@ def execute_pdml():
783844

784845
try:
785846
add_span_event(span, "Starting BeginTransaction")
786-
txn = api.begin_transaction(
787-
session=session.name,
788-
options=txn_options,
789-
metadata=self.metadata_with_request_id(
790-
self._next_nth_request,
791-
1,
792-
metadata,
793-
span,
794-
),
847+
call_metadata, error_augmenter = self.with_error_augmentation(
848+
self._next_nth_request,
849+
1,
850+
metadata,
851+
span,
795852
)
853+
with error_augmenter:
854+
txn = api.begin_transaction(
855+
session=session.name,
856+
options=txn_options,
857+
metadata=call_metadata,
858+
)
796859

797860
txn_selector = TransactionSelector(id=txn.id)
798861

@@ -2060,5 +2123,10 @@ def _retry_on_aborted(func, retry_config):
20602123
:type retry_config: Retry
20612124
:param retry_config: retry object with the settings to be used
20622125
"""
2063-
retry = retry_config.with_predicate(if_exception_type(Aborted))
2126+
2127+
def _is_aborted(exc):
2128+
"""Check if exception is Aborted."""
2129+
return isinstance(exc, Aborted)
2130+
2131+
retry = retry_config.with_predicate(_is_aborted)
20642132
return retry(func)

‎google/cloud/spanner_v1/exceptions.py‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
# Copyright 2026 Google LLC All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Cloud Spanner exception utilities with request ID support."""
16+
17+
from google.api_core.exceptions import GoogleAPICallError
18+
19+
20+
def wrap_with_request_id(error, request_id=None):
21+
"""Add request ID information to a GoogleAPICallError.
22+
23+
This function adds request_id as an attribute to the exception,
24+
preserving the original exception type for exception handling compatibility.
25+
The request_id is also appended to the error message so it appears in logs.
26+
27+
Args:
28+
error: The error to augment. If not a GoogleAPICallError, returns as-is
29+
request_id (str): The request ID to include
30+
31+
Returns:
32+
The original error with request_id attribute added and message updated
33+
(if GoogleAPICallError and request_id is provided), otherwise returns
34+
the original error unchanged.
35+
"""
36+
if isinstance(error, GoogleAPICallError) and request_id:
37+
# Add request_id as an attribute for programmatic access
38+
error.request_id = request_id
39+
# Modify the message to include request_id so it appears in logs
40+
if hasattr(error, "message") and error.message:
41+
error.message = f"{error.message}, request_id = {request_id}"
42+
return error

‎google/cloud/spanner_v1/pool.py‎

Lines changed: 20 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -259,15 +259,17 @@ def bind(self, database):
259259
f"Creating {request.session_count} sessions",
260260
span_event_attributes,
261261
)
262-
resp = api.batch_create_sessions(
263-
request=request,
264-
metadata=database.metadata_with_request_id(
265-
database._next_nth_request,
266-
1,
267-
metadata,
268-
span,
269-
),
262+
call_metadata, error_augmenter = database.with_error_augmentation(
263+
database._next_nth_request,
264+
1,
265+
metadata,
266+
span,
270267
)
268+
with error_augmenter:
269+
resp = api.batch_create_sessions(
270+
request=request,
271+
metadata=call_metadata,
272+
)
271273

272274
add_span_event(
273275
span,
@@ -570,15 +572,17 @@ def bind(self, database):
570572
) as span, MetricsCapture():
571573
returned_session_count = 0
572574
while returned_session_count < self.size:
573-
resp = api.batch_create_sessions(
574-
request=request,
575-
metadata=database.metadata_with_request_id(
576-
database._next_nth_request,
577-
1,
578-
metadata,
579-
span,
580-
),
575+
call_metadata, error_augmenter = database.with_error_augmentation(
576+
database._next_nth_request,
577+
1,
578+
metadata,
579+
span,
581580
)
581+
with error_augmenter:
582+
resp = api.batch_create_sessions(
583+
request=request,
584+
metadata=call_metadata,
585+
)
582586

583587
add_span_event(
584588
span,

‎google/cloud/spanner_v1/request_id_header.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,16 @@ def with_request_id(
4646
if span:
4747
span.set_attribute(X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR, req_id)
4848

49+
return all_metadata, req_id
50+
51+
52+
def with_request_id_metadata_only(
53+
client_id, channel_id, nth_request, attempt, other_metadata=[], span=None
54+
):
55+
"""Return metadata with request ID header, discarding the request ID value."""
56+
all_metadata, _ = with_request_id(
57+
client_id, channel_id, nth_request, attempt, other_metadata, span
58+
)
4959
return all_metadata
5060

5161

0 commit comments

Comments
 (0)