Support row_alias and col_aliases in INSERT statement for mysql…

… and generic dialects (#1136)
apache · Mar 8, 2024 · 11899fd · 11899fd
1 parent 5da66ad
commit 11899fd
Show file tree

Hide file tree

Showing 4 changed files with 148 additions and 4 deletions.
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -1759,6 +1759,8 @@ pub enum Statement {
         replace_into: bool,
         /// Only for mysql
         priority: Option<MysqlInsertPriority>,
+        /// Only for mysql
+        insert_alias: Option<InsertAliases>,
     },
     /// ```sql
     /// INSTALL
@@ -2773,6 +2775,7 @@ impl fmt::Display for Statement {
                 returning,
                 replace_into,
                 priority,
+                insert_alias,
             } => {
                 let table_name = if let Some(alias) = table_alias {
                     format!("{table_name} AS {alias}")
@@ -2822,6 +2825,16 @@ impl fmt::Display for Statement {
                     write!(f, "DEFAULT VALUES")?;
                 }
 
+                if let Some(insert_alias) = insert_alias {
+                    write!(f, " AS {0}", insert_alias.row_alias)?;
+
+                    if let Some(col_aliases) = &insert_alias.col_aliases {
+                        if !col_aliases.is_empty() {
+                            write!(f, " ({})", display_comma_separated(col_aliases))?;
+                        }
+                    }
+                }
+
                 if let Some(on) = on {
                     write!(f, "{on}")?;
                 }
@@ -4194,6 +4207,14 @@ pub enum OnInsert {
     OnConflict(OnConflict),
 }
 
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct InsertAliases {
+    pub row_alias: ObjectName,
+    pub col_aliases: Option<Vec<Ident>>,
+}
+
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -8419,6 +8419,19 @@ impl<'a> Parser<'a> {
                     (columns, partitioned, after_columns, source)
                 };
 
+            let insert_alias = if dialect_of!(self is MySqlDialect | GenericDialect)
+                && self.parse_keyword(Keyword::AS)
+            {
+                let row_alias = self.parse_object_name(false)?;
+                let col_aliases = Some(self.parse_parenthesized_column_list(Optional, false)?);
+                Some(InsertAliases {
+                    row_alias,
+                    col_aliases,
+                })
+            } else {
+                None
+            };
+
             let on = if self.parse_keyword(Keyword::ON) {
                 if self.parse_keyword(Keyword::CONFLICT) {
                     let conflict_target =
@@ -8488,6 +8501,7 @@ impl<'a> Parser<'a> {
                 returning,
                 replace_into,
                 priority,
+                insert_alias,
             })
         }
     }

diff --git a/tests/sqlparser_mysql.rs b/tests/sqlparser_mysql.rs
@@ -19,7 +19,7 @@ use matches::assert_matches;
 use sqlparser::ast::MysqlInsertPriority::{Delayed, HighPriority, LowPriority};
 use sqlparser::ast::*;
 use sqlparser::dialect::{GenericDialect, MySqlDialect};
-use sqlparser::parser::ParserOptions;
+use sqlparser::parser::{ParserError, ParserOptions};
 use sqlparser::tokenizer::Token;
 use test_utils::*;
 
@@ -1330,6 +1330,112 @@ fn parse_priority_insert() {
     }
 }
 
+#[test]
+fn parse_insert_as() {
+    let sql = r"INSERT INTO `table` (`date`) VALUES ('2024-01-01') AS `alias`";
+    match mysql_and_generic().verified_stmt(sql) {
+        Statement::Insert {
+            table_name,
+            columns,
+            source,
+            insert_alias,
+            ..
+        } => {
+            assert_eq!(
+                ObjectName(vec![Ident::with_quote('`', "table")]),
+                table_name
+            );
+            assert_eq!(vec![Ident::with_quote('`', "date")], columns);
+            let insert_alias = insert_alias.unwrap();
+
+            assert_eq!(
+                ObjectName(vec![Ident::with_quote('`', "alias")]),
+                insert_alias.row_alias
+            );
+            assert_eq!(Some(vec![]), insert_alias.col_aliases);
+            assert_eq!(
+                Some(Box::new(Query {
+                    with: None,
+                    body: Box::new(SetExpr::Values(Values {
+                        explicit_row: false,
+                        rows: vec![vec![Expr::Value(Value::SingleQuotedString(
+                            "2024-01-01".to_string()
+                        ))]]
+                    })),
+                    order_by: vec![],
+                    limit: None,
+                    limit_by: vec![],
+                    offset: None,
+                    fetch: None,
+                    locks: vec![],
+                    for_clause: None,
+                })),
+                source
+            );
+        }
+        _ => unreachable!(),
+    }
+
+    let sql = r"INSERT INTO `table` (`date`) VALUES ('2024-01-01') AS `alias` ()";
+    assert!(matches!(
+        mysql_and_generic().parse_sql_statements(sql),
+        Err(ParserError::ParserError(_))
+    ));
+
+    let sql = r"INSERT INTO `table` (`id`, `date`) VALUES (1, '2024-01-01') AS `alias` (`mek_id`, `mek_date`)";
+    match mysql_and_generic().verified_stmt(sql) {
+        Statement::Insert {
+            table_name,
+            columns,
+            source,
+            insert_alias,
+            ..
+        } => {
+            assert_eq!(
+                ObjectName(vec![Ident::with_quote('`', "table")]),
+                table_name
+            );
+            assert_eq!(
+                vec![Ident::with_quote('`', "id"), Ident::with_quote('`', "date")],
+                columns
+            );
+            let insert_alias = insert_alias.unwrap();
+            assert_eq!(
+                ObjectName(vec![Ident::with_quote('`', "alias")]),
+                insert_alias.row_alias
+            );
+            assert_eq!(
+                Some(vec![
+                    Ident::with_quote('`', "mek_id"),
+                    Ident::with_quote('`', "mek_date")
+                ]),
+                insert_alias.col_aliases
+            );
+            assert_eq!(
+                Some(Box::new(Query {
+                    with: None,
+                    body: Box::new(SetExpr::Values(Values {
+                        explicit_row: false,
+                        rows: vec![vec![
+                            Expr::Value(number("1")),
+                            Expr::Value(Value::SingleQuotedString("2024-01-01".to_string()))
+                        ]]
+                    })),
+                    order_by: vec![],
+                    limit: None,
+                    limit_by: vec![],
+                    offset: None,
+                    fetch: None,
+                    locks: vec![],
+                    for_clause: None,
+                })),
+                source
+            );
+        }
+        _ => unreachable!(),
+    }
+}
+
 #[test]
 fn parse_replace_insert() {
     let sql = r"REPLACE DELAYED INTO tasks (title, priority) VALUES ('Test Some Inserts', 1)";

diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
@@ -3764,7 +3764,8 @@ fn test_simple_postgres_insert_with_alias() {
             on: None,
             returning: None,
             replace_into: false,
-            priority: None
+            priority: None,
+            insert_alias: None
         }
     )
 }
@@ -3830,7 +3831,8 @@ fn test_simple_postgres_insert_with_alias() {
             on: None,
             returning: None,
             replace_into: false,
-            priority: None
+            priority: None,
+            insert_alias: None
         }
     )
 }
@@ -3892,7 +3894,8 @@ fn test_simple_insert_with_quoted_alias() {
             on: None,
             returning: None,
             replace_into: false,
-            priority: None
+            priority: None,
+            insert_alias: None,
         }
     )
 }