Moved MetricCapture out of Tracer logic

lszinv · lszinv · commit c07edcaf3b8c · 2025-01-30T21:53:29.000-08:00
diff --git a/google/cloud/spanner_v1/batch.py b/google/cloud/spanner_v1/batch.py
@@ -32,6 +32,7 @@
 from google.cloud.spanner_v1._helpers import _retry_on_aborted_exception
 from google.cloud.spanner_v1._helpers import _check_rst_stream_error
 from google.api_core.exceptions import InternalServerError
+from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
 import time
 
 DEFAULT_RETRY_TIMEOUT_SECS = 30
@@ -226,7 +227,7 @@ def commit(
             self._session,
             trace_attributes,
             observability_options=observability_options,
-        ):
+        ), MetricsCapture():
             method = functools.partial(
                 api.commit,
                 request=request,
@@ -348,7 +349,7 @@ def batch_write(self, request_options=None, exclude_txn_from_change_streams=Fals
             self._session,
             trace_attributes,
             observability_options=observability_options,
-        ):
+        ), MetricsCapture():
             method = functools.partial(
                 api.batch_write,
                 request=request,
diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py
@@ -72,6 +72,7 @@
     get_current_span,
     trace_call,
 )
+from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
 
 
 SPANNER_DATA_SCOPE = "https://www.googleapis.com/auth/spanner.data"
@@ -702,7 +703,7 @@ def execute_pdml():
             with trace_call(
                 "CloudSpanner.Database.execute_partitioned_pdml",
                 observability_options=self.observability_options,
-            ) as span:
+            ) as span, MetricsCapture():
                 with SessionCheckout(self._pool) as session:
                     add_span_event(span, "Starting BeginTransaction")
                     txn = api.begin_transaction(
@@ -897,7 +898,7 @@ def run_in_transaction(self, func, *args, **kw):
         with trace_call(
             "CloudSpanner.Database.run_in_transaction",
             observability_options=observability_options,
-        ):
+        ), MetricsCapture():
             # Sanity check: Is there a transaction already running?
             # If there is, then raise a red flag. Otherwise, mark that this one
             # is running.
@@ -1489,7 +1490,7 @@ def generate_read_batches(
             f"CloudSpanner.{type(self).__name__}.generate_read_batches",
             extra_attributes=dict(table=table, columns=columns),
             observability_options=self.observability_options,
-        ):
+        ), MetricsCapture():
             partitions = self._get_snapshot().partition_read(
                 table=table,
                 columns=columns,
@@ -1540,7 +1541,7 @@ def process_read_batch(
         with trace_call(
             f"CloudSpanner.{type(self).__name__}.process_read_batch",
             observability_options=observability_options,
-        ):
+        ), MetricsCapture():
             kwargs = copy.deepcopy(batch["read"])
             keyset_dict = kwargs.pop("keyset")
             kwargs["keyset"] = KeySet._from_dict(keyset_dict)
@@ -1625,7 +1626,7 @@ def generate_query_batches(
             f"CloudSpanner.{type(self).__name__}.generate_query_batches",
             extra_attributes=dict(sql=sql),
             observability_options=self.observability_options,
-        ):
+        ), MetricsCapture():
             partitions = self._get_snapshot().partition_query(
                 sql=sql,
                 params=params,
@@ -1681,7 +1682,7 @@ def process_query_batch(
         with trace_call(
             f"CloudSpanner.{type(self).__name__}.process_query_batch",
             observability_options=self.observability_options,
-        ):
+        ), MetricsCapture():
             return self._get_snapshot().execute_sql(
                 partition=batch["partition"],
                 **batch["query"],
@@ -1746,7 +1747,7 @@ def run_partitioned_query(
             f"CloudSpanner.${type(self).__name__}.run_partitioned_query",
             extra_attributes=dict(sql=sql),
             observability_options=self.observability_options,
-        ):
+        ), MetricsCapture():
             partitions = list(
                 self.generate_query_batches(
                     sql,
diff --git a/google/cloud/spanner_v1/merged_result_set.py b/google/cloud/spanner_v1/merged_result_set.py
@@ -18,6 +18,7 @@
 from threading import Lock, Event
 
 from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
+from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
 
 if TYPE_CHECKING:
     from google.cloud.spanner_v1.database import BatchSnapshot
@@ -45,7 +46,7 @@ def run(self):
         with trace_call(
             "CloudSpanner.PartitionExecutor.run",
             observability_options=observability_options,
-        ):
+        ), MetricsCapture():
             self.__run()
 
     def __run(self):
diff --git a/google/cloud/spanner_v1/pool.py b/google/cloud/spanner_v1/pool.py
@@ -32,6 +32,8 @@
 )
 from warnings import warn
 
+from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
+
 _NOW = datetime.datetime.utcnow  # unit tests may replace
 
 
@@ -242,7 +244,7 @@ def bind(self, database):
         with trace_call(
             "CloudSpanner.FixedPool.BatchCreateSessions",
             observability_options=observability_options,
-        ) as span:
+        ) as span, MetricsCapture():
             returned_session_count = 0
             while not self._sessions.full():
                 request.session_count = requested_session_count - self._sessions.qsize()
@@ -552,7 +554,7 @@ def bind(self, database):
         with trace_call(
             "CloudSpanner.PingingPool.BatchCreateSessions",
             observability_options=observability_options,
-        ) as span:
+        ) as span, MetricsCapture():
             returned_session_count = 0
             while returned_session_count < self.size:
                 resp = api.batch_create_sessions(
diff --git a/google/cloud/spanner_v1/session.py b/google/cloud/spanner_v1/session.py
@@ -40,6 +40,8 @@
 from google.cloud.spanner_v1.snapshot import Snapshot
 from google.cloud.spanner_v1.transaction import Transaction
 
+from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
+
 
 DEFAULT_RETRY_TIMEOUT_SECS = 30
 """Default timeout used by :meth:`Session.run_in_transaction`."""
@@ -165,7 +167,7 @@ def create(self):
             self,
             self._labels,
             observability_options=observability_options,
-        ):
+        ), MetricsCapture():
             session_pb = api.create_session(
                 request=request,
                 metadata=metadata,
@@ -205,7 +207,7 @@ def exists(self):
         observability_options = getattr(self._database, "observability_options", None)
         with trace_call(
             "CloudSpanner.GetSession", self, observability_options=observability_options
-        ) as span:
+        ) as span, MetricsCapture():
             try:
                 api.get_session(name=self.name, metadata=metadata)
                 if span:
@@ -248,7 +250,7 @@ def delete(self):
                 "session.name": self.name,
             },
             observability_options=observability_options,
-        ):
+        ), MetricsCapture():
             api.delete_session(name=self.name, metadata=metadata)
 
     def ping(self):
@@ -467,7 +469,7 @@ def run_in_transaction(self, func, *args, **kw):
             "CloudSpanner.Session.run_in_transaction",
             self,
             observability_options=observability_options,
-        ) as span:
+        ) as span, MetricsCapture():
             while True:
                 if self._transaction is None:
                     txn = self.transaction()
diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py
@@ -43,6 +43,8 @@
 from google.cloud.spanner_v1.streamed import StreamedResultSet
 from google.cloud.spanner_v1 import RequestOptions
 
+from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
+
 _STREAM_RESUMPTION_INTERNAL_ERROR_MESSAGES = (
     "RST_STREAM",
     "Received unexpected EOS on DATA frame from server",
@@ -96,7 +98,7 @@ def _restart_on_unavailable(
                     session,
                     attributes,
                     observability_options=observability_options,
-                ):
+                ), MetricsCapture():
                     iterator = method(request=request)
             for item in iterator:
                 item_buffer.append(item)
@@ -119,7 +121,7 @@ def _restart_on_unavailable(
                 session,
                 attributes,
                 observability_options=observability_options,
-            ):
+            ), MetricsCapture():
                 request.resume_token = resume_token
                 if transaction is not None:
                     transaction_selector = transaction._make_txn_selector()
@@ -139,7 +141,7 @@ def _restart_on_unavailable(
                 session,
                 attributes,
                 observability_options=observability_options,
-            ):
+            ), MetricsCapture():
                 request.resume_token = resume_token
                 if transaction is not None:
                     transaction_selector = transaction._make_txn_selector()
@@ -689,7 +691,7 @@ def partition_read(
             self._session,
             extra_attributes=trace_attributes,
             observability_options=getattr(database, "observability_options", None),
-        ):
+        ), MetricsCapture():
             method = functools.partial(
                 api.partition_read,
                 request=request,
@@ -792,7 +794,7 @@ def partition_query(
             self._session,
             trace_attributes,
             observability_options=getattr(database, "observability_options", None),
-        ):
+        ), MetricsCapture():
             method = functools.partial(
                 api.partition_query,
                 request=request,
@@ -938,7 +940,7 @@ def begin(self):
             f"CloudSpanner.{type(self).__name__}.begin",
             self._session,
             observability_options=getattr(database, "observability_options", None),
-        ):
+        ), MetricsCapture():
             method = functools.partial(
                 api.begin_transaction,
                 session=self._session.name,
diff --git a/google/cloud/spanner_v1/transaction.py b/google/cloud/spanner_v1/transaction.py
@@ -34,6 +34,7 @@
 from google.cloud.spanner_v1.batch import _BatchBase
 from google.cloud.spanner_v1._opentelemetry_tracing import add_span_event, trace_call
 from google.cloud.spanner_v1 import RequestOptions
+from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
 from google.api_core import gapic_v1
 from google.api_core.exceptions import InternalServerError
 from dataclasses import dataclass
@@ -118,7 +119,7 @@ def _execute_request(
         request.transaction = transaction
         with trace_call(
             trace_name, session, attributes, observability_options=observability_options
-        ):
+        ), MetricsCapture():
             method = functools.partial(method, request=request)
             response = _retry(
                 method,
@@ -160,7 +161,7 @@ def begin(self):
             f"CloudSpanner.{type(self).__name__}.begin",
             self._session,
             observability_options=observability_options,
-        ) as span:
+        ) as span, MetricsCapture():
             method = functools.partial(
                 api.begin_transaction,
                 session=self._session.name,
@@ -202,7 +203,7 @@ def rollback(self):
                 f"CloudSpanner.{type(self).__name__}.rollback",
                 self._session,
                 observability_options=observability_options,
-            ):
+            ), MetricsCapture():
                 method = functools.partial(
                     api.rollback,
                     session=self._session.name,
@@ -250,7 +251,7 @@ def commit(
             self._session,
             trace_attributes,
             observability_options,
-        ) as span:
+        ) as span, MetricsCapture():
             self._check_state()
             if self._transaction_id is None and len(self._mutations) > 0:
                 self.begin()
diff --git a/testing/constraints-3.7.txt b/testing/constraints-3.7.txt
@@ -17,3 +17,4 @@ protobuf==3.20.2
 deprecated==1.2.14
 grpc-interceptor==0.15.4
 google-cloud-monitoring==2.16.0
+mmh3==4.1.0
diff --git a/tests/unit/test_metrics.py b/tests/unit/test_metrics.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
 from unittest.mock import MagicMock
 from google.api_core.exceptions import ServiceUnavailable
 from google.cloud.spanner_v1.client import Client
@@ -21,6 +22,10 @@
     SpannerMetricsTracerFactory,
 )
 
+pytest.importorskip("opentelemetry")
+# Skip if semconv attributes are not present, as tracing wont' be enabled either
+# pytest.importorskip("opentelemetry.semconv.attributes.otel_attributes")
+
 
 def test_metrics_emission_with_failure_attempt(monkeypatch):
     monkeypatch.setenv("SPANNER_ENABLE_BUILTIN_METRICS", "true")
