Support multiple PARTITION statements in ALTER TABLE ADD statement (apache#1011)

Co-authored-by: Chris A <[email protected]>
Co-authored-by: Andrew Lamb <[email protected]>

Author: 3 people

src/ast/ddl.rs

@@ -69,7 +69,7 @@ pub enum AlterTableOperation {
     /// Add Partitions
     AddPartitions {
         if_not_exists: bool,
-        new_partitions: Vec<Expr>,
+        new_partitions: Vec<Partition>,
     },
     DropPartitions {
         partitions: Vec<Expr>,
@@ -119,8 +119,8 @@ impl fmt::Display for AlterTableOperation {
                 new_partitions,
             } => write!(
                 f,
-                "ADD{ine} PARTITION ({})",
-                display_comma_separated(new_partitions),
+                "ADD{ine} {}",
+                display_separated(new_partitions, " "),
                 ine = if *if_not_exists { " IF NOT EXISTS" } else { "" }
             ),
             AlterTableOperation::AddConstraint(c) => write!(f, "ADD {c}"),
@@ -771,3 +771,21 @@ impl fmt::Display for UserDefinedTypeCompositeAttributeDef {
         Ok(())
     }
 }
+
+/// PARTITION statement used in ALTER TABLE et al. such as in Hive SQL
+#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct Partition {
+    pub partitions: Vec<Expr>,
+}
+
+impl fmt::Display for Partition {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(
+            f,
+            "PARTITION ({})",
+            display_comma_separated(&self.partitions)
+        )
+    }
+}

src/ast/mod.rs

@@ -31,8 +31,9 @@ pub use self::data_type::{
 pub use self::dcl::{AlterRoleOperation, ResetConfig, RoleOption, SetConfigValue};
 pub use self::ddl::{
     AlterColumnOperation, AlterIndexOperation, AlterTableOperation, ColumnDef, ColumnOption,
-    ColumnOptionDef, GeneratedAs, IndexType, KeyOrIndexDisplay, ProcedureParam, ReferentialAction,
-    TableConstraint, UserDefinedTypeCompositeAttributeDef, UserDefinedTypeRepresentation,
+    ColumnOptionDef, GeneratedAs, IndexType, KeyOrIndexDisplay, Partition, ProcedureParam,
+    ReferentialAction, TableConstraint, UserDefinedTypeCompositeAttributeDef,
+    UserDefinedTypeRepresentation,
 };
 pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{

src/dialect/generic.rs

@@ -18,6 +18,10 @@ use crate::dialect::Dialect;
 pub struct GenericDialect;
 
 impl Dialect for GenericDialect {
+    fn is_delimited_identifier_start(&self, ch: char) -> bool {
+        ch == '"' || ch == '`'
+    }
+
     fn is_identifier_start(&self, ch: char) -> bool {
         ch.is_alphabetic() || ch == '_' || ch == '#' || ch == '@'
     }

src/parser/mod.rs

@@ -4195,20 +4195,32 @@ impl<'a> Parser<'a> {
         Ok(SqlOption { name, value })
     }
 
+    pub fn parse_partition(&mut self) -> Result<Partition, ParserError> {
+        self.expect_token(&Token::LParen)?;
+        let partitions = self.parse_comma_separated(Parser::parse_expr)?;
+        self.expect_token(&Token::RParen)?;
+        Ok(Partition { partitions })
+    }
+
     pub fn parse_alter_table_operation(&mut self) -> Result<AlterTableOperation, ParserError> {
         let operation = if self.parse_keyword(Keyword::ADD) {
             if let Some(constraint) = self.parse_optional_table_constraint()? {
                 AlterTableOperation::AddConstraint(constraint)
             } else {
                 let if_not_exists =
                     self.parse_keywords(&[Keyword::IF, Keyword::NOT, Keyword::EXISTS]);
-                if self.parse_keyword(Keyword::PARTITION) {
-                    self.expect_token(&Token::LParen)?;
-                    let partitions = self.parse_comma_separated(Parser::parse_expr)?;
-                    self.expect_token(&Token::RParen)?;
+                let mut new_partitions = vec![];
+                loop {
+                    if self.parse_keyword(Keyword::PARTITION) {
+                        new_partitions.push(self.parse_partition()?);
+                    } else {
+                        break;
+                    }
+                }
+                if !new_partitions.is_empty() {
                     AlterTableOperation::AddPartitions {
                         if_not_exists,
-                        new_partitions: partitions,
+                        new_partitions,
                     }
                 } else {
                     let column_keyword = self.parse_keyword(Keyword::COLUMN);

tests/sqlparser_hive.rs

@@ -128,6 +128,12 @@ fn test_add_partition() {
     hive().verified_stmt(add);
 }
 
+#[test]
+fn test_add_multiple_partitions() {
+    let add = "ALTER TABLE db.table ADD IF NOT EXISTS PARTITION (`a` = 'asdf', `b` = 2) PARTITION (`a` = 'asdh', `b` = 3)";
+    hive().verified_stmt(add);
+}
+
 #[test]
 fn test_drop_partition() {
     let drop = "ALTER TABLE db.table DROP PARTITION (a = 1)";