Support cache/uncache table syntax (#670)

* feat: support cache/uncache table syntax

* fix: support the full cache table syntax

src/ast/mod.rs

@@ -1430,6 +1430,25 @@ pub enum Statement {
         // Specifies the actions to perform when values match or do not match.
         clauses: Vec<MergeClause>,
     },
+    /// CACHE [ FLAG ] TABLE <table_name> [ OPTIONS('K1' = 'V1', 'K2' = V2) ] [ AS ] [ <query> ]
+    /// Based on Spark SQL,see <https://docs.databricks.com/spark/latest/spark-sql/language-manual/sql-ref-syntax-aux-cache-cache-table.html>
+    Cache {
+        // Table flag
+        table_flag: Option<ObjectName>,
+        // Table name
+        table_name: ObjectName,
+        has_as: bool,
+        // Table confs
+        options: Vec<SqlOption>,
+        // Cache table as a Query
+        query: Option<Query>,
+    },
+    /// UNCACHE TABLE [ IF EXISTS ]  <table_name>
+    UNCache {
+        // Table name
+        table_name: ObjectName,
+        if_exists: bool,
+    },
 }
 
 impl fmt::Display for Statement {
@@ -2397,6 +2416,53 @@ impl fmt::Display for Statement {
                 write!(f, "ON {} ", on)?;
                 write!(f, "{}", display_separated(clauses, " "))
             }
+            Statement::Cache {
+                table_name,
+                table_flag,
+                has_as,
+                options,
+                query,
+            } => {
+                if table_flag.is_some() {
+                    write!(
+                        f,
+                        "CACHE {table_flag} TABLE {table_name}",
+                        table_flag = table_flag.clone().unwrap(),
+                        table_name = table_name,
+                    )?;
+                } else {
+                    write!(f, "CACHE TABLE {table_name}", table_name = table_name,)?;
+                }
+
+                if !options.is_empty() {
+                    write!(f, " OPTIONS({})", display_comma_separated(options))?;
+                }
+
+                let has_query = query.is_some();
+                if *has_as && has_query {
+                    write!(f, " AS {query}", query = query.clone().unwrap())
+                } else if !has_as && has_query {
+                    write!(f, " {query}", query = query.clone().unwrap())
+                } else if *has_as && !has_query {
+                    write!(f, " AS")
+                } else {
+                    Ok(())
+                }
+            }
+            Statement::UNCache {
+                table_name,
+                if_exists,
+            } => {
+                if *if_exists {
+                    write!(
+                        f,
+                        "UNCACHE TABLE IF EXISTS {table_name}",
+                        table_name = table_name
+                    )
+                } else {
+                    write!(f, "UNCACHE TABLE {table_name}", table_name = table_name)
+                }
+            }
         }
     }
 }

src/keywords.rs

@@ -383,6 +383,7 @@ define_keywords!(
     OPEN,
     OPERATOR,
     OPTION,
+    OPTIONS,
     OR,
     ORC,
     ORDER,
@@ -554,6 +555,7 @@ define_keywords!(
     TYPE,
     UESCAPE,
     UNBOUNDED,
+    UNCACHE,
     UNCOMMITTED,
     UNION,
     UNIQUE,

src/parser.rs

@@ -170,12 +170,14 @@ impl<'a> Parser<'a> {
                 Keyword::TRUNCATE => Ok(self.parse_truncate()?),
                 Keyword::MSCK => Ok(self.parse_msck()?),
                 Keyword::CREATE => Ok(self.parse_create()?),
+                Keyword::CACHE => Ok(self.parse_cache_table()?),
                 Keyword::DROP => Ok(self.parse_drop()?),
                 Keyword::DISCARD => Ok(self.parse_discard()?),
                 Keyword::DECLARE => Ok(self.parse_declare()?),
                 Keyword::FETCH => Ok(self.parse_fetch_statement()?),
                 Keyword::DELETE => Ok(self.parse_delete()?),
                 Keyword::INSERT => Ok(self.parse_insert()?),
+                Keyword::UNCACHE => Ok(self.parse_uncache_table()?),
                 Keyword::UPDATE => Ok(self.parse_update()?),
                 Keyword::ALTER => Ok(self.parse_alter()?),
                 Keyword::COPY => Ok(self.parse_copy()?),
@@ -1907,6 +1909,115 @@ impl<'a> Parser<'a> {
         }
     }
 
+    /// Parse a CACHE TABLE statement
+    pub fn parse_cache_table(&mut self) -> Result<Statement, ParserError> {
+        let (mut table_flag, mut options, mut has_as, mut query) = (None, vec![], false, None);
+        if self.parse_keyword(Keyword::TABLE) {
+            let table_name = self.parse_object_name()?;
+            if self.peek_token() != Token::EOF {
+                if let Token::Word(word) = self.peek_token() {
+                    if word.keyword == Keyword::OPTIONS {
+                        options = self.parse_options(Keyword::OPTIONS)?
+                    }
+                };
+
+                if self.peek_token() != Token::EOF {
+                    let (a, q) = self.parse_as_query()?;
+                    has_as = a;
+                    query = Some(q);
+                }
+
+                Ok(Statement::Cache {
+                    table_flag,
+                    table_name,
+                    has_as,
+                    options,
+                    query,
+                })
+            } else {
+                Ok(Statement::Cache {
+                    table_flag,
+                    table_name,
+                    has_as,
+                    options,
+                    query,
+                })
+            }
+        } else {
+            table_flag = Some(self.parse_object_name()?);
+            if self.parse_keyword(Keyword::TABLE) {
+                let table_name = self.parse_object_name()?;
+                if self.peek_token() != Token::EOF {
+                    if let Token::Word(word) = self.peek_token() {
+                        if word.keyword == Keyword::OPTIONS {
+                            options = self.parse_options(Keyword::OPTIONS)?
+                        }
+                    };
+
+                    if self.peek_token() != Token::EOF {
+                        let (a, q) = self.parse_as_query()?;
+                        has_as = a;
+                        query = Some(q);
+                    }
+
+                    Ok(Statement::Cache {
+                        table_flag,
+                        table_name,
+                        has_as,
+                        options,
+                        query,
+                    })
+                } else {
+                    Ok(Statement::Cache {
+                        table_flag,
+                        table_name,
+                        has_as,
+                        options,
+                        query,
+                    })
+                }
+            } else {
+                if self.peek_token() == Token::EOF {
+                    self.prev_token();
+                }
+                self.expected("a `TABLE` keyword", self.peek_token())
+            }
+        }
+    }
+
+    /// Parse 'AS' before as query,such as `WITH XXX AS SELECT XXX` oer `CACHE TABLE AS SELECT XXX`
+    pub fn parse_as_query(&mut self) -> Result<(bool, Query), ParserError> {
+        match self.peek_token() {
+            Token::Word(word) => match word.keyword {
+                Keyword::AS => {
+                    self.next_token();
+                    Ok((true, self.parse_query()?))
+                }
+                _ => Ok((false, self.parse_query()?)),
+            },
+            _ => self.expected("a QUERY statement", self.peek_token()),
+        }
+    }
+
+    /// Parse a UNCACHE TABLE statement
+    pub fn parse_uncache_table(&mut self) -> Result<Statement, ParserError> {
+        let has_table = self.parse_keyword(Keyword::TABLE);
+        if has_table {
+            let if_exists = self.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);
+            let table_name = self.parse_object_name()?;
+            if self.peek_token() == Token::EOF {
+                Ok(Statement::UNCache {
+                    table_name,
+                    if_exists,
+                })
+            } else {
+                self.expected("an `EOF`", self.peek_token())
+            }
+        } else {
+            self.expected("a `TABLE` keyword", self.peek_token())
+        }
+    }
+
     /// SQLite-specific `CREATE VIRTUAL TABLE`
     pub fn parse_create_virtual_table(&mut self) -> Result<Statement, ParserError> {
         self.expect_keyword(Keyword::TABLE)?;