Feat: Added Operation and GFE Metrics

Author: lszinv

google/cloud/spanner_v1/_opentelemetry_tracing.py

@@ -33,6 +33,8 @@
 except ImportError:
     HAS_OPENTELEMETRY_INSTALLED = False
 
+from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
+
 TRACER_NAME = "cloud.google.com/python/spanner"
 TRACER_VERSION = gapic_version.__version__
 extended_tracing_globally_disabled = (
@@ -107,26 +109,27 @@ def trace_call(name, session=None, extra_attributes=None, observability_options=
     with tracer.start_as_current_span(
         name, kind=trace.SpanKind.CLIENT, attributes=attributes
     ) as span:
-        try:
-            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))
+        with MetricsCapture():
+            try:
+                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():

google/cloud/spanner_v1/metrics/constants.py

@@ -1,4 +1,5 @@
-# Copyright 2025 Google LLC
+# -*- coding: utf-8 -*-
+# Copyright 2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

google/cloud/spanner_v1/metrics/metrics_capture.py

@@ -0,0 +1,67 @@
+# -*- coding: utf-8 -*-
+# Copyright 2025 Google LLC All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This module provides functionality for capturing metrics in Cloud Spanner operations.
+
+It includes a context manager class, MetricsCapture, which automatically handles the
+start and completion of metrics tracing for a given operation. This ensures that metrics
+are consistently recorded for Cloud Spanner operations, facilitating observability and
+performance monitoring.
+"""
+
+from .spanner_metrics_tracer_factory import SpannerMetricsTracerFactory
+
+
+class MetricsCapture:
+    """Context manager for capturing metrics in Cloud Spanner operations.
+
+    This class provides a context manager interface to automatically handle
+    the start and completion of metrics tracing for a given operation.
+    """
+
+    def __enter__(self):
+        """Enter the runtime context related to this object.
+
+        This method initializes a new metrics tracer for the operation and
+        records the start of the operation.
+
+        Returns:
+            MetricsCapture: The instance of the context manager.
+        """
+        factory = SpannerMetricsTracerFactory()
+
+        # Define a new metrics tracer for the new operation
+        SpannerMetricsTracerFactory.current_metrics_tracer = (
+            factory.create_metrics_tracer()
+        )
+        SpannerMetricsTracerFactory.current_metrics_tracer.record_operation_start()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Exit the runtime context related to this object.
+
+        This method records the completion of the operation. If an exception
+        occurred, it will be propagated after the metrics are recorded.
+
+        Args:
+            exc_type (Type[BaseException]): The exception type.
+            exc_value (BaseException): The exception value.
+            traceback (TracebackType): The traceback object.
+
+        Returns:
+            bool: False to propagate the exception if any occurred.
+        """
+        SpannerMetricsTracerFactory.current_metrics_tracer.record_operation_completion()
+        return False  # Propagate the exception if any

google/cloud/spanner_v1/metrics/metrics_exporter.py

@@ -1,3 +1,4 @@
+# -*- coding: utf-8 -*-
 # Copyright 2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

google/cloud/spanner_v1/metrics/metrics_gfe_tracer.py

@@ -0,0 +1,87 @@
+# -*- coding: utf-8 -*-
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+try:
+    from opentelemetry.metrics import Counter, Histogram, get_meter_provider
+
+    HAS_OPENTELEMETRY_INSTALLED = True
+except ImportError:  # pragma: NO COVER
+    HAS_OPENTELEMETRY_INSTALLED = False
+from typing import List
+import re
+from google.cloud.spanner_v1.metrics.constants import (
+    BUILT_IN_METRICS_METER_NAME,
+    SPANNER_SERVICE_NAME,
+    METRIC_NAME_GFE_LATENCY,
+    METRIC_NAME_GFE_MISSING_HEADER_COUNT,
+)
+from google.cloud.spanner_v1 import __version__
+
+
+class MetricsGfeTracer:
+    _instrument_gfe_latency: Histogram
+    _instrument_gfe_missing_header_count: Counter
+    enabled = False
+
+    def __init__(self) -> None:
+        self._create_gfe_metric_instruments()
+
+    def record_gfe_metrics(self, metadata: List):
+        if not self.enabled:
+            return
+
+        gfe_headers = [
+            header.value
+            for header in metadata
+            if header.key == "server-timing" and header.value.startswith("gfe")
+        ]
+
+        if len(gfe_headers) == 0:
+            self.record_gfe_missing_header_count()
+            return
+
+        for gfe_header in gfe_headers:
+            match = re.search(r"dur=(\d+)", gfe_header)
+            if match:
+                duration = int(match.group(1))
+                self.record_gfe_latency(duration)
+
+    def record_gfe_latency(self, latency: int) -> None:
+        print("real called")
+        if not self.enabled:
+            return
+        self._instrument_gfe_latency.record(amount=latency)
+
+    def record_gfe_missing_header_count(self) -> None:
+        if not self.enabled:
+            return
+        self._instrument_gfe_missing_header_count.add(amount=1)
+
+    def _create_gfe_metric_instruments(self):
+        meter_provider = get_meter_provider()
+        meter = meter_provider.get_meter(
+            name=BUILT_IN_METRICS_METER_NAME, version=__version__
+        )
+        self._instrument_gfe_latency = meter.create_histogram(
+            name=f"{SPANNER_SERVICE_NAME}/{METRIC_NAME_GFE_LATENCY}",
+            unit="ms",
+            description="GFE Latency.",
+        )
+        self._instrument_gfe_missing_header_count = meter.create_counter(
+            name=f"{SPANNER_SERVICE_NAME}/{METRIC_NAME_GFE_MISSING_HEADER_COUNT}",
+            unit="1",
+            description="GFE missing header count.",
+        )

google/cloud/spanner_v1/metrics/metrics_interceptor.py

@@ -1,4 +1,5 @@
-# Copyright 2025 Google LLC
+# -*- coding: utf-8 -*-
+# Copyright 2025 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -19,6 +20,7 @@
     GOOGLE_CLOUD_RESOURCE_KEY,
     SPANNER_METHOD_PREFIX,
 )
+
 from typing import Dict
 from .spanner_metrics_tracer_factory import SpannerMetricsTracerFactory
 import re
@@ -142,4 +144,10 @@ def intercept(self, invoked_method, request_or_iterator, call_details):
         response = invoked_method(request_or_iterator, call_details)
         SpannerMetricsTracerFactory.current_metrics_tracer.record_attempt_completion()
 
+        # Process and send GFE metrics if enabled
+        if SpannerMetricsTracerFactory.metrics_gfe_tracer.enabled:
+            metadata = response.initial_metadata()
+            SpannerMetricsTracerFactory.metrics_gfe_tracer.record_gfe_metrics(metadata)
+        else:
+            print("disabled")
         return response

google/cloud/spanner_v1/metrics/metrics_tracer.py

@@ -56,7 +56,7 @@ class MetricAttemptTracer:
     direct_path_used: bool
     status: str
 
-    def __init__(self):
+    def __init__(self) -> None:
         """
         Initialize a MetricAttemptTracer instance with default values.
 

google/cloud/spanner_v1/metrics/spanner_metrics_tracer_factory.py

@@ -35,6 +35,7 @@
     HAS_OPENTELEMETRY_INSTALLED = False
 
 from .metrics_tracer import MetricsTracer
+from .metrics_gfe_tracer import MetricsGfeTracer
 from google.cloud.spanner_v1 import __version__
 from uuid import uuid4
 
@@ -44,6 +45,7 @@ class SpannerMetricsTracerFactory(MetricsTracerFactory):
 
     _metrics_tracer_factory: "SpannerMetricsTracerFactory" = None
     current_metrics_tracer: MetricsTracer = None
+    metrics_gfe_tracer: MetricsGfeTracer = MetricsGfeTracer()
 
     def __new__(cls, enabled: bool = True) -> "SpannerMetricsTracerFactory":
         """Create a new instance of SpannerMetricsTracerFactory if it doesn't already exist."""

tests/unit/test_metrics_capture.py

@@ -0,0 +1,32 @@
+import pytest
+from unittest import mock
+from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
+from google.cloud.spanner_v1.metrics.metrics_tracer_factory import MetricsTracerFactory
+
+
+@pytest.fixture
+def mock_tracer_factory():
+    with mock.patch.object(
+        MetricsTracerFactory, "create_metrics_tracer"
+    ) as mock_create:
+        yield mock_create
+
+
+def test_metrics_capture_enter(mock_tracer_factory):
+    mock_tracer = mock.Mock()
+    mock_tracer_factory.return_value = mock_tracer
+
+    with MetricsCapture() as capture:
+        assert capture is not None
+        mock_tracer_factory.assert_called_once()
+        mock_tracer.record_operation_start.assert_called_once()
+
+
+def test_metrics_capture_exit(mock_tracer_factory):
+    mock_tracer = mock.Mock()
+    mock_tracer_factory.return_value = mock_tracer
+
+    with MetricsCapture():
+        pass
+
+    mock_tracer.record_operation_completion.assert_called_once()

tests/unit/test_metrics_exporter.py

@@ -1,4 +1,4 @@
-# Copyright 2016 Google LLC All rights reserved.
+# Copyright 2025 Google LLC All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.