From 9848a9065a7a0bfe975c9452f06fe5a620b9bc35 Mon Sep 17 00:00:00 2001 From: Bruno Roy Date: Thu, 10 Sep 2026 17:22:40 -0700 Subject: [PATCH 1/4] [Arroyo] Port streaming SQL dialect to sqlparser 0.62 Preserve PostgreSQL-style expressions and Arroyo capabilities using dialect hooks instead of additional concrete dialect checks. Keep Arroyo type identity so generated columns do not require STORED; cover this behavior with a PostgreSQL comparison regression. Ported-from: 32b9bbe958cf44efc01d22c296e4ac1b51cd960c Ported-from: 83877ac25d20876cd515c31f04d14ba6b9d5d5d5 --- src/dialect/arroyo.rs | 142 ++++++++++++++++++++++++++++++++++++++ src/dialect/bigquery.rs | 4 ++ src/dialect/generic.rs | 16 +++++ src/dialect/mod.rs | 38 ++++++++++ src/dialect/postgresql.rs | 28 ++++++++ src/dialect/redshift.rs | 8 +++ src/parser/mod.rs | 21 +++--- tests/sqlparser_arroyo.rs | 50 ++++++++++++++ 8 files changed, 294 insertions(+), 13 deletions(-) create mode 100644 src/dialect/arroyo.rs create mode 100644 tests/sqlparser_arroyo.rs diff --git a/src/dialect/arroyo.rs b/src/dialect/arroyo.rs new file mode 100644 index 0000000000..a2b7f15358 --- /dev/null +++ b/src/dialect/arroyo.rs @@ -0,0 +1,142 @@ +// 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 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..da91f2fe54 100644 --- a/src/dialect/generic.rs +++ b/src/dialect/generic.rs @@ -24,6 +24,22 @@ use crate::dialect::Dialect; pub struct GenericDialect; impl Dialect for GenericDialect { + 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..1f0dff49e4 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,41 @@ pub trait Dialect: Debug + Any { 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 +1883,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/parser/mod.rs b/src/parser/mod.rs index 668c520e5e..55faf9a2ac 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), @@ -12220,7 +12216,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 +16115,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..fefec5add5 --- /dev/null +++ b/tests/sqlparser_arroyo.rs @@ -0,0 +1,50 @@ +// 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::{BinaryOperator, Expr}; +use sqlparser::dialect::{dialect_from_str, ArroyoDialect, 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()); +} From 40f51724fb88c18606e6b457d642b66daf6ba6ac Mon Sep 17 00:00:00 2001 From: Bruno Roy Date: Thu, 10 Sep 2026 17:26:53 -0700 Subject: [PATCH 2/4] [Arroyo] Port WATERMARK table constraints Preserve optional watermark expressions, constraint spans and formatting. Gate the syntax with a dialect capability, enabled for Arroyo and Generic; cover missing syntax, named constraints and PostgreSQL rejection. Ported-from: 2279ac9d852b05187129fbbf82ea48fad33261dc --- src/ast/spans.rs | 6 +++++ src/ast/table_constraints.rs | 18 +++++++++++++ src/dialect/arroyo.rs | 4 +++ src/dialect/generic.rs | 4 +++ src/dialect/mod.rs | 5 ++++ src/keywords.rs | 1 + src/parser/mod.rs | 25 ++++++++++++++++++ tests/sqlparser_arroyo.rs | 50 ++++++++++++++++++++++++++++++++++-- 8 files changed, 111 insertions(+), 2 deletions(-) diff --git a/src/ast/spans.rs b/src/ast/spans.rs index 0dc834ba03..4344c23c7f 100644 --- a/src/ast/spans.rs +++ b/src/ast/spans.rs @@ -650,6 +650,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)), } } } 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 index a2b7f15358..66912196f7 100644 --- a/src/dialect/arroyo.rs +++ b/src/dialect/arroyo.rs @@ -20,6 +20,10 @@ use crate::parser::{Parser, ParserError}; pub struct ArroyoDialect {} impl Dialect for ArroyoDialect { + fn supports_watermark_constraint(&self) -> bool { + true + } + fn identifier_quote_style(&self, _identifier: &str) -> Option { Some('"') } diff --git a/src/dialect/generic.rs b/src/dialect/generic.rs index da91f2fe54..16ce871f4e 100644 --- a/src/dialect/generic.rs +++ b/src/dialect/generic.rs @@ -24,6 +24,10 @@ use crate::dialect::Dialect; pub struct GenericDialect; impl Dialect for GenericDialect { + fn supports_watermark_constraint(&self) -> bool { + true + } + fn supports_array_overlap_operator(&self) -> bool { true } diff --git a/src/dialect/mod.rs b/src/dialect/mod.rs index 1f0dff49e4..08375389cc 100644 --- a/src/dialect/mod.rs +++ b/src/dialect/mod.rs @@ -1150,6 +1150,11 @@ pub trait Dialect: Debug + Any { 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 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 55faf9a2ac..efc1eaf1e3 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -9922,6 +9922,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) diff --git a/tests/sqlparser_arroyo.rs b/tests/sqlparser_arroyo.rs index fefec5add5..38c22ffebf 100644 --- a/tests/sqlparser_arroyo.rs +++ b/tests/sqlparser_arroyo.rs @@ -10,8 +10,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use sqlparser::ast::{BinaryOperator, Expr}; -use sqlparser::dialect::{dialect_from_str, ArroyoDialect, PostgreSqlDialect}; +use sqlparser::ast::{BinaryOperator, Expr, Spanned, Statement, TableConstraint}; +use sqlparser::dialect::{dialect_from_str, ArroyoDialect, GenericDialect, PostgreSqlDialect}; use sqlparser::parser::Parser; use sqlparser::test_utils::TestedDialects; @@ -48,3 +48,49 @@ fn generated_columns_do_not_require_stored() { 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}"); + } +} From 8a28842f4ee76ec5f78a8bab200ea25c88489bb1 Mon Sep 17 00:00:00 2001 From: Bruno Roy Date: Thu, 10 Sep 2026 17:29:35 -0700 Subject: [PATCH 3/4] [Arroyo] Port METADATA FROM column options Retain the metadata key and its source span, with a capability gate checked before consuming tokens. Cover escaped keys, round trips, malformed input and rejection by PostgreSQL. Ported-from: 8c1c36b2acab1ec3082805f7210787ac266bd144 --- src/ast/ddl.rs | 6 ++++++ src/ast/spans.rs | 1 + src/dialect/arroyo.rs | 4 ++++ src/dialect/generic.rs | 4 ++++ src/dialect/mod.rs | 5 +++++ src/parser/mod.rs | 10 ++++++++++ tests/sqlparser_arroyo.rs | 42 ++++++++++++++++++++++++++++++++++++++- 7 files changed, 71 insertions(+), 1 deletion(-) diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs index 67aefb3928..1ac69d9c2c 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, diff --git a/src/ast/spans.rs b/src/ast/spans.rs index 4344c23c7f..1c437552b0 100644 --- a/src/ast/spans.rs +++ b/src/ast/spans.rs @@ -838,6 +838,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/dialect/arroyo.rs b/src/dialect/arroyo.rs index 66912196f7..b8b54cc03f 100644 --- a/src/dialect/arroyo.rs +++ b/src/dialect/arroyo.rs @@ -20,6 +20,10 @@ use crate::parser::{Parser, ParserError}; pub struct ArroyoDialect {} impl Dialect for ArroyoDialect { + fn supports_metadata_column_option(&self) -> bool { + true + } + fn supports_watermark_constraint(&self) -> bool { true } diff --git a/src/dialect/generic.rs b/src/dialect/generic.rs index 16ce871f4e..26eeae9f22 100644 --- a/src/dialect/generic.rs +++ b/src/dialect/generic.rs @@ -24,6 +24,10 @@ use crate::dialect::Dialect; pub struct GenericDialect; impl Dialect for GenericDialect { + fn supports_metadata_column_option(&self) -> bool { + true + } + fn supports_watermark_constraint(&self) -> bool { true } diff --git a/src/dialect/mod.rs b/src/dialect/mod.rs index 08375389cc..43623e6e62 100644 --- a/src/dialect/mod.rs +++ b/src/dialect/mod.rs @@ -1150,6 +1150,11 @@ pub trait Dialect: Debug + Any { 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 diff --git a/src/parser/mod.rs b/src/parser/mod.rs index efc1eaf1e3..5c6efb8afb 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -9262,6 +9262,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) { diff --git a/tests/sqlparser_arroyo.rs b/tests/sqlparser_arroyo.rs index 38c22ffebf..1c15eeba6b 100644 --- a/tests/sqlparser_arroyo.rs +++ b/tests/sqlparser_arroyo.rs @@ -10,7 +10,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use sqlparser::ast::{BinaryOperator, Expr, Spanned, Statement, TableConstraint}; +use sqlparser::ast::{BinaryOperator, ColumnOption, Expr, Spanned, Statement, TableConstraint}; use sqlparser::dialect::{dialect_from_str, ArroyoDialect, GenericDialect, PostgreSqlDialect}; use sqlparser::parser::Parser; use sqlparser::test_utils::TestedDialects; @@ -94,3 +94,43 @@ fn invalid_watermark_constraints() { 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}"); + } +} From 887b4eaed88d5bae9c7749ca4d78b8ac61679665 Mon Sep 17 00:00:00 2001 From: Bruno Roy Date: Thu, 10 Sep 2026 17:31:10 -0700 Subject: [PATCH 4/4] [Arroyo] Port connector PARTITIONED BY expressions Retain the arroyo_partitions AST field and builder conversions for Iceberg transforms and identity partitions. Keep the existing clause position after table options, preserving upstream Hive partition column parsing. Include source spans and regression tests for round trips, builder conversions, dialect boundaries and malformed clauses. Ported-from: 7086ac7dd6070c49b9924c8e3bf575cd5d7b15cc --- src/ast/ddl.rs | 24 +++++- src/ast/helpers/stmt_create_table.rs | 11 +++ src/ast/spans.rs | 4 +- src/dialect/arroyo.rs | 4 + src/dialect/generic.rs | 4 + src/dialect/mod.rs | 6 ++ src/parser/mod.rs | 14 ++++ tests/sqlparser_arroyo.rs | 113 ++++++++++++++++++++++++++- tests/sqlparser_duckdb.rs | 1 + tests/sqlparser_mssql.rs | 2 + tests/sqlparser_postgres.rs | 1 + 11 files changed, 179 insertions(+), 5 deletions(-) diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs index 1ac69d9c2c..29c0cfbac6 100644 --- a/src/ast/ddl.rs +++ b/src/ast/ddl.rs @@ -3066,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 { @@ -3256,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 1c437552b0..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)), ) } } diff --git a/src/dialect/arroyo.rs b/src/dialect/arroyo.rs index b8b54cc03f..40c02a2691 100644 --- a/src/dialect/arroyo.rs +++ b/src/dialect/arroyo.rs @@ -20,6 +20,10 @@ use crate::parser::{Parser, ParserError}; pub struct ArroyoDialect {} impl Dialect for ArroyoDialect { + fn supports_partitioned_by_expressions(&self) -> bool { + true + } + fn supports_metadata_column_option(&self) -> bool { true } diff --git a/src/dialect/generic.rs b/src/dialect/generic.rs index 26eeae9f22..e4fd4f71d6 100644 --- a/src/dialect/generic.rs +++ b/src/dialect/generic.rs @@ -24,6 +24,10 @@ 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 } diff --git a/src/dialect/mod.rs b/src/dialect/mod.rs index 43623e6e62..7843a2224f 100644 --- a/src/dialect/mod.rs +++ b/src/dialect/mod.rs @@ -1150,6 +1150,12 @@ 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 diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 5c6efb8afb..25efb3f6b7 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -8571,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 { @@ -8653,6 +8666,7 @@ impl<'a> Parser<'a> { .diststyle(diststyle) .distkey(distkey) .sortkey(sortkey) + .arroyo_partitions(arroyo_partitions) .build()) } diff --git a/tests/sqlparser_arroyo.rs b/tests/sqlparser_arroyo.rs index 1c15eeba6b..aa04e6d65e 100644 --- a/tests/sqlparser_arroyo.rs +++ b/tests/sqlparser_arroyo.rs @@ -10,8 +10,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -use sqlparser::ast::{BinaryOperator, ColumnOption, Expr, Spanned, Statement, TableConstraint}; -use sqlparser::dialect::{dialect_from_str, ArroyoDialect, GenericDialect, PostgreSqlDialect}; +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; @@ -134,3 +139,107 @@ fn invalid_metadata_fields() { 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, } );