Add support for WATERMARK FOR table constraint syntax

Author: mwylde

src/ast/ddl.rs

@@ -988,6 +988,21 @@ pub enum TableConstraint {
         /// Referred column identifier list.
         columns: Vec<Ident>,
     },
+    /// 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<Expr>,
+    },
 }
 
 impl fmt::Display for TableConstraint {
@@ -1115,6 +1130,16 @@ impl fmt::Display for TableConstraint {
 
                 Ok(())
             }
+            Self::Watermark {
+                column_name,
+                watermark_expr,
+            } => {
+                write!(f, "WATERMARK FOR {}", column_name)?;
+                if let Some(expr) = watermark_expr {
+                    write!(f, " AS {}", expr)?;
+                }
+                Ok(())
+            }
         }
     }
 }

src/ast/spans.rs

@@ -695,6 +695,9 @@ impl Spanned for TableConstraint {
                     .map(|i| i.span)
                     .chain(columns.iter().map(|i| i.span)),
             ),
+            TableConstraint::Watermark { column_name, watermark_expr } => {
+                union_spans(watermark_expr.iter().map(|e| e.span()).chain(core::iter::once(column_name.span)))
+            }
         }
     }
 }

src/keywords.rs

@@ -952,6 +952,7 @@ define_keywords!(
     VOLUME,
     WAREHOUSE,
     WAREHOUSES,
+    WATERMARK,
     WEEK,
     WEEKS,
     WHEN,

src/parser/mod.rs

@@ -7566,6 +7566,35 @@ impl<'a> Parser<'a> {
                     columns,
                 }))
             }
+            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(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)

tests/sqlparser_arroyo.rs

@@ -0,0 +1,83 @@
+// 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.
+
+#![warn(clippy::all)]
+
+use sqlparser::ast::{BinaryOperator, Expr, Ident, Statement, TableConstraint};
+use sqlparser::dialect::ArroyoDialect;
+use sqlparser::parser::Parser;
+use sqlparser::test_utils;
+use sqlparser::tokenizer::{Location, Span};
+
+#[test]
+fn test_watermark_with_expr() {
+    let sql = "CREATE TABLE orders (
+        customer_id INT,
+        order_id INT,
+        date_string TEXT,
+        timestamp TIMESTAMP GENERATED ALWAYS AS (CAST(date_string as TIMESTAMP)),
+        WATERMARK FOR timestamp AS timestamp + 5
+    ) WITH (
+        connector = 'kafka',
+        format = 'json',
+        type = 'source',
+        bootstrap_servers = 'localhost:9092',
+        topic = 'order_topic'
+    )";
+
+    let parse = Parser::parse_sql(&ArroyoDialect {}, sql).unwrap();
+    let Statement::CreateTable(ct) = parse.get(0).unwrap() else {
+        panic!("not create table")
+    };
+
+    assert_eq!(
+        ct.constraints,
+        vec![TableConstraint::Watermark {
+            column_name: Ident::new("timestamp"),
+            watermark_expr: Some(Expr::BinaryOp {
+                left: Box::new(Expr::Identifier(Ident::new("timestamp"))),
+                op: BinaryOperator::Plus,
+                right: Box::new(Expr::Value(test_utils::number("5").with_span(Span::new(
+                    Location::new(5, 4), Location::new(5, 10)
+                )))),
+            }),
+        }]
+    );
+}
+
+#[test]
+fn test_watermark_without_expr() {
+    let sql = "CREATE TABLE users (
+        customer_id INT,
+        timestamp TIMESTAMP,
+        WATERMARK FOR timestamp
+    ) WITH (
+        connector = 'kafka',
+        format = 'json',
+        type = 'source',
+        bootstrap_servers = 'localhost:9092',
+        topic = 'order_topic'
+    )";
+
+    let parse = Parser::parse_sql(&ArroyoDialect {}, sql).unwrap();
+    let Statement::CreateTable(ct) = parse.get(0).unwrap() else {
+        panic!("not create table")
+    };
+
+    assert_eq!(
+        ct.constraints,
+        vec![TableConstraint::Watermark {
+            column_name: Ident::new("timestamp"),
+            watermark_expr: None,
+        }]
+    );
+}