From 510aaba4639174b3067440c060aae50e2f2c2a8f Mon Sep 17 00:00:00 2001
From: Christopher Hotchkiss <christopher.hotchkiss@gmail.com>
Date: Sat, 28 Aug 2021 12:51:01 -0400
Subject: [PATCH] Added support for parsing a primary key #2

---
 src/engine/objects/parse_tree.rs             |  1 +
 src/engine/sql_parser/create/create_table.rs | 27 ++++++++++++++++++--
 tests/sql_create_table.rs                    |  6 +++--
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/engine/objects/parse_tree.rs b/src/engine/objects/parse_tree.rs
index 361b107..8c647bb 100644
--- a/src/engine/objects/parse_tree.rs
+++ b/src/engine/objects/parse_tree.rs
@@ -18,6 +18,7 @@ pub struct RawColumn {
     pub name: String,
     pub sql_type: String,
     pub null: bool,
+    pub primary_key: bool,
 }
 
 #[derive(Clone, Debug, PartialEq)]
diff --git a/src/engine/sql_parser/create/create_table.rs b/src/engine/sql_parser/create/create_table.rs
index 90a4db1..d11cfb7 100644
--- a/src/engine/sql_parser/create/create_table.rs
+++ b/src/engine/sql_parser/create/create_table.rs
@@ -57,7 +57,7 @@ fn match_columns<'a, E: ParseError<&'a str> + ContextError<&'a str>>(
 fn match_column_attribute<'a, E: ParseError<&'a str> + ContextError<&'a str>>(
     input: &'a str,
 ) -> IResult<&'a str, RawColumn, E> {
-    let (input, (_, name, _, sql_type, _, is_null, _)) = tuple((
+    let (input, (_, name, _, sql_type, _, is_null, _, is_primary_key, _)) = tuple((
         maybe_take_whitespace,
         parse_sql_identifier,
         take_whitespace,
@@ -65,6 +65,8 @@ fn match_column_attribute<'a, E: ParseError<&'a str> + ContextError<&'a str>>(
         maybe_take_whitespace,
         is_null,
         maybe_take_whitespace,
+        is_primary_key,
+        maybe_take_whitespace,
     ))(input)?;
     Ok((
         input,
@@ -72,6 +74,7 @@ fn match_column_attribute<'a, E: ParseError<&'a str> + ContextError<&'a str>>(
             name: name.to_string(),
             sql_type: sql_type.to_string(),
             null: is_null,
+            primary_key: is_primary_key,
         },
     ))
 }
@@ -90,6 +93,16 @@ fn is_null<'a, E: ParseError<&'a str> + ContextError<&'a str>>(
     Ok((input, true))
 }
 
+fn is_primary_key<'a, E: ParseError<&'a str> + ContextError<&'a str>>(
+    input: &'a str,
+) -> IResult<&'a str, bool, E> {
+    let (input, primary_key) = opt(match_primary_key)(input)?;
+    match primary_key {
+        Some(()) => Ok((input, true)),
+        None => Ok((input, false)),
+    }
+}
+
 fn match_not_null<'a, E: ParseError<&'a str> + ContextError<&'a str>>(
     input: &'a str,
 ) -> IResult<&'a str, (), E> {
@@ -105,6 +118,14 @@ fn match_null<'a, E: ParseError<&'a str> + ContextError<&'a str>>(
     Ok((input, ()))
 }
 
+fn match_primary_key<'a, E: ParseError<&'a str> + ContextError<&'a str>>(
+    input: &'a str,
+) -> IResult<&'a str, (), E> {
+    let (input, (_, _, _)) =
+        tuple((tag_no_case("primary"), take_whitespace, tag_no_case("key")))(input)?;
+    Ok((input, ()))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -112,7 +133,7 @@ mod tests {
 
     #[test]
     fn test_simple_table() -> Result<(), Box<dyn std::error::Error>> {
-        let test_string = "create table foo (bar text, baz text not null)";
+        let test_string = "create table foo (bar text primary key, baz text not null)";
 
         let (_, result) = parse_create_table::<VerboseError<&str>>(test_string)?;
 
@@ -128,11 +149,13 @@ mod tests {
                 name: "bar".to_string(),
                 sql_type: "text".to_string(),
                 null: true,
+                primary_key: true,
             },
             RawColumn {
                 name: "baz".to_string(),
                 sql_type: "text".to_string(),
                 null: false,
+                primary_key: false,
             },
         ];
         assert_eq!(columns, result.provided_columns);
diff --git a/tests/sql_create_table.rs b/tests/sql_create_table.rs
index ee40222..84955ce 100644
--- a/tests/sql_create_table.rs
+++ b/tests/sql_create_table.rs
@@ -6,12 +6,14 @@ async fn create_table_with_nullable() -> Result<(), Box<dyn std::error::Error>>
     client
         .batch_execute(
             "create table foo (
-            bar text, 
+            bar text primary key, 
             baz text not null, 
             another text null
         )",
         )
         .await?;
 
-    common::_request_shutdown(request_shutdown).await
+    common::_request_shutdown(request_shutdown).await?;
+    
+    Ok(())
 }