| 
72 | 72 |     get_current_span,  | 
73 | 73 |     trace_call,  | 
74 | 74 | )  | 
 | 75 | +from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture  | 
75 | 76 | 
 
  | 
76 | 77 | 
 
  | 
77 | 78 | SPANNER_DATA_SCOPE = "https://www.googleapis.com/auth/spanner.data"  | 
@@ -702,7 +703,7 @@ def execute_pdml():  | 
702 | 703 |             with trace_call(  | 
703 | 704 |                 "CloudSpanner.Database.execute_partitioned_pdml",  | 
704 | 705 |                 observability_options=self.observability_options,  | 
705 |  | -            ) as span:  | 
 | 706 | +            ) as span, MetricsCapture():  | 
706 | 707 |                 with SessionCheckout(self._pool) as session:  | 
707 | 708 |                     add_span_event(span, "Starting BeginTransaction")  | 
708 | 709 |                     txn = api.begin_transaction(  | 
@@ -897,7 +898,7 @@ def run_in_transaction(self, func, *args, **kw):  | 
897 | 898 |         with trace_call(  | 
898 | 899 |             "CloudSpanner.Database.run_in_transaction",  | 
899 | 900 |             observability_options=observability_options,  | 
900 |  | -        ):  | 
 | 901 | +        ), MetricsCapture():  | 
901 | 902 |             # Sanity check: Is there a transaction already running?  | 
902 | 903 |             # If there is, then raise a red flag. Otherwise, mark that this one  | 
903 | 904 |             # is running.  | 
@@ -1489,7 +1490,7 @@ def generate_read_batches(  | 
1489 | 1490 |             f"CloudSpanner.{type(self).__name__}.generate_read_batches",  | 
1490 | 1491 |             extra_attributes=dict(table=table, columns=columns),  | 
1491 | 1492 |             observability_options=self.observability_options,  | 
1492 |  | -        ):  | 
 | 1493 | +        ), MetricsCapture():  | 
1493 | 1494 |             partitions = self._get_snapshot().partition_read(  | 
1494 | 1495 |                 table=table,  | 
1495 | 1496 |                 columns=columns,  | 
@@ -1540,7 +1541,7 @@ def process_read_batch(  | 
1540 | 1541 |         with trace_call(  | 
1541 | 1542 |             f"CloudSpanner.{type(self).__name__}.process_read_batch",  | 
1542 | 1543 |             observability_options=observability_options,  | 
1543 |  | -        ):  | 
 | 1544 | +        ), MetricsCapture():  | 
1544 | 1545 |             kwargs = copy.deepcopy(batch["read"])  | 
1545 | 1546 |             keyset_dict = kwargs.pop("keyset")  | 
1546 | 1547 |             kwargs["keyset"] = KeySet._from_dict(keyset_dict)  | 
@@ -1625,7 +1626,7 @@ def generate_query_batches(  | 
1625 | 1626 |             f"CloudSpanner.{type(self).__name__}.generate_query_batches",  | 
1626 | 1627 |             extra_attributes=dict(sql=sql),  | 
1627 | 1628 |             observability_options=self.observability_options,  | 
1628 |  | -        ):  | 
 | 1629 | +        ), MetricsCapture():  | 
1629 | 1630 |             partitions = self._get_snapshot().partition_query(  | 
1630 | 1631 |                 sql=sql,  | 
1631 | 1632 |                 params=params,  | 
@@ -1681,7 +1682,7 @@ def process_query_batch(  | 
1681 | 1682 |         with trace_call(  | 
1682 | 1683 |             f"CloudSpanner.{type(self).__name__}.process_query_batch",  | 
1683 | 1684 |             observability_options=self.observability_options,  | 
1684 |  | -        ):  | 
 | 1685 | +        ), MetricsCapture():  | 
1685 | 1686 |             return self._get_snapshot().execute_sql(  | 
1686 | 1687 |                 partition=batch["partition"],  | 
1687 | 1688 |                 **batch["query"],  | 
@@ -1746,7 +1747,7 @@ def run_partitioned_query(  | 
1746 | 1747 |             f"CloudSpanner.${type(self).__name__}.run_partitioned_query",  | 
1747 | 1748 |             extra_attributes=dict(sql=sql),  | 
1748 | 1749 |             observability_options=self.observability_options,  | 
1749 |  | -        ):  | 
 | 1750 | +        ), MetricsCapture():  | 
1750 | 1751 |             partitions = list(  | 
1751 | 1752 |                 self.generate_query_batches(  | 
1752 | 1753 |                     sql,  | 
 | 
0 commit comments