chore(spanner): trace server streaming queries and reads in ResultSet

src/spanner/src/batch_read_only_transaction.rs

@@ -406,6 +406,7 @@ impl Partition {
             operation: StreamOperation::Query(req.clone()),
             channel_hint,
             gax_options,
+            method_name: "ExecuteStreamingSql",
         })
         .await
     }
@@ -437,6 +438,7 @@ impl Partition {
             operation: StreamOperation::Read(req.clone()),
             channel_hint,
             gax_options,
+            method_name: "StreamingRead",
         })
         .await
     }

src/spanner/src/client.rs

@@ -126,15 +126,15 @@ macro_rules! define_idempotent_rpc {
             channel_hint: usize,
             o11y: &crate::observability::Observability,
         ) -> crate::Result<$response_type> {
-            o11y.trace_operation($canonical_name, || async move {
+            o11y.trace_operation(
+                $canonical_name,
                 self.get_channel(channel_hint)
                     .inner
                     .$method()
                     .with_request(request)
                     .with_options(apply_request_defaults(options))
-                    .send()
-                    .await
-            })
+                    .send(),
+            )
             .await
         }
     };

src/spanner/src/observability/metrics.rs

@@ -14,6 +14,8 @@
 
 use crate::omni::InstanceType;
 #[cfg(feature = "_experimental-builtin-metrics")]
+use std::sync::Arc;
+#[cfg(feature = "_experimental-builtin-metrics")]
 use std::time::Duration;
 #[cfg(feature = "_experimental-builtin-metrics")]
 use std::time::Instant;
@@ -73,10 +75,10 @@ impl SpannerMetrics {
 }
 
 #[cfg(feature = "_experimental-builtin-metrics")]
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub(crate) struct Observability {
-    pub(crate) metrics: Option<SpannerMetrics>,
-    _meter_provider: Option<SdkMeterProvider>,
+    pub(crate) metrics: Option<Arc<SpannerMetrics>>,
+    _meter_provider: Option<Arc<SdkMeterProvider>>,
 }
 
 #[cfg(feature = "_experimental-builtin-metrics")]
@@ -101,8 +103,8 @@ impl Observability {
         }
 
         let project_id = match project_id {
-            Some(id) => id,
-            None => return Self::disabled(),
+            Some(id) if !id.is_empty() => id,
+            _ => return Self::disabled(),
         };
 
         // Create the Google Cloud Monitoring client using the same config
@@ -137,22 +139,24 @@ impl Observability {
         let metrics = SpannerMetrics::new(meter);
 
         Self {
-            metrics: Some(metrics),
-            _meter_provider: Some(meter_provider),
+            metrics: Some(Arc::new(metrics)),
+            _meter_provider: Some(Arc::new(meter_provider)),
         }
     }
 
-    pub(crate) async fn trace_operation<F, Fut, T>(
+    pub(crate) async fn trace_operation<Fut, T>(
         &self,
         method: &'static str,
-        f: F,
+        fut: Fut,
     ) -> crate::Result<T>
     where
-        F: FnOnce() -> Fut,
         Fut: std::future::Future<Output = crate::Result<T>>,
     {
+        if self.metrics.is_none() {
+            return fut.await;
+        }
         let start_time = Instant::now();
-        let result = f().await;
+        let result = fut.await;
         let elapsed = start_time.elapsed();
         self.record_operation(method, elapsed, &result);
         result
@@ -278,7 +282,7 @@ pub(crate) fn parse_server_timing(header_val: &str) -> ServerTimings {
 }
 
 #[cfg(not(feature = "_experimental-builtin-metrics"))]
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub(crate) struct Observability;
 
 #[cfg(not(feature = "_experimental-builtin-metrics"))]
@@ -296,22 +300,79 @@ impl Observability {
         Self
     }
 
-    pub(crate) async fn trace_operation<F, Fut, T>(
+    #[inline(always)]
+    pub(crate) async fn trace_operation<Fut, T>(
         &self,
         _method: &'static str,
-        f: F,
+        fut: Fut,
     ) -> crate::Result<T>
     where
-        F: FnOnce() -> Fut,
         Fut: std::future::Future<Output = crate::Result<T>>,
     {
-        f().await
+        fut.await
     }
 }
 
 #[cfg(all(test, feature = "_experimental-builtin-metrics"))]
 mod tests {
     use super::*;
+    use opentelemetry_sdk::metrics::InMemoryMetricExporter;
+
+    #[test]
+    fn test_observability_disabled() {
+        let o11y = Observability::disabled();
+        assert!(o11y.metrics.is_none());
+    }
+
+    #[test]
+    fn test_result_to_status_str() {
+        let ok_res: crate::Result<()> = Ok(());
+        assert_eq!(result_to_status_str(&ok_res), "OK");
+
+        let status_pd = google_cloud_gax::error::rpc::Status::default()
+            .set_code(google_cloud_gax::error::rpc::Code::PermissionDenied);
+        let err_pd: crate::Result<()> = Err(crate::Error::service(status_pd));
+        assert_eq!(result_to_status_str(&err_pd), "PERMISSION_DENIED");
+    }
+
+    #[test]
+    fn test_spanner_metrics_record_operation_and_attempt() {
+        let exporter = InMemoryMetricExporter::default();
+        let reader = opentelemetry_sdk::metrics::PeriodicReader::builder(exporter.clone()).build();
+        let provider = SdkMeterProvider::builder().with_reader(reader).build();
+        let meter = provider.meter("cloud.google.com/rust");
+        let metrics = SpannerMetrics::new(meter);
+        let o11y = Observability {
+            metrics: Some(Arc::new(metrics)),
+            _meter_provider: Some(Arc::new(provider.clone())),
+        };
+
+        let ok_res: crate::Result<()> = Ok(());
+        o11y.record_operation("ExecuteSql", Duration::from_millis(50), &ok_res);
+        o11y.record_attempt(
+            "ExecuteSql",
+            Duration::from_millis(40),
+            &ok_res,
+            Some(12.5),
+            Some(5.0),
+        );
+
+        provider.force_flush().expect("force_flush failed");
+
+        let finished = exporter
+            .get_finished_metrics()
+            .expect("get_finished_metrics");
+        assert!(!finished.is_empty());
+    }
+
+    #[tokio::test]
+    async fn test_trace_operation_success() {
+        let o11y = Observability::disabled();
+        let result = o11y
+            .trace_operation("ExecuteSql", async { Ok::<i32, crate::Error>(42) })
+            .await;
+        assert_eq!(result.expect("trace_operation result"), 42);
+    }
 
     #[test]
     fn test_parse_server_timing() {

src/spanner/src/read_only_transaction.rs

@@ -1072,7 +1072,7 @@ fn merge_request_options(
 
 /// Helper macro to execute a streaming SQL or streaming read RPC with retry logic.
 macro_rules! execute_stream_with_retry {
-    ($self:expr, $request:ident, $gax_options:ident, $rpc_method:ident, $operation_variant:path) => {{
+    ($self:expr, $request:ident, $gax_options:ident, $rpc_method:ident, $operation_variant:path, $method_name:expr) => {{
         let stream = match $self
             .client
             .spanner
@@ -1123,7 +1123,7 @@ macro_rules! execute_stream_with_retry {
             }
         };
 
-        ResultSet::create(ResultSetParams {
+        Box::pin(ResultSet::create(ResultSetParams {
             stream,
             transaction_selector: Some($self.transaction_selector.clone()),
             precommit_token_tracker: $self.precommit_token_tracker.clone(),
@@ -1133,7 +1133,8 @@ macro_rules! execute_stream_with_retry {
             operation: $operation_variant($request),
             channel_hint: $self.channel_hint,
             gax_options: $gax_options,
-        })
+            method_name: $method_name,
+        }))
         .await
     }};
 }
@@ -1158,7 +1159,8 @@ impl ReadContext {
             request,
             gax_options,
             execute_streaming_sql,
-            StreamOperation::Query
+            StreamOperation::Query,
+            "ExecuteStreamingSql"
         )
     }
 
@@ -1179,7 +1181,8 @@ impl ReadContext {
             request,
             gax_options,
             streaming_read,
-            StreamOperation::Read
+            StreamOperation::Read,
+            "StreamingRead"
         )
     }
 }
@@ -3493,7 +3496,8 @@ pub(crate) mod tests {
                 request,
                 gax_options,
                 execute_streaming_sql,
-                StreamOperation::Query
+                StreamOperation::Query,
+                "ExecuteStreamingSql"
             )
         }
 

src/spanner/src/result_set.rs

@@ -98,6 +98,7 @@ pub(crate) struct ResultSetParams {
     pub operation: StreamOperation,
     pub channel_hint: usize,
     pub gax_options: GaxRequestOptions,
+    pub method_name: &'static str,
 }
 
 // The maximum number of PartialResultSets to buffer without a resume token.
@@ -111,8 +112,11 @@ const DEFAULT_ATTEMPT_LIMIT: u32 = 10;
 impl ResultSet {
     /// Creates a new result set asynchronously, waiting for the first chunk to arrive.
     pub(crate) async fn create(params: ResultSetParams) -> crate::Result<Self> {
+        let method_name = params.method_name;
         let mut result_set = Self::new(params);
-        result_set.init_stream().await?;
+        let o11y = result_set.client.o11y.clone();
+        let fut = Box::pin(result_set.init_stream());
+        o11y.trace_operation(method_name, fut).await?;
         Ok(result_set)
     }
 
@@ -128,6 +132,7 @@ impl ResultSet {
             operation,
             channel_hint,
             gax_options,
+            method_name: _,
         } = params;
 
         let gax_options = Self::apply_defaults(gax_options);
@@ -1871,6 +1876,7 @@ pub(crate) mod tests {
             operation: StreamOperation::Query(req),
             channel_hint: 0,
             gax_options: GaxRequestOptions::default(),
+            method_name: "ExecuteStreamingSql",
         })
         .await?;