From 5d7e0993285888240d79aeb79736b757ad0677ea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Thu, 8 Oct 2026 18:11:07 +0200 Subject: [PATCH] fix(spanner): eliminate duplicate metrics capture, measure full streaming latency, and cache resource attributes Eliminate redundant operation metrics capture across unary RPCs and streaming queries, measure streaming query latency accurately across full stream lifetimes, and optimize per-RPC metrics overhead through cached resource attributes. Problems Fixed: - Unary RPCs were capturing operation metrics twice: trace_call unconditionally entered MetricsCapture, while caller methods across Database, Session, Snapshot, Transaction, and Batch also wrapped their calls in MetricsCapture. This doubled operation counts and latencies per RPC, creating one legitimate record and one phantom record without resource labels. - Streaming operations (execute_sql and read) exited MetricsCapture immediately after receiving the response iterator instead of waiting for rows to finish streaming, resulting in near-zero operation latencies and duplicate operation records on stream retries. - Every RPC traversed the object hierarchy (session -> database -> instance -> client) to extract project, instance, and database strings, followed by dictionary allocations and repeated setter calls on every tracer instance. Key Changes: - Decoupled tracing from metrics capture: Removed MetricsCapture from trace_call so tracing strictly manages spans, while callers handle operation metrics once with full resource attributes. - Measured full streaming duration: Restructured _restart_on_unavailable so that an operation spans the full lifetime of the stream, recording completion in a finally block when the stream finishes, fails, or is cancelled. Stream retries are recorded as attempts of the same logical operation. - Cached resource attributes on Database: Database._resource_info pre-merges project, instance, and database labels with client-level attributes once into a cached dictionary. create_metrics_tracer simply clones this dictionary, eliminating per-RPC object traversal and setter calls. - Consolidated session resource resolution: Added _resource_info to _SessionWrapper in _helpers.py and delegated to _resource_info_from_database, eliminating duplicate property definitions across Snapshot, Batch, MutationGroups, Session, and SessionPool. - Short-circuited metrics capture when OpenTelemetry is not installed to bypass factory instantiation and ContextVar operations. - Added comprehensive unit and in-memory mock server tests verifying metrics accuracy, caching behavior, fallback paths, and stream resumption. --- .../google/cloud/spanner_v1/_async/batch.py | 20 - .../google/cloud/spanner_v1/_async/client.py | 10 +- .../cloud/spanner_v1/_async/database.py | 34 +- .../google/cloud/spanner_v1/_async/pool.py | 9 +- .../google/cloud/spanner_v1/_async/session.py | 12 +- .../cloud/spanner_v1/_async/snapshot.py | 78 +-- .../google/cloud/spanner_v1/_helpers.py | 23 + .../spanner_v1/_opentelemetry_tracing.py | 60 ++- .../google/cloud/spanner_v1/batch.py | 20 - .../google/cloud/spanner_v1/client.py | 23 +- .../google/cloud/spanner_v1/database.py | 30 +- .../spanner_v1/metrics/metrics_capture.py | 63 ++- .../spanner_v1/metrics/metrics_tracer.py | 21 + .../metrics/metrics_tracer_factory.py | 65 ++- .../metrics/spanner_metrics_tracer_factory.py | 15 +- .../google/cloud/spanner_v1/pool.py | 9 +- .../google/cloud/spanner_v1/session.py | 12 +- .../google/cloud/spanner_v1/snapshot.py | 73 +-- .../mockserver_tests/test_frontend_metrics.py | 453 +++++++++++++++++ .../tests/unit/_async/test_batch.py | 54 ++ .../tests/unit/_async/test_client_extra.py | 35 ++ .../tests/unit/_async/test_database_extra.py | 34 +- .../tests/unit/_async/test_pool.py | 21 + .../tests/unit/_async/test_session.py | 38 ++ .../tests/unit/_async/test_snapshot.py | 466 ++++++++++++++++++ .../tests/unit/test__opentelemetry_tracing.py | 49 ++ .../tests/unit/test_batch.py | 54 ++ .../tests/unit/test_client.py | 54 ++ .../tests/unit/test_database.py | 38 ++ .../tests/unit/test_metrics_capture.py | 123 +++++ .../tests/unit/test_metrics_tracer.py | 18 +- .../tests/unit/test_metrics_tracer_factory.py | 104 ++++ .../tests/unit/test_pool.py | 21 + .../tests/unit/test_session.py | 37 ++ .../tests/unit/test_snapshot.py | 406 ++++++++++++++- .../test_spanner_metrics_tracer_factory.py | 38 ++ 36 files changed, 2394 insertions(+), 226 deletions(-) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/batch.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/batch.py index 966519265d50..1faa0493e8e9 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/batch.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/batch.py @@ -78,16 +78,6 @@ def __init__(self, session, client_context=None): self.commit_stats: Optional[CommitResponse.CommitStats] = None self._client_context = _validate_client_context(client_context) - @property - def _resource_info(self): - """Resource information for metrics labels.""" - database = self._session._database - return { - "project": database._instance._client.project, - "instance": database._instance.instance_id, - "database": database.database_id, - } - def insert(self, table, columns, values): """Insert one or more new table rows. @@ -406,16 +396,6 @@ def __init__(self, session, client_context=None): self.committed: bool = False self._client_context = _validate_client_context(client_context) - @property - def _resource_info(self): - """Resource information for metrics labels.""" - database = self._session._database - return { - "project": database._instance._client.project, - "instance": database._instance.instance_id, - "database": database.database_id, - } - def group(self): """Returns a new `MutationGroup` to which mutations can be added.""" mutation_group = BatchWriteRequest.MutationGroup() diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/client.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/client.py index a18bbc9c605a..26364aa7fad0 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/client.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/client.py @@ -138,7 +138,7 @@ def _get_spanner_log_client_options_env(): return os.getenv(LOG_CLIENT_OPTIONS_ENV_VAR, "false").lower() == "true" -def _initialize_metrics(project, credentials): +def _initialize_metrics(project, credentials, emulator_host=None): """ Initializes the Spanner built-in metrics. @@ -151,7 +151,11 @@ def _initialize_metrics(project, credentials): if not _metrics_monitor_initialized: meter_provider = metrics.NoOpMeterProvider() try: - if not _get_spanner_emulator_host(): + if ( + not _get_spanner_emulator_host() + and not emulator_host + and not isinstance(credentials, AnonymousCredentials) + ): meter_provider = MeterProvider( metric_readers=[ PeriodicExportingMetricReader( @@ -417,7 +421,7 @@ def __init__( and not disable_builtin_metrics and HAS_GOOGLE_CLOUD_MONITORING_INSTALLED ): - _initialize_metrics(project, credentials) + _initialize_metrics(project, credentials, self._emulator_host) else: SpannerMetricsTracerFactory(enabled=False) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/database.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/database.py index c12769172c23..0ec52039253a 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/database.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/database.py @@ -93,6 +93,9 @@ trace_call, ) from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture +from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( + SpannerMetricsTracerFactory, +) from google.cloud.spanner_v1.table import Table SPANNER_DATA_SCOPE = "https://www.googleapis.com/auth/spanner.data" @@ -240,15 +243,28 @@ def __init__( @property def _resource_info(self): """Resource information for metrics labels.""" - return { - "project": ( - self._instance._client.project - if self._instance and self._instance._client - else None - ), - "instance": self._instance.instance_id if self._instance else None, - "database": self.database_id, - } + cached = getattr(self, "_cached_resource_info", None) + if cached is None: + instance = self._instance + client = getattr(instance, "_client", None) if instance else None + project = getattr(client, "project", None) if client else None + instance_id = getattr(instance, "instance_id", None) if instance else None + database_id = self.database_id + try: + factory = SpannerMetricsTracerFactory() + cached = factory.create_resource_info( + project=project, + instance=instance_id, + database=database_id, + ) + except Exception: + cached = { + "project": project, + "instance": instance_id, + "database": database_id, + } + self._cached_resource_info = cached + return cached @classmethod def from_pb(cls, database_pb, instance, pool=None): diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/pool.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/pool.py index 9bdd0d3fcd91..b43c4f1bc56c 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/pool.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/pool.py @@ -27,6 +27,7 @@ from google.cloud.spanner_v1._helpers import ( _metadata_with_leader_aware_routing, _metadata_with_prefix, + _resource_info_from_database, ) from google.cloud.spanner_v1._opentelemetry_tracing import ( add_span_event, @@ -104,13 +105,7 @@ def __init__(self, labels=None, database_role=None): @property def _resource_info(self): """Resource information for metrics labels.""" - if self._database is None: - return None - return { - "project": self._database._instance._client.project, - "instance": self._database._instance.instance_id, - "database": self._database.database_id, - } + return _resource_info_from_database(self._database) @property def labels(self): diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/session.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/session.py index b5911ad343b1..1bbfa1a1c136 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/session.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/session.py @@ -32,6 +32,7 @@ _get_retry_delay, _metadata_with_leader_aware_routing, _metadata_with_prefix, + _resource_info_from_database, ) from google.cloud.spanner_v1._opentelemetry_tracing import ( add_span_event, @@ -94,11 +95,7 @@ def __init__(self, database, labels=None, database_role=None, is_multiplexed=Fal @property def _resource_info(self): """Resource information for metrics labels.""" - return { - "project": self._database._instance._client.project, - "instance": self._database._instance.instance_id, - "database": self._database.database_id, - } + return _resource_info_from_database(self._database) def __lt__(self, other): return self._session_id < other._session_id @@ -352,7 +349,10 @@ async def ping(self): metadata = _metadata_with_prefix(database.name) nth_request = database._next_nth_request - with trace_call("CloudSpanner.Session.ping", self) as span: + with ( + trace_call("CloudSpanner.Session.ping", self) as span, + MetricsCapture(self._resource_info), + ): call_metadata, error_augmenter = database.with_error_augmentation( nth_request, 1, metadata, span ) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py index e0d73c7509e7..5e9bd6948939 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py @@ -45,6 +45,9 @@ ) from google.cloud.spanner_v1._opentelemetry_tracing import add_span_event, trace_call from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture +from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( + SpannerMetricsTracerFactory, +) from google.cloud.spanner_v1.types import MultiplexedSessionPrecommitToken from google.cloud.spanner_v1.types.mutation import Mutation from google.cloud.spanner_v1.types.result_set import PartialResultSet, ResultSet @@ -120,21 +123,25 @@ async def _restart_on_unavailable( current_request_id = None stream_finished = False + metrics_tracer = None + factory = SpannerMetricsTracerFactory() + if factory.enabled: + metrics_tracer = factory.create_metrics_tracer(resource_info) + if metrics_tracer: + metrics_tracer.record_operation_start() + try: while True: try: # Get results iterator. if iterator is None: - with ( - trace_call( - trace_name, - session, - attributes, - observability_options=observability_options, - metadata=metadata, - ) as span, - MetricsCapture(resource_info), - ): + with trace_call( + trace_name, + session, + attributes, + observability_options=observability_options, + metadata=metadata, + ) as span: ( call_metadata, current_request_id, @@ -144,11 +151,25 @@ async def _restart_on_unavailable( metadata, span, ) - iterator = await CrossSync.run_if_async( - method, - request=request, - metadata=call_metadata, - ) + token = None + if metrics_tracer is not None: + token = SpannerMetricsTracerFactory.set_current_tracer( + metrics_tracer + ) + try: + iterator = await CrossSync.run_if_async( + method, + request=request, + metadata=call_metadata, + ) + finally: + if token is not None: + try: + SpannerMetricsTracerFactory.reset_current_tracer( + token + ) + except ValueError: + pass # Add items from iterator to buffer. item: PartialResultSet @@ -224,11 +245,18 @@ async def _restart_on_unavailable( if stream_finished: break finally: - if iterator is not None and hasattr(iterator, "cancel"): - try: - iterator.cancel() - except Exception: - pass + try: + if metrics_tracer is not None: + try: + metrics_tracer.record_operation_completion() + except Exception: + pass + finally: + if iterator is not None and hasattr(iterator, "cancel"): + try: + iterator.cancel() + except Exception: + pass class _SnapshotBase(_SessionWrapper): @@ -262,16 +290,6 @@ def __init__(self, session, client_context=None): # This is used to prevent the "Transaction has not begun" race condition. self._transaction_begin_event: CrossSync.Event = CrossSync.Event() - @property - def _resource_info(self): - """Resource information for metrics labels.""" - database = self._session._database - return { - "project": database._instance._client.project, - "instance": database._instance.instance_id, - "database": database.database_id, - } - @CrossSync.convert async def _wait_for_transaction_begin(self) -> None: """Claims the inline-begin for this request, or waits for it to complete. diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py index 369dc64f32bc..f3386e3aeda6 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py @@ -742,6 +742,22 @@ def _parse_interval(value_pb): _STRUCT_TYPE_CODE = int(TypeCode.STRUCT) +def _resource_info_from_database(database): + """Resource information dictionary from database instance or fallback.""" + if database is None: + return None + resource_info = getattr(database, "_resource_info", None) + if resource_info is not None: + return resource_info + instance = getattr(database, "_instance", None) + client = getattr(instance, "_client", None) if instance else None + return { + "project": getattr(client, "project", None), + "instance": getattr(instance, "instance_id", None), + "database": getattr(database, "database_id", None), + } + + class _SessionWrapper(object): """Base class for objects wrapping a session. @@ -752,6 +768,13 @@ class _SessionWrapper(object): def __init__(self, session): self._session = session + @property + def _resource_info(self): + """Resource information for metrics labels.""" + session = self._session + database = getattr(session, "_database", None) if session is not None else None + return _resource_info_from_database(database) + def _append_routing_headers(metadata): """Appends routing and backend-specific headers to the metadata.""" diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_opentelemetry_tracing.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_opentelemetry_tracing.py index ddca0656eb0b..ed315ac6f80d 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_opentelemetry_tracing.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_opentelemetry_tracing.py @@ -27,7 +27,6 @@ _metadata_with_span_context, ) from google.cloud.spanner_v1.gapic_version import __version__ as gapic_version -from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture from google.cloud.spanner_v1.services.spanner.client import SpannerClient TRACER_NAME = "cloud.google.com/python/spanner" @@ -98,16 +97,14 @@ def trace_call( ): child_span = trace.NonRecordingSpan(current_span.get_span_context()) with trace.use_span(child_span): - with MetricsCapture(): - if enable_end_to_end_tracing: - _metadata_with_span_context(metadata) - yield child_span + if enable_end_to_end_tracing: + _metadata_with_span_context(metadata) + yield child_span return - with MetricsCapture(): - if enable_end_to_end_tracing: - _metadata_with_span_context(metadata) - yield trace.INVALID_SPAN + if enable_end_to_end_tracing: + _metadata_with_span_context(metadata) + yield trace.INVALID_SPAN return # Slow path: resolve attributes and configure span when tracing is active. @@ -158,29 +155,28 @@ def trace_call( with tracer.start_as_current_span( name, kind=trace.SpanKind.CLIENT, attributes=attributes ) as span: - with MetricsCapture(): - try: - if enable_end_to_end_tracing: - _metadata_with_span_context(metadata) - yield span - except Exception as error: - span.set_status(Status(StatusCode.ERROR, str(error))) - # OpenTelemetry-Python imposes invoking span.record_exception on __exit__ - # on any exception. We should file a bug later on with them to only - # invoke .record_exception if not already invoked, hence we should not - # invoke .record_exception on our own else we shall have 2 exceptions. - raise - else: - # All spans still have set_status available even if for example - # NonRecordingSpan doesn't have "_status". - absent_span_status = getattr(span, "_status", None) is None - if absent_span_status or span._status.status_code == StatusCode.UNSET: - # OpenTelemetry-Python only allows a status change - # if the current code is UNSET or ERROR. At the end - # of the generator's consumption, only set it to OK - # it wasn't previously set otherwise. - # https://github.com/googleapis/python-spanner/issues/1246 - span.set_status(Status(StatusCode.OK)) + try: + if enable_end_to_end_tracing: + _metadata_with_span_context(metadata) + yield span + except Exception as error: + span.set_status(Status(StatusCode.ERROR, str(error))) + # OpenTelemetry-Python imposes invoking span.record_exception on __exit__ + # on any exception. We should file a bug later on with them to only + # invoke .record_exception if not already invoked, hence we should not + # invoke .record_exception on our own else we shall have 2 exceptions. + raise + else: + # All spans still have set_status available even if for example + # NonRecordingSpan doesn't have "_status". + absent_span_status = getattr(span, "_status", None) is None + if absent_span_status or span._status.status_code == StatusCode.UNSET: + # OpenTelemetry-Python only allows a status change + # if the current code is UNSET or ERROR. At the end + # of the generator's consumption, only set it to OK + # it wasn't previously set otherwise. + # https://github.com/googleapis/python-spanner/issues/1246 + span.set_status(Status(StatusCode.OK)) def get_current_span(): diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/batch.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/batch.py index c0ff1cc1a613..19567e3bb4c7 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/batch.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/batch.py @@ -69,16 +69,6 @@ def __init__(self, session, client_context=None): self.commit_stats: Optional[CommitResponse.CommitStats] = None self._client_context = _validate_client_context(client_context) - @property - def _resource_info(self): - """Resource information for metrics labels.""" - database = self._session._database - return { - "project": database._instance._client.project, - "instance": database._instance.instance_id, - "database": database.database_id, - } - def insert(self, table, columns, values): """Insert one or more new table rows. @@ -352,16 +342,6 @@ def __init__(self, session, client_context=None): self.committed: bool = False self._client_context = _validate_client_context(client_context) - @property - def _resource_info(self): - """Resource information for metrics labels.""" - database = self._session._database - return { - "project": database._instance._client.project, - "instance": database._instance.instance_id, - "database": database.database_id, - } - def group(self): """Returns a new `MutationGroup` to which mutations can be added.""" mutation_group = BatchWriteRequest.MutationGroup() diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/client.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/client.py index 7ae7f58a0e0d..781603f92078 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/client.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/client.py @@ -119,7 +119,7 @@ def _get_spanner_log_client_options_env(): return os.getenv(LOG_CLIENT_OPTIONS_ENV_VAR, "false").lower() == "true" -def _initialize_metrics(project, credentials): +def _initialize_metrics(project, credentials, emulator_host=None): """Initializes the Spanner built-in metrics. This function sets up the OpenTelemetry MeterProvider and the SpannerMetricsTracerFactory. @@ -130,7 +130,11 @@ def _initialize_metrics(project, credentials): if not _metrics_monitor_initialized: meter_provider = metrics.NoOpMeterProvider() try: - if not _get_spanner_emulator_host(): + if ( + not _get_spanner_emulator_host() + and (not emulator_host) + and (not isinstance(credentials, AnonymousCredentials)) + ): meter_provider = MeterProvider( metric_readers=[ PeriodicExportingMetricReader( @@ -320,9 +324,7 @@ def __init__( raise ValueError( "Both username and password must be specified for Omni authentication" ) - from google.cloud.spanner_v1.omni.credentials import ( - SpannerOmniCredentials, - ) + from google.cloud.spanner_v1.omni.credentials import SpannerOmniCredentials if has_username and has_password: credentials = SpannerOmniCredentials( @@ -339,11 +341,10 @@ def __init__( disable_builtin_metrics = True elif isinstance(credentials, AnonymousCredentials): self._emulator_host = self._client_options.api_endpoint - else: - if username is not None or password is not None: - raise ValueError( - "username and password can only be used when instance_type='omni'." - ) + elif username is not None or password is not None: + raise ValueError( + "username and password can only be used when instance_type='omni'." + ) super(Client, self).__init__( project=project, credentials=credentials, @@ -366,7 +367,7 @@ def __init__( and (not disable_builtin_metrics) and HAS_GOOGLE_CLOUD_MONITORING_INSTALLED ): - _initialize_metrics(project, credentials) + _initialize_metrics(project, credentials, self._emulator_host) else: SpannerMetricsTracerFactory(enabled=False) self._route_to_leader_enabled = route_to_leader_enabled diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/database.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/database.py index 0508e0ae23c8..31022b9e88ec 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/database.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/database.py @@ -65,6 +65,9 @@ from google.cloud.spanner_v1.keyset import KeySet from google.cloud.spanner_v1.merged_result_set import MergedResultSet from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture +from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( + SpannerMetricsTracerFactory, +) from google.cloud.spanner_v1.pool import BurstyPool from google.cloud.spanner_v1.services.spanner.client import ( SpannerClient as SpannerClient, @@ -208,13 +211,26 @@ def __init__( @property def _resource_info(self): """Resource information for metrics labels.""" - return { - "project": self._instance._client.project - if self._instance and self._instance._client - else None, - "instance": self._instance.instance_id if self._instance else None, - "database": self.database_id, - } + cached = getattr(self, "_cached_resource_info", None) + if cached is None: + instance = self._instance + client = getattr(instance, "_client", None) if instance else None + project = getattr(client, "project", None) if client else None + instance_id = getattr(instance, "instance_id", None) if instance else None + database_id = self.database_id + try: + factory = SpannerMetricsTracerFactory() + cached = factory.create_resource_info( + project=project, instance=instance_id, database=database_id + ) + except Exception: + cached = { + "project": project, + "instance": instance_id, + "database": database_id, + } + self._cached_resource_info = cached + return cached @classmethod def from_pb(cls, database_pb, instance, pool=None): diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_capture.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_capture.py index 77c1f86c4dee..815dd24fcfa7 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_capture.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_capture.py @@ -22,7 +22,10 @@ from contextvars import Token -from .spanner_metrics_tracer_factory import SpannerMetricsTracerFactory +from .spanner_metrics_tracer_factory import ( + HAS_OPENTELEMETRY_INSTALLED, + SpannerMetricsTracerFactory, +) class MetricsCapture: @@ -32,7 +35,7 @@ class MetricsCapture: the start and completion of metrics tracing for a given operation. """ - _token: Token + _token: Token = None """Token to reset the context variable after the operation completes.""" def __init__(self, resource_info: dict = None): @@ -42,6 +45,8 @@ def __init__(self, resource_info: dict = None): resource_info (dict): Optional dictionary containing project, instance and database info. """ self._resource_info = resource_info + self._token = None + self._tracer = None def __enter__(self): """Enter the runtime context related to this object. @@ -52,26 +57,29 @@ def __enter__(self): Returns: MetricsCapture: The instance of the context manager. """ + if not HAS_OPENTELEMETRY_INSTALLED: + return self + # Short circuit out if metrics are disabled factory = SpannerMetricsTracerFactory() if not factory.enabled: return self - # Define a new metrics tracer for the new operation - # Set the context var and keep the token for reset - tracer = factory.create_metrics_tracer() - - if tracer and self._resource_info: - if "project" in self._resource_info: - tracer.set_project(self._resource_info["project"]) - if "instance" in self._resource_info: - tracer.set_instance(self._resource_info["instance"]) - if "database" in self._resource_info: - tracer.set_database(self._resource_info["database"]) - + tracer = factory.create_metrics_tracer(self._resource_info) + self._tracer = tracer self._token = SpannerMetricsTracerFactory.set_current_tracer(tracer) if tracer: - tracer.record_operation_start() + try: + tracer.record_operation_start() + except BaseException: + if self._token is not None: + try: + SpannerMetricsTracerFactory.reset_current_tracer(self._token) + except ValueError: + pass + self._token = None + self._tracer = None + raise return self def __exit__(self, exc_type, exc_value, traceback): @@ -88,15 +96,22 @@ def __exit__(self, exc_type, exc_value, traceback): Returns: bool: False to propagate the exception if any occurred. """ - # Short circuit out if metrics are disable - if not SpannerMetricsTracerFactory().enabled: + token = self._token + if token is None: return False - tracer = SpannerMetricsTracerFactory.get_current_tracer() - if tracer: - tracer.record_operation_completion() - - # Reset the context var using the token - if getattr(self, "_token", None): - SpannerMetricsTracerFactory.reset_current_tracer(self._token) + try: + tracer = self._tracer or SpannerMetricsTracerFactory.get_current_tracer() + if tracer: + try: + tracer.record_operation_completion() + except Exception: + pass + finally: + try: + SpannerMetricsTracerFactory.reset_current_tracer(token) + except ValueError: + pass + self._token = None + self._tracer = None return False # Propagate the exception if any diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_tracer.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_tracer.py index 6106fa6e18b0..16399c77701c 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_tracer.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_tracer.py @@ -587,6 +587,27 @@ def _create_attempt_otel_attributes(self) -> dict: attributes[METRIC_LABEL_KEY_STATUS] = self.current_op.current_attempt.status return attributes + def set_resource_info( + self, resource_info: Optional[Dict[str, str]] + ) -> "MetricsTracer": + """Set resource attributes (project, instance, database) for metrics tracing. + + :param resource_info: A dictionary containing project, instance, and database information. + :return: This instance of MetricsTracer for method chaining. + """ + if resource_info: + client_attributes = self._client_attributes + project = resource_info.get("project") + if project and MONITORED_RES_LABEL_KEY_PROJECT not in client_attributes: + client_attributes[MONITORED_RES_LABEL_KEY_PROJECT] = project + instance = resource_info.get("instance") + if instance and MONITORED_RES_LABEL_KEY_INSTANCE not in client_attributes: + client_attributes[MONITORED_RES_LABEL_KEY_INSTANCE] = instance + database = resource_info.get("database") + if database and METRIC_LABEL_KEY_DATABASE not in client_attributes: + client_attributes[METRIC_LABEL_KEY_DATABASE] = database + return self + def set_project(self, project: str) -> "MetricsTracer": """ Set the project attribute for metrics tracing. diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_tracer_factory.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_tracer_factory.py index 0363c08898c0..95c5ac21d33f 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_tracer_factory.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_tracer_factory.py @@ -14,7 +14,7 @@ """Factory for creating MetricTracer instances, facilitating metrics collection and tracing.""" -from typing import Dict +from typing import Dict, Optional from google.cloud.spanner_v1.metrics.constants import ( BUILT_IN_METRICS_METER_NAME, @@ -48,6 +48,12 @@ from google.cloud.spanner_v1.gapic_version import __version__ +class _ResourceInfo(dict): + """Resource info dictionary containing pre-merged client attributes for metrics.""" + + _client_attributes: Optional[Dict[str, str]] = None + + class MetricsTracerFactory: """Factory class for creating MetricTracer instances. This class facilitates the creation of MetricTracer objects, which are responsible for collecting and tracing metrics.""" @@ -250,34 +256,79 @@ def enable_direct_path(self, enable: bool = False) -> "MetricsTracerFactory": self._client_attributes[METRIC_LABEL_KEY_DIRECT_PATH_ENABLED] = enable return self - def create_metrics_tracer(self) -> MetricsTracer: - """ - Create and return a MetricsTracer instance with default settings and client attributes. + def create_metrics_tracer( + self, resource_info: Optional[Dict[str, str]] = None + ) -> Optional[MetricsTracer]: + """Create and return a MetricsTracer instance with default settings and client attributes. This method initializes a MetricsTracer instance with default settings for metrics tracing, including metrics tracing enabled if OpenTelemetry is installed and the direct path disabled by default. - It also sets the client attributes based on the factory's configuration. + It also sets the client attributes based on the factory's configuration and optional resource info. + + Args: + resource_info (dict, optional): A dictionary of resource information (project, instance, database). Returns: - MetricsTracer: A MetricsTracer instance with default settings and client attributes. + MetricsTracer: A MetricsTracer instance with default settings and client attributes, or None if OpenTelemetry is not installed. """ if not HAS_OPENTELEMETRY_INSTALLED: return None + pre_merged = getattr(resource_info, "_client_attributes", None) + if isinstance(pre_merged, dict): + client_attributes = pre_merged.copy() + else: + client_attributes = self._client_attributes.copy() + metrics_tracer = MetricsTracer( enabled=self.enabled and HAS_OPENTELEMETRY_INSTALLED, instrument_attempt_latency=self._instrument_attempt_latency, instrument_attempt_counter=self._instrument_attempt_counter, instrument_operation_latency=self._instrument_operation_latency, instrument_operation_counter=self._instrument_operation_counter, - client_attributes=self._client_attributes.copy(), + client_attributes=client_attributes, instrument_gfe_latency=self._instrument_gfe_latency, instrument_gfe_connectivity_error_count=self._instrument_gfe_connectivity_error_count, instrument_afe_latency=self._instrument_afe_latency, instrument_afe_connectivity_error_count=self._instrument_afe_connectivity_error_count, ) + if resource_info and not isinstance(pre_merged, dict): + metrics_tracer.set_resource_info(resource_info) return metrics_tracer + def create_resource_info( + self, + project: Optional[str] = None, + instance: Optional[str] = None, + database: Optional[str] = None, + ) -> _ResourceInfo: + """Create a resource info dictionary with pre-merged client attributes. + + Args: + project (str, optional): The project ID. + instance (str, optional): The instance ID. + database (str, optional): The database ID. + + Returns: + _ResourceInfo: A dictionary with project, instance, and database keys, + carrying pre-merged client attributes for metrics tracing. + """ + info = _ResourceInfo( + project=project, + instance=instance, + database=database, + ) + if HAS_OPENTELEMETRY_INSTALLED: + client_attributes = self._client_attributes.copy() + if project and MONITORED_RES_LABEL_KEY_PROJECT not in client_attributes: + client_attributes[MONITORED_RES_LABEL_KEY_PROJECT] = project + if instance and MONITORED_RES_LABEL_KEY_INSTANCE not in client_attributes: + client_attributes[MONITORED_RES_LABEL_KEY_INSTANCE] = instance + if database and METRIC_LABEL_KEY_DATABASE not in client_attributes: + client_attributes[METRIC_LABEL_KEY_DATABASE] = database + info._client_attributes = client_attributes + return info + def _create_metric_instruments(self, service_name: str) -> None: """ Creates and sets up metric instruments for the given service name. diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/spanner_metrics_tracer_factory.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/spanner_metrics_tracer_factory.py index 5342d6474689..8ee1df2305f9 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/spanner_metrics_tracer_factory.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/spanner_metrics_tracer_factory.py @@ -18,6 +18,7 @@ import contextvars import logging import os +from typing import Optional from .constants import SPANNER_SERVICE_NAME from .metrics_tracer_factory import MetricsTracerFactory @@ -51,7 +52,7 @@ class SpannerMetricsTracerFactory(MetricsTracerFactory): "current_metrics_tracer", default=None ) - def __new__(cls, enabled: bool = True) -> "SpannerMetricsTracerFactory": + def __new__(cls, enabled: Optional[bool] = None) -> "SpannerMetricsTracerFactory": """ Create a new instance of SpannerMetricsTracerFactory if it doesn't already exist. @@ -60,14 +61,17 @@ def __new__(cls, enabled: bool = True) -> "SpannerMetricsTracerFactory": if it hasn't been created yet. Args: - enabled (bool): A flag indicating whether metrics tracing is enabled. Defaults to True. + enabled (bool, optional): A flag indicating whether metrics tracing is enabled. + Defaults to None (which defaults to True upon initial creation, and preserves + existing state on subsequent calls). Returns: SpannerMetricsTracerFactory: The singleton instance of SpannerMetricsTracerFactory. """ if cls._metrics_tracer_factory is None: + init_enabled = True if enabled is None else enabled cls._metrics_tracer_factory = MetricsTracerFactory( - enabled, SPANNER_SERVICE_NAME + init_enabled, SPANNER_SERVICE_NAME ) if not HAS_OPENTELEMETRY_INSTALLED: return cls._metrics_tracer_factory @@ -80,9 +84,8 @@ def __new__(cls, enabled: bool = True) -> "SpannerMetricsTracerFactory": cls._generate_client_hash(client_uid) ) cls._metrics_tracer_factory.set_location(_get_cloud_region()) - - if cls._metrics_tracer_factory.enabled != enabled: - cls._metrics_tracer_factory.enabled = enabled + elif enabled is not None and cls._metrics_tracer_factory.enabled != enabled: + cls._metrics_tracer_factory.enabled = enabled return cls._metrics_tracer_factory diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/pool.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/pool.py index 74dbfa9e60ae..f22ad3846462 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/pool.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/pool.py @@ -28,6 +28,7 @@ from google.cloud.spanner_v1._helpers import ( _metadata_with_leader_aware_routing, _metadata_with_prefix, + _resource_info_from_database, ) from google.cloud.spanner_v1._opentelemetry_tracing import ( add_span_event, @@ -94,13 +95,7 @@ def __init__(self, labels=None, database_role=None): @property def _resource_info(self): """Resource information for metrics labels.""" - if self._database is None: - return None - return { - "project": self._database._instance._client.project, - "instance": self._database._instance.instance_id, - "database": self._database.database_id, - } + return _resource_info_from_database(self._database) @property def labels(self): diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/session.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/session.py index 618d1cc947a8..ae2124f313fa 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/session.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/session.py @@ -31,6 +31,7 @@ _get_retry_delay, _metadata_with_leader_aware_routing, _metadata_with_prefix, + _resource_info_from_database, ) from google.cloud.spanner_v1._opentelemetry_tracing import ( add_span_event, @@ -85,11 +86,7 @@ def __init__(self, database, labels=None, database_role=None, is_multiplexed=Fal @property def _resource_info(self): """Resource information for metrics labels.""" - return { - "project": self._database._instance._client.project, - "instance": self._database._instance.instance_id, - "database": self._database.database_id, - } + return _resource_info_from_database(self._database) def __lt__(self, other): return self._session_id < other._session_id @@ -310,7 +307,10 @@ def ping(self): api = database.spanner_api metadata = _metadata_with_prefix(database.name) nth_request = database._next_nth_request - with trace_call("CloudSpanner.Session.ping", self) as span: + with ( + trace_call("CloudSpanner.Session.ping", self) as span, + MetricsCapture(self._resource_info), + ): call_metadata, error_augmenter = database.with_error_augmentation( nth_request, 1, metadata, span ) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py index 91eafced6ca6..aed6a94ab574 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/snapshot.py @@ -47,6 +47,9 @@ ) from google.cloud.spanner_v1._opentelemetry_tracing import add_span_event, trace_call from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture +from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( + SpannerMetricsTracerFactory, +) from google.cloud.spanner_v1.streamed import StreamedResultSet from google.cloud.spanner_v1.types import MultiplexedSessionPrecommitToken from google.cloud.spanner_v1.types.mutation import Mutation @@ -114,28 +117,45 @@ def _restart_on_unavailable( nth_request = getattr(request_id_manager, "_next_nth_request", 0) current_request_id = None stream_finished = False + metrics_tracer = None + factory = SpannerMetricsTracerFactory() + if factory.enabled: + metrics_tracer = factory.create_metrics_tracer(resource_info) + if metrics_tracer: + metrics_tracer.record_operation_start() try: while True: try: if iterator is None: - with ( - trace_call( - trace_name, - session, - attributes, - observability_options=observability_options, - metadata=metadata, - ) as span, - MetricsCapture(resource_info), - ): + with trace_call( + trace_name, + session, + attributes, + observability_options=observability_options, + metadata=metadata, + ) as span: call_metadata, current_request_id = ( request_id_manager.metadata_and_request_id( nth_request, attempt, metadata, span ) ) - iterator = CrossSync._Sync_Impl.run_if_async( - method, request=request, metadata=call_metadata - ) + token = None + if metrics_tracer is not None: + token = SpannerMetricsTracerFactory.set_current_tracer( + metrics_tracer + ) + try: + iterator = CrossSync._Sync_Impl.run_if_async( + method, request=request, metadata=call_metadata + ) + finally: + if token is not None: + try: + SpannerMetricsTracerFactory.reset_current_tracer( + token + ) + except ValueError: + pass item: PartialResultSet for item in iterator: item_buffer.append(item) @@ -196,11 +216,18 @@ def _restart_on_unavailable( if stream_finished: break finally: - if iterator is not None and hasattr(iterator, "cancel"): - try: - iterator.cancel() - except Exception: - pass + try: + if metrics_tracer is not None: + try: + metrics_tracer.record_operation_completion() + except Exception: + pass + finally: + if iterator is not None and hasattr(iterator, "cancel"): + try: + iterator.cancel() + except Exception: + pass class _SnapshotBase(_SessionWrapper): @@ -228,16 +255,6 @@ def __init__(self, session, client_context=None): CrossSync._Sync_Impl.Event() ) - @property - def _resource_info(self): - """Resource information for metrics labels.""" - database = self._session._database - return { - "project": database._instance._client.project, - "instance": database._instance.instance_id, - "database": database.database_id, - } - def _wait_for_transaction_begin(self) -> None: """Claims the inline-begin for this request, or waits for it to complete. diff --git a/packages/google-cloud-spanner/tests/mockserver_tests/test_frontend_metrics.py b/packages/google-cloud-spanner/tests/mockserver_tests/test_frontend_metrics.py index 980ca3d1e380..efbe2f66d580 100644 --- a/packages/google-cloud-spanner/tests/mockserver_tests/test_frontend_metrics.py +++ b/packages/google-cloud-spanner/tests/mockserver_tests/test_frontend_metrics.py @@ -217,6 +217,208 @@ def test_gfe_missing_header_count_exported(self): finally: pass + def test_operation_metrics_no_duplicates(self): + add_select1_result() + reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[reader]) + + with ( + mock.patch( + "google.cloud.spanner_v1.metrics.metrics_tracer_factory.get_meter_provider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client._get_spanner_emulator_host", + return_value=None, + ), + ): + client = Client( + project="p", + credentials=AnonymousCredentials(), + client_options=ClientOptions( + api_endpoint="localhost:" + str(MockServerTestBase.port), + ), + ) + instance = client.instance("test-instance") + database = instance.database( + "test-database", + pool=FixedSizePool(size=10), + enable_interceptors_in_tests=True, + ) + database._interceptors.append(MetricsInterceptor()) + database._spanner_api = None # Force recreation with the new interceptor + + with database.snapshot() as snapshot: + results = snapshot.execute_sql("select 1") + list(results) + + metric_data = reader.get_metrics_data() + self.assertIsNotNone(metric_data) + metrics = { + metric.name: metric + for resource_metric in metric_data.resource_metrics + for scope_metric in resource_metric.scope_metrics + for metric in scope_metric.metrics + } + + self.assertIn("operation_count", metrics) + self.assertIn("operation_latencies", metrics) + + operation_count_metric = metrics["operation_count"] + streaming_points = [ + point + for point in operation_count_metric.data.data_points + if point.attributes.get("method") == "Spanner.ExecuteStreamingSql" + ] + self.assertEqual( + len(streaming_points), + 1, + f"Expected exactly 1 ExecuteStreamingSql data point, found {len(streaming_points)}", + ) + self.assertEqual(streaming_points[0].value, 1) + self.assertEqual(streaming_points[0].attributes.get("project_id"), "p") + self.assertEqual( + streaming_points[0].attributes.get("instance_id"), "test-instance" + ) + self.assertEqual( + streaming_points[0].attributes.get("database"), "test-database" + ) + + for point in operation_count_metric.data.data_points: + self.assertEqual(point.attributes.get("project_id"), "p") + self.assertEqual(point.attributes.get("instance_id"), "test-instance") + self.assertEqual(point.attributes.get("database"), "test-database") + + operation_latencies_metric = metrics["operation_latencies"] + streaming_latency_points = [ + point + for point in operation_latencies_metric.data.data_points + if point.attributes.get("method") == "Spanner.ExecuteStreamingSql" + ] + self.assertEqual( + len(streaming_latency_points), + 1, + f"Expected exactly 1 ExecuteStreamingSql latency point, found {len(streaming_latency_points)}", + ) + self.assertEqual(streaming_latency_points[0].count, 1) + self.assertEqual(streaming_latency_points[0].attributes.get("project_id"), "p") + self.assertEqual( + streaming_latency_points[0].attributes.get("instance_id"), "test-instance" + ) + self.assertEqual( + streaming_latency_points[0].attributes.get("database"), "test-database" + ) + + def test_operation_metrics_disabled_no_metrics_exported(self): + add_select1_result() + reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[reader]) + + with ( + mock.patch( + "google.cloud.spanner_v1.metrics.metrics_tracer_factory.get_meter_provider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client._get_spanner_emulator_host", + return_value=None, + ), + ): + SpannerMetricsTracerFactory(enabled=False) + try: + client = Client( + project="p", + credentials=AnonymousCredentials(), + client_options=ClientOptions( + api_endpoint="localhost:" + str(MockServerTestBase.port), + ), + ) + instance = client.instance("test-instance") + database = instance.database( + "test-database", + pool=FixedSizePool(size=10), + enable_interceptors_in_tests=True, + ) + database._interceptors.append(MetricsInterceptor()) + database._spanner_api = None + + with database.snapshot() as snapshot: + results = snapshot.execute_sql("select 1") + rows = list(results) + self.assertEqual(len(rows), 1) + finally: + SpannerMetricsTracerFactory(enabled=True) + + metric_data = reader.get_metrics_data() + if metric_data is not None: + metrics = { + metric.name: metric + for resource_metric in metric_data.resource_metrics + for scope_metric in resource_metric.scope_metrics + for metric in scope_metric.metrics + } + if "operation_count" in metrics: + self.assertEqual(len(metrics["operation_count"].data.data_points), 0) + if "operation_latencies" in metrics: + self.assertEqual( + len(metrics["operation_latencies"].data.data_points), 0 + ) + + def test_operation_metrics_completion_error_does_not_fail_query(self): + from google.cloud.spanner_v1.metrics.metrics_tracer import MetricsTracer + + add_select1_result() + reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[reader]) + + with ( + mock.patch( + "google.cloud.spanner_v1.metrics.metrics_tracer_factory.get_meter_provider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client._get_spanner_emulator_host", + return_value=None, + ), + mock.patch.object( + MetricsTracer, + "record_operation_completion", + side_effect=RuntimeError("telemetry export error"), + ), + ): + client = Client( + project="p", + credentials=AnonymousCredentials(), + client_options=ClientOptions( + api_endpoint="localhost:" + str(MockServerTestBase.port), + ), + ) + instance = client.instance("test-instance") + database = instance.database( + "test-database", + pool=FixedSizePool(size=10), + enable_interceptors_in_tests=True, + ) + database._interceptors.append(MetricsInterceptor()) + database._spanner_api = None + + with database.snapshot() as snapshot: + results = snapshot.execute_sql("select 1") + rows = list(results) + self.assertEqual(len(rows), 1) + class TestFrontendMetricsAsyncIntegration(AsyncMockServerTestBase): def setUp(self): @@ -450,3 +652,254 @@ async def test_async_gfe_missing_header_count_exported(self): ] self.assertTrue(len(streaming_points) > 0) self.assertGreaterEqual(streaming_points[0].value, 1) + + async def test_async_operation_metrics_no_duplicates(self): + add_select1_result() + reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[reader]) + + with ( + mock.patch( + "google.cloud.spanner_v1.metrics.metrics_tracer_factory.get_meter_provider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1._async.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client._get_spanner_emulator_host", + return_value=None, + ), + mock.patch( + "google.cloud.spanner_v1._async.client._get_spanner_emulator_host", + return_value=None, + ), + ): + client = AsyncClient( + project="p", + credentials=AnonymousCredentials(), + client_options=ClientOptions( + api_endpoint="localhost:" + str(AsyncMockServerTestBase.port), + ), + ) + instance = client.instance("test-instance") + database = await instance.database( + "test-database", + pool=AsyncFixedSizePool(size=10), + ) + async with grpc.aio.insecure_channel( + "localhost:" + str(AsyncMockServerTestBase.port), + ) as channel: + transport = SpannerGrpcAsyncIOTransport( + channel=channel, + metrics_interceptor=MetricsInterceptor(), + ) + database._spanner_api = SpannerAsyncClient( + client_info=client._client_info, + transport=transport, + ) + + async with database.snapshot() as snapshot: + results = await snapshot.execute_sql("select 1") + async for _ in results: + pass + + metric_data = reader.get_metrics_data() + self.assertIsNotNone(metric_data) + metrics = { + metric.name: metric + for resource_metric in metric_data.resource_metrics + for scope_metric in resource_metric.scope_metrics + for metric in scope_metric.metrics + } + + self.assertIn("operation_count", metrics) + self.assertIn("operation_latencies", metrics) + + operation_count_metric = metrics["operation_count"] + streaming_points = [ + point + for point in operation_count_metric.data.data_points + if point.attributes.get("method") == "Spanner.ExecuteStreamingSql" + ] + self.assertEqual( + len(streaming_points), + 1, + f"Expected exactly 1 ExecuteStreamingSql data point, found {len(streaming_points)}", + ) + self.assertEqual(streaming_points[0].value, 1) + self.assertEqual(streaming_points[0].attributes.get("project_id"), "p") + self.assertEqual( + streaming_points[0].attributes.get("instance_id"), "test-instance" + ) + self.assertEqual( + streaming_points[0].attributes.get("database"), "test-database" + ) + + for point in operation_count_metric.data.data_points: + self.assertEqual(point.attributes.get("project_id"), "p") + self.assertEqual(point.attributes.get("instance_id"), "test-instance") + self.assertEqual(point.attributes.get("database"), "test-database") + + operation_latencies_metric = metrics["operation_latencies"] + streaming_latency_points = [ + point + for point in operation_latencies_metric.data.data_points + if point.attributes.get("method") == "Spanner.ExecuteStreamingSql" + ] + self.assertEqual( + len(streaming_latency_points), + 1, + f"Expected exactly 1 ExecuteStreamingSql latency point, found {len(streaming_latency_points)}", + ) + self.assertEqual(streaming_latency_points[0].count, 1) + self.assertEqual(streaming_latency_points[0].attributes.get("project_id"), "p") + self.assertEqual( + streaming_latency_points[0].attributes.get("instance_id"), "test-instance" + ) + self.assertEqual( + streaming_latency_points[0].attributes.get("database"), "test-database" + ) + + async def test_async_operation_metrics_disabled_no_metrics_exported(self): + add_select1_result() + reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[reader]) + + with ( + mock.patch( + "google.cloud.spanner_v1.metrics.metrics_tracer_factory.get_meter_provider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1._async.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client._get_spanner_emulator_host", + return_value=None, + ), + mock.patch( + "google.cloud.spanner_v1._async.client._get_spanner_emulator_host", + return_value=None, + ), + ): + SpannerMetricsTracerFactory(enabled=False) + try: + client = AsyncClient( + project="p", + credentials=AnonymousCredentials(), + client_options=ClientOptions( + api_endpoint="localhost:" + str(AsyncMockServerTestBase.port), + ), + ) + instance = client.instance("test-instance") + database = await instance.database( + "test-database", + pool=AsyncFixedSizePool(size=10), + ) + async with grpc.aio.insecure_channel( + "localhost:" + str(AsyncMockServerTestBase.port), + ) as channel: + transport = SpannerGrpcAsyncIOTransport( + channel=channel, + metrics_interceptor=MetricsInterceptor(), + ) + database._spanner_api = SpannerAsyncClient( + client_info=client._client_info, + transport=transport, + ) + + async with database.snapshot() as snapshot: + results = await snapshot.execute_sql("select 1") + async for _ in results: + pass + finally: + SpannerMetricsTracerFactory(enabled=True) + + metric_data = reader.get_metrics_data() + if metric_data is not None: + metrics = { + metric.name: metric + for resource_metric in metric_data.resource_metrics + for scope_metric in resource_metric.scope_metrics + for metric in scope_metric.metrics + } + if "operation_count" in metrics: + self.assertEqual(len(metrics["operation_count"].data.data_points), 0) + if "operation_latencies" in metrics: + self.assertEqual( + len(metrics["operation_latencies"].data.data_points), 0 + ) + + async def test_async_operation_metrics_completion_error_does_not_fail_query(self): + from google.cloud.spanner_v1.metrics.metrics_tracer import MetricsTracer + + add_select1_result() + reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[reader]) + + with ( + mock.patch( + "google.cloud.spanner_v1.metrics.metrics_tracer_factory.get_meter_provider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1._async.client.MeterProvider", + return_value=meter_provider, + ), + mock.patch( + "google.cloud.spanner_v1.client._get_spanner_emulator_host", + return_value=None, + ), + mock.patch( + "google.cloud.spanner_v1._async.client._get_spanner_emulator_host", + return_value=None, + ), + mock.patch.object( + MetricsTracer, + "record_operation_completion", + side_effect=RuntimeError("telemetry export error"), + ), + ): + client = AsyncClient( + project="p", + credentials=AnonymousCredentials(), + client_options=ClientOptions( + api_endpoint="localhost:" + str(AsyncMockServerTestBase.port), + ), + ) + instance = client.instance("test-instance") + database = await instance.database( + "test-database", + pool=AsyncFixedSizePool(size=10), + ) + async with grpc.aio.insecure_channel( + "localhost:" + str(AsyncMockServerTestBase.port), + ) as channel: + transport = SpannerGrpcAsyncIOTransport( + channel=channel, + metrics_interceptor=MetricsInterceptor(), + ) + database._spanner_api = SpannerAsyncClient( + client_info=client._client_info, + transport=transport, + ) + + async with database.snapshot() as snapshot: + results = await snapshot.execute_sql("select 1") + async for _ in results: + pass diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_batch.py b/packages/google-cloud-spanner/tests/unit/_async/test_batch.py index 12b465be877e..b5a5a958f718 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_batch.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_batch.py @@ -50,6 +50,33 @@ def test_ctor(self): self.assertIs(base._session, session) self.assertEqual(len(base._mutations), 0) + def test_resource_info_with_database(self): + session = mock.Mock() + mock_db = mock.Mock(spec=["_resource_info"]) + mock_db._resource_info = {"project": "p", "instance": "i", "database": "d"} + session._database = mock_db + base = self._make_one(session) + self.assertEqual(base._resource_info, mock_db._resource_info) + + def test_resource_info_fallback(self): + session = mock.Mock() + mock_db = mock.Mock(spec=["_instance", "database_id"]) + mock_db._instance = mock.Mock(spec=["_client", "instance_id"]) + mock_db._instance.instance_id = "i" + mock_db._instance._client = mock.Mock(spec=["project"]) + mock_db._instance._client.project = "p" + mock_db.database_id = "d" + session._database = mock_db + base = self._make_one(session) + self.assertEqual( + base._resource_info, + {"project": "p", "instance": "i", "database": "d"}, + ) + + def test_resource_info_none(self): + base = self._make_one(None) + self.assertIsNone(base._resource_info) + def test_insert(self): session = mock.Mock() base = self._make_one(session) @@ -259,6 +286,33 @@ def _make_session(self): return session + def test_resource_info_with_database(self): + session = mock.Mock() + mock_db = mock.Mock(spec=["_resource_info"]) + mock_db._resource_info = {"project": "p", "instance": "i", "database": "d"} + session._database = mock_db + groups = self._make_one(session) + self.assertEqual(groups._resource_info, mock_db._resource_info) + + def test_resource_info_fallback(self): + session = mock.Mock() + mock_db = mock.Mock(spec=["_instance", "database_id"]) + mock_db._instance = mock.Mock(spec=["_client", "instance_id"]) + mock_db._instance.instance_id = "i" + mock_db._instance._client = mock.Mock(spec=["project"]) + mock_db._instance._client.project = "p" + mock_db.database_id = "d" + session._database = mock_db + groups = self._make_one(session) + self.assertEqual( + groups._resource_info, + {"project": "p", "instance": "i", "database": "d"}, + ) + + def test_resource_info_none(self): + groups = self._make_one(None) + self.assertIsNone(groups._resource_info) + async def test_batch_write_ok(self): session = self._make_session() diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_client_extra.py b/packages/google-cloud-spanner/tests/unit/_async/test_client_extra.py index 6dbba5bacc25..cca47a952ea8 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_client_extra.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_client_extra.py @@ -240,3 +240,38 @@ def test_initialize_metrics_emulator_branch(self): ) as set_mock: MUT._initialize_metrics("project", self.credentials) set_mock.assert_called_once() + + def test_initialize_metrics_emulator_host_parameter(self): + from google.cloud.spanner_v1._async import client as MUT + + with mock.patch( + "google.cloud.spanner_v1._async.client._metrics_monitor_initialized", + False, + ): + with mock.patch( + "google.cloud.spanner_v1._async.client.metrics.set_meter_provider" + ) as set_mock: + MUT._initialize_metrics( + "project", + self.credentials, + emulator_host="localhost:9010", + ) + set_mock.assert_called_once() + + def test_initialize_metrics_anonymous_credentials(self): + from google.auth.credentials import AnonymousCredentials + + from google.cloud.spanner_v1._async import client as MUT + + with mock.patch( + "google.cloud.spanner_v1._async.client._metrics_monitor_initialized", + False, + ): + with mock.patch( + "google.cloud.spanner_v1._async.client.metrics.set_meter_provider" + ) as set_mock: + MUT._initialize_metrics( + "project", + AnonymousCredentials(), + ) + set_mock.assert_called_once() diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_database_extra.py b/packages/google-cloud-spanner/tests/unit/_async/test_database_extra.py index ea4d6643605c..272ec504cbc5 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_database_extra.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_database_extra.py @@ -470,7 +470,7 @@ async def test_batch_commit_options_coverage(self): # coverage for line 1603, 1655 db = Database("db", self.instance) await db._pool.bind(db) - batch = db.batch() + _ = db.batch() mock_session = mock.MagicMock() mock_session.name = "projects/p/instances/i/databases/db/sessions/s" @@ -511,3 +511,35 @@ async def test_batch_commit_options_coverage(self): sc = SnapshotCheckout(db) _ = sc._resource_info # line 1655 + + def test_resource_info(self): + db = Database("db", self.instance) + res_info = db._resource_info + self.assertEqual(res_info["database"], "db") + self.assertEqual(res_info["project"], self.instance._client.project) + self.assertEqual(res_info["instance"], "i") + self.assertIsNotNone(getattr(res_info, "_client_attributes", None)) + self.assertEqual(res_info._client_attributes["database"], "db") + self.assertEqual( + res_info._client_attributes["project_id"], self.instance._client.project + ) + self.assertEqual(res_info._client_attributes["instance_id"], "i") + # Caching check: second access returns identical instance + self.assertIs(db._resource_info, res_info) + + def test_resource_info_fallback_on_exception(self): + db = Database("db", self.instance) + from google.cloud.spanner_v1.metrics.metrics_tracer_factory import ( + MetricsTracerFactory, + ) + + with mock.patch.object( + MetricsTracerFactory, + "create_resource_info", + side_effect=RuntimeError("Boom"), + ): + res_info = db._resource_info + self.assertEqual(res_info["database"], "db") + self.assertEqual(res_info["project"], self.instance._client.project) + self.assertEqual(res_info["instance"], "i") + self.assertIsNone(getattr(res_info, "_client_attributes", None)) diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_pool.py b/packages/google-cloud-spanner/tests/unit/_async/test_pool.py index 01f584abe57a..b8ff662e37ed 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_pool.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_pool.py @@ -151,6 +151,27 @@ def test_resource_info_unbound(self): pool = self._make_one() self.assertIsNone(pool._resource_info) + def test_resource_info_bound(self): + pool = self._make_one() + mock_db = mock.Mock(spec=["_resource_info"]) + mock_db._resource_info = {"project": "p", "instance": "i", "database": "d"} + pool._database = mock_db + self.assertEqual(pool._resource_info, mock_db._resource_info) + + def test_resource_info_fallback(self): + pool = self._make_one() + mock_db = mock.Mock(spec=["_instance", "database_id"]) + mock_db._instance = mock.Mock(spec=["_client", "instance_id"]) + mock_db._instance.instance_id = "i" + mock_db._instance._client = mock.Mock(spec=["project"]) + mock_db._instance._client.project = "p" + mock_db.database_id = "d" + pool._database = mock_db + self.assertEqual( + pool._resource_info, + {"project": "p", "instance": "i", "database": "d"}, + ) + def test_session_deprecated(self): import warnings diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_session.py b/packages/google-cloud-spanner/tests/unit/_async/test_session.py index 3159b55a7ec9..8cebc5a7807b 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_session.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_session.py @@ -706,6 +706,44 @@ async def test_ping_hit(self, mock_region): attributes=dict(self.BASE_ATTRIBUTES, x_goog_spanner_request_id=req_id), ) + @mock.patch( + "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", + return_value="global", + ) + @CrossSync.pytest + async def test_ping_invokes_metrics_capture(self, mock_region): + gax_api = self._make_spanner_api() + gax_api.execute_sql.return_value = "1" + database = self._make_database() + database.spanner_api = gax_api + session = self._make_one(database) + session._session_id = self.SESSION_ID + + with mock.patch( + "google.cloud.spanner_v1._async.session.MetricsCapture" + ) as mock_capture: + await session.ping() + mock_capture.assert_called_once_with(session._resource_info) + + def test_resource_info_fallback_when_database_lacks_resource_info(self): + database = mock.Mock(spec=["_instance", "database_id"]) + database.database_id = "test-db" + instance = mock.Mock(spec=["_client", "instance_id"]) + instance.instance_id = "test-inst" + client = mock.Mock(spec=["project"]) + client.project = "test-proj" + instance._client = client + database._instance = instance + + session = self._make_one(database) + self.assertEqual( + session._resource_info, + {"project": "test-proj", "instance": "test-inst", "database": "test-db"}, + ) + + session._database = None + self.assertIsNone(session._resource_info) + @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", return_value="global", diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py b/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py index 39926148d9de..4126fccf0e6d 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_snapshot.py @@ -21,6 +21,7 @@ from google.api_core.exceptions import ( InternalServerError, InvalidArgument, + NotFound, ServiceUnavailable, ) @@ -69,6 +70,27 @@ def _make_snapshot(self, *args, **kwargs): s._lock = CrossSync.Lock() return s + def test_resource_info_fallback_when_database_lacks_resource_info(self): + database = mock.Mock(spec=["_instance", "database_id"]) + database.database_id = "test-db" + instance = mock.Mock(spec=["_client", "instance_id"]) + instance.instance_id = "test-inst" + client = mock.Mock(spec=["project"]) + client.project = "test-proj" + instance._client = client + database._instance = instance + + session = mock.Mock() + session._database = database + snapshot = self._make_snapshot(session, multi_use=True) + self.assertEqual( + snapshot._resource_info, + {"project": "test-proj", "instance": "test-inst", "database": "test-db"}, + ) + + snapshot._session = None + self.assertIsNone(snapshot._resource_info) + async def test_read_errors(self): snapshot = self._make_snapshot(_Session(), multi_use=False) snapshot._read_request_count = 1 @@ -503,6 +525,450 @@ async def test_restart_on_unavailable_retry_before_last(self): mock_drain.assert_called_once_with(stream_two) stream_two.cancel.assert_not_called() + async def test_restart_on_unavailable_metrics_tracer_lifecycle(self): + from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable + from google.cloud.spanner_v1.metrics.metrics_tracer_factory import ( + MetricsTracerFactory, + ) + from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( + SpannerMetricsTracerFactory, + ) + + mock_tracer = mock.Mock() + + item = mock.Mock( + value=0, + resume_token=b"", + metadata=None, + precommit_token=None, + last=False, + _pb=None, + spec=["value", "resume_token", "metadata", "precommit_token", "last"], + ) + raw = _MockIterator(item) + tracer_during_call = [] + + def restart(*args, **kwargs): + tracer_during_call.append(SpannerMetricsTracerFactory.get_current_tracer()) + return raw + + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + session = _Session(database) + snapshot = self._make_snapshot(session, multi_use=True) + resource_info = {"project": "p", "instance": "i", "database": "d"} + + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=mock_tracer + ) as mock_create_tracer: + resumable = _restart_on_unavailable( + restart, + request, + session=session, + transaction=snapshot, + request_id_manager=session._database, + resource_info=resource_info, + ) + + # Not started before iteration begins (lazy generator) + mock_tracer.record_operation_start.assert_not_called() + + items = [] + async for it in resumable: + # During row iteration, ambient contextvar must NOT be contaminated + self.assertIsNone(SpannerMetricsTracerFactory.get_current_tracer()) + items.append(it) + + self.assertEqual(items, [item]) + # Tracer factory was called with resource info + mock_create_tracer.assert_called_once_with(resource_info) + # Operation was started once + mock_tracer.record_operation_start.assert_called_once() + # Contextvar was active during RPC invocation + self.assertEqual(tracer_during_call, [mock_tracer]) + # Operation completion recorded once + mock_tracer.record_operation_completion.assert_called_once() + # Context is clean after completion + self.assertIsNone(SpannerMetricsTracerFactory.get_current_tracer()) + + async def test_restart_on_unavailable_metrics_tracer_spans_retries(self): + from google.api_core.exceptions import ServiceUnavailable + + from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable + from google.cloud.spanner_v1.metrics.metrics_tracer_factory import ( + MetricsTracerFactory, + ) + from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( + SpannerMetricsTracerFactory, + ) + + mock_tracer = mock.Mock() + + item_one = mock.Mock( + value=0, + resume_token=b"token", + metadata=None, + precommit_token=None, + last=False, + _pb=None, + spec=["value", "resume_token", "metadata", "precommit_token", "last"], + ) + item_two = mock.Mock( + value=1, + resume_token=b"", + metadata=None, + precommit_token=None, + last=False, + _pb=None, + spec=["value", "resume_token", "metadata", "precommit_token", "last"], + ) + before = _MockIterator( + item_one, fail_after=True, error=ServiceUnavailable("testing") + ) + after = _MockIterator(item_two) + tracers_during_calls = [] + + def restart(*args, **kwargs): + tracers_during_calls.append( + SpannerMetricsTracerFactory.get_current_tracer() + ) + if len(tracers_during_calls) == 1: + return before + return after + + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + session = _Session(database) + snapshot = self._make_snapshot(session, multi_use=True) + resource_info = {"project": "p", "instance": "i", "database": "d"} + + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=mock_tracer + ): + resumable = _restart_on_unavailable( + restart, + request, + session=session, + transaction=snapshot, + request_id_manager=session._database, + resource_info=resource_info, + ) + + items = [] + async for it in resumable: + items.append(it) + + self.assertEqual(items, [item_one, item_two]) + self.assertEqual(len(tracers_during_calls), 2) + # Both attempts used the same metrics tracer + self.assertEqual(tracers_during_calls, [mock_tracer, mock_tracer]) + mock_tracer.record_operation_start.assert_called_once() + mock_tracer.record_operation_completion.assert_called_once() + self.assertIsNone(SpannerMetricsTracerFactory.get_current_tracer()) + + async def test_restart_on_unavailable_early_break_no_context_error(self): + from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable + from google.cloud.spanner_v1.metrics.metrics_tracer_factory import ( + MetricsTracerFactory, + ) + from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( + SpannerMetricsTracerFactory, + ) + + mock_tracer = mock.Mock() + + item_one = mock.Mock( + value=0, + resume_token=b"", + metadata=None, + precommit_token=None, + last=False, + _pb=None, + spec=["value", "resume_token", "metadata", "precommit_token", "last"], + ) + item_two = mock.Mock( + value=1, + resume_token=b"", + metadata=None, + precommit_token=None, + last=False, + _pb=None, + spec=["value", "resume_token", "metadata", "precommit_token", "last"], + ) + raw = _MockIterator(item_one, item_two) + restart = mock.Mock(return_value=raw) + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + session = _Session(database) + snapshot = self._make_snapshot(session, multi_use=True) + resource_info = {"project": "p", "instance": "i", "database": "d"} + + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=mock_tracer + ): + resumable = _restart_on_unavailable( + restart, + request, + session=session, + transaction=snapshot, + request_id_manager=session._database, + resource_info=resource_info, + ) + + # Consume only the first item and break early. + # In prior implementation, breaking early caused ValueError: + # Token was created in a different Context during aclose() cleanup. + async for it in resumable: + self.assertEqual(it, item_one) + break + await resumable.aclose() + + # Must invoke operation completion on early termination/generator cleanup + mock_tracer.record_operation_completion.assert_called_once() + self.assertIsNone(SpannerMetricsTracerFactory.get_current_tracer()) + + async def test_restart_on_unavailable_real_factory_early_break_no_context_error( + self, + ): + # End-to-end test using real SpannerMetricsTracerFactory and contextvars + from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable + from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( + SpannerMetricsTracerFactory, + ) + + item = mock.Mock( + value=0, + resume_token=b"", + metadata=None, + precommit_token=None, + last=False, + _pb=None, + spec=["value", "resume_token", "metadata", "precommit_token", "last"], + ) + raw = _MockIterator(item) + restart = mock.Mock(return_value=raw) + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + session = _Session(database) + snapshot = self._make_snapshot(session, multi_use=True) + resource_info = {"project": "p", "instance": "i", "database": "d"} + + resumable = _restart_on_unavailable( + restart, + request, + session=session, + transaction=snapshot, + request_id_manager=session._database, + resource_info=resource_info, + ) + + # Breaking early must not raise ValueError during aclose() + async for it in resumable: + self.assertEqual(it, item) + break + await resumable.aclose() + + self.assertIsNone(SpannerMetricsTracerFactory.get_current_tracer()) + + async def test_restart_on_unavailable_metrics_tracer_partial_resource_info_and_disabled( + self, + ): + from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable + from google.cloud.spanner_v1.metrics.metrics_tracer_factory import ( + MetricsTracerFactory, + ) + from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( + SpannerMetricsTracerFactory, + ) + + mock_tracer = mock.Mock() + item = mock.Mock( + value=0, + resume_token=b"", + metadata=None, + precommit_token=None, + last=False, + _pb=None, + spec=["value", "resume_token", "metadata", "precommit_token", "last"], + ) + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + session = _Session(database) + snapshot = self._make_snapshot(session, multi_use=True) + + # 1. resource_info is None + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=mock_tracer + ): + resumable = _restart_on_unavailable( + mock.Mock(return_value=_MockIterator(item)), + request, + session=session, + transaction=snapshot, + request_id_manager=session._database, + resource_info=None, + ) + async for it in resumable: + pass + mock_tracer.record_operation_start.assert_called_once() + mock_tracer.record_operation_completion.assert_called_once() + + # 2. partial resource_info (only project, only instance, only database) + for partial_info in [{"project": "p"}, {"instance": "i"}, {"database": "d"}]: + mock_tracer.reset_mock() + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=mock_tracer + ) as mock_create_tracer: + resumable = _restart_on_unavailable( + mock.Mock(return_value=_MockIterator(item)), + request, + session=session, + transaction=snapshot, + request_id_manager=session._database, + resource_info=partial_info, + ) + async for it in resumable: + pass + mock_create_tracer.assert_called_once_with(partial_info) + + # 3. metrics factory disabled + SpannerMetricsTracerFactory(enabled=False) + try: + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer" + ) as mock_create: + resumable = _restart_on_unavailable( + mock.Mock(return_value=_MockIterator(item)), + request, + session=session, + transaction=snapshot, + request_id_manager=session._database, + resource_info={"project": "p"}, + ) + items = [] + async for it in resumable: + items.append(it) + self.assertEqual(items, [item]) + mock_create.assert_not_called() + self.assertIsNone(SpannerMetricsTracerFactory.get_current_tracer()) + finally: + SpannerMetricsTracerFactory(enabled=True) + + # 4. create_metrics_tracer returns None + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=None + ): + resumable = _restart_on_unavailable( + mock.Mock(return_value=_MockIterator(item)), + request, + session=session, + transaction=snapshot, + request_id_manager=session._database, + resource_info={"project": "p"}, + ) + items = [] + async for it in resumable: + items.append(it) + self.assertEqual(items, [item]) + self.assertIsNone(SpannerMetricsTracerFactory.get_current_tracer()) + + async def test_restart_on_unavailable_metrics_completion_error_handled(self): + from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable + from google.cloud.spanner_v1.metrics.metrics_tracer_factory import ( + MetricsTracerFactory, + ) + + mock_tracer = mock.Mock() + mock_tracer.record_operation_completion.side_effect = RuntimeError( + "metrics export failure" + ) + + item = mock.Mock( + value=0, + resume_token=b"", + metadata=None, + precommit_token=None, + last=False, + _pb=None, + spec=["value", "resume_token", "metadata", "precommit_token", "last"], + ) + mock_iterator = _MockIterator(item) + mock_iterator.cancel = mock.Mock() + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + session = _Session(database) + snapshot = self._make_snapshot(session, multi_use=True) + + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=mock_tracer + ): + resumable = _restart_on_unavailable( + mock.Mock(return_value=mock_iterator), + request, + session=session, + transaction=snapshot, + request_id_manager=session._database, + resource_info={"project": "p"}, + ) + items = [] + async for it in resumable: + items.append(it) + break + await resumable.aclose() + + self.assertEqual(items, [item]) + mock_tracer.record_operation_completion.assert_called_once() + mock_iterator.cancel.assert_called_once() + + async def test_restart_on_unavailable_metrics_completion_error_preserves_query_error( + self, + ): + from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable + from google.cloud.spanner_v1.metrics.metrics_tracer_factory import ( + MetricsTracerFactory, + ) + + mock_tracer = mock.Mock() + mock_tracer.record_operation_completion.side_effect = RuntimeError( + "metrics export failure" + ) + + item = mock.Mock( + value=0, + resume_token=b"", + metadata=None, + precommit_token=None, + last=False, + _pb=None, + spec=["value", "resume_token", "metadata", "precommit_token", "last"], + ) + mock_iterator = _MockIterator( + item, fail_after=True, error=NotFound("table not found") + ) + mock_iterator.cancel = mock.Mock() + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + session = _Session(database) + snapshot = self._make_snapshot(session, multi_use=True) + + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=mock_tracer + ): + resumable = _restart_on_unavailable( + mock.Mock(return_value=mock_iterator), + request, + session=session, + transaction=snapshot, + request_id_manager=session._database, + resource_info={"project": "p"}, + ) + with self.assertRaises(NotFound): + async for _ in resumable: + pass + + mock_tracer.record_operation_completion.assert_called_once() + mock_iterator.cancel.assert_called_once() + async def test_execute_sql_ok(self): database = _Database() fields = [StructType.Field(name="col", type_=Type(code=TypeCode.STRING))] diff --git a/packages/google-cloud-spanner/tests/unit/test__opentelemetry_tracing.py b/packages/google-cloud-spanner/tests/unit/test__opentelemetry_tracing.py index 8d04a85720bc..2aea6deb9c34 100644 --- a/packages/google-cloud-spanner/tests/unit/test__opentelemetry_tracing.py +++ b/packages/google-cloud-spanner/tests/unit/test__opentelemetry_tracing.py @@ -676,3 +676,52 @@ def test_add_span_event(self): span = mock.Mock() _opentelemetry_tracing.add_span_event(span, "test_event", {"attr": "val"}) span.add_event.assert_called_once_with("test_event", {"attr": "val"}) + + @mock.patch("google.cloud.spanner_v1.metrics.metrics_capture.MetricsCapture") + def test_trace_call_does_not_invoke_metrics_capture(self, mock_metrics_capture): + from opentelemetry import trace + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + + # 1. No-op path (no ambient span) + with _opentelemetry_tracing.trace_call( + "CloudSpanner.TestNoOpSpan", + _make_session(), + observability_options={"tracer_provider": trace.NoOpTracerProvider()}, + ) as noop_span: + self.assertIs(noop_span, trace.INVALID_SPAN) + mock_metrics_capture.assert_not_called() + + # 2. No-op path with ambient valid span context (sub-branch 1a) + valid_span_context = trace.SpanContext( + trace_id=0x12345678123456781234567812345678, + span_id=0x1234567812345678, + is_remote=False, + trace_flags=trace.TraceFlags(trace.TraceFlags.SAMPLED), + ) + ambient_span = trace.NonRecordingSpan(valid_span_context) + with trace.use_span(ambient_span): + with _opentelemetry_tracing.trace_call( + "CloudSpanner.TestNoOpSpanWithAmbient", + _make_session(), + observability_options={"tracer_provider": trace.NoOpTracerProvider()}, + ) as child_span: + self.assertIsInstance(child_span, trace.NonRecordingSpan) + mock_metrics_capture.assert_not_called() + + # 3. Active span path + tracer_provider = TracerProvider() + exporter = InMemorySpanExporter() + tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + observability_options = dict(tracer_provider=tracer_provider) + + with _opentelemetry_tracing.trace_call( + "CloudSpanner.TestActiveSpan", + _make_session(), + observability_options=observability_options, + ): + pass + mock_metrics_capture.assert_not_called() diff --git a/packages/google-cloud-spanner/tests/unit/test_batch.py b/packages/google-cloud-spanner/tests/unit/test_batch.py index 247bdb58ab6c..b463eadfc659 100644 --- a/packages/google-cloud-spanner/tests/unit/test_batch.py +++ b/packages/google-cloud-spanner/tests/unit/test_batch.py @@ -104,6 +104,33 @@ def test_ctor(self): self.assertIs(base._session, session) self.assertEqual(len(base._mutations), 0) + def test_resource_info_with_database(self): + session = mock.Mock() + mock_db = mock.Mock(spec=["_resource_info"]) + mock_db._resource_info = {"project": "p", "instance": "i", "database": "d"} + session._database = mock_db + base = self._make_one(session) + self.assertEqual(base._resource_info, mock_db._resource_info) + + def test_resource_info_fallback(self): + session = mock.Mock() + mock_db = mock.Mock(spec=["_instance", "database_id"]) + mock_db._instance = mock.Mock(spec=["_client", "instance_id"]) + mock_db._instance.instance_id = "i" + mock_db._instance._client = mock.Mock(spec=["project"]) + mock_db._instance._client.project = "p" + mock_db.database_id = "d" + session._database = mock_db + base = self._make_one(session) + self.assertEqual( + base._resource_info, + {"project": "p", "instance": "i", "database": "d"}, + ) + + def test_resource_info_none(self): + base = self._make_one(None) + self.assertIsNone(base._resource_info) + def test_insert(self): session = _Session() base = self._make_one(session) @@ -667,6 +694,33 @@ def test_ctor(self): groups = self._make_one(session) self.assertIs(groups._session, session) + def test_resource_info_with_database(self): + session = mock.Mock() + mock_db = mock.Mock(spec=["_resource_info"]) + mock_db._resource_info = {"project": "p", "instance": "i", "database": "d"} + session._database = mock_db + groups = self._make_one(session) + self.assertEqual(groups._resource_info, mock_db._resource_info) + + def test_resource_info_fallback(self): + session = mock.Mock() + mock_db = mock.Mock(spec=["_instance", "database_id"]) + mock_db._instance = mock.Mock(spec=["_client", "instance_id"]) + mock_db._instance.instance_id = "i" + mock_db._instance._client = mock.Mock(spec=["project"]) + mock_db._instance._client.project = "p" + mock_db.database_id = "d" + session._database = mock_db + groups = self._make_one(session) + self.assertEqual( + groups._resource_info, + {"project": "p", "instance": "i", "database": "d"}, + ) + + def test_resource_info_none(self): + groups = self._make_one(None) + self.assertIsNone(groups._resource_info) + @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", return_value="global", diff --git a/packages/google-cloud-spanner/tests/unit/test_client.py b/packages/google-cloud-spanner/tests/unit/test_client.py index 241b4efa02c3..c8d4a9895d20 100644 --- a/packages/google-cloud-spanner/tests/unit/test_client.py +++ b/packages/google-cloud-spanner/tests/unit/test_client.py @@ -1111,3 +1111,57 @@ def test_constructor_w_omni_explicit_credentials_instance(self): credentials=creds, ) self.assertIs(client._credentials, creds) + + +class TestInitializeMetrics(unittest.TestCase): + def test_initialize_metrics_emulator_branch(self): + from google.cloud.spanner_v1 import client as MUT + + with mock.patch( + "google.cloud.spanner_v1.client._get_spanner_emulator_host", + return_value="localhost", + ): + with mock.patch( + "google.cloud.spanner_v1.client._metrics_monitor_initialized", + False, + ): + with mock.patch( + "google.cloud.spanner_v1.client.metrics.set_meter_provider" + ) as set_mock: + MUT._initialize_metrics("project", mock.Mock()) + set_mock.assert_called_once() + + def test_initialize_metrics_emulator_host_parameter(self): + from google.cloud.spanner_v1 import client as MUT + + with mock.patch( + "google.cloud.spanner_v1.client._metrics_monitor_initialized", + False, + ): + with mock.patch( + "google.cloud.spanner_v1.client.metrics.set_meter_provider" + ) as set_mock: + MUT._initialize_metrics( + "project", + mock.Mock(), + emulator_host="localhost:9010", + ) + set_mock.assert_called_once() + + def test_initialize_metrics_anonymous_credentials(self): + from google.auth.credentials import AnonymousCredentials + + from google.cloud.spanner_v1 import client as MUT + + with mock.patch( + "google.cloud.spanner_v1.client._metrics_monitor_initialized", + False, + ): + with mock.patch( + "google.cloud.spanner_v1.client.metrics.set_meter_provider" + ) as set_mock: + MUT._initialize_metrics( + "project", + AnonymousCredentials(), + ) + set_mock.assert_called_once() diff --git a/packages/google-cloud-spanner/tests/unit/test_database.py b/packages/google-cloud-spanner/tests/unit/test_database.py index b63e19a61baf..a0f6043d9407 100644 --- a/packages/google-cloud-spanner/tests/unit/test_database.py +++ b/packages/google-cloud-spanner/tests/unit/test_database.py @@ -141,6 +141,44 @@ def test_ctor_defaults(self): self.assertIsNone(database.database_role) self.assertTrue(database._route_to_leader_enabled, True) + def test_resource_info(self): + client = _Client() + instance = _Instance(self.INSTANCE_NAME, client=client) + database = self._make_one(self.DATABASE_ID, instance) + + res_info = database._resource_info + self.assertEqual(res_info["project"], self.PROJECT_ID) + self.assertEqual(res_info["instance"], self.INSTANCE_ID) + self.assertEqual(res_info["database"], self.DATABASE_ID) + self.assertIsNotNone(getattr(res_info, "_client_attributes", None)) + self.assertEqual(res_info._client_attributes["project_id"], self.PROJECT_ID) + self.assertEqual(res_info._client_attributes["instance_id"], self.INSTANCE_ID) + self.assertEqual(res_info._client_attributes["database"], self.DATABASE_ID) + + # Caching check: second access returns identical instance + self.assertIs(database._resource_info, res_info) + + def test_resource_info_fallback_on_exception(self): + client = _Client() + instance = _Instance(self.INSTANCE_NAME, client=client) + database = self._make_one(self.DATABASE_ID, instance) + database._cached_resource_info = None + + from google.cloud.spanner_v1.metrics.metrics_tracer_factory import ( + MetricsTracerFactory, + ) + + with mock.patch.object( + MetricsTracerFactory, + "create_resource_info", + side_effect=RuntimeError("Boom"), + ): + res_info = database._resource_info + self.assertEqual(res_info["project"], self.PROJECT_ID) + self.assertEqual(res_info["instance"], self.INSTANCE_ID) + self.assertEqual(res_info["database"], self.DATABASE_ID) + self.assertIsNone(getattr(res_info, "_client_attributes", None)) + def test_ctor_w_explicit_pool(self): instance = _Instance(self.INSTANCE_NAME) pool = _Pool() diff --git a/packages/google-cloud-spanner/tests/unit/test_metrics_capture.py b/packages/google-cloud-spanner/tests/unit/test_metrics_capture.py index 1bd1c19f9bab..c5cc0e7aae2e 100644 --- a/packages/google-cloud-spanner/tests/unit/test_metrics_capture.py +++ b/packages/google-cloud-spanner/tests/unit/test_metrics_capture.py @@ -50,3 +50,126 @@ def test_metrics_capture_exit(mock_tracer_factory): pass mock_tracer.record_operation_completion.assert_called_once() + + +def test_metrics_capture_with_resource_info(mock_tracer_factory): + mock_tracer = mock.Mock() + mock_tracer_factory.return_value = mock_tracer + resource_info = { + "project": "test-project", + "instance": "test-instance", + "database": "test-database", + } + + with MetricsCapture(resource_info): + pass + + mock_tracer_factory.assert_called_once_with(resource_info) + mock_tracer.record_operation_completion.assert_called_once() + + +def test_metrics_capture_retains_tracer_when_context_cleared(mock_tracer_factory): + mock_tracer = mock.Mock() + mock_tracer_factory.return_value = mock_tracer + + with MetricsCapture() as capture: + # Simulate inner code modifying or clearing current tracer + assert capture._tracer is mock_tracer + + mock_tracer.record_operation_completion.assert_called_once() + + +def test_metrics_capture_handles_value_error_on_reset(mock_tracer_factory): + mock_tracer = mock.Mock() + mock_tracer_factory.return_value = mock_tracer + + with mock.patch.object( + SpannerMetricsTracerFactory, + "reset_current_tracer", + side_effect=ValueError("Token was created in a different Context"), + ): + # Exiting context must not raise ValueError + with MetricsCapture(): + pass + + mock_tracer.record_operation_completion.assert_called_once() + + +def test_metrics_capture_disabled(): + SpannerMetricsTracerFactory(enabled=False) + try: + with MetricsCapture() as capture: + assert capture._token is None + finally: + SpannerMetricsTracerFactory(enabled=True) + + +def test_metrics_capture_exit_without_token(): + capture = MetricsCapture() + assert capture.__exit__(None, None, None) is False + + +def test_metrics_capture_partial_resource_info(mock_tracer_factory): + mock_tracer = mock.Mock() + mock_tracer_factory.return_value = mock_tracer + + with MetricsCapture({"project": "p"}): + pass + mock_tracer_factory.assert_called_once_with({"project": "p"}) + + mock_tracer_factory.reset_mock() + mock_tracer.reset_mock() + with MetricsCapture({"instance": "i"}): + pass + mock_tracer_factory.assert_called_once_with({"instance": "i"}) + + mock_tracer_factory.reset_mock() + mock_tracer.reset_mock() + with MetricsCapture({"database": "d"}): + pass + mock_tracer_factory.assert_called_once_with({"database": "d"}) + + +def test_metrics_capture_tracer_is_none(mock_tracer_factory): + mock_tracer_factory.return_value = None + with MetricsCapture(): + pass + + +def test_metrics_capture_exit_suppresses_metrics_tracer_exception( + mock_tracer_factory, +): + mock_tracer = mock.Mock() + mock_tracer.record_operation_completion.side_effect = RuntimeError( + "metrics recording failed" + ) + mock_tracer_factory.return_value = mock_tracer + + # Exception from record_operation_completion must be caught and not propagated + with MetricsCapture(): + pass + mock_tracer.record_operation_completion.assert_called_once() + + +def test_metrics_capture_without_opentelemetry(): + with mock.patch( + "google.cloud.spanner_v1.metrics.metrics_capture.HAS_OPENTELEMETRY_INSTALLED", + False, + ): + with MetricsCapture() as capture: + assert capture._token is None + + +def test_metrics_capture_enter_resets_token_on_start_failure(mock_tracer_factory): + mock_tracer = mock.Mock() + mock_tracer.record_operation_start.side_effect = RuntimeError("start failed") + mock_tracer_factory.return_value = mock_tracer + + capture = MetricsCapture() + with pytest.raises(RuntimeError): + with capture: + pass + + assert capture._token is None + assert capture._tracer is None + assert SpannerMetricsTracerFactory.get_current_tracer() is None diff --git a/packages/google-cloud-spanner/tests/unit/test_metrics_tracer.py b/packages/google-cloud-spanner/tests/unit/test_metrics_tracer.py index c645905c40ee..3b3cc084f872 100644 --- a/packages/google-cloud-spanner/tests/unit/test_metrics_tracer.py +++ b/packages/google-cloud-spanner/tests/unit/test_metrics_tracer.py @@ -175,11 +175,27 @@ def test_set_location(metrics_tracer): metrics_tracer.set_location("test_location") assert metrics_tracer.client_attributes["location"] == "test_location" - # Ensure it does not overwrite metrics_tracer.set_location("new_location") assert metrics_tracer.client_attributes["location"] == "test_location" +def test_set_resource_info(metrics_tracer): + resource_info = { + "project": "test_project", + "instance": "test_instance", + "database": "test_database", + } + metrics_tracer.set_resource_info(resource_info) + assert metrics_tracer.client_attributes["project_id"] == "test_project" + assert metrics_tracer.client_attributes["instance_id"] == "test_instance" + assert metrics_tracer.client_attributes["database"] == "test_database" + + # None and empty resource_info are safe no-ops + metrics_tracer.set_resource_info(None) + metrics_tracer.set_resource_info({}) + assert metrics_tracer.client_attributes["project_id"] == "test_project" + + def test_set_client_hash(metrics_tracer): metrics_tracer.set_client_hash("test_hash") assert metrics_tracer.client_attributes["client_hash"] == "test_hash" diff --git a/packages/google-cloud-spanner/tests/unit/test_metrics_tracer_factory.py b/packages/google-cloud-spanner/tests/unit/test_metrics_tracer_factory.py index ec352683241e..f2b408342ff0 100644 --- a/packages/google-cloud-spanner/tests/unit/test_metrics_tracer_factory.py +++ b/packages/google-cloud-spanner/tests/unit/test_metrics_tracer_factory.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest import mock + import pytest from google.cloud.spanner_v1.metrics.metrics_tracer import MetricsTracer @@ -50,3 +52,105 @@ def test_client_attributes(metrics_tracer_factory): attributes = metrics_tracer_factory.client_attributes assert attributes["project_id"] == "test_project" assert attributes["instance_id"] == "test_instance" + + +def test_create_metrics_tracer_with_resource_info(): + factory = MetricsTracerFactory( + enabled=True, + service_name="test_service", + ) + resource_info = { + "project": "custom_project", + "instance": "custom_instance", + "database": "custom_database", + } + tracer = factory.create_metrics_tracer(resource_info) + assert isinstance(tracer, MetricsTracer) + assert tracer.client_attributes["project_id"] == "custom_project" + assert tracer.client_attributes["instance_id"] == "custom_instance" + assert tracer.client_attributes["database"] == "custom_database" + + +def test_create_resource_info_and_tracer(): + factory = MetricsTracerFactory( + enabled=True, + service_name="test_service", + ) + resource_info = factory.create_resource_info( + project="pre_proj", + instance="pre_inst", + database="pre_db", + ) + assert resource_info["project"] == "pre_proj" + assert resource_info["instance"] == "pre_inst" + assert resource_info["database"] == "pre_db" + assert getattr(resource_info, "_client_attributes") is not None + assert resource_info._client_attributes["project_id"] == "pre_proj" + assert resource_info._client_attributes["instance_id"] == "pre_inst" + assert resource_info._client_attributes["database"] == "pre_db" + + tracer = factory.create_metrics_tracer(resource_info) + assert isinstance(tracer, MetricsTracer) + assert tracer.client_attributes["project_id"] == "pre_proj" + assert tracer.client_attributes["instance_id"] == "pre_inst" + assert tracer.client_attributes["database"] == "pre_db" + + +def test_create_resource_info_partial_arguments(): + factory = MetricsTracerFactory( + enabled=True, + service_name="test_service", + ) + resource_info = factory.create_resource_info( + project="pre_proj", + instance=None, + database=None, + ) + assert resource_info["project"] == "pre_proj" + assert resource_info["instance"] is None + assert resource_info["database"] is None + assert "project_id" in resource_info._client_attributes + assert "instance_id" not in resource_info._client_attributes + assert "database" not in resource_info._client_attributes + + +def test_create_metrics_tracer_with_mock_resource_info(): + factory = MetricsTracerFactory( + enabled=True, + service_name="test_service", + ) + # MagicMock dynamically returns a child Mock for _client_attributes + mock_info = mock.MagicMock() + mock_info.get.side_effect = lambda k: {"project": "mock_p"}.get(k) + tracer = factory.create_metrics_tracer(mock_info) + assert isinstance(tracer, MetricsTracer) + # client_attributes must be a dict, not a MagicMock + assert isinstance(tracer.client_attributes, dict) + assert tracer.client_attributes["project_id"] == "mock_p" + + +def test_create_metrics_tracer_without_opentelemetry(): + factory = MetricsTracerFactory(enabled=True, service_name="test_service") + with mock.patch( + "google.cloud.spanner_v1.metrics.metrics_tracer_factory.HAS_OPENTELEMETRY_INSTALLED", + False, + ): + assert factory.create_metrics_tracer() is None + res_info = factory.create_resource_info(project="p", instance="i", database="d") + assert res_info["project"] == "p" + assert getattr(res_info, "_client_attributes", None) is None + + +def test_create_resource_info_does_not_overwrite_existing_attributes(): + factory = MetricsTracerFactory(enabled=True, service_name="test_service") + factory.set_project("factory_proj") + factory.set_instance("factory_inst") + factory.set_database("factory_db") + resource_info = factory.create_resource_info( + project="new_proj", + instance="new_inst", + database="new_db", + ) + assert resource_info._client_attributes["project_id"] == "factory_proj" + assert resource_info._client_attributes["instance_id"] == "factory_inst" + assert resource_info._client_attributes["database"] == "factory_db" diff --git a/packages/google-cloud-spanner/tests/unit/test_pool.py b/packages/google-cloud-spanner/tests/unit/test_pool.py index 1a4642fa98d8..829de710a7da 100644 --- a/packages/google-cloud-spanner/tests/unit/test_pool.py +++ b/packages/google-cloud-spanner/tests/unit/test_pool.py @@ -229,6 +229,27 @@ def test_resource_info_unbound(self): pool = self._make_one() self.assertIsNone(pool._resource_info) + def test_resource_info_bound(self): + pool = self._make_one() + mock_db = mock.Mock(spec=["_resource_info"]) + mock_db._resource_info = {"project": "p", "instance": "i", "database": "d"} + pool._database = mock_db + self.assertEqual(pool._resource_info, mock_db._resource_info) + + def test_resource_info_fallback(self): + pool = self._make_one() + mock_db = mock.Mock(spec=["_instance", "database_id"]) + mock_db._instance = mock.Mock(spec=["_client", "instance_id"]) + mock_db._instance.instance_id = "i" + mock_db._instance._client = mock.Mock(spec=["project"]) + mock_db._instance._client.project = "p" + mock_db.database_id = "d" + pool._database = mock_db + self.assertEqual( + pool._resource_info, + {"project": "p", "instance": "i", "database": "d"}, + ) + def test__new_session_wo_labels(self): pool = self._make_one() database = pool._database = build_database() diff --git a/packages/google-cloud-spanner/tests/unit/test_session.py b/packages/google-cloud-spanner/tests/unit/test_session.py index 6ed3950e0a28..eb45d1b7369f 100644 --- a/packages/google-cloud-spanner/tests/unit/test_session.py +++ b/packages/google-cloud-spanner/tests/unit/test_session.py @@ -671,6 +671,43 @@ def test_ping_hit(self, mock_region): attributes=dict(self.BASE_ATTRIBUTES, x_goog_spanner_request_id=req_id), ) + @mock.patch( + "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", + return_value="global", + ) + def test_ping_invokes_metrics_capture(self, mock_region): + gax_api = self._make_spanner_api() + gax_api.execute_sql.return_value = "1" + database = self._make_database() + database.spanner_api = gax_api + session = self._make_one(database) + session._session_id = self.SESSION_ID + + with mock.patch( + "google.cloud.spanner_v1.session.MetricsCapture" + ) as mock_capture: + session.ping() + mock_capture.assert_called_once_with(session._resource_info) + + def test_resource_info_fallback_when_database_lacks_resource_info(self): + database = mock.Mock(spec=["_instance", "database_id"]) + database.database_id = "test-db" + instance = mock.Mock(spec=["_client", "instance_id"]) + instance.instance_id = "test-inst" + client = mock.Mock(spec=["project"]) + client.project = "test-proj" + instance._client = client + database._instance = instance + + session = self._make_one(database) + self.assertEqual( + session._resource_info, + {"project": "test-proj", "instance": "test-inst", "database": "test-db"}, + ) + + session._database = None + self.assertIsNone(session._resource_info) + @mock.patch( "google.cloud.spanner_v1._opentelemetry_tracing._get_cloud_region", return_value="global", diff --git a/packages/google-cloud-spanner/tests/unit/test_snapshot.py b/packages/google-cloud-spanner/tests/unit/test_snapshot.py index bfc52ee4461d..e938625a675e 100644 --- a/packages/google-cloud-spanner/tests/unit/test_snapshot.py +++ b/packages/google-cloud-spanner/tests/unit/test_snapshot.py @@ -17,7 +17,7 @@ import mock from google.api_core import gapic_v1 -from google.api_core.exceptions import Aborted, InternalServerError +from google.api_core.exceptions import Aborted, InternalServerError, NotFound from google.api_core.retry import Retry from google.cloud.spanner_admin_database_v1 import Database @@ -902,6 +902,389 @@ def test_iteration_w_multiple_span_creation(self, mock_region): ), ) + def test_restart_on_unavailable_metrics_tracer_lifecycle(self): + from google.cloud.spanner_v1.metrics.metrics_tracer_factory import ( + MetricsTracerFactory, + ) + from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( + SpannerMetricsTracerFactory, + ) + from google.cloud.spanner_v1.snapshot import _restart_on_unavailable + + mock_tracer = mock.Mock() + + item = self._make_item(0) + raw = _MockIterator(item) + tracer_during_call = [] + + def restart(*args, **kwargs): + tracer_during_call.append(SpannerMetricsTracerFactory.get_current_tracer()) + return raw + + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + database.spanner_api = build_spanner_api() + session = _Session(database) + derived = _build_snapshot_derived(session) + resource_info = {"project": "p", "instance": "i", "database": "d"} + + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=mock_tracer + ) as mock_create_tracer: + resumable = _restart_on_unavailable( + restart, + request, + session=session, + transaction=derived, + request_id_manager=session._database, + resource_info=resource_info, + ) + + # Not started before iteration begins (lazy generator) + mock_tracer.record_operation_start.assert_not_called() + + items = [] + for it in resumable: + # During row iteration, ambient contextvar must NOT be contaminated + self.assertIsNone(SpannerMetricsTracerFactory.get_current_tracer()) + items.append(it) + + self.assertEqual(items, [item]) + # Tracer factory was called with resource info + mock_create_tracer.assert_called_once_with(resource_info) + # Operation was started once + mock_tracer.record_operation_start.assert_called_once() + # Contextvar was active during RPC invocation + self.assertEqual(tracer_during_call, [mock_tracer]) + # Operation completion recorded once + mock_tracer.record_operation_completion.assert_called_once() + # Context is clean after completion + self.assertIsNone(SpannerMetricsTracerFactory.get_current_tracer()) + + def test_restart_on_unavailable_metrics_tracer_spans_retries(self): + from google.api_core.exceptions import ServiceUnavailable + + from google.cloud.spanner_v1.metrics.metrics_tracer_factory import ( + MetricsTracerFactory, + ) + from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( + SpannerMetricsTracerFactory, + ) + from google.cloud.spanner_v1.snapshot import _restart_on_unavailable + + mock_tracer = mock.Mock() + + FIRST = (self._make_item(0), self._make_item(1, resume_token=RESUME_TOKEN)) + SECOND = (self._make_item(2),) + before = _MockIterator( + *FIRST, fail_after=True, error=ServiceUnavailable("testing") + ) + after = _MockIterator(*SECOND) + tracers_during_calls = [] + + def restart(*args, **kwargs): + tracers_during_calls.append( + SpannerMetricsTracerFactory.get_current_tracer() + ) + if len(tracers_during_calls) == 1: + return before + return after + + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + database.spanner_api = build_spanner_api() + session = _Session(database) + derived = _build_snapshot_derived(session) + resource_info = {"project": "p", "instance": "i", "database": "d"} + + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=mock_tracer + ): + resumable = _restart_on_unavailable( + restart, + request, + session=session, + transaction=derived, + request_id_manager=session._database, + resource_info=resource_info, + ) + + items = list(resumable) + self.assertEqual(items, list(FIRST + SECOND)) + self.assertEqual(len(tracers_during_calls), 2) + # Both attempts used the same metrics tracer + self.assertEqual(tracers_during_calls, [mock_tracer, mock_tracer]) + mock_tracer.record_operation_start.assert_called_once() + mock_tracer.record_operation_completion.assert_called_once() + self.assertIsNone(SpannerMetricsTracerFactory.get_current_tracer()) + + def test_restart_on_unavailable_early_break_no_context_error(self): + from google.cloud.spanner_v1.metrics.metrics_tracer_factory import ( + MetricsTracerFactory, + ) + from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( + SpannerMetricsTracerFactory, + ) + from google.cloud.spanner_v1.snapshot import _restart_on_unavailable + + mock_tracer = mock.Mock() + + FIRST = (self._make_item(0), self._make_item(1)) + raw = _MockIterator(*FIRST) + restart = mock.Mock(return_value=raw) + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + database.spanner_api = build_spanner_api() + session = _Session(database) + derived = _build_snapshot_derived(session) + resource_info = {"project": "p", "instance": "i", "database": "d"} + + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=mock_tracer + ): + resumable = _restart_on_unavailable( + restart, + request, + session=session, + transaction=derived, + request_id_manager=session._database, + resource_info=resource_info, + ) + + for it in resumable: + self.assertEqual(it, FIRST[0]) + break + resumable.close() + + mock_tracer.record_operation_completion.assert_called_once() + self.assertIsNone(SpannerMetricsTracerFactory.get_current_tracer()) + + def test_restart_on_unavailable_real_factory_early_break_no_context_error(self): + from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( + SpannerMetricsTracerFactory, + ) + from google.cloud.spanner_v1.snapshot import _restart_on_unavailable + + FIRST = (self._make_item(0), self._make_item(1)) + raw = _MockIterator(*FIRST) + restart = mock.Mock(return_value=raw) + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + database.spanner_api = build_spanner_api() + session = _Session(database) + derived = _build_snapshot_derived(session) + resource_info = {"project": "p", "instance": "i", "database": "d"} + + resumable = _restart_on_unavailable( + restart, + request, + session=session, + transaction=derived, + request_id_manager=session._database, + resource_info=resource_info, + ) + + for it in resumable: + self.assertEqual(it, FIRST[0]) + break + resumable.close() + + self.assertIsNone(SpannerMetricsTracerFactory.get_current_tracer()) + + def test_restart_on_unavailable_metrics_tracer_partial_resource_info_and_disabled( + self, + ): + from google.cloud.spanner_v1.metrics.metrics_tracer_factory import ( + MetricsTracerFactory, + ) + from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import ( + SpannerMetricsTracerFactory, + ) + from google.cloud.spanner_v1.snapshot import _restart_on_unavailable + + mock_tracer = mock.Mock() + item = self._make_item(0) + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + database.spanner_api = build_spanner_api() + session = _Session(database) + derived = _build_snapshot_derived(session) + + # 1. resource_info is None + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=mock_tracer + ): + resumable = _restart_on_unavailable( + mock.Mock(return_value=_MockIterator(item)), + request, + session=session, + transaction=derived, + request_id_manager=session._database, + resource_info=None, + ) + for it in resumable: + pass + mock_tracer.record_operation_start.assert_called_once() + mock_tracer.record_operation_completion.assert_called_once() + + # 2. partial resource_info (only project, only instance, only database) + for partial_info in [{"project": "p"}, {"instance": "i"}, {"database": "d"}]: + mock_tracer.reset_mock() + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=mock_tracer + ) as mock_create_tracer: + resumable = _restart_on_unavailable( + mock.Mock(return_value=_MockIterator(item)), + request, + session=session, + transaction=derived, + request_id_manager=session._database, + resource_info=partial_info, + ) + for it in resumable: + pass + mock_create_tracer.assert_called_once_with(partial_info) + + # 3. metrics factory disabled + SpannerMetricsTracerFactory(enabled=False) + try: + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer" + ) as mock_create: + resumable = _restart_on_unavailable( + mock.Mock(return_value=_MockIterator(item)), + request, + session=session, + transaction=derived, + request_id_manager=session._database, + resource_info={"project": "p"}, + ) + items = [] + for it in resumable: + items.append(it) + self.assertEqual(items, [item]) + mock_create.assert_not_called() + self.assertIsNone(SpannerMetricsTracerFactory.get_current_tracer()) + finally: + SpannerMetricsTracerFactory(enabled=True) + + # 4. create_metrics_tracer returns None + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=None + ): + resumable = _restart_on_unavailable( + mock.Mock(return_value=_MockIterator(item)), + request, + session=session, + transaction=derived, + request_id_manager=session._database, + resource_info={"project": "p"}, + ) + items = [] + for it in resumable: + items.append(it) + self.assertEqual(items, [item]) + self.assertIsNone(SpannerMetricsTracerFactory.get_current_tracer()) + + def test_restart_on_unavailable_metrics_completion_error_handled(self): + from google.cloud.spanner_v1.metrics.metrics_tracer_factory import ( + MetricsTracerFactory, + ) + from google.cloud.spanner_v1.snapshot import _restart_on_unavailable + + mock_tracer = mock.Mock() + mock_tracer.record_operation_completion.side_effect = RuntimeError( + "metrics export failure" + ) + + item = mock.Mock( + value=0, + resume_token=b"", + metadata=None, + precommit_token=None, + last=False, + _pb=None, + spec=["value", "resume_token", "metadata", "precommit_token", "last"], + ) + mock_iterator = _MockIterator(item) + mock_iterator.cancel = mock.Mock() + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + database.spanner_api = build_spanner_api() + session = _Session(database) + derived = _build_snapshot_derived(session) + + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=mock_tracer + ): + resumable = _restart_on_unavailable( + mock.Mock(return_value=mock_iterator), + request, + session=session, + transaction=derived, + request_id_manager=session._database, + resource_info={"project": "p"}, + ) + items = [] + for it in resumable: + items.append(it) + break + resumable.close() + + self.assertEqual(items, [item]) + mock_tracer.record_operation_completion.assert_called_once() + mock_iterator.cancel.assert_called_once() + + def test_restart_on_unavailable_metrics_completion_error_preserves_query_error( + self, + ): + from google.cloud.spanner_v1.metrics.metrics_tracer_factory import ( + MetricsTracerFactory, + ) + from google.cloud.spanner_v1.snapshot import _restart_on_unavailable + + mock_tracer = mock.Mock() + mock_tracer.record_operation_completion.side_effect = RuntimeError( + "metrics export failure" + ) + + item = mock.Mock( + value=0, + resume_token=b"", + metadata=None, + precommit_token=None, + last=False, + _pb=None, + spec=["value", "resume_token", "metadata", "precommit_token", "last"], + ) + mock_iterator = _MockIterator( + item, fail_after=True, error=NotFound("table not found") + ) + mock_iterator.cancel = mock.Mock() + request = mock.Mock(test="test", spec=["test", "resume_token"]) + database = _Database() + database.spanner_api = build_spanner_api() + session = _Session(database) + derived = _build_snapshot_derived(session) + + with mock.patch.object( + MetricsTracerFactory, "create_metrics_tracer", return_value=mock_tracer + ): + resumable = _restart_on_unavailable( + mock.Mock(return_value=mock_iterator), + request, + session=session, + transaction=derived, + request_id_manager=session._database, + resource_info={"project": "p"}, + ) + with self.assertRaises(NotFound): + for _ in resumable: + pass + + mock_tracer.record_operation_completion.assert_called_once() + mock_iterator.cancel.assert_called_once() + class Test_SnapshotBase(OpenTelemetryBase): def test_ctor(self): @@ -924,6 +1307,27 @@ def test_ctor(self): self.assertNoSpans() + def test_resource_info_fallback_when_database_lacks_resource_info(self): + database = mock.Mock(spec=["_instance", "database_id"]) + database.database_id = "test-db" + instance = mock.Mock(spec=["_client", "instance_id"]) + instance.instance_id = "test-inst" + client = mock.Mock(spec=["project"]) + client.project = "test-proj" + instance._client = client + database._instance = instance + + session = mock.Mock() + session._database = database + derived = _build_snapshot_derived(session=session) + self.assertEqual( + derived._resource_info, + {"project": "test-proj", "instance": "test-inst", "database": "test-db"}, + ) + + derived._session = None + self.assertIsNone(derived._resource_info) + def test__wait_for_transaction_begin_claims_inline_begin(self): derived = _build_snapshot_derived(multi_use=True) diff --git a/packages/google-cloud-spanner/tests/unit/test_spanner_metrics_tracer_factory.py b/packages/google-cloud-spanner/tests/unit/test_spanner_metrics_tracer_factory.py index f8db3de3dbfa..2b76e01ada01 100644 --- a/packages/google-cloud-spanner/tests/unit/test_spanner_metrics_tracer_factory.py +++ b/packages/google-cloud-spanner/tests/unit/test_spanner_metrics_tracer_factory.py @@ -46,3 +46,41 @@ def test_get_instance_config(self): def test_get_client_name(self): client_name = SpannerMetricsTracerFactory._get_client_name() assert isinstance(client_name, str) + + def test_singleton_enabled_toggle(self): + factory = SpannerMetricsTracerFactory(enabled=True) + assert factory.enabled is True + factory_disabled = SpannerMetricsTracerFactory(enabled=False) + assert factory_disabled is factory + assert factory.enabled is False + factory_re_enabled = SpannerMetricsTracerFactory(enabled=True) + assert factory_re_enabled is factory + assert factory.enabled is True + + def test_generate_client_hash_empty(self): + assert SpannerMetricsTracerFactory._generate_client_hash("") == "000000" + + def test_generate_client_uid_exception(self, monkeypatch): + import os + + monkeypatch.setattr( + os, "uname", lambda: (_ for _ in ()).throw(RuntimeError("uname fail")) + ) + assert SpannerMetricsTracerFactory._generate_client_uid() == "" + + def test_context_var_methods(self): + mock_tracer = object() + token = SpannerMetricsTracerFactory.set_current_tracer(mock_tracer) + try: + assert SpannerMetricsTracerFactory.get_current_tracer() is mock_tracer + finally: + SpannerMetricsTracerFactory.reset_current_tracer(token) + assert SpannerMetricsTracerFactory.get_current_tracer() is None + + def test_without_opentelemetry(self, monkeypatch): + import google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory as mod + + monkeypatch.setattr(mod, "HAS_OPENTELEMETRY_INSTALLED", False) + SpannerMetricsTracerFactory._metrics_tracer_factory = None + factory = SpannerMetricsTracerFactory(enabled=True) + assert factory is not None