Change ALTER TABLE constraints parsing

- merge PrimaryKey and UniqueKey variants
- support `CHECK` constraints, removing the separate `Key` struct
- make `CONSTRAINT constraint_name` optional
- remove `KEY` without qualifiers (wasn't parsed and there doesn't
  appear to be such a thing)
- change `UNIQUE KEY` -> `UNIQUE`
- change `REMOVE CONSTRAINT` -> `DROP CONSTRAINT` and note its parsing
  is not implemented

Spec:
- ANSI SQL: see <table constraint definition> in https://jakewheat.github.io/sql-overview/sql-2011-foundation-grammar.html#_11_6_table_constraint_definition
- Postgres: look for "and table_constraint is:" in https://www.postgresql.org/docs/11/sql-altertable.html

src/sqlast/ddl.rs

@@ -0,0 +1,86 @@
+//! AST types specific to CREATE/ALTER variants of `SQLStatement`
+//! (commonly referred to as Data Definition Language, or DDL)
+use super::{ASTNode, SQLIdent, SQLObjectName};
+
+/// An `ALTER TABLE` (`SQLStatement::SQLAlterTable`) operation
+#[derive(Debug, Clone, PartialEq)]
+pub enum AlterOperation {
+    /// `ADD <table_constraint>`
+    AddConstraint(TableConstraint),
+    /// TODO: implement `DROP CONSTRAINT name`
+    DropConstraint { name: SQLIdent },
+}
+
+impl ToString for AlterOperation {
+    fn to_string(&self) -> String {
+        match self {
+            AlterOperation::AddConstraint(constraint) => format!("ADD {}", constraint.to_string()),
+            AlterOperation::DropConstraint { name } => format!("DROP CONSTRAINT {}", name),
+        }
+    }
+}
+
+/// A table-level constraint, specified in a `CREATE TABLE` or an
+/// `ALTER TABLE ADD <constraint>` statement.
+#[derive(Debug, Clone, PartialEq)]
+pub enum TableConstraint {
+    /// `[ CONSTRAINT <name> ] { PRIMARY KEY | UNIQUE } (<columns>)`
+    Unique {
+        name: Option<SQLIdent>,
+        columns: Vec<SQLIdent>,
+        /// Whether this is a `PRIMARY KEY` or just a `UNIQUE` constraint
+        is_primary: bool,
+    },
+    /// A referential integrity constraint (`[ CONSTRAINT <name> ] FOREIGN KEY (<columns>)
+    /// REFERENCES <foreign_table> (<referred_columns>)`)
+    ForeignKey {
+        name: Option<SQLIdent>,
+        columns: Vec<SQLIdent>,
+        foreign_table: SQLObjectName,
+        referred_columns: Vec<SQLIdent>,
+    },
+    /// `[ CONSTRAINT <name> ] CHECK (<expr>)`
+    Check {
+        name: Option<SQLIdent>,
+        expr: Box<ASTNode>,
+    },
+}
+
+impl ToString for TableConstraint {
+    fn to_string(&self) -> String {
+        fn format_constraint_name(name: &Option<SQLIdent>) -> String {
+            name.as_ref()
+                .map(|name| format!("CONSTRAINT {} ", name))
+                .unwrap_or_default()
+        }
+        match self {
+            TableConstraint::Unique {
+                name,
+                columns,
+                is_primary,
+            } => format!(
+                "{}{} ({})",
+                format_constraint_name(name),
+                if *is_primary { "PRIMARY KEY" } else { "UNIQUE" },
+                columns.join(", ")
+            ),
+            TableConstraint::ForeignKey {
+                name,
+                columns,
+                foreign_table,
+                referred_columns,
+            } => format!(
+                "{}FOREIGN KEY ({}) REFERENCES {}({})",
+                format_constraint_name(name),
+                columns.join(", "),
+                foreign_table.to_string(),
+                referred_columns.join(", ")
+            ),
+            TableConstraint::Check { name, expr } => format!(
+                "{}CHECK ({})",
+                format_constraint_name(name),
+                expr.to_string()
+            ),
+        }
+    }
+}

src/sqlast/mod.rs

@@ -14,18 +14,18 @@
 
 //! SQL Abstract Syntax Tree (AST) types
 
+mod ddl;
 mod query;
 mod sql_operator;
 mod sqltype;
-mod table_key;
 mod value;
 
+pub use self::ddl::{AlterOperation, TableConstraint};
 pub use self::query::{
     Cte, Fetch, Join, JoinConstraint, JoinOperator, SQLOrderByExpr, SQLQuery, SQLSelect,
     SQLSelectItem, SQLSetExpr, SQLSetOperator, TableFactor,
 };
 pub use self::sqltype::SQLType;
-pub use self::table_key::{AlterOperation, Key, TableKey};
 pub use self::value::Value;
 
 pub use self::sql_operator::SQLOperator;

src/sqlast/table_key.rs

@@ -1,61 +0,0 @@
-use super::{SQLIdent, SQLObjectName};
-
-#[derive(Debug, Clone, PartialEq)]
-pub enum AlterOperation {
-    AddConstraint(TableKey),
-    RemoveConstraint { name: SQLIdent },
-}
-
-impl ToString for AlterOperation {
-    fn to_string(&self) -> String {
-        match self {
-            AlterOperation::AddConstraint(table_key) => {
-                format!("ADD CONSTRAINT {}", table_key.to_string())
-            }
-            AlterOperation::RemoveConstraint { name } => format!("REMOVE CONSTRAINT {}", name),
-        }
-    }
-}
-
-#[derive(Debug, Clone, PartialEq)]
-pub struct Key {
-    pub name: SQLIdent,
-    pub columns: Vec<SQLIdent>,
-}
-
-#[derive(Debug, Clone, PartialEq)]
-pub enum TableKey {
-    PrimaryKey(Key),
-    UniqueKey(Key),
-    Key(Key),
-    ForeignKey {
-        key: Key,
-        foreign_table: SQLObjectName,
-        referred_columns: Vec<SQLIdent>,
-    },
-}
-
-impl ToString for TableKey {
-    fn to_string(&self) -> String {
-        match self {
-            TableKey::PrimaryKey(ref key) => {
-                format!("{} PRIMARY KEY ({})", key.name, key.columns.join(", "))
-            }
-            TableKey::UniqueKey(ref key) => {
-                format!("{} UNIQUE KEY ({})", key.name, key.columns.join(", "))
-            }
-            TableKey::Key(ref key) => format!("{} KEY ({})", key.name, key.columns.join(", ")),
-            TableKey::ForeignKey {
-                key,
-                foreign_table,
-                referred_columns,
-            } => format!(
-                "{} FOREIGN KEY ({}) REFERENCES {}({})",
-                key.name,
-                key.columns.join(", "),
-                foreign_table.to_string(),
-                referred_columns.join(", ")
-            ),
-        }
-    }
-}

src/sqlparser.rs

@@ -912,33 +912,45 @@ impl Parser {
         Ok(columns)
     }
 
-    pub fn parse_table_key(&mut self, constraint_name: SQLIdent) -> Result<TableKey, ParserError> {
-        let is_primary_key = self.parse_keywords(vec!["PRIMARY", "KEY"]);
-        let is_unique_key = self.parse_keywords(vec!["UNIQUE", "KEY"]);
-        let is_foreign_key = self.parse_keywords(vec!["FOREIGN", "KEY"]);
-        let column_names = self.parse_parenthesized_column_list(Mandatory)?;
-        let key = Key {
-            name: constraint_name,
-            columns: column_names,
-        };
-        if is_primary_key {
-            Ok(TableKey::PrimaryKey(key))
-        } else if is_unique_key {
-            Ok(TableKey::UniqueKey(key))
-        } else if is_foreign_key {
-            self.expect_keyword("REFERENCES")?;
-            let foreign_table = self.parse_object_name()?;
-            let referred_columns = self.parse_parenthesized_column_list(Mandatory)?;
-            Ok(TableKey::ForeignKey {
-                key,
-                foreign_table,
-                referred_columns,
-            })
+    pub fn parse_table_constraint(&mut self) -> Result<TableConstraint, ParserError> {
+        let name = if self.parse_keyword("CONSTRAINT") {
+            Some(self.parse_identifier()?)
         } else {
-            parser_err!(format!(
-                "Expecting primary key, unique key, or foreign key, found: {:?}",
-                self.peek_token()
-            ))
+            None
+        };
+        match self.next_token() {
+            Some(Token::SQLWord(ref k)) if k.keyword == "PRIMARY" || k.keyword == "UNIQUE" => {
+                let is_primary = k.keyword == "PRIMARY";
+                if is_primary {
+                    self.expect_keyword("KEY")?;
+                }
+                let columns = self.parse_parenthesized_column_list(Mandatory)?;
+                Ok(TableConstraint::Unique {
+                    name,
+                    columns,
+                    is_primary,
+                })
+            }
+            Some(Token::SQLWord(ref k)) if k.keyword == "FOREIGN" => {
+                self.expect_keyword("KEY")?;
+                let columns = self.parse_parenthesized_column_list(Mandatory)?;
+                self.expect_keyword("REFERENCES")?;
+                let foreign_table = self.parse_object_name()?;
+                let referred_columns = self.parse_parenthesized_column_list(Mandatory)?;
+                Ok(TableConstraint::ForeignKey {
+                    name,
+                    columns,
+                    foreign_table,
+                    referred_columns,
+                })
+            }
+            Some(Token::SQLWord(ref k)) if k.keyword == "CHECK" => {
+                self.expect_token(&Token::LParen)?;
+                let expr = Box::new(self.parse_expr()?);
+                self.expect_token(&Token::RParen)?;
+                Ok(TableConstraint::Check { name, expr })
+            }
+            _ => self.expected("PRIMARY, UNIQUE, or FOREIGN", self.peek_token()),
         }
     }
 
@@ -947,13 +959,7 @@ impl Parser {
         let _ = self.parse_keyword("ONLY");
         let table_name = self.parse_object_name()?;
         let operation = if self.parse_keyword("ADD") {
-            if self.parse_keyword("CONSTRAINT") {
-                let constraint_name = self.parse_identifier()?;
-                let table_key = self.parse_table_key(constraint_name)?;
-                AlterOperation::AddConstraint(table_key)
-            } else {
-                return self.expected("CONSTRAINT after ADD", self.peek_token());
-            }
+            AlterOperation::AddConstraint(self.parse_table_constraint()?)
         } else {
             return self.expected("ADD after ALTER TABLE", self.peek_token());
         };

tests/sqlparser_common.rs

@@ -736,26 +736,31 @@ fn parse_create_external_table() {
 }
 
 #[test]
-fn parse_alter_table_constraint_primary_key() {
-    let sql = "ALTER TABLE bazaar.address \
-               ADD CONSTRAINT address_pkey PRIMARY KEY (address_id)";
-    match verified_stmt(sql) {
-        SQLStatement::SQLAlterTable { name, .. } => {
-            assert_eq!(name.to_string(), "bazaar.address");
-        }
-        _ => unreachable!(),
-    }
-}
+fn parse_alter_table_constraints() {
+    check_one("CONSTRAINT address_pkey PRIMARY KEY (address_id)");
+    check_one("CONSTRAINT uk_task UNIQUE (report_date, task_id)");
+    check_one(
+        "CONSTRAINT customer_address_id_fkey FOREIGN KEY (address_id) \
+         REFERENCES public.address(address_id)",
+    );
+    check_one("CONSTRAINT ck CHECK (rtrim(ltrim(REF_CODE)) <> '')");
 
-#[test]
-fn parse_alter_table_constraint_foreign_key() {
-    let sql = "ALTER TABLE public.customer \
-        ADD CONSTRAINT customer_address_id_fkey FOREIGN KEY (address_id) REFERENCES public.address(address_id)";
-    match verified_stmt(sql) {
-        SQLStatement::SQLAlterTable { name, .. } => {
-            assert_eq!(name.to_string(), "public.customer");
+    check_one("PRIMARY KEY (foo, bar)");
+    check_one("UNIQUE (id)");
+    check_one("FOREIGN KEY (foo, bar) REFERENCES AnotherTable(foo, bar)");
+    check_one("CHECK (end_date > start_date OR end_date IS NULL)");
+
+    fn check_one(constraint_text: &str) {
+        match verified_stmt(&format!("ALTER TABLE tab ADD {}", constraint_text)) {
+            SQLStatement::SQLAlterTable {
+                name,
+                operation: AlterOperation::AddConstraint(constraint),
+            } => {
+                assert_eq!("tab", name.to_string());
+                assert_eq!(constraint_text, constraint.to_string());
+            }
+            _ => unreachable!(),
         }
-        _ => unreachable!(),
     }
 }