Port Arroyo SQL extensions to sqlparser 0.62

src/ast/ddl.rs

@@ -1939,6 +1939,9 @@ pub enum ColumnOption {
     Comment(String),
     /// `ON UPDATE <expr>` 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 }`
     /// <https://docs.aws.amazon.com/redshift/latest/dg/r_CREATE_TABLE_NEW.html>
     pub backup: Option<bool>,
+    /// Arroyo connector partition expressions, following table options:
+    /// `WITH (...) PARTITIONED BY (hour(ts), bucket(32, id), region)`.
+    pub arroyo_partitions: Option<Vec<Expr>>,
 }
 
 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}'")?;

src/ast/helpers/stmt_create_table.rs

@@ -183,6 +183,8 @@ pub struct CreateTableBuilder {
     pub sortkey: Option<Vec<Expr>>,
     /// Redshift `BACKUP` option.
     pub backup: Option<bool>,
+    /// Arroyo connector partition expressions.
+    pub arroyo_partitions: Option<Vec<Expr>>,
 }
 
 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<Vec<Expr>>) -> 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<CreateTable> for CreateTableBuilder {
             distkey: table.distkey,
             sortkey: table.sortkey,
             backup: table.backup,
+            arroyo_partitions: table.arroyo_partitions,
         }
     }
 }

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(),

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<Expr>,
+    },
 }
 
 impl From<UniqueConstraint> 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(())
+            }
         }
     }
 }

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<char> {
+        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<Result<u8, ParserError>> {
+        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
+    }
+}

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<Result<Statement, ParserError>> {
         if parser.parse_keyword(Keyword::BEGIN) {
             if parser.peek_keyword(Keyword::TRANSACTION)

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 == '`'
     }