From 7fef56f90d78500cf5a3c81ea0d904eb3645fee6 Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Sun, 22 Oct 2023 13:07:40 -0400
Subject: [PATCH] Implement select with where clauses

---
 src/abstract_tree/select.rs | 52 +++++++++++++++++++++++++++++++++----
 src/value/table.rs          | 41 +++--------------------------
 2 files changed, 50 insertions(+), 43 deletions(-)

diff --git a/src/abstract_tree/select.rs b/src/abstract_tree/select.rs
index ebe4377..bb96abc 100644
--- a/src/abstract_tree/select.rs
+++ b/src/abstract_tree/select.rs
@@ -1,6 +1,6 @@
 use serde::{Deserialize, Serialize};
 
-use crate::{AbstractTree, Expression, Identifier, Item, Value};
+use crate::{AbstractTree, Expression, Identifier, Item, Result, Table, Value, VariableMap};
 
 #[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, PartialOrd, Ord)]
 pub struct Select {
@@ -52,9 +52,9 @@ impl AbstractTree for Select {
         })
     }
 
-    fn run(&self, source: &str, context: &mut crate::VariableMap) -> crate::Result<crate::Value> {
+    fn run(&self, source: &str, context: &mut VariableMap) -> Result<Value> {
         let value = self.expression.run(source, context)?;
-        let table = value.as_table()?;
+        let old_table = value.as_table()?;
         let column_names = if self.identifiers.len() > 0 {
             self.identifiers
                 .iter()
@@ -62,9 +62,51 @@ impl AbstractTree for Select {
                 .map(|identifierier| identifierier.take_inner())
                 .collect()
         } else {
-            table.column_names().clone()
+            old_table.headers().clone()
         };
-        let new_table = table.select(&column_names);
+
+        let mut new_table = Table::new(column_names.to_vec());
+
+        for row in old_table.rows() {
+            let mut new_row = Vec::new();
+            let mut row_context = VariableMap::new();
+
+            for (i, value) in row.iter().enumerate() {
+                let column_name = old_table.headers().get(i).unwrap();
+
+                row_context.set_value(column_name.clone(), value.clone())?;
+
+                let new_table_column_index =
+                    new_table
+                        .headers()
+                        .iter()
+                        .enumerate()
+                        .find_map(|(index, new_column_name)| {
+                            if new_column_name == column_name {
+                                Some(index)
+                            } else {
+                                None
+                            }
+                        });
+
+                if let Some(index) = new_table_column_index {
+                    while new_row.len() < index + 1 {
+                        new_row.push(Value::Empty);
+                    }
+                    new_row[index] = value.clone();
+                }
+            }
+
+            if let Some(where_clause) = &self.item {
+                let should_include = where_clause.run(source, &mut row_context)?.as_boolean()?;
+
+                if should_include {
+                    new_table.insert(new_row)?;
+                }
+            } else {
+                new_table.insert(new_row)?;
+            }
+        }
 
         Ok(Value::Table(new_table))
     }
diff --git a/src/value/table.rs b/src/value/table.rs
index 331d887..2cafaee 100644
--- a/src/value/table.rs
+++ b/src/value/table.rs
@@ -1,4 +1,4 @@
-use crate::{Error, Result, Value, VariableMap};
+use crate::{Error, Item, Result, Value, VariableMap};
 use comfy_table::{Cell, Color, ContentArrangement, Table as ComfyTable};
 use serde::{Deserialize, Serialize};
 use std::{
@@ -42,7 +42,7 @@ impl Table {
         self.rows.reserve(additional);
     }
 
-    pub fn column_names(&self) -> &Vec<String> {
+    pub fn headers(&self) -> &Vec<String> {
         &self.headers
     }
 
@@ -86,46 +86,11 @@ impl Table {
     }
 
     pub fn get(&self, value: &Value) -> Option<&Vec<Value>> {
-        let primary_key = self.column_names().get(self.primary_key_index)?;
+        let primary_key = self.headers().get(self.primary_key_index)?;
 
         self.get_where(primary_key, value)
     }
 
-    pub fn select(&self, column_names: &[String]) -> Table {
-        let mut new_table = Table::new(column_names.to_vec());
-
-        for row in &self.rows {
-            let mut new_row = Vec::new();
-
-            for (i, value) in row.iter().enumerate() {
-                let column_name = self.headers.get(i).unwrap();
-                let new_table_column_index =
-                    new_table
-                        .headers
-                        .iter()
-                        .enumerate()
-                        .find_map(|(index, new_column_name)| {
-                            if new_column_name == column_name {
-                                Some(index)
-                            } else {
-                                None
-                            }
-                        });
-
-                if let Some(index) = new_table_column_index {
-                    while new_row.len() < index + 1 {
-                        new_row.push(Value::Empty);
-                    }
-                    new_row[index] = value.clone();
-                }
-            }
-
-            new_table.insert(new_row).unwrap();
-        }
-
-        new_table
-    }
-
     pub fn get_where(&self, column_name: &str, expected: &Value) -> Option<&Vec<Value>> {
         let column_index = self.get_column_index(column_name)?;