diff --git a/examples/cli.rs b/examples/cli.rs index 3c4299b209..51a63a7ea1 100644 --- a/examples/cli.rs +++ b/examples/cli.rs @@ -71,9 +71,9 @@ $ cargo run --example cli - [--dialectname] .expect("failed to read from stdin"); String::from_utf8(buf).expect("stdin content wasn't valid utf8") } else { - println!("Parsing from file '{}' using {:?}", &filename, dialect); + println!("Parsing from file '{}' using {:?}", filename, dialect); fs::read_to_string(&filename) - .unwrap_or_else(|_| panic!("Unable to read the file {}", &filename)) + .unwrap_or_else(|_| panic!("Unable to read the file {}", filename)) }; let without_bom = if contents.chars().next().unwrap() as u64 != 0xfeff { contents.as_str() diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs index 67aefb3928..99760a425e 100644 --- a/src/ast/ddl.rs +++ b/src/ast/ddl.rs @@ -1939,6 +1939,20 @@ pub enum ColumnOption { Comment(String), /// `ON UPDATE ` column option OnUpdate(Expr), + /// `METADATA FROM 'key'` + /// + /// A special type of column that gets its value from metadata + /// associated with the record. + /// + /// Example: + /// ```sql + /// CREATE TABLE logs ( + /// id TEXT, + /// kafka_topic STRING METADATA FROM 'topic', + /// log TEXT + /// ) + /// ``` + MetadataField(String, Span), /// `Generated`s are modifiers that follow a column definition in a `CREATE /// TABLE` statement. Generated { @@ -2085,6 +2099,9 @@ impl fmt::Display for ColumnOption { Collation(n) => write!(f, "COLLATE {n}"), Comment(v) => write!(f, "COMMENT '{}'", escape_single_quote_string(v)), OnUpdate(expr) => write!(f, "ON UPDATE {expr}"), + MetadataField(key, _) => { + write!(f, "METADATA FROM '{}'", escape_single_quote_string(key)) + } Generated { generated_as, sequence_options, @@ -3060,6 +3077,10 @@ pub struct CreateTable { /// Redshift `BACKUP` option: `BACKUP { YES | NO }` /// pub backup: Option, + /// Arroyo-specific: Iceberg partition transforms + /// Syntax: PARTITIONED BY (hour(ts), bucket(32, id), truncate(8, color)) + /// + pub arroyo_partitions: Option>, } impl fmt::Display for CreateTable { @@ -3250,8 +3271,25 @@ impl fmt::Display for CreateTable { if let Some(cluster_by) = self.cluster_by.as_ref() { write!(f, " CLUSTER BY {cluster_by}")?; } - if let options @ CreateTableOptions::Options(_) = &self.table_options { - write!(f, " {options}")?; + // Connector partitions are parsed after table options. Keep `OPTIONS` + // before them when both are present so GenericDialect does not reparse + // the partition expressions as Hive partition columns. + if self.arroyo_partitions.is_some() { + if let options @ CreateTableOptions::Options(_) = &self.table_options { + write!(f, " {options}")?; + } + } + if let Some(partitions) = &self.arroyo_partitions { + write!( + f, + " PARTITIONED BY ({})", + display_comma_separated(partitions) + )?; + } + if self.arroyo_partitions.is_none() { + if let options @ CreateTableOptions::Options(_) = &self.table_options { + write!(f, " {options}")?; + } } if let Some(external_volume) = self.external_volume.as_ref() { write!(f, " EXTERNAL_VOLUME='{external_volume}'")?; @@ -4698,7 +4736,7 @@ impl fmt::Display for AlterTable { if self.only { write!(f, "ONLY ")?; } - write!(f, "{} ", &self.name)?; + write!(f, "{} ", self.name)?; if let Some(cluster) = &self.on_cluster { write!(f, "ON CLUSTER {cluster} ")?; } diff --git a/src/ast/helpers/stmt_create_table.rs b/src/ast/helpers/stmt_create_table.rs index ab2feb6930..b9f2d37606 100644 --- a/src/ast/helpers/stmt_create_table.rs +++ b/src/ast/helpers/stmt_create_table.rs @@ -183,6 +183,8 @@ pub struct CreateTableBuilder { pub sortkey: Option>, /// Redshift `BACKUP` option. pub backup: Option, + /// Arroyo connector partition expressions. + pub arroyo_partitions: Option>, } impl CreateTableBuilder { @@ -248,6 +250,7 @@ impl CreateTableBuilder { distkey: None, sortkey: None, backup: None, + arroyo_partitions: None, } } /// Set `OR REPLACE` for the CREATE TABLE statement. @@ -556,6 +559,12 @@ impl CreateTableBuilder { self.backup = backup; self } + /// Set Arroyo connector partition expressions. + pub fn arroyo_partitions(mut self, partitions: Option>) -> Self { + self.arroyo_partitions = partitions; + self + } + /// Consume the builder and produce a `CreateTable`. pub fn build(self) -> CreateTable { CreateTable { @@ -618,6 +627,7 @@ impl CreateTableBuilder { distkey: self.distkey, sortkey: self.sortkey, backup: self.backup, + arroyo_partitions: self.arroyo_partitions, } } } @@ -699,6 +709,7 @@ impl From for CreateTableBuilder { distkey: table.distkey, sortkey: table.sortkey, backup: table.backup, + arroyo_partitions: table.arroyo_partitions, } } } diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 886bea26d5..b007f5da55 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -11748,7 +11748,7 @@ impl fmt::Display for AlterUser { let has_props = !self.set_props.options.is_empty(); if has_props { write!(f, " SET")?; - write!(f, " {}", &self.set_props)?; + write!(f, " {}", self.set_props)?; } if !self.unset_props.is_empty() { write!(f, " UNSET {}", display_comma_separated(&self.unset_props))?; diff --git a/src/ast/query.rs b/src/ast/query.rs index bbdd7540af..eecfb0490d 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -3520,7 +3520,7 @@ pub struct LockClause { impl fmt::Display for LockClause { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "FOR {}", &self.lock_type)?; + write!(f, "FOR {}", self.lock_type)?; if let Some(ref of) = self.of { write!(f, " OF {of}")?; } diff --git a/src/ast/spans.rs b/src/ast/spans.rs index 0dc834ba03..c1fa11137e 100644 --- a/src/ast/spans.rs +++ b/src/ast/spans.rs @@ -604,6 +604,7 @@ impl Spanned for CreateTable { distkey: _, sortkey: _, backup: _, + arroyo_partitions, } = self; union_spans( @@ -614,7 +615,8 @@ impl Spanned for CreateTable { .chain(query.iter().map(|i| i.span())) .chain(clone.iter().map(|i| i.span())) .chain(partition_of.iter().map(|i| i.span())) - .chain(for_values.iter().map(|i| i.span())), + .chain(for_values.iter().map(|i| i.span())) + .chain(arroyo_partitions.iter().flatten().map(Spanned::span)), ) } } @@ -650,6 +652,12 @@ impl Spanned for TableConstraint { TableConstraint::FulltextOrSpatial(constraint) => constraint.span(), TableConstraint::PrimaryKeyUsingIndex(constraint) | TableConstraint::UniqueUsingIndex(constraint) => constraint.span(), + TableConstraint::Watermark { + column_name, + watermark_expr, + } => column_name + .span + .union_opt(&watermark_expr.as_ref().map(Spanned::span)), } } } @@ -832,6 +840,7 @@ impl Spanned for ColumnOption { ColumnOption::Collation(object_name) => object_name.span(), ColumnOption::Comment(_) => Span::empty(), ColumnOption::OnUpdate(expr) => expr.span(), + ColumnOption::MetadataField(_, span) => *span, ColumnOption::Generated { .. } => Span::empty(), ColumnOption::Options(vec) => union_spans(vec.iter().map(|i| i.span())), ColumnOption::Identity(..) => Span::empty(), diff --git a/src/ast/table_constraints.rs b/src/ast/table_constraints.rs index 9ba196a81e..e313959f81 100644 --- a/src/ast/table_constraints.rs +++ b/src/ast/table_constraints.rs @@ -117,6 +117,21 @@ pub enum TableConstraint { /// /// [1]: https://www.postgresql.org/docs/current/sql-altertable.html UniqueUsingIndex(ConstraintUsingIndex), + /// Arroyo specific: Watermark definition for streaming tables + /// Syntax: + /// ```sql + /// WATERMARK FOR timestamp AS timestamp - INTERVAL '5 seconds' + /// ``` + /// or without an expression + /// ```sql + /// WATERMARK FOR timestamp + /// ``` + Watermark { + /// Column name to be used for the watermark + column_name: Ident, + /// Optional watermark expression + watermark_expr: Option, + }, } impl From for TableConstraint { @@ -166,6 +181,16 @@ impl fmt::Display for TableConstraint { TableConstraint::FulltextOrSpatial(constraint) => constraint.fmt(f), TableConstraint::PrimaryKeyUsingIndex(c) => c.fmt_with_keyword(f, "PRIMARY KEY"), TableConstraint::UniqueUsingIndex(c) => c.fmt_with_keyword(f, "UNIQUE"), + TableConstraint::Watermark { + column_name, + watermark_expr, + } => { + write!(f, "WATERMARK FOR {column_name}")?; + if let Some(expr) = watermark_expr { + write!(f, " AS {expr}")?; + } + Ok(()) + } } } } diff --git a/src/dialect/arroyo.rs b/src/dialect/arroyo.rs new file mode 100644 index 0000000000..fba4facb39 --- /dev/null +++ b/src/dialect/arroyo.rs @@ -0,0 +1,223 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use log::debug; + +use crate::dialect::{Dialect, Precedence}; +use crate::keywords::Keyword; +use crate::parser::{Parser, ParserError}; +use crate::tokenizer::Token; + +/// A [`Dialect`] for [Arroyo](https://www.arroyo.dev/) +/// This is based on the Postgres dialect +/// +/// Arroyo adds its own streaming SQL syntax extensions. +#[derive(Debug)] +pub struct ArroyoDialect {} + +const PERIOD_PREC: u8 = 200; +const DOUBLE_COLON_PREC: u8 = 140; +const BRACKET_PREC: u8 = 130; +const COLLATE_PREC: u8 = 120; +const AT_TZ_PREC: u8 = 110; +const CARET_PREC: u8 = 100; +const MUL_DIV_MOD_OP_PREC: u8 = 90; +const PLUS_MINUS_PREC: u8 = 80; +// there's no XOR operator in PostgreSQL, but support it here to avoid breaking tests +const XOR_PREC: u8 = 75; +const PG_OTHER_PREC: u8 = 70; +const BETWEEN_LIKE_PREC: u8 = 60; +const EQ_PREC: u8 = 50; +const IS_PREC: u8 = 40; +const NOT_PREC: u8 = 30; +const AND_PREC: u8 = 20; +const OR_PREC: u8 = 10; + +impl Dialect for ArroyoDialect { + fn identifier_quote_style(&self, _identifier: &str) -> Option { + Some('"') + } + + fn is_delimited_identifier_start(&self, ch: char) -> bool { + ch == '"' // Postgres does not support backticks to quote identifiers + } + + fn is_identifier_start(&self, ch: char) -> bool { + // See https://www.postgresql.org/docs/11/sql-syntax-lexical.html#SQL-SYNTAX-IDENTIFIERS + // We don't yet support identifiers beginning with "letters with + // diacritical marks" + ch.is_alphabetic() || ch == '_' + } + + fn is_identifier_part(&self, ch: char) -> bool { + ch.is_alphabetic() || ch.is_ascii_digit() || ch == '$' || ch == '_' + } + + /// See + fn is_custom_operator_part(&self, ch: char) -> bool { + matches!( + ch, + '+' | '-' + | '*' + | '/' + | '<' + | '>' + | '=' + | '~' + | '!' + | '@' + | '#' + | '%' + | '^' + | '&' + | '|' + | '`' + | '?' + ) + } + + fn get_next_precedence(&self, parser: &Parser) -> Option> { + let token = parser.peek_token_ref(); + debug!("get_next_precedence() {token:?}"); + + // we only return some custom value here when the behaviour (not merely the numeric value) differs + // from the default implementation + match &token.token { + Token::Word(w) + if w.keyword == Keyword::COLLATE && !parser.in_column_definition_state() => + { + Some(Ok(COLLATE_PREC)) + } + Token::LBracket => Some(Ok(BRACKET_PREC)), + Token::Arrow + | Token::LongArrow + | Token::HashArrow + | Token::HashLongArrow + | Token::AtArrow + | Token::ArrowAt + | Token::HashMinus + | Token::AtQuestion + | Token::AtAt + | Token::Question + | Token::QuestionAnd + | Token::QuestionPipe + | Token::ExclamationMark + | Token::Overlap + | Token::CaretAt + | Token::StringConcat + | Token::Sharp + | Token::ShiftRight + | Token::ShiftLeft + | Token::CustomBinaryOperator(_) => Some(Ok(PG_OTHER_PREC)), + // lowest prec to prevent it from turning into a binary op + Token::Colon => Some(Ok(self.prec_unknown())), + _ => None, + } + } + + fn prec_value(&self, prec: Precedence) -> u8 { + match prec { + Precedence::Period => PERIOD_PREC, + Precedence::DoubleColon => DOUBLE_COLON_PREC, + Precedence::AtTz => AT_TZ_PREC, + Precedence::MulDivModOp => MUL_DIV_MOD_OP_PREC, + Precedence::PlusMinus => PLUS_MINUS_PREC, + Precedence::Xor => XOR_PREC, + Precedence::Ampersand => PG_OTHER_PREC, + Precedence::Caret => CARET_PREC, + Precedence::Pipe => PG_OTHER_PREC, + Precedence::Colon => PG_OTHER_PREC, + Precedence::Between => BETWEEN_LIKE_PREC, + Precedence::Eq => EQ_PREC, + Precedence::Like => BETWEEN_LIKE_PREC, + Precedence::Is => IS_PREC, + Precedence::PgOther => PG_OTHER_PREC, + Precedence::UnaryNot => NOT_PREC, + Precedence::And => AND_PREC, + Precedence::Or => OR_PREC, + } + } + + fn supports_unicode_string_literal(&self) -> bool { + true + } + + fn supports_filter_during_aggregation(&self) -> bool { + true + } + + fn supports_group_by_expr(&self) -> bool { + true + } + + fn allow_extract_custom(&self) -> bool { + true + } + + fn allow_extract_single_quotes(&self) -> bool { + true + } + + /// see + fn supports_factorial_operator(&self) -> bool { + true + } + + fn supports_bitwise_shift_operators(&self) -> bool { + true + } + + /// see + fn supports_comment_on(&self) -> bool { + true + } + + /// Return true if the dialect supports empty projections in SELECT statements + /// + /// Example + /// ```sql + /// SELECT from table_name + /// ``` + fn supports_empty_projections(&self) -> bool { + true + } + + fn supports_nested_comments(&self) -> bool { + true + } + + fn supports_string_escape_constant(&self) -> bool { + true + } + + fn supports_numeric_literal_underscores(&self) -> bool { + true + } + + /// See: + fn supports_array_typedef_with_brackets(&self) -> bool { + true + } + + fn supports_geometric_types(&self) -> bool { + true + } + + // arroyo-specific features + fn supports_struct_literal(&self) -> bool { + true + } + + fn supports_insert_table_alias(&self) -> bool { + true + } +} diff --git a/src/dialect/mod.rs b/src/dialect/mod.rs index 6ab6cb15e4..3e485892c6 100644 --- a/src/dialect/mod.rs +++ b/src/dialect/mod.rs @@ -16,6 +16,7 @@ // under the License. mod ansi; +mod arroyo; mod bigquery; mod clickhouse; mod databricks; @@ -40,6 +41,7 @@ use core::str::Chars; use log::debug; pub use self::ansi::AnsiDialect; +pub use self::arroyo::ArroyoDialect; pub use self::bigquery::BigQueryDialect; pub use self::clickhouse::ClickHouseDialect; pub use self::databricks::DatabricksDialect; diff --git a/src/keywords.rs b/src/keywords.rs index 4fc8f72d1d..cf45cac9f1 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -1155,6 +1155,7 @@ define_keywords!( WAITFOR, WAREHOUSE, WAREHOUSES, + WATERMARK, WEEK, WEEKS, WHEN, diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 668c520e5e..8565163b02 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -1848,7 +1848,7 @@ impl<'a> Parser<'a> { | tok @ Token::PGSquareRoot | tok @ Token::PGCubeRoot | tok @ Token::AtSign - if dialect_is!(dialect is PostgreSqlDialect) => + if dialect_is!(dialect is PostgreSqlDialect | ArroyoDialect) => { let op = match tok { Token::DoubleExclamationMark => UnaryOperator::PGPrefixFactorial, @@ -1898,7 +1898,7 @@ impl<'a> Parser<'a> { ), }) } - Token::EscapedStringLiteral(_) if dialect_is!(dialect is PostgreSqlDialect | GenericDialect) => + Token::EscapedStringLiteral(_) if dialect_is!(dialect is PostgreSqlDialect | GenericDialect | ArroyoDialect) => { self.prev_token(); Ok(Expr::Value(self.parse_value()?)) @@ -3762,7 +3762,7 @@ impl<'a> Parser<'a> { Token::Caret => { // In PostgreSQL, ^ stands for the exponentiation operation, // and # stands for XOR. See https://www.postgresql.org/docs/current/functions-math.html - if dialect_is!(dialect is PostgreSqlDialect) { + if dialect_is!(dialect is PostgreSqlDialect | ArroyoDialect) { Some(BinaryOperator::PGExp) } else { Some(BinaryOperator::BitwiseXor) @@ -3779,19 +3779,19 @@ impl<'a> Parser<'a> { Token::ShiftRight if dialect.supports_bitwise_shift_operators() => { Some(BinaryOperator::PGBitwiseShiftRight) } - Token::Sharp if dialect_is!(dialect is PostgreSqlDialect | RedshiftSqlDialect) => { + Token::Sharp if dialect_is!(dialect is PostgreSqlDialect | RedshiftSqlDialect | ArroyoDialect) => { Some(BinaryOperator::PGBitwiseXor) } Token::Overlap if dialect_is!(dialect is PostgreSqlDialect | RedshiftSqlDialect) => { Some(BinaryOperator::PGOverlap) } - Token::Overlap if dialect_is!(dialect is PostgreSqlDialect | GenericDialect) => { + Token::Overlap if dialect_is!(dialect is PostgreSqlDialect | ArroyoDialect | GenericDialect) => { Some(BinaryOperator::PGOverlap) } Token::Overlap if dialect.supports_double_ampersand_operator() => { Some(BinaryOperator::And) } - Token::CaretAt if dialect_is!(dialect is PostgreSqlDialect | GenericDialect) => { + Token::CaretAt if dialect_is!(dialect is PostgreSqlDialect | ArroyoDialect | GenericDialect) => { Some(BinaryOperator::PGStartsWith) } Token::Tilde => Some(BinaryOperator::PGRegexMatch), @@ -4744,7 +4744,7 @@ impl<'a> Parser<'a> { if self.parse_keyword(expected) { Ok(self.get_current_token().clone()) } else { - self.expected_ref(format!("{:?}", &expected).as_str(), self.peek_token_ref()) + self.expected_ref(format!("{:?}", expected).as_str(), self.peek_token_ref()) } } @@ -4757,7 +4757,7 @@ impl<'a> Parser<'a> { if self.parse_keyword(expected) { Ok(()) } else { - self.expected_ref(format!("{:?}", &expected).as_str(), self.peek_token_ref()) + self.expected_ref(format!("{:?}", expected).as_str(), self.peek_token_ref()) } } @@ -8575,6 +8575,18 @@ impl<'a> Parser<'a> { create_table_config.partition_by }; + // Parse Arroyo-specific PARTITIONED BY for Iceberg + let arroyo_partitions = if dialect_of!(self is ArroyoDialect | GenericDialect) + && self.parse_keywords(&[Keyword::PARTITIONED, Keyword::BY]) + { + self.expect_token(&Token::LParen)?; + let partitions = self.parse_comma_separated(Parser::parse_expr)?; + self.expect_token(&Token::RParen)?; + Some(partitions) + } else { + None + }; + let on_commit = if self.parse_keywords(&[Keyword::ON, Keyword::COMMIT]) { Some(self.parse_create_table_on_commit()?) } else { @@ -8657,6 +8669,7 @@ impl<'a> Parser<'a> { .diststyle(diststyle) .distkey(distkey) .sortkey(sortkey) + .arroyo_partitions(arroyo_partitions) .build()) } @@ -9266,6 +9279,17 @@ impl<'a> Parser<'a> { Ok(Some(ColumnOption::Null)) } else if self.parse_keyword(Keyword::DEFAULT) { Ok(Some(ColumnOption::Default(self.parse_expr()?))) + } else if self.parse_keywords(&[Keyword::METADATA, Keyword::FROM]) + && dialect_of!(self is ArroyoDialect | GenericDialect) + { + // Parse metadata field syntax: METADATA FROM 'key' + let next_token = self.next_token(); + match next_token.token { + Token::SingleQuotedString(value) => { + Ok(Some(ColumnOption::MetadataField(value, next_token.span))) + } + _ => self.expected("string literal for metadata key", next_token), + } } else if dialect_of!(self is ClickHouseDialect| GenericDialect) && self.parse_keyword(Keyword::MATERIALIZED) { @@ -9926,6 +9950,33 @@ impl<'a> Parser<'a> { .into(), )) } + Token::Word(w) + if w.keyword == Keyword::WATERMARK + && dialect_of!(self is ArroyoDialect | GenericDialect) => + { + if let Some(name) = name { + return self.expected( + "WATERMARK option without constraint name", + TokenWithSpan { + token: Token::make_keyword(&name.to_string()), + span: next_token.span, + }, + ); + } + self.expect_keyword_is(Keyword::FOR)?; + let column_name = self.parse_identifier()?; + + // The AS keyword and expression are optional + let watermark_expr = if self.parse_keyword(Keyword::AS) { + Some(self.parse_expr()?) + } else { + None + }; + Ok(Some(TableConstraint::Watermark { + column_name, + watermark_expr, + })) + } _ => { if name.is_some() { self.expected("PRIMARY, UNIQUE, FOREIGN, or CHECK", next_token) @@ -12220,7 +12271,7 @@ impl<'a> Parser<'a> { }) => Ok(value), Token::SingleQuotedString(s) => Ok(s), Token::DoubleQuotedString(s) => Ok(s), - Token::EscapedStringLiteral(s) if dialect_of!(self is PostgreSqlDialect | GenericDialect) => { + Token::EscapedStringLiteral(s) if dialect_of!(self is PostgreSqlDialect | ArroyoDialect | GenericDialect) => { Ok(s) } Token::UnicodeStringLiteral(s) => Ok(s), @@ -13324,16 +13375,14 @@ impl<'a> Parser<'a> { "Trailing period in identifier".to_string(), ))? } - token => { - return Err(ParserError::ParserError(format!( - "Unexpected token following period in identifier: {token}" - )))? - } + token => Err(ParserError::ParserError(format!( + "Unexpected token following period in identifier: {token}" + )))?, } } Token::EOF => break, token => { - return Err(ParserError::ParserError(format!( + Err(ParserError::ParserError(format!( "Unexpected token in identifier: {token}" )))?; } @@ -16119,7 +16168,7 @@ impl<'a> Parser<'a> { alias, sample: None, }) - } else if dialect_of!(self is BigQueryDialect | PostgreSqlDialect | GenericDialect) + } else if dialect_of!(self is BigQueryDialect | PostgreSqlDialect | ArroyoDialect | GenericDialect) && self.parse_keyword(Keyword::UNNEST) { self.expect_token(&Token::LParen)?; @@ -16592,7 +16641,7 @@ impl<'a> Parser<'a> { where_clause = Some(self.parse_expr()?); } else { let tok = self.peek_token_ref(); - return parser_err!( + parser_err!( format!( "Expected one of DIMENSIONS, METRICS, FACTS or WHERE, got {}", tok.token diff --git a/tests/sqlparser_arroyo.rs b/tests/sqlparser_arroyo.rs new file mode 100644 index 0000000000..d1944939e5 --- /dev/null +++ b/tests/sqlparser_arroyo.rs @@ -0,0 +1,242 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use sqlparser::ast::helpers::stmt_create_table::CreateTableBuilder; +use sqlparser::ast::{ + BinaryOperator, ColumnOption, Expr, HiveDistributionStyle, Spanned, Statement, TableConstraint, +}; +use sqlparser::dialect::{ArroyoDialect, GenericDialect, HiveDialect, PostgreSqlDialect}; +use sqlparser::parser::Parser; +use sqlparser::test_utils::TestedDialects; + +fn arroyo() -> TestedDialects { + TestedDialects::new(vec![Box::new(ArroyoDialect {})]) +} + +#[test] +fn postgres_expressions_and_struct_types() { + for sql in [ + "SELECT 2 ^ 3, 7 # 2, 1 << 2, 8 >> 1", + "SELECT |/9, ||/27, @(-1), !!5", + "SELECT ARRAY[1] && ARRAY[2], 'hello' ^@ 'he'", + r#"SELECT E'hello\nworld', U&'hello', "quoted""#, + "SELECT * FROM UNNEST(ARRAY[1, 2]) WITH ORDINALITY AS t (value, ordinal)", + "CREATE TABLE events (payload STRUCT)", + "INSERT INTO events AS e SELECT * FROM source", + ] { + arroyo().verified_stmt(sql); + } + + let Expr::BinaryOp { op, .. } = arroyo().verified_expr("2 ^ 3") else { + panic!("expected binary expression"); + }; + assert_eq!(op, BinaryOperator::PGExp); + assert!(Parser::parse_sql(&ArroyoDialect {}, "SELECT 1_000").is_ok()); +} + +#[test] +fn generated_columns_do_not_require_stored() { + let sql = + "CREATE TABLE events (raw TEXT, ts TIMESTAMP GENERATED ALWAYS AS (CAST(raw AS TIMESTAMP)))"; + arroyo().verified_stmt(sql); + assert!(Parser::parse_sql(&PostgreSqlDialect {}, sql).is_err()); +} + +#[test] +fn watermark_constraints_round_trip() { + let dialects = TestedDialects::new(vec![Box::new(ArroyoDialect {}), Box::new(GenericDialect)]); + for expression in [None, Some("ts - INTERVAL '5 seconds'")] { + let suffix = expression.map(|e| format!(" AS {e}")).unwrap_or_default(); + let sql = format!("CREATE TABLE events (ts TIMESTAMP, WATERMARK FOR ts{suffix}) WITH (connector = 'kafka')"); + let Statement::CreateTable(table) = dialects.verified_stmt(&sql) else { + panic!("expected CREATE TABLE"); + }; + let [constraint @ TableConstraint::Watermark { + column_name, + watermark_expr, + }] = table.constraints.as_slice() + else { + panic!("expected one watermark constraint"); + }; + assert_eq!(column_name.value, "ts"); + assert_eq!( + *watermark_expr, + expression.map(|e| arroyo().verified_expr(e)) + ); + assert_eq!( + constraint.span(), + column_name + .span + .union_opt(&watermark_expr.as_ref().map(Spanned::span)) + ); + assert!(Parser::parse_sql(&PostgreSqlDialect {}, &sql).is_err()); + } + arroyo().verified_stmt( + r#"CREATE TABLE events ("event time" TIMESTAMP, WATERMARK FOR "event time")"#, + ); +} + +#[test] +fn invalid_watermark_constraints() { + for sql in [ + "CREATE TABLE events (ts TIMESTAMP, WATERMARK ts)", + "CREATE TABLE events (ts TIMESTAMP, WATERMARK FOR)", + "CREATE TABLE events (ts TIMESTAMP, WATERMARK FOR ts AS)", + "CREATE TABLE events (ts TIMESTAMP, CONSTRAINT wm WATERMARK FOR ts)", + ] { + assert!(Parser::parse_sql(&ArroyoDialect {}, sql).is_err(), "{sql}"); + } +} + +#[test] +fn metadata_fields_round_trip() { + let dialects = TestedDialects::new(vec![Box::new(ArroyoDialect {}), Box::new(GenericDialect)]); + for (literal, key) in [("'topic'", "topic"), ("'it''s a key'", "it's a key")] { + let sql = format!( + "CREATE TABLE logs (id INT, topic TEXT METADATA FROM {literal} NOT NULL, payload TEXT)" + ); + dialects.verified_stmt(&sql); + // The round-trip helper intentionally discards token spans. + let Statement::CreateTable(table) = Parser::parse_sql(&ArroyoDialect {}, &sql) + .unwrap() + .remove(0) + else { + panic!("expected CREATE TABLE"); + }; + assert_eq!(table.columns.len(), 3); + let options = &table.columns[1].options; + let ColumnOption::MetadataField(actual, span) = &options[0].option else { + panic!("expected metadata field"); + }; + assert_eq!(actual, key); + assert_eq!(options[0].span(), *span); + assert_eq!(span.end.column - span.start.column, literal.len() as u64); + assert_eq!(options[1].option, ColumnOption::NotNull); + assert!(Parser::parse_sql(&PostgreSqlDialect {}, &sql).is_err()); + } +} + +#[test] +fn invalid_metadata_fields() { + for sql in [ + "CREATE TABLE logs (topic TEXT METADATA)", + "CREATE TABLE logs (topic TEXT METADATA FROM)", + "CREATE TABLE logs (topic TEXT METADATA FROM topic)", + "CREATE TABLE logs (topic TEXT METADATA FROM 42)", + ] { + assert!(Parser::parse_sql(&ArroyoDialect {}, sql).is_err(), "{sql}"); + } +} + +#[test] +fn connector_partition_expressions_round_trip() { + let dialects = TestedDialects::new(vec![Box::new(ArroyoDialect {}), Box::new(GenericDialect)]); + for expressions in [ + vec!["hour(ts)", "bucket(32, id)", "truncate(8, color)"], + vec!["day(ts)"], + vec!["color"], + ] { + let sql = format!( + "CREATE TABLE ice (ts TIMESTAMP, id INT, color TEXT) WITH (connector = 'iceberg') PARTITIONED BY ({})", + expressions.join(", ") + ); + let statement = dialects.verified_stmt(&sql); + let Statement::CreateTable(table) = &statement else { + panic!("expected CREATE TABLE"); + }; + assert_eq!( + table.arroyo_partitions, + Some( + expressions + .iter() + .map(|e| arroyo().verified_expr(e)) + .collect() + ) + ); + assert_eq!(table.hive_distribution, HiveDistributionStyle::NONE); + let builder = CreateTableBuilder::try_from(statement.clone()).unwrap(); + assert_eq!(builder.build(), *table); + let rebuilt = CreateTableBuilder::from(table.clone()) + .arroyo_partitions(None) + .arroyo_partitions(table.arroyo_partitions.clone()) + .build(); + assert_eq!(rebuilt, *table); + assert!(Parser::parse_sql(&PostgreSqlDialect {}, &sql).is_err()); + } +} + +#[test] +fn connector_partitions_preserve_source_span() { + let sql = r#"CREATE TABLE ice (color TEXT) +WITH (connector = 'iceberg') +PARTITIONED BY (color)"#; + let Statement::CreateTable(table) = + Parser::parse_sql(&ArroyoDialect {}, sql).unwrap().remove(0) + else { + panic!("expected CREATE TABLE"); + }; + let partitions = table.arroyo_partitions.as_ref().unwrap(); + assert_eq!(partitions[0].span().start.line, 3); + assert_eq!(table.span().end, partitions[0].span().end); +} + +#[test] +fn generic_options_precede_connector_partitions_when_formatted() { + let sql = "CREATE TABLE ice (id INT) OPTIONS(foo = 'bar') PARTITIONED BY (bucket(32, id))"; + let Statement::CreateTable(table) = + TestedDialects::new(vec![Box::new(GenericDialect)]).verified_stmt(sql) + else { + panic!("expected CREATE TABLE"); + }; + assert!(table.arroyo_partitions.is_some()); + assert_eq!(table.hive_distribution, HiveDistributionStyle::NONE); +} + +#[test] +fn hive_partition_columns_are_unchanged() { + let dialects = TestedDialects::new(vec![ + Box::new(ArroyoDialect {}), + Box::new(GenericDialect), + Box::new(HiveDialect {}), + ]); + for sql in [ + "CREATE TABLE events (id INT) PARTITIONED BY (region STRING)", + "CREATE TABLE events (id INT) PARTITIONED BY (region)", + ] { + let Statement::CreateTable(table) = dialects.verified_stmt(sql) else { + panic!("expected CREATE TABLE"); + }; + assert!(table.arroyo_partitions.is_none()); + assert!(matches!( + table.hive_distribution, + HiveDistributionStyle::PARTITIONED { .. } + )); + } + let Statement::CreateTable(table) = arroyo().verified_stmt("CREATE TABLE events (id INT)") + else { + panic!("expected CREATE TABLE"); + }; + assert!(table.arroyo_partitions.is_none()); +} + +#[test] +fn invalid_connector_partitions() { + for suffix in [ + "PARTITIONED BY ()", + "PARTITIONED BY hour(ts)", + "PARTITIONED BY (hour(ts),)", + "PARTITIONED BY (hour(ts)) PARTITIONED BY (ts)", + ] { + let sql = format!("CREATE TABLE ice (ts TIMESTAMP) WITH (connector = 'iceberg') {suffix}"); + assert!(Parser::parse_sql(&ArroyoDialect {}, &sql).is_err(), "{sql}"); + } +} diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index 221c88971a..36e953ac10 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -1719,7 +1719,7 @@ fn parse_json_ops_without_colon() { ]; for (str_op, op, dialects) in binary_ops { - let select = dialects.verified_only_select(&format!("SELECT a {} b", &str_op)); + let select = dialects.verified_only_select(&format!("SELECT a {} b", str_op)); assert_eq!( SelectItem::UnnamedExpr(Expr::BinaryOp { left: Box::new(Expr::Identifier(Ident::new("a"))), @@ -2415,7 +2415,7 @@ fn parse_bitwise_ops() { ]; for (str_op, op, dialects) in bitwise_ops { - let select = dialects.verified_only_select(&format!("SELECT a {} b", &str_op)); + let select = dialects.verified_only_select(&format!("SELECT a {} b", str_op)); assert_eq!( SelectItem::UnnamedExpr(Expr::BinaryOp { left: Box::new(Expr::Identifier(Ident::new("a"))), @@ -18560,7 +18560,7 @@ fn parse_generic_unary_ops() { ("+", UnaryOperator::Plus), ]; for (str_op, op) in unary_ops { - let select = verified_only_select(&format!("SELECT {}expr", &str_op)); + let select = verified_only_select(&format!("SELECT {}expr", str_op)); assert_eq!( UnnamedExpr(UnaryOp { op: *op, diff --git a/tests/sqlparser_duckdb.rs b/tests/sqlparser_duckdb.rs index df62685808..f23eec48a4 100644 --- a/tests/sqlparser_duckdb.rs +++ b/tests/sqlparser_duckdb.rs @@ -794,6 +794,7 @@ fn test_duckdb_union_datatype() { distkey: Default::default(), sortkey: Default::default(), backup: Default::default(), + arroyo_partitions: None, }), stmt ); diff --git a/tests/sqlparser_mssql.rs b/tests/sqlparser_mssql.rs index 1e053da78c..218319c58b 100644 --- a/tests/sqlparser_mssql.rs +++ b/tests/sqlparser_mssql.rs @@ -2013,6 +2013,7 @@ fn parse_create_table_with_valid_options() { distkey: None, sortkey: None, backup: None, + arroyo_partitions: None, }) ); } @@ -2187,6 +2188,7 @@ fn parse_create_table_with_identity_column() { distkey: None, sortkey: None, backup: None, + arroyo_partitions: None, }), ); } diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 86315b1ef9..2177665793 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -2461,7 +2461,7 @@ fn parse_pg_unary_ops() { ("@", UnaryOperator::PGAbs), ]; for (str_op, op) in pg_unary_ops { - let select = pg().verified_only_select(&format!("SELECT {}a", &str_op)); + let select = pg().verified_only_select(&format!("SELECT {}a", str_op)); assert_eq!( SelectItem::UnnamedExpr(Expr::UnaryOp { op: *op, @@ -2477,7 +2477,7 @@ fn parse_pg_postfix_factorial() { let postfix_factorial = &[("!", UnaryOperator::PGPostfixFactorial)]; for (str_op, op) in postfix_factorial { - let select = pg().verified_only_select(&format!("SELECT a{}", &str_op)); + let select = pg().verified_only_select(&format!("SELECT a{}", str_op)); assert_eq!( SelectItem::UnnamedExpr(Expr::UnaryOp { op: *op, @@ -6716,6 +6716,7 @@ fn parse_trigger_related_functions() { distkey: None, sortkey: None, backup: None, + arroyo_partitions: None, } );