diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs index 67aefb3928..29c0cfbac6 100644 --- a/src/ast/ddl.rs +++ b/src/ast/ddl.rs @@ -1939,6 +1939,9 @@ pub enum ColumnOption { Comment(String), /// `ON UPDATE ` column option OnUpdate(Expr), + /// `METADATA FROM 'key'`: a column populated from record metadata in Arroyo. + /// The span covers the quoted metadata key. + MetadataField(String, Span), /// `Generated`s are modifiers that follow a column definition in a `CREATE /// TABLE` statement. Generated { @@ -2085,6 +2088,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 +3066,9 @@ pub struct CreateTable { /// Redshift `BACKUP` option: `BACKUP { YES | NO }` /// pub backup: Option, + /// Arroyo connector partition expressions, following table options: + /// `WITH (...) PARTITIONED BY (hour(ts), bucket(32, id), region)`. + pub arroyo_partitions: Option>, } impl fmt::Display for CreateTable { @@ -3250,8 +3259,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}'")?; 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/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..b95d4a3ca9 100644 --- a/src/ast/table_constraints.rs +++ b/src/ast/table_constraints.rs @@ -117,6 +117,14 @@ pub enum TableConstraint { /// /// [1]: https://www.postgresql.org/docs/current/sql-altertable.html UniqueUsingIndex(ConstraintUsingIndex), + /// Arroyo's watermark definition for a streaming table: + /// `WATERMARK FOR timestamp [AS timestamp - INTERVAL '5 seconds']`. + Watermark { + /// Column to use for the watermark. + column_name: Ident, + /// Optional expression used to compute the watermark. + watermark_expr: Option, + }, } impl From for TableConstraint { @@ -166,6 +174,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..40c02a2691 --- /dev/null +++ b/src/dialect/arroyo.rs @@ -0,0 +1,154 @@ +// 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 crate::dialect::{Dialect, PostgreSqlDialect, Precedence}; +use crate::parser::{Parser, ParserError}; + +/// Arroyo's streaming SQL dialect: PostgreSQL-style expressions with +/// angle-bracket struct types and generated columns that do not require STORED. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] +pub struct ArroyoDialect {} + +impl Dialect for ArroyoDialect { + fn supports_partitioned_by_expressions(&self) -> bool { + true + } + + fn supports_metadata_column_option(&self) -> bool { + true + } + + fn supports_watermark_constraint(&self) -> bool { + true + } + + fn identifier_quote_style(&self, _identifier: &str) -> Option { + Some('"') + } + + fn is_delimited_identifier_start(&self, ch: char) -> bool { + ch == '"' + } + + fn is_identifier_start(&self, ch: char) -> bool { + ch.is_alphabetic() || ch == '_' + } + + fn is_identifier_part(&self, ch: char) -> bool { + ch.is_alphabetic() || ch.is_ascii_digit() || ch == '$' || ch == '_' + } + + fn is_custom_operator_part(&self, ch: char) -> bool { + PostgreSqlDialect {}.is_custom_operator_part(ch) + } + + fn get_next_precedence(&self, parser: &Parser) -> Option> { + PostgreSqlDialect {}.get_next_precedence(parser) + } + + fn prec_value(&self, prec: Precedence) -> u8 { + PostgreSqlDialect {}.prec_value(prec) + } + + fn supports_pg_math_prefix_operators(&self) -> bool { + true + } + + fn supports_caret_exponentiation(&self) -> bool { + true + } + + fn supports_sharp_bitwise_xor(&self) -> bool { + true + } + + fn supports_array_overlap_operator(&self) -> bool { + true + } + + fn supports_starts_with_operator(&self) -> bool { + true + } + + fn supports_escaped_string_literal(&self) -> bool { + true + } + + fn supports_unnest_table_factor(&self) -> bool { + true + } + + 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 + } + + fn supports_factorial_operator(&self) -> bool { + true + } + + fn supports_bitwise_shift_operators(&self) -> bool { + true + } + + fn supports_comment_on(&self) -> bool { + true + } + + 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 + } + + fn supports_array_typedef_with_brackets(&self) -> bool { + true + } + + fn supports_geometric_types(&self) -> bool { + true + } + + fn supports_struct_literal(&self) -> bool { + true + } + + fn supports_insert_table_alias(&self) -> bool { + true + } +} diff --git a/src/dialect/bigquery.rs b/src/dialect/bigquery.rs index 8fca515182..ee07046d67 100644 --- a/src/dialect/bigquery.rs +++ b/src/dialect/bigquery.rs @@ -47,6 +47,10 @@ const RESERVED_FOR_COLUMN_ALIAS: &[Keyword] = &[ pub struct BigQueryDialect; impl Dialect for BigQueryDialect { + fn supports_unnest_table_factor(&self) -> bool { + true + } + fn parse_statement(&self, parser: &mut Parser) -> Option> { if parser.parse_keyword(Keyword::BEGIN) { if parser.peek_keyword(Keyword::TRANSACTION) diff --git a/src/dialect/generic.rs b/src/dialect/generic.rs index 25f57e3d18..e4fd4f71d6 100644 --- a/src/dialect/generic.rs +++ b/src/dialect/generic.rs @@ -24,6 +24,34 @@ use crate::dialect::Dialect; pub struct GenericDialect; impl Dialect for GenericDialect { + fn supports_partitioned_by_expressions(&self) -> bool { + true + } + + fn supports_metadata_column_option(&self) -> bool { + true + } + + fn supports_watermark_constraint(&self) -> bool { + true + } + + fn supports_array_overlap_operator(&self) -> bool { + true + } + + fn supports_starts_with_operator(&self) -> bool { + true + } + + fn supports_escaped_string_literal(&self) -> bool { + true + } + + fn supports_unnest_table_factor(&self) -> bool { + true + } + fn is_delimited_identifier_start(&self, ch: char) -> bool { ch == '"' || ch == '`' } diff --git a/src/dialect/mod.rs b/src/dialect/mod.rs index 6ab6cb15e4..7843a2224f 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; @@ -1148,6 +1150,57 @@ pub trait Dialect: Debug + Any { false } + /// Supports connector partition expressions after table options: + /// `WITH (...) PARTITIONED BY (hour(ts), bucket(32, id))`. + fn supports_partitioned_by_expressions(&self) -> bool { + false + } + + /// Supports `METADATA FROM 'key'` column options. + fn supports_metadata_column_option(&self) -> bool { + false + } + + /// Supports `WATERMARK FOR column [AS expression]` table constraints. + fn supports_watermark_constraint(&self) -> bool { + false + } + + /// Supports PostgreSQL's prefix math operators: `!!`, `|/`, `||/`, and `@`. + fn supports_pg_math_prefix_operators(&self) -> bool { + false + } + + /// Interprets `^` as exponentiation rather than bitwise XOR. + fn supports_caret_exponentiation(&self) -> bool { + false + } + + /// Supports `#` as bitwise XOR. + fn supports_sharp_bitwise_xor(&self) -> bool { + false + } + + /// Supports `&&` as the array-overlap operator. + fn supports_array_overlap_operator(&self) -> bool { + false + } + + /// Supports `^@` as the string starts-with operator. + fn supports_starts_with_operator(&self) -> bool { + false + } + + /// Accepts an escaped string literal token as an expression or literal string. + fn supports_escaped_string_literal(&self) -> bool { + false + } + + /// Supports `UNNEST(...)` as a table factor, including ordinality and offsets. + fn supports_unnest_table_factor(&self) -> bool { + false + } + /// Returns true if the dialect supports nested comments /// e.g. `/* /* nested */ */` fn supports_nested_comments(&self) -> bool { @@ -1846,6 +1899,7 @@ pub fn dialect_from_str(dialect_name: impl AsRef) -> Option Some(Box::new(GenericDialect)), + "arroyo" => Some(Box::new(ArroyoDialect {})), "mysql" => Some(Box::new(MySqlDialect {})), "postgresql" | "postgres" => Some(Box::new(PostgreSqlDialect {})), "hive" => Some(Box::new(HiveDialect {})), diff --git a/src/dialect/postgresql.rs b/src/dialect/postgresql.rs index fda676eb26..07af97e485 100644 --- a/src/dialect/postgresql.rs +++ b/src/dialect/postgresql.rs @@ -59,6 +59,34 @@ const AND_PREC: u8 = 20; const OR_PREC: u8 = 10; impl Dialect for PostgreSqlDialect { + fn supports_pg_math_prefix_operators(&self) -> bool { + true + } + + fn supports_caret_exponentiation(&self) -> bool { + true + } + + fn supports_sharp_bitwise_xor(&self) -> bool { + true + } + + fn supports_array_overlap_operator(&self) -> bool { + true + } + + fn supports_starts_with_operator(&self) -> bool { + true + } + + fn supports_escaped_string_literal(&self) -> bool { + true + } + + fn supports_unnest_table_factor(&self) -> bool { + true + } + fn identifier_quote_style(&self, _identifier: &str) -> Option { Some('"') } diff --git a/src/dialect/redshift.rs b/src/dialect/redshift.rs index db5bc53a0d..c1dbde5913 100644 --- a/src/dialect/redshift.rs +++ b/src/dialect/redshift.rs @@ -33,6 +33,14 @@ pub struct RedshiftSqlDialect {} // in the Postgres dialect, the query will be parsed as an array, while in the Redshift dialect it will // be a json path impl Dialect for RedshiftSqlDialect { + fn supports_sharp_bitwise_xor(&self) -> bool { + true + } + + fn supports_array_overlap_operator(&self) -> bool { + true + } + /// Determine if a character starts a potential nested quoted identifier. /// Example: RedShift supports the following quote styles to all mean the same thing: /// ```sql 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..25efb3f6b7 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.supports_pg_math_prefix_operators() => { let op = match tok { Token::DoubleExclamationMark => UnaryOperator::PGPrefixFactorial, @@ -1898,8 +1898,7 @@ impl<'a> Parser<'a> { ), }) } - Token::EscapedStringLiteral(_) if dialect_is!(dialect is PostgreSqlDialect | GenericDialect) => - { + Token::EscapedStringLiteral(_) if dialect.supports_escaped_string_literal() => { self.prev_token(); Ok(Expr::Value(self.parse_value()?)) } @@ -3762,7 +3761,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.supports_caret_exponentiation() { Some(BinaryOperator::PGExp) } else { Some(BinaryOperator::BitwiseXor) @@ -3779,19 +3778,16 @@ 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.supports_sharp_bitwise_xor() => { 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.supports_array_overlap_operator() => { 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.supports_starts_with_operator() => { Some(BinaryOperator::PGStartsWith) } Token::Tilde => Some(BinaryOperator::PGRegexMatch), @@ -8575,6 +8571,19 @@ impl<'a> Parser<'a> { create_table_config.partition_by }; + // Keep Hive's PARTITIONED BY column definitions above distinct from + // connector partition expressions following WITH table options. + let arroyo_partitions = if self.dialect.supports_partitioned_by_expressions() + && 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 +8666,7 @@ impl<'a> Parser<'a> { .diststyle(diststyle) .distkey(distkey) .sortkey(sortkey) + .arroyo_partitions(arroyo_partitions) .build()) } @@ -9266,6 +9276,16 @@ 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.dialect.supports_metadata_column_option() + && self.parse_keywords(&[Keyword::METADATA, Keyword::FROM]) + { + 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 +9946,31 @@ impl<'a> Parser<'a> { .into(), )) } + Token::Word(w) + if w.keyword == Keyword::WATERMARK + && self.dialect.supports_watermark_constraint() => + { + 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()?; + 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 +12265,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 self.dialect.supports_escaped_string_literal() => { Ok(s) } Token::UnicodeStringLiteral(s) => Ok(s), @@ -16119,8 +16164,7 @@ impl<'a> Parser<'a> { alias, sample: None, }) - } else if dialect_of!(self is BigQueryDialect | PostgreSqlDialect | GenericDialect) - && self.parse_keyword(Keyword::UNNEST) + } else if self.dialect.supports_unnest_table_factor() && self.parse_keyword(Keyword::UNNEST) { self.expect_token(&Token::LParen)?; let array_exprs = self.parse_comma_separated(Parser::parse_expr)?; diff --git a/tests/sqlparser_arroyo.rs b/tests/sqlparser_arroyo.rs new file mode 100644 index 0000000000..aa04e6d65e --- /dev/null +++ b/tests/sqlparser_arroyo.rs @@ -0,0 +1,245 @@ +// 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::{ + dialect_from_str, 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()); + assert!(dialect_from_str("arroyo").unwrap().is::()); +} + +#[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_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..c9bb63f8ec 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -6716,6 +6716,7 @@ fn parse_trigger_related_functions() { distkey: None, sortkey: None, backup: None, + arroyo_partitions: None, } );