feat(driver): add try_collect for rows (#196)

Author: everpcpc

driver/src/conn.rs

@@ -76,7 +76,7 @@ pub trait Connection: DynClone + Send + Sync {
         let row = self.query_row("SELECT version()").await?;
         let version = match row {
             Some(row) => {
-                let (version,): (String,) = row.try_into()?;
+                let (version,): (String,) = row.try_into().map_err(Error::Parsing)?;
                 version
             }
             None => "".to_string(),

driver/src/flight_sql.rs

@@ -78,7 +78,7 @@ impl Connection for FlightSQLConnection {
             Ok(_) => None,
             Err(err) => Some(Err(err)),
         });
-        Ok(Box::pin(rows))
+        Ok(RowIterator::new(Box::pin(rows)))
     }
 
     async fn query_iter_ext(&self, sql: &str) -> Result<(Schema, RowProgressIterator)> {
@@ -92,7 +92,7 @@ impl Connection for FlightSQLConnection {
             .ok_or(Error::Protocol("Ticket is empty".to_string()))?;
         let flight_data = client.do_get(ticket.clone()).await?;
         let (schema, rows) = FlightSQLRows::try_from_flight_data(flight_data).await?;
-        Ok((schema, Box::pin(rows)))
+        Ok((schema, RowProgressIterator::new(Box::pin(rows))))
     }
 
     async fn stream_load(

driver/src/rest_api.rs

@@ -62,13 +62,13 @@ impl Connection for RestAPIConnection {
             Ok(_) => None,
             Err(err) => Some(Err(err)),
         });
-        Ok(Box::pin(rows))
+        Ok(RowIterator::new(Box::pin(rows)))
     }
 
     async fn query_iter_ext(&self, sql: &str) -> Result<(Schema, RowProgressIterator)> {
         let resp = self.client.query(sql).await?;
         let (schema, rows) = RestAPIRows::from_response(self.client.clone(), resp)?;
-        Ok((schema, Box::pin(rows)))
+        Ok((schema, RowProgressIterator::new(Box::pin(rows))))
     }
 
     async fn query_row(&self, sql: &str) -> Result<Option<Row>> {

driver/tests/driver/select_iter.rs

@@ -167,19 +167,17 @@ async fn select_iter_struct() {
     ];
 
     let sql_select = format!("SELECT * FROM `{}`", table);
-    let mut rows = conn.query_iter(&sql_select).await.unwrap();
-    let mut row_count = 0;
-    while let Some(row) = rows.next().await {
-        let v: RowResult = row.unwrap().try_into().unwrap();
-        let expected_row = &expected[row_count];
+    let rows = conn.query_iter(&sql_select).await.unwrap();
+    let results = rows.try_collect::<RowResult>().await.unwrap();
+    for (idx, v) in results.iter().enumerate() {
+        let expected_row = &expected[idx];
         assert_eq!(v.i64, expected_row.i64);
         assert_eq!(v.u64, expected_row.u64);
         assert_eq!(v.f64, expected_row.f64);
         assert_eq!(v.s, expected_row.s);
         assert_eq!(v.s2, expected_row.s2);
         assert_eq!(v.d, expected_row.d);
         assert_eq!(v.t, expected_row.t);
-        row_count += 1;
     }
 }
 

macros/src/from_row.rs

@@ -38,18 +38,18 @@ pub fn from_row_derive(tokens_input: TokenStream) -> TokenStream {
                                // it is safe to unwrap
                 let t = col_value.get_type();
                 <#field_type>::try_from(col_value)
-                    .map_err(|_| #path::Error::InvalidResponse(format!("failed converting column {} from type({:?}) to type({})", col_ix, t, std::any::type_name::<#field_type>())))?
+                    .map_err(|_| format!("failed converting column {} from type({:?}) to type({})", col_ix, t, std::any::type_name::<#field_type>()))?
             },
         }
     });
 
     let fields_count = struct_fields.named.len();
     let generated = quote! {
         impl #impl_generics TryFrom<#path::Row> for #struct_name #ty_generics #where_clause {
-            type Error = #path::Error;
-            fn try_from(row: #path::Row) -> #path::Result<Self> {
+            type Error = String;
+            fn try_from(row: #path::Row) -> #path::Result<Self, String> {
                 if #fields_count != row.len() {
-                    return Err(#path::Error::InvalidResponse(format!("row size mismatch: expected {} columns, got {}", #fields_count, row.len())));
+                    return Err(format!("row size mismatch: expected {} columns, got {}", #fields_count, row.len()));
                 }
                 let mut vals_iter = row.into_iter().enumerate();
                 Ok(#struct_name {

sql/Cargo.toml

@@ -16,6 +16,7 @@ flight-sql = ["dep:arrow-array", "dep:arrow-cast", "dep:arrow-flight", "dep:arro
 [dependencies]
 databend-client = { workspace = true }
 
+async-trait = "0.1.68"
 chrono = { version = "0.4.26", default-features = false }
 serde = { version = "1.0.164", default-features = false, features = ["derive"] }
 serde_json = { version = "1.0.97", default-features = false, features = ["std"] }

sql/src/from_row.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::error::{Error, Result};
+use crate::error::Result;
 use crate::rows::Row;
 use crate::value::Value;
 
@@ -29,15 +29,15 @@ macro_rules! impl_tuple_from_row {
         where
             $($Ti: TryFrom<Value>),+
         {
-            type Error = Error;
-            fn try_from(row: Row) -> Result<Self> {
+            type Error = String;
+            fn try_from(row: Row) -> Result<Self, String> {
                 // It is not possible yet to get the number of metavariable repetitions
                 // ref: https://github.com/rust-lang/lang-team/issues/28#issue-644523674
                 // This is a workaround
                 let expected_len = <[()]>::len(&[$(replace_expr!(($Ti) ())),*]);
 
                 if expected_len != row.len() {
-                    return Err(Error::InvalidResponse(format!("row size mismatch: expected {} columns, got {}", expected_len, row.len())));
+                    return Err(format!("row size mismatch: expected {} columns, got {}", expected_len, row.len()));
                 }
                 let mut vals_iter = row.into_iter().enumerate();
 
@@ -50,7 +50,7 @@ macro_rules! impl_tuple_from_row {
                                            // so it is safe to unwrap
                             let t = col_value.get_type();
                             $Ti::try_from(col_value)
-                                .map_err(|_| Error::InvalidResponse(format!("failed converting column {} from type({:?}) to type({})", col_ix, t, std::any::type_name::<$Ti>())))?
+                                .map_err(|_| format!("failed converting column {} from type({:?}) to type({})", col_ix, t, std::any::type_name::<$Ti>()))?
                         }
                     ,)+
                 ))

sql/src/lib.rs

@@ -21,7 +21,7 @@ pub mod value;
 #[doc(hidden)]
 pub mod _macro_internal {
     pub use crate::error::{Error, Result};
-    pub use crate::rows::Row;
+    pub use crate::rows::{Row, RowIterator};
     pub use crate::schema::Schema;
     pub use crate::value::Value;
 }

sql/src/rows.rs

@@ -13,9 +13,11 @@
 // limitations under the License.
 
 use std::pin::Pin;
+use std::task::Context;
+use std::task::Poll;
 
 use serde::Deserialize;
-use tokio_stream::Stream;
+use tokio_stream::{Stream, StreamExt};
 
 #[cfg(feature = "flight-sql")]
 use arrow::record_batch::RecordBatch;
@@ -24,9 +26,6 @@ use crate::error::{Error, Result};
 use crate::schema::SchemaRef;
 use crate::value::Value;
 
-pub type RowIterator = Pin<Box<dyn Stream<Item = Result<Row>> + Send>>;
-pub type RowProgressIterator = Pin<Box<dyn Stream<Item = Result<RowWithProgress>> + Send>>;
-
 #[derive(Clone, Debug)]
 pub enum RowWithProgress {
     Row(Row),
@@ -149,3 +148,57 @@ impl IntoIterator for Rows {
         self.0.into_iter()
     }
 }
+
+pub struct RowIterator(Pin<Box<dyn Stream<Item = Result<Row>> + Send>>);
+
+impl RowIterator {
+    pub fn new(it: Pin<Box<dyn Stream<Item = Result<Row>> + Send>>) -> Self {
+        Self(it)
+    }
+
+    pub async fn try_collect<T>(mut self) -> Result<Vec<T>>
+    where
+        T: TryFrom<Row>,
+        T::Error: std::fmt::Display,
+    {
+        let mut ret = Vec::new();
+        while let Some(row) = self.0.next().await {
+            let v = T::try_from(row?).map_err(|e| Error::Parsing(e.to_string()))?;
+            ret.push(v)
+        }
+        Ok(ret)
+    }
+}
+
+impl Stream for RowIterator {
+    type Item = Result<Row>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        Pin::new(&mut self.0).poll_next(cx)
+    }
+}
+
+pub struct RowProgressIterator(Pin<Box<dyn Stream<Item = Result<RowWithProgress>> + Send>>);
+
+impl RowProgressIterator {
+    pub fn new(it: Pin<Box<dyn Stream<Item = Result<RowWithProgress>> + Send>>) -> Self {
+        Self(it)
+    }
+
+    pub async fn filter_rows(self) -> RowIterator {
+        let rows = self.0.filter_map(|r| match r {
+            Ok(RowWithProgress::Row(r)) => Some(Ok(r)),
+            Ok(_) => None,
+            Err(err) => Some(Err(err)),
+        });
+        RowIterator(Box::pin(rows))
+    }
+}
+
+impl Stream for RowProgressIterator {
+    type Item = Result<RowWithProgress>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        Pin::new(&mut self.0).poll_next(cx)
+    }
+}