add create index and drop index support

Author: mashuai

src/ast/mod.rs

@@ -476,6 +476,15 @@ pub enum Statement {
         file_format: Option<FileFormat>,
         location: Option<String>,
     },
+    /// CREATE INDEX
+    CreateIndex {
+        /// index name
+        name: ObjectName,
+        table_name: ObjectName,
+        columns: Vec<Ident>,
+        unique: bool,
+        if_not_exists: bool,
+    },
     /// ALTER TABLE
     AlterTable {
         /// Table name
@@ -655,6 +664,28 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
+            Statement::CreateIndex {
+                name,
+                table_name,
+                columns,
+                unique,
+                if_not_exists,
+            } => {
+                write!(
+                    f,
+                    "CREATE{}INDEX{}{} ON {}({}",
+                    if *unique { " UNIQUE " } else { " " },
+                    if *if_not_exists {
+                        " IF NOT EXISTS "
+                    } else {
+                        " "
+                    },
+                    name,
+                    table_name,
+                    display_separated(columns, ",")
+                )?;
+                write!(f, ");")
+            }
             Statement::AlterTable { name, operation } => {
                 write!(f, "ALTER TABLE {} {}", name, operation)
             }
@@ -819,13 +850,15 @@ impl FromStr for FileFormat {
 pub enum ObjectType {
     Table,
     View,
+    Index,
 }
 
 impl fmt::Display for ObjectType {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         f.write_str(match self {
             ObjectType::Table => "TABLE",
             ObjectType::View => "VIEW",
+            ObjectType::Index => "INDEX",
         })
     }
 }

src/dialect/keywords.rs

@@ -420,7 +420,8 @@ define_keywords!(
     WORK,
     YEAR,
     ZONE,
-    END_EXEC = "END-EXEC"
+    END_EXEC = "END-EXEC",
+    INDEX
 );
 
 /// These keywords can't be used as a table alias, so that `FROM table_name alias`

src/parser.rs

@@ -855,13 +855,17 @@ impl Parser {
     pub fn parse_create(&mut self) -> Result<Statement, ParserError> {
         if self.parse_keyword("TABLE") {
             self.parse_create_table()
+        } else if self.parse_keyword("INDEX") {
+            self.parse_create_index(false)
+        } else if self.parse_keywords(vec!["UNIQUE", "INDEX"]) {
+            self.parse_create_index(true)
         } else if self.parse_keyword("MATERIALIZED") || self.parse_keyword("VIEW") {
             self.prev_token();
             self.parse_create_view()
         } else if self.parse_keyword("EXTERNAL") {
             self.parse_create_external_table()
         } else {
-            self.expected("TABLE or VIEW after CREATE", self.peek_token())
+            self.expected("TABLE, VIEW or INDEX after CREATE", self.peek_token())
         }
     }
 
@@ -912,8 +916,10 @@ impl Parser {
             ObjectType::Table
         } else if self.parse_keyword("VIEW") {
             ObjectType::View
+        } 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 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
@@ -932,6 +938,23 @@ impl Parser {
         })
     }
 
+    pub fn parse_create_index(&mut self, unique: bool) -> Result<Statement, ParserError> {
+        let if_not_exists = self.parse_keywords(vec!["IF", "NOT", "EXISTS"]);
+        let index_name = self.parse_object_name()?;
+        self.expect_keyword("ON")?;
+        let table_name = self.parse_object_name()?;
+        self.expect_token(&Token::LParen)?;
+        let columns = self.parse_comma_separated(Parser::parse_identifier)?;
+        self.expect_token(&Token::RParen)?;
+        Ok(Statement::CreateIndex {
+            name: index_name,
+            table_name,
+            columns,
+            unique,
+            if_not_exists,
+        })
+    }
+
     pub fn parse_create_table(&mut self) -> Result<Statement, ParserError> {
         let if_not_exists = self.parse_keywords(vec!["IF", "NOT", "EXISTS"]);
         let table_name = self.parse_object_name()?;

tests/sqlparser_common.rs

@@ -2631,6 +2631,44 @@ fn ensure_multiple_dialects_are_tested() {
     let _ = parse_sql_statements("SELECT @foo");
 }
 
+#[test]
+fn parse_create_index() {
+    let sql = "CREATE UNIQUE INDEX IF NOT EXISTS idx_name ON test(name,age);";
+    let ident_vec = vec![Ident::new("name"), Ident::new("age")];
+    match verified_stmt(sql) {
+        Statement::CreateIndex {
+            name,
+            table_name,
+            columns,
+            unique,
+            if_not_exists,
+        } => {
+            assert_eq!("idx_name", name.to_string());
+            assert_eq!("test", table_name.to_string());
+            assert_eq!(ident_vec, columns);
+            assert_eq!(true, unique);
+            assert_eq!(true, if_not_exists)
+        }
+        _ => unreachable!(),
+    }
+}
+#[test]
+fn parse_drop_index() {
+    let sql = "DROP INDEX idx_a";
+    match verified_stmt(sql) {
+        Statement::Drop {
+            names, object_type, ..
+        } => {
+            assert_eq!(
+                vec!["idx_a"],
+                names.iter().map(ToString::to_string).collect::<Vec<_>>()
+            );
+            assert_eq!(ObjectType::Index, object_type);
+        }
+        _ => unreachable!(),
+    }
+}
+
 fn parse_sql_statements(sql: &str) -> Result<Vec<Statement>, ParserError> {
     all_dialects().parse_sql_statements(sql)
 }