add SELECT INTO statement variation (apache#447)

mobuchowski · web-flow · commit d38c30e9ffd8 · 2022-03-30T16:01:55.000-04:00
Signed-off-by: Maciej Obuchowski &lt;obuchowski.maciej@gmail.com&gt;
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -36,8 +36,8 @@ pub use self::ddl::{
 pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{
     Cte, Fetch, Join, JoinConstraint, JoinOperator, LateralView, LockType, Offset, OffsetRows,
-    OrderByExpr, Query, Select, SelectItem, SetExpr, SetOperator, TableAlias, TableFactor,
-    TableWithJoins, Top, Values, With,
+    OrderByExpr, Query, Select, SelectInto, SelectItem, SetExpr, SetOperator, TableAlias,
+    TableFactor, TableWithJoins, Top, Values, With,
 };
 pub use self::value::{DateTimeField, TrimWhereField, Value};
 
diff --git a/src/ast/query.rs b/src/ast/query.rs
@@ -136,6 +136,8 @@ pub struct Select {
     pub top: Option<Top>,
     /// projection expressions
     pub projection: Vec<SelectItem>,
+    /// INTO
+    pub into: Option<SelectInto>,
     /// FROM
     pub from: Vec<TableWithJoins>,
     /// LATERAL VIEWs
@@ -161,6 +163,11 @@ impl fmt::Display for Select {
             write!(f, " {}", top)?;
         }
         write!(f, " {}", display_comma_separated(&self.projection))?;
+
+        if let Some(ref into) = self.into {
+            write!(f, " {}", into)?;
+        }
+
         if !self.from.is_empty() {
             write!(f, " FROM {}", display_comma_separated(&self.from))?;
         }
@@ -636,3 +643,20 @@ impl fmt::Display for Values {
         Ok(())
     }
 }
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct SelectInto {
+    pub temporary: bool,
+    pub unlogged: bool,
+    pub name: ObjectName,
+}
+
+impl fmt::Display for SelectInto {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let temporary = if self.temporary { " TEMPORARY" } else { "" };
+        let unlogged = if self.unlogged { " UNLOGGED" } else { "" };
+
+        write!(f, "INTO{}{} {}", temporary, unlogged, self.name)
+    }
+}
diff --git a/src/keywords.rs b/src/keywords.rs
@@ -499,6 +499,7 @@ define_keywords!(
     UNION,
     UNIQUE,
     UNKNOWN,
+    UNLOGGED,
     UNNEST,
     UNSIGNED,
     UPDATE,
@@ -600,4 +601,5 @@ pub const RESERVED_FOR_COLUMN_ALIAS: &[Keyword] = &[
     Keyword::DISTRIBUTE,
     // Reserved only as a column alias in the `SELECT` clause
     Keyword::FROM,
+    Keyword::INTO,
 ];
diff --git a/src/parser.rs b/src/parser.rs
@@ -2892,6 +2892,21 @@ impl<'a> Parser<'a> {
 
         let projection = self.parse_comma_separated(Parser::parse_select_item)?;
 
+        let into = if self.parse_keyword(Keyword::INTO) {
+            let temporary = self
+                .parse_one_of_keywords(&[Keyword::TEMP, Keyword::TEMPORARY])
+                .is_some();
+            let unlogged = self.parse_keyword(Keyword::UNLOGGED);
+            let name = self.parse_object_name()?;
+            Some(SelectInto {
+                temporary,
+                unlogged,
+                name,
+            })
+        } else {
+            None
+        };
+
         // Note that for keywords to be properly handled here, they need to be
         // added to `RESERVED_FOR_COLUMN_ALIAS` / `RESERVED_FOR_TABLE_ALIAS`,
         // otherwise they may be parsed as an alias as part of the `projection`
@@ -2973,6 +2988,7 @@ impl<'a> Parser<'a> {
             distinct,
             top,
             projection,
+            into,
             from,
             lateral_views,
             selection,
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
@@ -380,6 +380,32 @@ fn parse_select_all_distinct() {
     );
 }
 
+#[test]
+fn parse_select_into() {
+    let sql = "SELECT * INTO table0 FROM table1";
+    one_statement_parses_to(sql, "SELECT * INTO table0 FROM table1");
+    let select = verified_only_select(sql);
+    assert_eq!(
+        &SelectInto {
+            temporary: false,
+            unlogged: false,
+            name: ObjectName(vec![Ident::new("table0")])
+        },
+        only(&select.into)
+    );
+
+    let sql = "SELECT * INTO TEMPORARY UNLOGGED table0 FROM table1";
+    one_statement_parses_to(sql, "SELECT * INTO TEMPORARY UNLOGGED table0 FROM table1");
+
+    // Do not allow aliases here
+    let sql = "SELECT * INTO table0 asdf FROM table1";
+    let result = parse_sql_statements(sql);
+    assert_eq!(
+        ParserError::ParserError("Expected end of statement, found: asdf".to_string()),
+        result.unwrap_err()
+    )
+}
+
 #[test]
 fn parse_select_wildcard() {
     let sql = "SELECT * FROM foo";
@@ -4235,6 +4261,7 @@ fn parse_merge() {
                         distinct: false,
                         top: None,
                         projection: vec![SelectItem::Wildcard],
+                        into: None,
                         from: vec![TableWithJoins {
                             relation: TableFactor::Table {
                                 name: ObjectName(vec![Ident::new("s"), Ident::new("foo")]),
diff --git a/tests/sqlparser_mysql.rs b/tests/sqlparser_mysql.rs
@@ -305,6 +305,7 @@ fn parse_quote_identifiers_2() {
                     value: "quoted ` identifier".into(),
                     quote_style: Some('`'),
                 }))],
+                into: None,
                 from: vec![],
                 lateral_views: vec![],
                 selection: None,
@@ -732,6 +733,7 @@ fn parse_substring_in_select() {
                                 false
                             ))))
                         })],
+                        into: None,
                         from: vec![TableWithJoins {
                             relation: TableFactor::Table {
                                 name: ObjectName(vec![Ident {
diff --git a/tests/sqpparser_clickhouse.rs b/tests/sqpparser_clickhouse.rs
@@ -53,6 +53,7 @@ fn parse_map_access_expr() {
                     distinct: false,
                 })],
             })],
+            into: None,
             from: vec![TableWithJoins {
                 relation: Table {
                     name: ObjectName(vec![Ident::new("foos")]),
