feat: add requestID info in error exceptions (#1415) · googleapis/python-spanner@2c5eb96
@@ -25,7 +25,6 @@
25252626import google.auth.credentials
2727from google.api_core.retry import Retry
28-from google.api_core.retry import if_exception_type
2928from google.cloud.exceptions import NotFound
3029from google.api_core.exceptions import Aborted
3130from google.api_core import gapic_v1
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)
5960from google.cloud.spanner_v1.batch import Batch
6061from google.cloud.spanner_v1.batch import MutationGroups
@@ -496,6 +497,66 @@ def metadata_with_request_id(
496497span,
497498 )
498499500+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+499560def __eq__(self, other):
500561if not isinstance(other, self.__class__):
501562return NotImplemented
@@ -783,16 +844,18 @@ def execute_pdml():
783844784845try:
785846add_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+ )
796859797860txn_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)
20642132return retry(func)