From 0b93246cc51d526dafd5a4c26995456a266f02a1 Mon Sep 17 00:00:00 2001
From: Timothy Warren <twarren@nabancard.com>
Date: Thu, 11 Apr 2019 17:16:35 -0400
Subject: [PATCH] Regex and Integration tests

---
 Cargo.toml                |  2 ++
 src/lib.rs                |  3 ++
 src/query_builder.rs      | 27 ++++++++++++------
 tests/integration_test.rs | 59 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 83 insertions(+), 8 deletions(-)
 create mode 100644 tests/integration_test.rs

diff --git a/Cargo.toml b/Cargo.toml
index 65c087c..b60d06c 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -7,6 +7,8 @@ categories = ["database"]
 
 [dependencies]
 dotenv = "0.13.0"
+lazy_static = "1.3.0"
+regex = "1.1.5"
 serde_json = "1.0.39"
 
 [dependencies.pg]
diff --git a/src/lib.rs b/src/lib.rs
index a17f30b..adff51d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -7,6 +7,9 @@
 #![allow(dead_code)]
 #![allow(unused_variables)]
 
+#[macro_use]
+extern crate lazy_static;
+
 pub mod drivers;
 pub mod enums;
 pub mod query_builder;
diff --git a/src/query_builder.rs b/src/query_builder.rs
index 907ae78..c2b9cd6 100644
--- a/src/query_builder.rs
+++ b/src/query_builder.rs
@@ -10,6 +10,7 @@ use crate::enums::*;
 use crate::split_map_join;
 use crate::types::Wild;
 
+use regex::Regex;
 use query_state::QueryState;
 
 #[derive(Debug)]
@@ -52,14 +53,23 @@ impl QueryBuilder {
 
     /// Set the fields to select from the database as a string
     pub fn select(&mut self, fields: &str) -> &mut Self {
-        let fields = split_map_join(fields, ",", |s| s.trim().to_string());
+        lazy_static! {
+            static ref RE: Regex = Regex::new(r"(?i) as ").unwrap();
+        };
 
-        // Split identifiers on `As` keyword so they can be quoted properly
-        // @TODO split identifiers on `as` keyword (needs to be case-insensitive)
+        let fields = split_map_join(fields, ",", |s| {
+            if ! RE.is_match(s) {
+                return self.driver.quote_identifier(s.trim());
+            }
 
-        // Quote the identifiers (where there was an `as` keyword)
-
-        // Rejoin those identifiers
+            // Do a operation similar to split_map_join for the
+            // regex matches, quoting each identifier
+            RE.split(s)
+                .into_iter()
+                .map(|p| self.driver.quote_identifier(p))
+                .collect::<Vec<String>>()
+                .join(" as ")
+        });
 
         self.state.append_select_string(&fields);
 
@@ -175,7 +185,8 @@ impl QueryBuilder {
     /// Set a key and value for an insert or update query
     pub fn set(&mut self, key: &str, value: Wild) -> &mut Self {
         // @TODO figure a way to make this easier to use
-        self.state.append_set_array_keys(key).append_values(value);
+        let key = self.driver.quote_identifier(key);
+        self.state.append_set_array_keys(&key).append_values(value);
 
         self
     }
@@ -538,7 +549,7 @@ mod tests {
 
         qb.set("foo", Box::new("bar"));
 
-        assert_eq!(qb.state.get_set_array_keys()[0], "foo");
+        assert_eq!(qb.state.get_set_array_keys()[0], "\"foo\"");
         assert!(qb.state.get_values()[0].is::<&str>());
 
         // @TODO find a way to make this kind of operation much more ergonomic
diff --git a/tests/integration_test.rs b/tests/integration_test.rs
new file mode 100644
index 0000000..f7dbcc9
--- /dev/null
+++ b/tests/integration_test.rs
@@ -0,0 +1,59 @@
+use stringqb::query_builder::QueryBuilder;
+
+#[test]
+fn minimal_select_query() {
+    let mut qb = QueryBuilder::default();
+
+    let sql = qb.from("foo").get_compiled_select();
+
+    assert_eq!(sql, "SELECT *\nFROM \"foo\"");
+}
+
+#[test]
+fn select_keys_as_query() {
+    let mut qb = QueryBuilder::default();
+
+    let sql = qb.select("foo as bar, baz").from("a").get_compiled_select();
+
+    assert_eq!(sql, "SELECT \"foo\" as \"bar\",\"baz\"\nFROM \"a\"");
+}
+
+#[test]
+fn select_keys_query() {
+    let mut qb = QueryBuilder::default();
+
+    let sql = qb.select("a, b, c").from("foo").get_compiled_select();
+
+    assert_eq!(sql, "SELECT \"a\",\"b\",\"c\"\nFROM \"foo\"");
+}
+
+#[test]
+fn select_vec_keys_query() {
+    let mut qb = QueryBuilder::default();
+
+    let sql = qb
+        .select_vec(vec!["a", "b", "c"])
+        .from("foo")
+        .get_compiled_select();
+
+    assert_eq!(sql, "SELECT \"a\",\"b\",\"c\"\nFROM \"foo\"");
+}
+
+#[test]
+#[should_panic]
+fn select_without_from() {
+    let qb = QueryBuilder::default();
+
+    qb.get_compiled_select();
+}
+
+#[test]
+fn basic_insert_query() {
+    let mut qb = QueryBuilder::default();
+
+    qb.set("foo", Box::new("bar"));
+
+    let sql = qb.get_compiled_insert("foobar");
+
+    assert_eq!(sql, "INSERT INTO \"foobar\" (\"foo\")\nVALUES(?)");
+}