From 254ccfb4d87852edc9fbdbcf10672dd529afcaab Mon Sep 17 00:00:00 2001 From: Lukasz Stefaniak Date: Fri, 27 Oct 2023 20:52:47 +0200 Subject: [PATCH] snowflake: add support for LATERAL FLATTEN and similar (#1026) --- src/ast/query.rs | 23 +++++++++++++ src/parser/mod.rs | 19 ++++++++--- tests/sqlparser_common.rs | 65 +++++++++++++++++++++++++++++++++--- tests/sqlparser_snowflake.rs | 6 ++++ 4 files changed, 105 insertions(+), 8 deletions(-) diff --git a/src/ast/query.rs b/src/ast/query.rs index 4289b0bd..38128dff 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -695,6 +695,13 @@ pub enum TableFactor { expr: Expr, alias: Option, }, + /// `e.g. LATERAL FLATTEN()[ AS ]` + Function { + lateral: bool, + name: ObjectName, + args: Vec, + alias: Option, + }, /// ```sql /// SELECT * FROM UNNEST ([10,20,30]) as numbers WITH OFFSET; /// +---------+--------+ @@ -793,6 +800,22 @@ impl fmt::Display for TableFactor { } Ok(()) } + TableFactor::Function { + lateral, + name, + args, + alias, + } => { + if *lateral { + write!(f, "LATERAL ")?; + } + write!(f, "{name}")?; + write!(f, "({})", display_comma_separated(args))?; + if let Some(alias) = alias { + write!(f, " AS {alias}")?; + } + Ok(()) + } TableFactor::TableFunction { expr, alias } => { write!(f, "TABLE({expr})")?; if let Some(alias) = alias { diff --git a/src/parser/mod.rs b/src/parser/mod.rs index eb4ef68a..c3976c74 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -6555,11 +6555,21 @@ impl<'a> Parser<'a> { /// A table name or a parenthesized subquery, followed by optional `[AS] alias` pub fn parse_table_factor(&mut self) -> Result { if self.parse_keyword(Keyword::LATERAL) { - // LATERAL must always be followed by a subquery. - if !self.consume_token(&Token::LParen) { - self.expected("subquery after LATERAL", self.peek_token())?; + // LATERAL must always be followed by a subquery or table function. + if self.consume_token(&Token::LParen) { + self.parse_derived_table_factor(Lateral) + } else { + let name = self.parse_object_name()?; + self.expect_token(&Token::LParen)?; + let args = self.parse_optional_args()?; + let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?; + Ok(TableFactor::Function { + lateral: true, + name, + args, + alias, + }) } - self.parse_derived_table_factor(Lateral) } else if self.parse_keyword(Keyword::TABLE) { // parse table function (SELECT * FROM TABLE () [ AS ]) self.expect_token(&Token::LParen)?; @@ -6638,6 +6648,7 @@ impl<'a> Parser<'a> { match &mut table_and_joins.relation { TableFactor::Derived { alias, .. } | TableFactor::Table { alias, .. } + | TableFactor::Function { alias, .. } | TableFactor::UNNEST { alias, .. } | TableFactor::TableFunction { alias, .. } | TableFactor::Pivot { alias, .. } diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index fcad2455..befdf512 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -18,6 +18,8 @@ //! sqlparser regardless of the chosen dialect (i.e. it doesn't conflict with //! dialect-specific parsing rules). +extern crate core; + use matches::assert_matches; use sqlparser::ast::SelectItem::UnnamedExpr; use sqlparser::ast::TableFactor::{Pivot, Unpivot}; @@ -37,6 +39,9 @@ use test_utils::{ #[macro_use] mod test_utils; +#[cfg(test)] +use pretty_assertions::assert_eq; + #[test] fn parse_insert_values() { let row = vec![ @@ -5976,12 +5981,10 @@ fn lateral_derived() { chk(false); chk(true); - let sql = "SELECT * FROM customer LEFT JOIN LATERAL generate_series(1, customer.id)"; + let sql = "SELECT * FROM LATERAL UNNEST ([10,20,30]) as numbers WITH OFFSET;"; let res = parse_sql_statements(sql); assert_eq!( - ParserError::ParserError( - "Expected subquery after LATERAL, found: generate_series".to_string() - ), + ParserError::ParserError("Expected end of statement, found: WITH".to_string()), res.unwrap_err() ); @@ -5995,6 +5998,60 @@ fn lateral_derived() { ); } +#[test] +fn lateral_function() { + let sql = "SELECT * FROM customer LEFT JOIN LATERAL generate_series(1, customer.id)"; + let actual_select_only = verified_only_select(sql); + let expected = Select { + distinct: None, + top: None, + projection: vec![SelectItem::Wildcard(WildcardAdditionalOptions { + opt_exclude: None, + opt_except: None, + opt_rename: None, + opt_replace: None, + })], + into: None, + from: vec![TableWithJoins { + relation: TableFactor::Table { + name: ObjectName(vec![Ident { + value: "customer".to_string(), + quote_style: None, + }]), + alias: None, + args: None, + with_hints: vec![], + version: None, + partitions: vec![], + }, + joins: vec![Join { + relation: TableFactor::Function { + lateral: true, + name: ObjectName(vec!["generate_series".into()]), + args: vec![ + FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(number("1")))), + FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::CompoundIdentifier( + vec![Ident::new("customer"), Ident::new("id")], + ))), + ], + alias: None, + }, + join_operator: JoinOperator::LeftOuter(JoinConstraint::None), + }], + }], + lateral_views: vec![], + selection: None, + group_by: GroupByExpr::Expressions(vec![]), + cluster_by: vec![], + distribute_by: vec![], + sort_by: vec![], + having: None, + named_window: vec![], + qualify: None, + }; + assert_eq!(actual_select_only, expected); +} + #[test] fn parse_start_transaction() { match verified_stmt("START TRANSACTION READ ONLY, READ WRITE, ISOLATION LEVEL SERIALIZABLE") { diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs index f0a07797..2b7e34cc 100644 --- a/tests/sqlparser_snowflake.rs +++ b/tests/sqlparser_snowflake.rs @@ -176,6 +176,12 @@ fn parse_array() { ); } +#[test] +fn parse_lateral_flatten() { + snowflake().verified_only_select(r#"SELECT * FROM TABLE(FLATTEN(input => parse_json('{"a":1, "b":[77,88]}'), outer => true)) AS f"#); + snowflake().verified_only_select(r#"SELECT emp.employee_ID, emp.last_name, index, value AS project_name FROM employees AS emp, LATERAL FLATTEN(INPUT => emp.project_names) AS proj_names"#); +} + #[test] fn parse_json_using_colon() { let sql = "SELECT a:b FROM t";