Merge pull request #5 from MaterializeInc/peek-immediate

Support PEEK IMMEDIATE

Author: benesch

src/ast/mod.rs

@@ -452,8 +452,11 @@ pub enum Statement {
     Commit { chain: bool },
     /// ROLLBACK [ TRANSACTION | WORK ] [ AND [ NO ] CHAIN ]
     Rollback { chain: bool },
-    /// PEEK
-    Peek { name: ObjectName },
+    /// PEEK [ IMMEDIATE ]
+    Peek {
+        name: ObjectName,
+        immediate: bool,
+    },
     /// TAIL
     Tail { name: ObjectName },
     /// The mysql-ish `SHOW COLUMNS FROM`
@@ -692,7 +695,13 @@ impl fmt::Display for Statement {
             Statement::Rollback { chain } => {
                 write!(f, "ROLLBACK{}", if *chain { " AND CHAIN" } else { "" },)
             }
-            Statement::Peek { name } => write!(f, "PEEK {}", name),
+            Statement::Peek { name, immediate } => {
+                f.write_str("PEEK ")?;
+                if *immediate {
+                    f.write_str("IMMEDIATE ")?;
+                }
+                write!(f, "{}", name)
+            },
             Statement::Show { object_type } => {
                 use ObjectType::*;
                 write!(

src/ast/visit.rs

@@ -470,8 +470,8 @@ pub trait Visit<'ast> {
 
     fn visit_rollback(&mut self, _chain: bool) {}
 
-    fn visit_peek(&mut self, name: &'ast ObjectName) {
-        visit_peek(self, name)
+    fn visit_peek(&mut self, name: &'ast ObjectName, immediate: bool) {
+        visit_peek(self, name, immediate)
     }
 
     fn visit_show_columns(&mut self, name: &'ast ObjectName) {
@@ -559,8 +559,8 @@ pub fn visit_statement<'ast, V: Visit<'ast> + ?Sized>(visitor: &mut V, statement
         Statement::SetTransaction { modes } => visitor.visit_set_transaction(modes),
         Statement::Commit { chain } => visitor.visit_commit(*chain),
         Statement::Rollback { chain } => visitor.visit_rollback(*chain),
-        Statement::Peek { name } => {
-            visitor.visit_peek(name);
+        Statement::Peek { name, immediate } => {
+            visitor.visit_peek(name, *immediate);
         }
         Statement::Show { object_type } => visitor.visit_show(*object_type),
         Statement::ShowColumns { table_name } => visitor.visit_show_columns(table_name),
@@ -1397,7 +1397,7 @@ pub fn visit_transaction_mode<'ast, V: Visit<'ast> + ?Sized>(
     }
 }
 
-pub fn visit_peek<'ast, V: Visit<'ast> + ?Sized>(visitor: &mut V, name: &'ast ObjectName) {
+pub fn visit_peek<'ast, V: Visit<'ast> + ?Sized>(visitor: &mut V, name: &'ast ObjectName, _immediate: bool) {
     visitor.visit_object_name(name);
 }
 

src/dialect/keywords.rs

@@ -197,6 +197,7 @@ define_keywords!(
     HOUR,
     IDENTITY,
     IF,
+    IMMEDIATE,
     IN,
     INDICATOR,
     INNER,

src/parser.rs

@@ -135,6 +135,7 @@ impl Parser {
                     "COMMIT" => Ok(self.parse_commit()?),
                     "ROLLBACK" => Ok(self.parse_rollback()?),
                     "PEEK" => Ok(Statement::Peek {
+                        immediate: self.parse_keyword("IMMEDIATE"),
                         name: self.parse_object_name()?,
                     }),
                     "SHOW" => Ok(self.parse_show()?),

tests/sqlparser_common.rs

@@ -2354,8 +2354,18 @@ fn parse_drop_source() {
 fn parse_peek() {
     let sql = "PEEK foo.bar";
     match verified_stmt(sql) {
-        Statement::Peek { name } => {
+        Statement::Peek { name, immediate } => {
             assert_eq!("foo.bar", name.to_string());
+            assert!(!immediate);
+        }
+        _ => assert!(false),
+    }
+
+    let sql = "PEEK IMMEDIATE foo.bar";
+    match verified_stmt(sql) {
+        Statement::Peek { name, immediate } => {
+            assert_eq!("foo.bar", name.to_string());
+            assert!(immediate);
         }
         _ => assert!(false),
     }