Merge branch 'main' into feat/export_data_for_bigquery

Author: chenkovsky

src/ast/mod.rs

@@ -6780,6 +6780,7 @@ pub enum ActionCreateObjectType {
     OrganiationListing,
     ReplicationGroup,
     Role,
+    Schema,
     Share,
     User,
     Warehouse,
@@ -6801,6 +6802,7 @@ impl fmt::Display for ActionCreateObjectType {
             ActionCreateObjectType::OrganiationListing => write!(f, "ORGANIZATION LISTING"),
             ActionCreateObjectType::ReplicationGroup => write!(f, "REPLICATION GROUP"),
             ActionCreateObjectType::Role => write!(f, "ROLE"),
+            ActionCreateObjectType::Schema => write!(f, "SCHEMA"),
             ActionCreateObjectType::Share => write!(f, "SHARE"),
             ActionCreateObjectType::User => write!(f, "USER"),
             ActionCreateObjectType::Warehouse => write!(f, "WAREHOUSE"),
@@ -7038,6 +7040,8 @@ pub enum GrantObjects {
     AllMaterializedViewsInSchema { schemas: Vec<ObjectName> },
     /// Grant privileges on `ALL EXTERNAL TABLES IN SCHEMA <schema_name> [, ...]`
     AllExternalTablesInSchema { schemas: Vec<ObjectName> },
+    /// Grant privileges on `ALL FUNCTIONS IN SCHEMA <schema_name> [, ...]`
+    AllFunctionsInSchema { schemas: Vec<ObjectName> },
     /// Grant privileges on `FUTURE SCHEMAS IN DATABASE <database_name> [, ...]`
     FutureSchemasInDatabase { databases: Vec<ObjectName> },
     /// Grant privileges on `FUTURE TABLES IN SCHEMA <schema_name> [, ...]`
@@ -7158,6 +7162,13 @@ impl fmt::Display for GrantObjects {
                     display_comma_separated(schemas)
                 )
             }
+            GrantObjects::AllFunctionsInSchema { schemas } => {
+                write!(
+                    f,
+                    "ALL FUNCTIONS IN SCHEMA {}",
+                    display_comma_separated(schemas)
+                )
+            }
             GrantObjects::FutureSchemasInDatabase { databases } => {
                 write!(
                     f,
@@ -7833,6 +7844,7 @@ pub enum ObjectType {
     Stage,
     Type,
     User,
+    Stream,
 }
 
 impl fmt::Display for ObjectType {
@@ -7849,6 +7861,7 @@ impl fmt::Display for ObjectType {
             ObjectType::Stage => "STAGE",
             ObjectType::Type => "TYPE",
             ObjectType::User => "USER",
+            ObjectType::Stream => "STREAM",
         })
     }
 }

src/dialect/bigquery.rs

@@ -19,6 +19,7 @@ use crate::ast::Statement;
 use crate::dialect::Dialect;
 use crate::keywords::Keyword;
 use crate::parser::{Parser, ParserError};
+use crate::tokenizer::Token;
 
 /// These keywords are disallowed as column identifiers. Such that
 /// `SELECT 5 AS <col> FROM T` is rejected by BigQuery.
@@ -47,6 +48,13 @@ pub struct BigQueryDialect;
 impl Dialect for BigQueryDialect {
     fn parse_statement(&self, parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
         if parser.parse_keyword(Keyword::BEGIN) {
+            if parser.peek_keyword(Keyword::TRANSACTION)
+                || parser.peek_token_ref().token == Token::SemiColon
+                || parser.peek_token_ref().token == Token::EOF
+            {
+                parser.prev_token();
+                return None;
+            }
             return Some(parser.parse_begin_exception_end());
         }
 

src/keywords.rs

@@ -870,6 +870,7 @@ define_keywords!(
     STORAGE_SERIALIZATION_POLICY,
     STORED,
     STRAIGHT_JOIN,
+    STREAM,
     STRICT,
     STRING,
     STRUCT,

src/parser/mod.rs

@@ -6326,6 +6326,8 @@ impl<'a> Parser<'a> {
             ObjectType::Type
         } else if self.parse_keyword(Keyword::USER) {
             ObjectType::User
+        } else if self.parse_keyword(Keyword::STREAM) {
+            ObjectType::Stream
         } else if self.parse_keyword(Keyword::FUNCTION) {
             return self.parse_drop_function();
         } else if self.parse_keyword(Keyword::POLICY) {
@@ -14109,6 +14111,15 @@ impl<'a> Parser<'a> {
                 Some(GrantObjects::AllMaterializedViewsInSchema {
                     schemas: self.parse_comma_separated(|p| p.parse_object_name(false))?,
                 })
+            } else if self.parse_keywords(&[
+                Keyword::ALL,
+                Keyword::FUNCTIONS,
+                Keyword::IN,
+                Keyword::SCHEMA,
+            ]) {
+                Some(GrantObjects::AllFunctionsInSchema {
+                    schemas: self.parse_comma_separated(|p| p.parse_object_name(false))?,
+                })
             } else if self.parse_keywords(&[
                 Keyword::FUTURE,
                 Keyword::SCHEMAS,
@@ -14415,6 +14426,8 @@ impl<'a> Parser<'a> {
             Some(ActionCreateObjectType::Integration)
         } else if self.parse_keyword(Keyword::ROLE) {
             Some(ActionCreateObjectType::Role)
+        } else if self.parse_keyword(Keyword::SCHEMA) {
+            Some(ActionCreateObjectType::Schema)
         } else if self.parse_keyword(Keyword::SHARE) {
             Some(ActionCreateObjectType::Share)
         } else if self.parse_keyword(Keyword::USER) {

tests/sqlparser_bigquery.rs

@@ -2568,7 +2568,7 @@ fn test_struct_trailing_and_nested_bracket() {
 }
 
 #[test]
-fn test_export() {
+fn test_export_data() {
     bigquery().verified_stmt(concat!(
         "EXPORT DATA OPTIONS(",
         "uri = 'gs://bucket/folder/*', ",
@@ -2587,3 +2587,13 @@ fn test_export() {
         "SELECT field1, field2 FROM mydataset.table1 ORDER BY field1 LIMIT 10",
     ));
 }
+
+#[test]
+fn test_begin_transaction() {
+    bigquery().verified_stmt("BEGIN TRANSACTION");
+}
+
+#[test]
+fn test_begin_statement() {
+    bigquery().verified_stmt("BEGIN");
+}

tests/sqlparser_common.rs

@@ -9528,6 +9528,7 @@ fn parse_grant() {
     verified_stmt("GRANT SELECT ON ALL VIEWS IN SCHEMA db1.sc1 TO ROLE role1");
     verified_stmt("GRANT SELECT ON ALL MATERIALIZED VIEWS IN SCHEMA db1.sc1 TO ROLE role1");
     verified_stmt("GRANT SELECT ON ALL EXTERNAL TABLES IN SCHEMA db1.sc1 TO ROLE role1");
+    verified_stmt("GRANT USAGE ON ALL FUNCTIONS IN SCHEMA db1.sc1 TO ROLE role1");
     verified_stmt("GRANT USAGE ON SCHEMA sc1 TO a:b");
     verified_stmt("GRANT USAGE ON SCHEMA sc1 TO GROUP group1");
     verified_stmt("GRANT OWNERSHIP ON ALL TABLES IN SCHEMA DEV_STAS_ROGOZHIN TO ROLE ANALYST");
@@ -9551,6 +9552,7 @@ fn parse_grant() {
     verified_stmt("GRANT USAGE ON FUNCTION db1.sc1.foo(INT) TO ROLE role1");
     verified_stmt("GRANT ROLE role1 TO ROLE role2");
     verified_stmt("GRANT ROLE role1 TO USER user");
+    verified_stmt("GRANT CREATE SCHEMA ON DATABASE db1 TO ROLE role1");
 }
 
 #[test]
@@ -16357,3 +16359,21 @@ fn parse_create_user() {
         _ => unreachable!(),
     }
 }
+
+#[test]
+fn parse_drop_stream() {
+    let sql = "DROP STREAM s1";
+    match verified_stmt(sql) {
+        Statement::Drop {
+            names, object_type, ..
+        } => {
+            assert_eq!(
+                vec!["s1"],
+                names.iter().map(ToString::to_string).collect::<Vec<_>>()
+            );
+            assert_eq!(ObjectType::Stream, object_type);
+        }
+        _ => unreachable!(),
+    }
+    verified_stmt("DROP STREAM IF EXISTS s1");
+}