Merge pull request MaterializeInc#15 from umanwizard/flush

Parse `FLUSH` statements

Author: umanwizard

src/ast/mod.rs

@@ -531,6 +531,12 @@ pub enum Statement {
         url: String,
         with_options: Vec<SqlOption>,
     },
+    /// `FLUSH SOURCE`
+    FlushSource {
+        name: ObjectName
+    },
+    /// `FLUSH ALL SOURCES`
+    FlushAllSources,
     /// `CREATE VIEW`
     CreateView {
         /// View name
@@ -931,6 +937,8 @@ impl fmt::Display for Statement {
             }
             Statement::Tail { name } => write!(f, "TAIL {}", name),
             Statement::Explain { stage, query } => write!(f, "EXPLAIN {} FOR {}", stage, query),
+            Statement::FlushSource { name } => write!(f, "FLUSH SOURCE {}", name),
+            Statement::FlushAllSources => write!(f, "FLUSH ALL SOURCES"),
         }
     }
 }

src/ast/visit_macro.rs

@@ -562,6 +562,12 @@ macro_rules! make_visitor {
             fn visit_explain(&mut self, stage: &'ast $($mut)* Stage, query: &'ast $($mut)* Query) {
                 visit_explain(self, stage, query)
             }
+            fn visit_flush(&mut self, name: &'ast $($mut)* ObjectName) {
+                visit_flush(self, name)
+            }
+            fn visit_flush_all(&mut self) {
+                visit_flush_all(self)
+            }
         }
 
         pub fn visit_statement<'ast, V: $name<'ast> + ?Sized>(visitor: &mut V, statement: &'ast $($mut)* Statement) {
@@ -665,6 +671,8 @@ macro_rules! make_visitor {
                     visitor.visit_tail(name);
                 }
                 Statement::Explain { stage, query } => visitor.visit_explain(stage, query),
+                Statement::FlushSource { name } => visitor.visit_flush(name),
+                Statement::FlushAllSources => visitor.visit_flush_all(),
             }
         }
 
@@ -1595,6 +1603,12 @@ macro_rules! make_visitor {
         pub fn visit_explain<'ast, V: $name<'ast> + ?Sized>(visitor: &mut V, _stage: &'ast $($mut)* Stage, query: &'ast $($mut)* Query) {
             visitor.visit_query(query);
         }
+
+        pub fn visit_flush<'ast, V: $name<'ast> + ?Sized>(visitor: &mut V, name: &'ast $($mut)* ObjectName) {
+            visitor.visit_object_name(name);
+        }
+
+        pub fn visit_flush_all<'ast, V: $name<'ast> + ?Sized>(_visitor: &mut V) {}
     }
 }
 

src/dialect/keywords.rs

@@ -180,6 +180,7 @@ define_keywords!(
     FIRST_VALUE,
     FLOAT,
     FLOOR,
+    FLUSH,
     FOLLOWING,
     FOR,
     FOREIGN,

src/parser.rs

@@ -149,6 +149,7 @@ impl Parser {
                         name: self.parse_object_name()?,
                     }),
                     "EXPLAIN" => Ok(self.parse_explain()?),
+                    "FLUSH" => Ok(self.parse_flush()?),
                     _ => parser_err!(format!(
                         "Unexpected keyword {:?} at the beginning of a statement",
                         w.to_string()
@@ -2383,6 +2384,23 @@ impl Parser {
             query: Box::new(self.parse_query()?),
         })
     }
+
+    /// Parse a statement like `FLUSH SOURCE foo` or `FLUSH ALL SOURCES`,
+    /// assuming that the `FLUSH` token has already been consumed.
+    ///
+    /// This causes the source (or sources) to downgrade their capability(-ies),
+    /// promising not to send any new data for the current timestamp
+    pub fn parse_flush(&mut self) -> Result<Statement, ParserError> {
+        if self.parse_keywords(vec!["ALL", "SOURCES"]) {
+            Ok(Statement::FlushAllSources)
+        } else if self.parse_keyword("SOURCE") {
+            Ok(Statement::FlushSource {
+                name: self.parse_object_name()?
+            })
+        } else {
+            self.expected("ALL SOURCES or SOURCE", self.peek_token())?
+        }
+    }
 }
 
 impl Word {

tests/sqlparser_common.rs

@@ -3320,6 +3320,23 @@ fn parse_explain() {
     );
 }
 
+#[test]
+fn parse_flush() {
+    let ast = verified_stmt("FLUSH ALL SOURCES");
+    assert_eq!(
+        ast,
+        Statement::FlushAllSources,
+    );
+
+    let ast = verified_stmt("FLUSH SOURCE foo");
+    assert_eq!(
+        ast,
+        Statement::FlushSource {
+            name: ObjectName(vec![Ident::new("foo")])
+        }
+    );
+}
+
 fn parse_sql_statements(sql: &str) -> Result<Vec<Statement>, ParserError> {
     all_dialects().parse_sql_statements(sql)
 }