Snowflake: support trailing options in CREATE TABLE (apache#1931)

Author: yoavcloud

src/ast/helpers/stmt_create_table.rs

@@ -383,6 +383,26 @@ impl CreateTableBuilder {
         self
     }
 
+    /// Returns true if the statement has exactly one source of info on the schema of the new table.
+    /// This is Snowflake-specific, some dialects allow more than one source.
+    pub(crate) fn validate_schema_info(&self) -> bool {
+        let mut sources = 0;
+        if !self.columns.is_empty() {
+            sources += 1;
+        }
+        if self.query.is_some() {
+            sources += 1;
+        }
+        if self.like.is_some() {
+            sources += 1;
+        }
+        if self.clone.is_some() {
+            sources += 1;
+        }
+
+        sources == 1
+    }
+
     pub fn build(self) -> Statement {
         Statement::CreateTable(CreateTable {
             or_replace: self.or_replace,

src/dialect/bigquery.rs

@@ -144,4 +144,8 @@ impl Dialect for BigQueryDialect {
     fn supports_pipe_operator(&self) -> bool {
         true
     }
+
+    fn supports_create_table_multi_schema_info_sources(&self) -> bool {
+        true
+    }
 }

src/dialect/mod.rs

@@ -590,6 +590,13 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Returne true if the dialect supports specifying multiple options
+    /// in a `CREATE TABLE` statement for the structure of the new table. For example:
+    /// `CREATE TABLE t (a INT, b INT) AS SELECT 1 AS b, 2 AS a`
+    fn supports_create_table_multi_schema_info_sources(&self) -> bool {
+        false
+    }
+
     /// Dialect-specific infix parser override
     ///
     /// This method is called to parse the next infix expression.

src/dialect/snowflake.rs

@@ -555,17 +555,14 @@ pub fn parse_create_table(
                 Keyword::AS => {
                     let query = parser.parse_query()?;
                     builder = builder.query(Some(query));
-                    break;
                 }
                 Keyword::CLONE => {
                     let clone = parser.parse_object_name(false).ok();
                     builder = builder.clone_clause(clone);
-                    break;
                 }
                 Keyword::LIKE => {
                     let like = parser.parse_object_name(false).ok();
                     builder = builder.like(like);
-                    break;
                 }
                 Keyword::CLUSTER => {
                     parser.expect_keyword_is(Keyword::BY)?;
@@ -691,7 +688,7 @@ pub fn parse_create_table(
                 builder = builder.columns(columns).constraints(constraints);
             }
             Token::EOF => {
-                if builder.columns.is_empty() {
+                if !builder.validate_schema_info() {
                     return Err(ParserError::ParserError(
                         "unexpected end of input".to_string(),
                     ));
@@ -700,7 +697,7 @@ pub fn parse_create_table(
                 break;
             }
             Token::SemiColon => {
-                if builder.columns.is_empty() {
+                if !builder.validate_schema_info() {
                     return Err(ParserError::ParserError(
                         "unexpected end of input".to_string(),
                     ));

tests/sqlparser_common.rs

@@ -4347,8 +4347,9 @@ fn parse_create_table_as() {
     // BigQuery allows specifying table schema in CTAS
     // ANSI SQL and PostgreSQL let you only specify the list of columns
     // (without data types) in a CTAS, but we have yet to support that.
+    let dialects = all_dialects_where(|d| d.supports_create_table_multi_schema_info_sources());
     let sql = "CREATE TABLE t (a INT, b INT) AS SELECT 1 AS b, 2 AS a";
-    match verified_stmt(sql) {
+    match dialects.verified_stmt(sql) {
         Statement::CreateTable(CreateTable { columns, query, .. }) => {
             assert_eq!(columns.len(), 2);
             assert_eq!(columns[0].to_string(), "a INT".to_string());
@@ -4453,20 +4454,6 @@ fn parse_create_or_replace_table() {
         }
         _ => unreachable!(),
     }
-
-    let sql = "CREATE TABLE t (a INT, b INT) AS SELECT 1 AS b, 2 AS a";
-    match verified_stmt(sql) {
-        Statement::CreateTable(CreateTable { columns, query, .. }) => {
-            assert_eq!(columns.len(), 2);
-            assert_eq!(columns[0].to_string(), "a INT".to_string());
-            assert_eq!(columns[1].to_string(), "b INT".to_string());
-            assert_eq!(
-                query,
-                Some(Box::new(verified_query("SELECT 1 AS b, 2 AS a")))
-            );
-        }
-        _ => unreachable!(),
-    }
 }
 
 #[test]

tests/sqlparser_snowflake.rs

@@ -995,6 +995,51 @@ fn test_snowflake_create_iceberg_table_without_location() {
     );
 }
 
+#[test]
+fn test_snowflake_create_table_trailing_options() {
+    // Serialization to SQL assume that in `CREATE TABLE AS` the options come before the `AS (<query>)`
+    // but Snowflake supports also the other way around
+    snowflake()
+        .verified_stmt("CREATE TEMPORARY TABLE dst ON COMMIT PRESERVE ROWS AS (SELECT * FROM src)");
+    snowflake()
+        .parse_sql_statements(
+            "CREATE TEMPORARY TABLE dst AS (SELECT * FROM src) ON COMMIT PRESERVE ROWS",
+        )
+        .unwrap();
+
+    // Same for `CREATE TABLE LIKE|CLONE`:
+    snowflake().verified_stmt("CREATE TEMPORARY TABLE dst LIKE src ON COMMIT PRESERVE ROWS");
+    snowflake()
+        .parse_sql_statements("CREATE TEMPORARY TABLE dst ON COMMIT PRESERVE ROWS LIKE src")
+        .unwrap();
+
+    snowflake().verified_stmt("CREATE TEMPORARY TABLE dst CLONE src ON COMMIT PRESERVE ROWS");
+    snowflake()
+        .parse_sql_statements("CREATE TEMPORARY TABLE dst ON COMMIT PRESERVE ROWS CLONE src")
+        .unwrap();
+}
+
+#[test]
+fn test_snowflake_create_table_valid_schema_info() {
+    // Validate there's exactly one source of information on the schema of the new table
+    assert_eq!(
+        snowflake()
+            .parse_sql_statements("CREATE TABLE dst")
+            .is_err(),
+        true
+    );
+    assert_eq!(
+        snowflake().parse_sql_statements("CREATE OR REPLACE TEMP TABLE dst LIKE src AS (SELECT * FROM CUSTOMERS) ON COMMIT PRESERVE ROWS").is_err(),
+        true
+    );
+    assert_eq!(
+        snowflake()
+            .parse_sql_statements("CREATE OR REPLACE TEMP TABLE dst CLONE customers LIKE customer2")
+            .is_err(),
+        true
+    );
+}
+
 #[test]
 fn parse_sf_create_or_replace_view_with_comment_missing_equal() {
     assert!(snowflake_and_generic()