diff --git a/google/cloud/spanner_v1/metrics/metrics_interceptor.py b/google/cloud/spanner_v1/metrics/metrics_interceptor.py index 4b55056dab..f6d048ac98 100644 --- a/google/cloud/spanner_v1/metrics/metrics_interceptor.py +++ b/google/cloud/spanner_v1/metrics/metrics_interceptor.py @@ -14,9 +14,14 @@ """Interceptor for collecting Cloud Spanner metrics.""" +import logging + from grpc_interceptor import ClientInterceptor from .constants import ( GOOGLE_CLOUD_RESOURCE_KEY, + METRIC_LABEL_KEY_DATABASE, + MONITORED_RES_LABEL_KEY_INSTANCE, + MONITORED_RES_LABEL_KEY_PROJECT, SPANNER_METHOD_PREFIX, ) @@ -24,6 +29,8 @@ from .spanner_metrics_tracer_factory import SpannerMetricsTracerFactory import re +_logger = logging.getLogger(__name__) + class MetricsInterceptor(ClientInterceptor): """Interceptor that collects metrics for Cloud Spanner operations.""" @@ -92,27 +99,48 @@ def _set_metrics_tracer_attributes(self, resources: Dict[str, str]) -> None: """ Sets the metric tracer attributes based on the provided resources. - This method updates the current metric tracer's attributes with the project, instance, and database information extracted from the resources dictionary. If the current metric tracer is not set, the method does nothing. + This method updates the current metric tracer's attributes with the + project, instance, and database information extracted from the resources + dictionary. If the current metric tracer is not set, the method does + nothing. + + Before setting each attribute, any existing value for that key is removed + from the current tracer's client_attributes to ensure the set_* methods + (which use set-once semantics) can overwrite stale values from previous + operations. + + Additionally, this method updates the factory's client attributes for + project and instance to ensure these values are available for subsequent + operations. Database is not propagated to the factory because each + Spanner RPC may target a different database within the same instance. Args: - resources (Dict[str, str]): A dictionary containing project, instance, and database information. + resources (Dict[str, str]): A dictionary containing project, + instance, and database information. """ if SpannerMetricsTracerFactory.current_metrics_tracer is None: return if resources: + tracer = SpannerMetricsTracerFactory.current_metrics_tracer + factory = SpannerMetricsTracerFactory() + + # For each resource key, remove the existing value from + # client_attributes so the tracer's set_* method (which only + # writes if the key is absent) will accept the fresh value. if "project" in resources: - SpannerMetricsTracerFactory.current_metrics_tracer.set_project( - resources["project"] - ) + tracer.client_attributes.pop(MONITORED_RES_LABEL_KEY_PROJECT, None) + tracer.set_project(resources["project"]) + factory.set_project(resources["project"]) + if "instance" in resources: - SpannerMetricsTracerFactory.current_metrics_tracer.set_instance( - resources["instance"] - ) + tracer.client_attributes.pop(MONITORED_RES_LABEL_KEY_INSTANCE, None) + tracer.set_instance(resources["instance"]) + factory.set_instance(resources["instance"]) + if "database" in resources: - SpannerMetricsTracerFactory.current_metrics_tracer.set_database( - resources["database"] - ) + tracer.client_attributes.pop(METRIC_LABEL_KEY_DATABASE, None) + tracer.set_database(resources["database"]) def intercept(self, invoked_method, request_or_iterator, call_details): """Intercept gRPC calls to collect metrics. @@ -133,24 +161,33 @@ def intercept(self, invoked_method, request_or_iterator, call_details): return invoked_method(request_or_iterator, call_details) # Setup Metric Tracer attributes from call details - ## Extract Project / Instance / Databse from header information - resources = self._extract_resource_from_path(call_details.metadata) - self._set_metrics_tracer_attributes(resources) + try: + ## Extract Project / Instance / Database from header information + resources = self._extract_resource_from_path(call_details.metadata) + self._set_metrics_tracer_attributes(resources) - ## Format method to be be spanner. - method_name = self._remove_prefix( - call_details.method, SPANNER_METHOD_PREFIX - ).replace("/", ".") + ## Format method to be spanner. + method_name = self._remove_prefix( + call_details.method, SPANNER_METHOD_PREFIX + ).replace("/", ".") + + SpannerMetricsTracerFactory.current_metrics_tracer.set_method(method_name) + SpannerMetricsTracerFactory.current_metrics_tracer.record_attempt_start() + except Exception: + _logger.warning("Failed to set up metrics tracer attributes", exc_info=True) - SpannerMetricsTracerFactory.current_metrics_tracer.set_method(method_name) - SpannerMetricsTracerFactory.current_metrics_tracer.record_attempt_start() response = invoked_method(request_or_iterator, call_details) - SpannerMetricsTracerFactory.current_metrics_tracer.record_attempt_completion() - - # Process and send GFE metrics if enabled - if SpannerMetricsTracerFactory.current_metrics_tracer.gfe_enabled: - metadata = response.initial_metadata() - SpannerMetricsTracerFactory.current_metrics_trace.record_gfe_metrics( - metadata - ) + + try: + SpannerMetricsTracerFactory.current_metrics_tracer.record_attempt_completion() + + # Process and send GFE metrics if enabled + if SpannerMetricsTracerFactory.current_metrics_tracer.gfe_enabled: + metadata = response.initial_metadata() + SpannerMetricsTracerFactory.current_metrics_tracer.record_gfe_metrics( + metadata + ) + except Exception: + _logger.warning("Failed to record metrics", exc_info=True) + return response diff --git a/google/cloud/spanner_v1/metrics/spanner_metrics_tracer_factory.py b/google/cloud/spanner_v1/metrics/spanner_metrics_tracer_factory.py index 9566e61a28..89dae1fd0c 100644 --- a/google/cloud/spanner_v1/metrics/spanner_metrics_tracer_factory.py +++ b/google/cloud/spanner_v1/metrics/spanner_metrics_tracer_factory.py @@ -80,7 +80,7 @@ def __new__( cls._metrics_tracer_factory.gfe_enabled = gfe_enabled if cls._metrics_tracer_factory.enabled != enabled: - cls._metrics_tracer_factory.enabeld = enabled + cls._metrics_tracer_factory.enabled = enabled return cls._metrics_tracer_factory diff --git a/tests/unit/test_metrics_interceptor.py b/tests/unit/test_metrics_interceptor.py index e32003537f..285ad0f9cf 100644 --- a/tests/unit/test_metrics_interceptor.py +++ b/tests/unit/test_metrics_interceptor.py @@ -26,6 +26,25 @@ def interceptor(): return MetricsInterceptor() +@pytest.fixture +def clean_factory(): + """Return a factory with project_id and instance_id cleared for test isolation.""" + factory = SpannerMetricsTracerFactory() + factory._client_attributes.pop("instance_id", None) + factory._client_attributes.pop("project_id", None) + return factory + + +@pytest.fixture +def mock_tracer(clean_factory): + """Set up a clean MockMetricTracer and return the factory.""" + SpannerMetricsTracerFactory.current_metrics_tracer = MockMetricTracer() + return clean_factory + + +# --- Parsing tests --- + + def test_parse_resource_path_valid(interceptor): path = "projects/my_project/instances/my_instance/databases/my_database" expected = { @@ -57,8 +76,11 @@ def test_extract_resource_from_path(interceptor): assert interceptor._extract_resource_from_path(metadata) == expected -def test_set_metrics_tracer_attributes(interceptor): - SpannerMetricsTracerFactory.current_metrics_tracer = MockMetricTracer() +# --- _set_metrics_tracer_attributes tests --- + + +def test_set_metrics_tracer_attributes(interceptor, mock_tracer): + """Verify tracer receives project, instance, and database from resources.""" resources = { "project": "my_project", "instance": "my_instance", @@ -66,12 +88,164 @@ def test_set_metrics_tracer_attributes(interceptor): } interceptor._set_metrics_tracer_attributes(resources) - assert SpannerMetricsTracerFactory.current_metrics_tracer.project == "my_project" - assert SpannerMetricsTracerFactory.current_metrics_tracer.instance == "my_instance" - assert SpannerMetricsTracerFactory.current_metrics_tracer.database == "my_database" + + tracer = SpannerMetricsTracerFactory.current_metrics_tracer + assert tracer.client_attributes.get("project_id") == "my_project" + assert tracer.client_attributes.get("instance_id") == "my_instance" + assert tracer.client_attributes.get("database") == "my_database" + + +def test_set_metrics_tracer_attributes_updates_factory(interceptor, mock_tracer): + """Verify that the factory's client_attributes are updated with project and instance.""" + resources = { + "project": "test_project", + "instance": "test_instance", + "database": "test_database", + } + + interceptor._set_metrics_tracer_attributes(resources) + + # Verify factory attributes are updated + assert mock_tracer.client_attributes.get("instance_id") == "test_instance" + assert mock_tracer.client_attributes.get("project_id") == "test_project" + # Database should NOT be set in factory (it may vary per operation) + assert "database" not in mock_tracer.client_attributes + + +def test_set_metrics_tracer_attributes_no_tracer(interceptor, clean_factory): + """Verify that nothing happens when current_metrics_tracer is None.""" + SpannerMetricsTracerFactory.current_metrics_tracer = None + + resources = { + "project": "test_project", + "instance": "test_instance", + } + + interceptor._set_metrics_tracer_attributes(resources) + + # Factory should NOT be updated when current_metrics_tracer is None + assert "instance_id" not in clean_factory.client_attributes + assert "project_id" not in clean_factory.client_attributes + + +def test_set_metrics_tracer_attributes_empty_resources(interceptor, mock_tracer): + """Verify that nothing happens when resources is empty.""" + interceptor._set_metrics_tracer_attributes({}) + + assert "instance_id" not in mock_tracer.client_attributes + assert "project_id" not in mock_tracer.client_attributes + + +def test_set_metrics_tracer_attributes_none_resources(interceptor, mock_tracer): + """Verify that nothing happens when resources is None.""" + interceptor._set_metrics_tracer_attributes(None) + + assert "instance_id" not in mock_tracer.client_attributes + assert "project_id" not in mock_tracer.client_attributes + + +def test_set_metrics_tracer_attributes_partial_resources_project_only( + interceptor, mock_tracer +): + """Verify that only project is set when instance is missing.""" + resources = {"project": "test_project"} + + interceptor._set_metrics_tracer_attributes(resources) + + # Only project should be set + assert mock_tracer.client_attributes.get("project_id") == "test_project" + assert "instance_id" not in mock_tracer.client_attributes + tracer = SpannerMetricsTracerFactory.current_metrics_tracer + assert tracer.client_attributes.get("project_id") == "test_project" + assert "instance_id" not in tracer.client_attributes + + +def test_set_metrics_tracer_attributes_partial_resources_instance_only( + interceptor, mock_tracer +): + """Verify that only instance is set when project is missing.""" + resources = {"instance": "test_instance"} + + interceptor._set_metrics_tracer_attributes(resources) + + # Only instance should be set + assert mock_tracer.client_attributes.get("instance_id") == "test_instance" + assert "project_id" not in mock_tracer.client_attributes + tracer = SpannerMetricsTracerFactory.current_metrics_tracer + assert tracer.client_attributes.get("instance_id") == "test_instance" + assert "project_id" not in tracer.client_attributes + + +def test_set_metrics_tracer_attributes_partial_resources_database_only( + interceptor, mock_tracer +): + """Verify that database is set on tracer but NOT propagated to factory.""" + resources = {"database": "test_database"} + + interceptor._set_metrics_tracer_attributes(resources) + + tracer = SpannerMetricsTracerFactory.current_metrics_tracer + assert tracer.client_attributes.get("database") == "test_database" + # Factory should NOT have database, project, or instance + assert "database" not in mock_tracer.client_attributes + assert "project_id" not in mock_tracer.client_attributes + assert "instance_id" not in mock_tracer.client_attributes + + +def test_set_metrics_tracer_attributes_overwrites_stale_tracer_values(interceptor): + """Verify that request resource values replace stale tracer client_attributes.""" + stale_tracer = MockMetricTracer() + # Directly populate to simulate stale values from factory copy + stale_tracer.client_attributes["project_id"] = "stale_project" + stale_tracer.client_attributes["instance_id"] = "stale_instance" + stale_tracer.client_attributes["database"] = "stale_database" + + SpannerMetricsTracerFactory.current_metrics_tracer = stale_tracer + resources = { + "project": "fresh_project", + "instance": "fresh_instance", + "database": "fresh_database", + } + + interceptor._set_metrics_tracer_attributes(resources) + + assert stale_tracer.client_attributes.get("project_id") == "fresh_project" + assert stale_tracer.client_attributes.get("instance_id") == "fresh_instance" + assert stale_tracer.client_attributes.get("database") == "fresh_database" + + +def test_new_tracer_inherits_factory_attributes(interceptor, mock_tracer): + """ + Integration test: Verify that a new tracer created after + _set_metrics_tracer_attributes inherits project and instance from factory. + + This is the core test for the bug fix - ensuring that subsequent operations + get tracers with the correct attributes. + """ + resources = { + "project": "inherited_project", + "instance": "inherited_instance", + "database": "db1", + } + interceptor._set_metrics_tracer_attributes(resources) + + # Simulate second operation: create a new tracer from factory + new_tracer = mock_tracer.create_metrics_tracer() + + if new_tracer is None: + pytest.skip("OpenTelemetry not installed; cannot verify tracer inheritance") + + assert new_tracer.client_attributes.get("project_id") == "inherited_project" + assert new_tracer.client_attributes.get("instance_id") == "inherited_instance" + # Database should NOT be inherited (it's per-operation) + assert "database" not in new_tracer.client_attributes + + +# --- intercept tests --- def test_intercept_with_tracer(interceptor): + """Verify intercept records metrics and invokes the gRPC method.""" SpannerMetricsTracerFactory.current_metrics_tracer = MockMetricTracer() SpannerMetricsTracerFactory.current_metrics_tracer.record_attempt_start = ( MagicMock() @@ -102,21 +276,66 @@ def test_intercept_with_tracer(interceptor): mock_invoked_method.assert_called_once_with("request", call_details) +def test_intercept_no_tracer(interceptor): + """Verify that intercept returns directly when current_metrics_tracer is None.""" + SpannerMetricsTracerFactory.current_metrics_tracer = None + + invoked_response = MagicMock() + mock_invoked_method = MagicMock(return_value=invoked_response) + call_details = MagicMock(method="spanner.someMethod", metadata=[]) + + response = interceptor.intercept(mock_invoked_method, "request", call_details) + assert response == invoked_response + mock_invoked_method.assert_called_once_with("request", call_details) + + +def test_intercept_factory_disabled(interceptor): + """Verify that intercept returns directly when factory is disabled.""" + SpannerMetricsTracerFactory.current_metrics_tracer = MockMetricTracer() + factory = SpannerMetricsTracerFactory() + original_enabled = factory.enabled + + try: + factory.enabled = False + + invoked_response = MagicMock() + mock_invoked_method = MagicMock(return_value=invoked_response) + call_details = MagicMock(method="spanner.someMethod", metadata=[]) + + response = interceptor.intercept(mock_invoked_method, "request", call_details) + assert response == invoked_response + mock_invoked_method.assert_called_once_with("request", call_details) + finally: + factory.enabled = original_enabled + + +# --- MockMetricTracer --- + + class MockMetricTracer: + """Mock that mirrors MetricsTracer's set-once semantics. + + The real MetricsTracer.set_project/set_instance/set_database only write + if the key is absent from _client_attributes. This guard is critical for + testing the pop-then-set pattern in _set_metrics_tracer_attributes. + """ + def __init__(self): - self.project = None - self.instance = None - self.database = None self.method = None + self.client_attributes = {} + self.gfe_enabled = False def set_project(self, project): - self.project = project + if "project_id" not in self.client_attributes: + self.client_attributes["project_id"] = project def set_instance(self, instance): - self.instance = instance + if "instance_id" not in self.client_attributes: + self.client_attributes["instance_id"] = instance def set_database(self, database): - self.database = database + if "database" not in self.client_attributes: + self.client_attributes["database"] = database def set_method(self, method): self.method = method