Skip to content

Commit

Permalink
fix: support the full cache table syntax
Browse files Browse the repository at this point in the history
  • Loading branch information
@francis-du
francis-du committed Oct 14, 2022
1 parent f76cfe6 commit 66c68c9
Show file tree
Hide file tree
Showing 4 changed files with 313 additions and 57 deletions.
71 changes: 48 additions & 23 deletions src/ast/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1430,18 +1430,24 @@ pub enum Statement {
// Specifies the actions to perform when values match or do not match.
clauses: Vec<MergeClause>,
},
/// CACHE TABLE <table_name> [[AS] [QUERY]]
/// 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,
alias: bool,
has_as: bool,
// Table confs
options: Vec<SqlOption>,
// Cache table as a Query
query: Option<Query>,
},
/// UNCACHE TABLE <table_name>
/// UNCACHE TABLE [ IF EXISTS ] <table_name>
UNCache {
// Table name
table_name: ObjectName,
if_exists: bool,
},
}

Expand Down Expand Up @@ -2412,31 +2418,50 @@ impl fmt::Display for Statement {
}
Statement::Cache {
table_name,
alias,
table_flag,
has_as,
options,
query,
} => {
if query.is_some() {
if *alias {
write!(
f,
"CACHE TABLE {table_name} AS {query}",
table_name = table_name,
query = query.clone().unwrap()
)
} else {
write!(
f,
"CACHE TABLE {table_name} {query}",
table_name = table_name,
query = query.clone().unwrap()
)
}
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 {
write!(f, "CACHE TABLE {table_name}", table_name = table_name)
Ok(())
}
}
Statement::UNCache { table_name } => {
write!(f, "UNCACHE TABLE {table_name}", table_name = table_name)
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)
}
}
}
}
Expand Down
1 change: 1 addition & 0 deletions src/keywords.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -383,6 +383,7 @@ define_keywords!(
OPEN,
OPERATOR,
OPTION,
OPTIONS,
OR,
ORC,
ORDER,
Expand Down
106 changes: 82 additions & 24 deletions src/parser.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1909,49 +1909,107 @@ impl<'a> Parser<'a> {

/// Parse a CACHE TABLE statement
pub fn parse_cache_table(&mut self) -> Result<Statement, ParserError> {
let has_table = self.parse_keyword(Keyword::TABLE);
if has_table {
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 {
match self.peek_token() {
Token::Word(word) => match word.keyword {
Keyword::AS => {
self.next_token();
Ok(Statement::Cache {
table_name,
alias: true,
query: Some(self.parse_query()?),
})
}
_ => Ok(Statement::Cache {
table_name,
alias: false,
query: Some(self.parse_query()?),
}),
},
_ => self.expected("a QUERY statement", self.peek_token()),
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,
alias: false,
query: None,
has_as,
options,
query,
})
}
} else {
self.expected("a `TABLE` keyword", self.peek_token())
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 })
Ok(Statement::UNCache {
table_name,
if_exists,
})
} else {
self.expected("an `EOF` keyword", self.peek_token())
self.expected("an `EOF`", self.peek_token())
}
} else {
self.expected("a `TABLE` keyword", self.peek_token())
Expand Down

0 comments on commit 66c68c9

Please sign in to comment.