feat: implement select iterator utility function

gluesql · Oct 15, 2023 · e0521ca · e0521ca
1 parent a509e78
commit e0521ca
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 46 deletions.
diff --git a/core/src/executor/execute.rs b/core/src/executor/execute.rs
@@ -53,6 +53,39 @@ pub enum Payload {
     ShowVariable(PayloadVariable),
 }
 
+impl Payload {
+    /// Exports `select` payloads as an [`std::iter::Iterator`].
+    ///
+    /// The items of the Iterator are `HashMap<Column, Value>`, and they are borrowed by default.
+    /// If ownership is required, you need to acquire them directly.
+    ///
+    /// - Some: [`Payload::Select`], [`Payload::SelectMap`]
+    /// - None: otherwise
+    pub fn select(&self) -> Option<impl Iterator<Item = HashMap<&str, &Value>>> {
+        #[derive(iter_enum::Iterator)]
+        enum Iter<I1, I2> {
+            Schema(I1),
+            Schemaless(I2),
+        }
+
+        Some(match self {
+            Payload::Select { labels, rows } => Iter::Schema(rows.iter().map(move |row| {
+                labels
+                    .iter()
+                    .zip(row.iter())
+                    .map(|(label, value)| (label.as_str(), value))
+                    .collect::<HashMap<_, _>>()
+            })),
+            Payload::SelectMap(rows) => Iter::Schemaless(rows.iter().map(|row| {
+                row.iter()
+                    .map(|(k, v)| (k.as_str(), v))
+                    .collect::<HashMap<_, _>>()
+            })),
+            _ => return None,
+        })
+    }
+}
+
 #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
 pub enum PayloadVariable {
     Tables(Vec<String>),

diff --git a/pkg/rust/examples/hello_world.rs b/pkg/rust/examples/hello_world.rs
@@ -1,56 +1,14 @@
 #[cfg(feature = "sled-storage")]
 mod hello_world {
     use {
-        gluesql::{
-            prelude::{Glue, Payload},
-            sled_storage::SledStorage,
-        },
+        gluesql::{prelude::Glue, sled_storage::SledStorage},
         std::fs,
     };
 
-    struct GreetTable {
-        rows: Vec<GreetRow>,
-    }
     struct GreetRow {
         name: String,
     }
 
-    impl TryFrom<Payload> for GreetTable {
-        type Error = &'static str;
-
-        fn try_from(payload: Payload) -> Result<Self, Self::Error> {
-            match payload {
-                Payload::Select { labels: _, rows } => {
-                    let rows = rows
-                        .into_iter()
-                        .map(|mut row| {
-                            let name = row.remove(0);
-                            GreetRow {
-                                name: String::from(name),
-                            }
-                        })
-                        .collect::<Vec<_>>();
-
-                    Ok(Self { rows })
-                }
-                Payload::SelectMap(rows) => {
-                    let rows = rows
-                        .into_iter()
-                        .map(|row| {
-                            let name = row.get("name").unwrap();
-                            GreetRow {
-                                name: String::from(name),
-                            }
-                        })
-                        .collect::<Vec<_>>();
-
-                    Ok(Self { rows })
-                }
-                _ => Err("unexpected payload, expected a select query result"),
-            }
-        }
-    }
-
     pub async fn run() {
         /*
             Initiate a connection
@@ -92,10 +50,22 @@ mod hello_world {
         */
         assert_eq!(result.len(), 1);
 
-        let table: GreetTable = result.remove(0).try_into().unwrap();
-        assert_eq!(table.rows.len(), 1);
+        let payload = result.remove(0);
+
+        let rows = payload
+            .select()
+            .unwrap()
+            .map(|map| {
+                let name = *map.get("name").unwrap();
+                let name = name.into();
+
+                GreetRow { name }
+            })
+            .collect::<Vec<_>>();
+
+        assert_eq!(rows.len(), 1);
 
-        println!("Hello {}!", table.rows[0].name); // Will always output "Hello World!"
+        println!("Hello {}!", rows[0].name); // Will always output "Hello World!"
     }
 }