Merge pull request #14 from wangandi/indexing

Support CREATE/DROP INDEX

Author: wangandi

src/ast/mod.rs

@@ -499,6 +499,15 @@ pub enum Statement {
         file_format: Option<FileFormat>,
         location: Option<String>,
     },
+    /// `CREATE INDEX`
+    CreateIndex {
+        /// Index name
+        name: Ident,
+        /// `ON` table or view name
+        on_name: ObjectName,
+        /// Expressions that form part of the index key
+        key_parts: Vec<Expr>,
+    },
     /// `ALTER TABLE`
     AlterTable {
         /// Table name
@@ -759,6 +768,20 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
+            Statement::CreateIndex {
+                name,
+                on_name,
+                key_parts,
+            } => {
+                write!(
+                    f,
+                    "CREATE INDEX {} ON {} ({})",
+                    name,
+                    on_name,
+                    display_comma_separated(key_parts),
+                )?;
+                Ok(())
+            }
             Statement::AlterTable { name, operation } => {
                 write!(f, "ALTER TABLE {} {}", name, operation)
             }
@@ -797,6 +820,7 @@ impl fmt::Display for Statement {
                         View => "VIEWS",
                         Source => "SOURCES",
                         Sink => "SINKS",
+                        Index => unreachable!(),
                     }
                 )
             }
@@ -963,6 +987,7 @@ pub enum ObjectType {
     View,
     Source,
     Sink,
+    Index,
 }
 
 impl fmt::Display for ObjectType {
@@ -972,6 +997,7 @@ impl fmt::Display for ObjectType {
             ObjectType::View => "VIEW",
             ObjectType::Source => "SOURCE",
             ObjectType::Sink => "SINK",
+            ObjectType::Index => "INDEX",
         })
     }
 }

src/ast/visit_macro.rs

@@ -382,6 +382,15 @@ macro_rules! make_visitor {
                 visit_create_view(self, name, columns, query, materialized, with_options)
             }
 
+            fn visit_create_index(
+                &mut self,
+                name: &'ast $($mut)* Ident,
+                on_name: &'ast $($mut)* ObjectName,
+                key_parts: &'ast $($mut)* Vec<Expr>
+            ){
+                visit_create_index(self, name, on_name, key_parts)
+            }
+
             fn visit_create_table(
                 &mut self,
                 name: &'ast $($mut)* ObjectName,
@@ -602,6 +611,11 @@ macro_rules! make_visitor {
                     materialized,
                     with_options,
                 } => visitor.visit_create_view(name, columns, query, *materialized, with_options),
+                Statement::CreateIndex {
+                    name,
+                    on_name,
+                    key_parts,
+                } => visitor.visit_create_index(name, on_name, key_parts),
                 Statement::Drop {
                     object_type,
                     if_exists,
@@ -1288,6 +1302,19 @@ macro_rules! make_visitor {
             visitor.visit_query(query);
         }
 
+        pub fn visit_create_index<'ast, V: $name<'ast> + ?Sized>(
+            visitor: &mut V,
+            name: &'ast $($mut)* Ident,
+            on_name: &'ast $($mut)* ObjectName,
+            key_parts: &'ast $($mut)* Vec<Expr>,
+        ) {
+            visitor.visit_ident(name);
+            visitor.visit_object_name(on_name);
+            for key_part in key_parts {
+                visitor.visit_expr(key_part);
+            }
+        }
+
         pub fn visit_create_table<'ast, V: $name<'ast> + ?Sized>(
             visitor: &mut V,
             name: &'ast $($mut)* ObjectName,

src/dialect/keywords.rs

@@ -203,6 +203,7 @@ define_keywords!(
     IF,
     IMMEDIATE,
     IN,
+    INDEX,
     INDICATOR,
     INNER,
     INOUT,

src/parser.rs

@@ -1055,9 +1055,11 @@ impl Parser {
             self.parse_create_sink()
         } else if self.parse_keyword("EXTERNAL") {
             self.parse_create_external_table()
+        } else if self.parse_keyword("INDEX") {
+            self.parse_create_index()
         } else {
             self.expected(
-                "TABLE, VIEW, SOURCE, or SINK after CREATE",
+                "TABLE, VIEW, SOURCE, SINK, or INDEX after CREATE",
                 self.peek_token(),
             )
         }
@@ -1163,6 +1165,20 @@ impl Parser {
         })
     }
 
+    pub fn parse_create_index(&mut self) -> Result<Statement, ParserError> {
+        let name = self.parse_identifier()?;
+        self.expect_keyword("ON")?;
+        let on_name = self.parse_object_name()?;
+        self.expect_token(&Token::LParen)?;
+        let key_parts = self.parse_comma_separated(Parser::parse_expr)?;
+        self.expect_token(&Token::RParen)?;
+        Ok(Statement::CreateIndex {
+            name,
+            on_name,
+            key_parts,
+        })
+    }
+
     pub fn parse_drop(&mut self) -> Result<Statement, ParserError> {
         let object_type = if self.parse_keyword("TABLE") {
             ObjectType::Table
@@ -1172,8 +1188,13 @@ impl Parser {
             ObjectType::Source
         } else if self.parse_keywords(vec!["SINK"]) {
             ObjectType::Sink
+        } else if self.parse_keyword("INDEX") {
+            ObjectType::Index
         } else {
-            return self.expected("TABLE or VIEW after DROP", self.peek_token());
+            return self.expected(
+                "TABLE, VIEW, SOURCE, SINK, or INDEX after DROP",
+                self.peek_token(),
+            );
         };
         // Many dialects support the non standard `IF EXISTS` clause and allow
         // specifying multiple objects to delete in a single statement

tests/sqlparser_common.rs

@@ -1872,7 +1872,9 @@ fn parse_show_objects() {
 fn parse_show_create_view() {
     assert_eq!(
         verified_stmt("SHOW CREATE VIEW foo"),
-        Statement::ShowCreateView { view_name: ObjectName(vec!["foo".into()]) }
+        Statement::ShowCreateView {
+            view_name: ObjectName(vec!["foo".into()])
+        }
     )
 }
 
@@ -2711,6 +2713,103 @@ fn parse_create_sink() {
     }
 }
 
+#[test]
+fn parse_create_index() {
+    let sql = "CREATE INDEX foo ON myschema.bar (a, b)";
+    match verified_stmt(sql) {
+        Statement::CreateIndex {
+            name,
+            on_name,
+            key_parts,
+        } => {
+            assert_eq!("foo", name.to_string());
+            assert_eq!("myschema.bar", on_name.to_string());
+            assert_eq!(
+                key_parts,
+                vec![
+                    Expr::Identifier("a".to_string()),
+                    Expr::Identifier("b".to_string())
+                ]
+            );
+        }
+        _ => assert!(false),
+    }
+
+    let sql = "CREATE INDEX fizz ON baz (ascii(x), a IS NOT NULL, (EXISTS (SELECT y FROM boop WHERE boop.z = z)), delta)";
+    match verified_stmt(sql) {
+        Statement::CreateIndex {
+            name,
+            on_name,
+            key_parts,
+        } => {
+            assert_eq!("fizz", name.to_string());
+            assert_eq!("baz", on_name.to_string());
+            assert_matches!(key_parts[0], Expr::Function(..));
+            assert_eq!(
+                key_parts[1],
+                Expr::IsNotNull(Box::new(Expr::Identifier("a".to_string())))
+            );
+            if let Expr::Nested(expr) = &key_parts[2] {
+                assert_matches!(**expr, Expr::Exists(..));
+            } else {
+                assert!(false);
+            }
+            assert_eq!(key_parts[3], Expr::Identifier("delta".to_string()));
+        }
+        _ => assert!(false),
+    }
+
+    let sql = "CREATE INDEX ind ON tab ((col + 1))";
+    match verified_stmt(sql) {
+        Statement::CreateIndex {
+            name,
+            on_name,
+            key_parts,
+        } => {
+            assert_eq!("ind", name.to_string());
+            assert_eq!("tab", on_name.to_string());
+            assert_eq!(
+                key_parts,
+                vec![Expr::Nested(Box::new(Expr::BinaryOp {
+                    left: Box::new(Expr::Identifier("col".to_string())),
+                    op: BinaryOperator::Plus,
+                    right: Box::new(Expr::Value(Value::Number("1".to_string())))
+                }))],
+            );
+        }
+        _ => assert!(false),
+    }
+    let sql = "CREATE INDEX qualifiers ON no_parentheses (alpha.omega)";
+    match verified_stmt(sql) {
+        Statement::CreateIndex {
+            name,
+            on_name,
+            key_parts,
+        } => {
+            assert_eq!("qualifiers", name.to_string());
+            assert_eq!("no_parentheses", on_name.to_string());
+            assert_eq!(
+                key_parts,
+                vec![Expr::CompoundIdentifier(vec![
+                    "alpha".to_string(),
+                    "omega".to_string()
+                ])],
+            );
+        }
+        _ => assert!(false),
+    }
+}
+
+#[test]
+fn parse_invalid_create_index() {
+    // Index names should not have a schema in front of it
+    let res = parse_sql_statements("CREATE INDEX myschema.ind ON foo(b)");
+    assert_eq!(
+        ParserError::ParserError("Expected ON, found: .".to_string()),
+        res.unwrap_err(),
+    );
+}
+
 #[test]
 fn parse_drop_table() {
     let sql = "DROP TABLE foo";
@@ -2803,6 +2902,28 @@ fn parse_drop_source() {
     }
 }
 
+#[test]
+fn parse_drop_index() {
+    let sql = "DROP INDEX IF EXISTS myschema.myindex";
+    match verified_stmt(sql) {
+        Statement::Drop {
+            object_type,
+            if_exists,
+            names,
+            cascade,
+        } => {
+            assert_eq!(true, if_exists);
+            assert_eq!(
+                vec!["myschema.myindex"],
+                names.iter().map(|n| n.to_string()).collect::<Vec<_>>()
+            );
+            assert_eq!(false, cascade);
+            assert_eq!(ObjectType::Index, object_type);
+        }
+        _ => unreachable!(),
+    }
+}
+
 #[test]
 fn parse_peek() {
     let sql = "PEEK foo.bar";