GeezerVison is showing github.com — feat: add requestID info in error exceptions (#1415) · googleapis/python-spanner@2c5eb96
This repository was archived by the owner on Jun 8, 2026. It is now read-only.
File treeExpand file treeCollapse file tree Open diff view settings Expand file treeCollapse file tree Open diff view settings Collapse file Copy file name to clipboardExpand all lines: google/cloud/spanner_v1/__init__.py +3Lines changed: 3 additions & 0 deletions | Original file line number | Diff line number | Diff line change |
|---|
| 65 | 65 | from .types.type import TypeCode | 66 | 66 | from .data_types import JsonObject, Interval | 67 | 67 | from .transaction import BatchTransactionId, DefaultTransactionOptions | | 68 | +from .exceptions import wrap_with_request_id | 68 | 69 |
| 69 | 70 | from google.cloud.spanner_v1 import param_types | 70 | 71 | from google.cloud.spanner_v1.client import Client | | 88 | 89 | # google.cloud.spanner_v1 | 89 | 90 | "__version__", | 90 | 91 | "param_types", | | 92 | + # google.cloud.spanner_v1.exceptions | | 93 | + "wrap_with_request_id", | 91 | 94 | # google.cloud.spanner_v1.client | 92 | 95 | "Client", | 93 | 96 | # google.cloud.spanner_v1.keyset | |
Collapse file Copy file name to clipboardExpand all lines: google/cloud/spanner_v1/_helpers.py +68-3Lines changed: 68 additions & 3 deletions | Original file line number | Diff line number | Diff line change |
|---|
| 22 | 22 | import threading | 23 | 23 | import logging | 24 | 24 | import uuid | | 25 | +from contextlib import contextmanager | 25 | 26 |
| 26 | 27 | from google.protobuf.struct_pb2 import ListValue | 27 | 28 | from google.protobuf.struct_pb2 import Value | | 34 | 35 | from google.cloud.spanner_v1.types import ExecuteSqlRequest | 35 | 36 | from google.cloud.spanner_v1.types import TransactionOptions | 36 | 37 | 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 | +) | 38 | 42 | from google.cloud.spanner_v1.types import TypeCode | | 43 | +from google.cloud.spanner_v1.exceptions import wrap_with_request_id | 39 | 44 |
| 40 | 45 | from google.rpc.error_details_pb2 import RetryInfo | 41 | 46 |
| @@ -612,9 +617,11 @@ def _retry( | 612 | 617 | try: | 613 | 618 | return func() | 614 | 619 | except Exception as exc: | 615 | | - if ( | | 620 | + is_allowed = ( | 616 | 621 | 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: | 618 | 625 | if ( | 619 | 626 | allowed_exceptions is not None | 620 | 627 | and allowed_exceptions[exc.__class__] is not None | @@ -767,9 +774,67 @@ def reset(self): | 767 | 774 |
| 768 | 775 |
| 769 | 776 | 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 | + """ | 770 | 804 | return with_request_id(*args, **kwargs) | 771 | 805 |
| 772 | 806 |
| | 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 | + | 773 | 838 | def _merge_Transaction_Options( | 774 | 839 | defaultTransactionOptions: TransactionOptions, | 775 | 840 | mergeTransactionOptions: TransactionOptions, | |
Collapse file Copy file name to clipboardExpand all lines: google/cloud/spanner_v1/batch.py +12-10Lines changed: 12 additions & 10 deletions | Original file line number | Diff line number | Diff line change |
|---|
@@ -252,20 +252,22 @@ def wrapped_method(): | 252 | 252 | max_commit_delay=max_commit_delay, | 253 | 253 | request_options=request_options, | 254 | 254 | ) | | 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 | + ) | 255 | 264 | commit_method = functools.partial( | 256 | 265 | api.commit, | 257 | 266 | 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, | 267 | 268 | ) | 268 | | - return commit_method() | | 269 | + with error_augmenter: | | 270 | + return commit_method() | 269 | 271 |
| 270 | 272 | response = _retry_on_aborted_exception( | 271 | 273 | wrapped_method, | |
Collapse file Copy file name to clipboardExpand all lines: google/cloud/spanner_v1/database.py +79-11Lines changed: 79 additions & 11 deletions | Original file line number | Diff line number | Diff line change |
|---|
| 25 | 25 |
| 26 | 26 | import google.auth.credentials | 27 | 27 | from google.api_core.retry import Retry | 28 | | -from google.api_core.retry import if_exception_type | 29 | 28 | from google.cloud.exceptions import NotFound | 30 | 29 | from google.api_core.exceptions import Aborted | 31 | 30 | from google.api_core import gapic_v1 | | 55 | 54 | _metadata_with_prefix, | 56 | 55 | _metadata_with_leader_aware_routing, | 57 | 56 | _metadata_with_request_id, | | 57 | + _augment_errors_with_request_id, | | 58 | + _metadata_with_request_id_and_req_id, | 58 | 59 | ) | 59 | 60 | from google.cloud.spanner_v1.batch import Batch | 60 | 61 | from google.cloud.spanner_v1.batch import MutationGroups | @@ -496,6 +497,66 @@ def metadata_with_request_id( | 496 | 497 | span, | 497 | 498 | ) | 498 | 499 |
| | 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 | + | 499 | 560 | def __eq__(self, other): | 500 | 561 | if not isinstance(other, self.__class__): | 501 | 562 | return NotImplemented | @@ -783,16 +844,18 @@ def execute_pdml(): | 783 | 844 |
| 784 | 845 | try: | 785 | 846 | 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, | 795 | 852 | ) | | 853 | + with error_augmenter: | | 854 | + txn = api.begin_transaction( | | 855 | + session=session.name, | | 856 | + options=txn_options, | | 857 | + metadata=call_metadata, | | 858 | + ) | 796 | 859 |
| 797 | 860 | txn_selector = TransactionSelector(id=txn.id) | 798 | 861 |
| @@ -2060,5 +2123,10 @@ def _retry_on_aborted(func, retry_config): | 2060 | 2123 | :type retry_config: Retry | 2061 | 2124 | :param retry_config: retry object with the settings to be used | 2062 | 2125 | """ | 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) | 2064 | 2132 | return retry(func) |
Collapse file Copy file name to clipboard +42Lines changed: 42 additions & 0 deletions | Original file line number | Diff line number | Diff line change |
|---|
| | 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 |
Collapse file Copy file name to clipboardExpand all lines: google/cloud/spanner_v1/pool.py +20-16Lines changed: 20 additions & 16 deletions | Original file line number | Diff line number | Diff line change |
|---|
@@ -259,15 +259,17 @@ def bind(self, database): | 259 | 259 | f"Creating {request.session_count} sessions", | 260 | 260 | span_event_attributes, | 261 | 261 | ) | 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, | 270 | 267 | ) | | 268 | + with error_augmenter: | | 269 | + resp = api.batch_create_sessions( | | 270 | + request=request, | | 271 | + metadata=call_metadata, | | 272 | + ) | 271 | 273 |
| 272 | 274 | add_span_event( | 273 | 275 | span, | @@ -570,15 +572,17 @@ def bind(self, database): | 570 | 572 | ) as span, MetricsCapture(): | 571 | 573 | returned_session_count = 0 | 572 | 574 | 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, | 581 | 580 | ) | | 581 | + with error_augmenter: | | 582 | + resp = api.batch_create_sessions( | | 583 | + request=request, | | 584 | + metadata=call_metadata, | | 585 | + ) | 582 | 586 |
| 583 | 587 | add_span_event( | 584 | 588 | span, | |
Collapse file Copy file name to clipboardExpand all lines: google/cloud/spanner_v1/request_id_header.py +10Lines changed: 10 additions & 0 deletions | Original file line number | Diff line number | Diff line change |
|---|
@@ -46,6 +46,16 @@ def with_request_id( | 46 | 46 | if span: | 47 | 47 | span.set_attribute(X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR, req_id) | 48 | 48 |
| | 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 | + ) | 49 | 59 | return all_metadata | 50 | 60 |
| 51 | 61 |
| |
|