Merge pull request #407 from danielnorberg/dano/spanner-query-stats

[spanner] Expose query stats in RowIterator

Author: yoshidan

spanner/src/reader.rs

@@ -5,7 +5,7 @@ use prost_types::{value::Kind, Value};
 
 use google_cloud_gax::grpc::{Code, Response, Status, Streaming};
 use google_cloud_googleapis::spanner::v1::struct_type::Field;
-use google_cloud_googleapis::spanner::v1::{ExecuteSqlRequest, PartialResultSet, ReadRequest, ResultSetMetadata};
+use google_cloud_googleapis::spanner::v1::{ExecuteSqlRequest, PartialResultSet, ReadRequest, ResultSetMetadata, ResultSetStats};
 
 use crate::row::Row;
 use crate::session::SessionHandle;
@@ -186,6 +186,7 @@ where
     reader: T,
     rs: ResultSet,
     reader_option: Option<CallOptions>,
+    stats: Option<ResultSetStats>,
 }
 
 impl<'a, T> RowIterator<'a, T>
@@ -210,6 +211,7 @@ where
             reader,
             rs,
             reader_option: None,
+            stats: None,
         })
     }
 
@@ -241,6 +243,10 @@ where
                 if !result_set.resume_token.is_empty() {
                     self.reader.update_token(result_set.resume_token);
                 }
+                // Capture stats if present (only sent with the last response)
+                if result_set.stats.is_some() {
+                    self.stats = result_set.stats;
+                }
                 self.rs
                     .add(result_set.metadata, result_set.values, result_set.chunked_value)
             }
@@ -262,6 +268,13 @@ where
         None
     }
 
+    /// Returns query execution statistics if available.
+    /// Stats are only available after all rows have been consumed and only when
+    /// the query was executed with a QueryMode that includes stats (Profile, WithStats, or WithPlanAndStats).
+    pub fn stats(&self) -> Option<&ResultSetStats> {
+        self.stats.as_ref()
+    }
+
     /// next returns the next result.
     /// Its second return value is None if there are no more results.
     pub async fn next(&mut self) -> Result<Option<Row>, Status> {

spanner/src/transaction.rs

@@ -132,8 +132,6 @@ impl Transaction {
 
     /// query executes a query against the database. It returns a RowIterator for
     /// retrieving the resulting rows.
-    ///
-    /// query returns only row data, without a query plan or execution statistics.
     pub async fn query_with_option(
         &mut self,
         statement: Statement,