mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-10-31 15:17:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Support DECLARE syntax for snowflake and bigquery (#1122)

Browse source
This commit is contained in:
Ifeanyi Ubah 2024-02-29 13:34:00 +01:00 • committed by GitHub
parent f75bb4be20
commit 57113a9344
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 741 additions and 68 deletions
Download patch file Download diff file Expand all files Collapse all files

269
src/ast/mod.rs
View file

@ -1422,6 +1422,211 @@ pub enum Password {
NullPassword,
}
/// Represents an expression assignment within a variable `DECLARE` statement.
///
/// Examples:
/// ```sql
/// DECLARE variable_name := 42
/// DECLARE variable_name DEFAULT 42
/// ```
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum DeclareAssignment {
/// Plain expression specified.
Expr(Box<Expr>),
/// Expression assigned via the `DEFAULT` keyword
Default(Box<Expr>),
/// Expression assigned via the `:=` syntax
///
/// Example:
/// ```sql
/// DECLARE variable_name := 42;
/// ```
DuckAssignment(Box<Expr>),
/// Expression via the `FOR` keyword
///
/// Example:
/// ```sql
/// DECLARE c1 CURSOR FOR res
/// ```
For(Box<Expr>),
}
impl fmt::Display for DeclareAssignment {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
DeclareAssignment::Expr(expr) => {
write!(f, "{expr}")
}
DeclareAssignment::Default(expr) => {
write!(f, "DEFAULT {expr}")
}
DeclareAssignment::DuckAssignment(expr) => {
write!(f, ":= {expr}")
}
DeclareAssignment::For(expr) => {
write!(f, "FOR {expr}")
}
}
}
}
/// Represents the type of a `DECLARE` statement.
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum DeclareType {
/// Cursor variable type. e.g. [Snowflake] [Postgres]
///
/// [Snowflake]: https://docs.snowflake.com/en/developer-guide/snowflake-scripting/cursors#declaring-a-cursor
/// [Postgres]: https://www.postgresql.org/docs/current/plpgsql-cursors.html
Cursor,
/// Result set variable type. [Snowflake]
///
/// Syntax:
/// ```text
/// <resultset_name> RESULTSET [ { DEFAULT | := } ( <query> ) ] ;
/// ```
/// [Snowflake]: https://docs.snowflake.com/en/sql-reference/snowflake-scripting/declare#resultset-declaration-syntax
ResultSet,
/// Exception declaration syntax. [Snowflake]
///
/// Syntax:
/// ```text
/// <exception_name> EXCEPTION [ ( <exception_number> , '<exception_message>' ) ] ;
/// ```
/// [Snowflake]: https://docs.snowflake.com/en/sql-reference/snowflake-scripting/declare#exception-declaration-syntax
Exception,
}
impl fmt::Display for DeclareType {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
DeclareType::Cursor => {
write!(f, "CURSOR")
}
DeclareType::ResultSet => {
write!(f, "RESULTSET")
}
DeclareType::Exception => {
write!(f, "EXCEPTION")
}
}
}
}
/// A `DECLARE` statement.
/// [Postgres] [Snowflake] [BigQuery]
///
/// Examples:
/// ```sql
/// DECLARE variable_name := 42
/// DECLARE liahona CURSOR FOR SELECT * FROM films;
/// ```
///
/// [Postgres]: https://www.postgresql.org/docs/current/sql-declare.html
/// [Snowflake]: https://docs.snowflake.com/en/sql-reference/snowflake-scripting/declare
/// [BigQuery]: https://cloud.google.com/bigquery/docs/reference/standard-sql/procedural-language#declare
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct Declare {
/// The name(s) being declared.
/// Example: `DECLARE a, b, c DEFAULT 42;
pub names: Vec<Ident>,
/// Data-type assigned to the declared variable.
/// Example: `DECLARE x INT64 DEFAULT 42;
pub data_type: Option<DataType>,
/// Expression being assigned to the declared variable.
pub assignment: Option<DeclareAssignment>,
/// Represents the type of the declared variable.
pub declare_type: Option<DeclareType>,
/// Causes the cursor to return data in binary rather than in text format.
pub binary: Option<bool>,
/// None = Not specified
/// Some(true) = INSENSITIVE
/// Some(false) = ASENSITIVE
pub sensitive: Option<bool>,
/// None = Not specified
/// Some(true) = SCROLL
/// Some(false) = NO SCROLL
pub scroll: Option<bool>,
/// None = Not specified
/// Some(true) = WITH HOLD, specifies that the cursor can continue to be used after the transaction that created it successfully commits
/// Some(false) = WITHOUT HOLD, specifies that the cursor cannot be used outside of the transaction that created it
pub hold: Option<bool>,
/// `FOR <query>` clause in a CURSOR declaration.
pub for_query: Option<Box<Query>>,
}
impl fmt::Display for Declare {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let Declare {
names,
data_type,
assignment,
declare_type,
binary,
sensitive,
scroll,
hold,
for_query,
} = self;
write!(f, "{}", display_comma_separated(names))?;
if let Some(true) = binary {
write!(f, " BINARY")?;
}
if let Some(sensitive) = sensitive {
if *sensitive {
write!(f, " INSENSITIVE")?;
} else {
write!(f, " ASENSITIVE")?;
}
}
if let Some(scroll) = scroll {
if *scroll {
write!(f, " SCROLL")?;
} else {
write!(f, " NO SCROLL")?;
}
}
if let Some(declare_type) = declare_type {
write!(f, " {declare_type}")?;
}
if let Some(hold) = hold {
if *hold {
write!(f, " WITH HOLD")?;
} else {
write!(f, " WITHOUT HOLD")?;
}
}
if let Some(query) = for_query {
write!(f, " FOR {query}")?;
}
if let Some(data_type) = data_type {
write!(f, " {data_type}")?;
}
if let Some(expr) = assignment {
write!(f, " {expr}")?;
}
Ok(())
}
}
/// Sql options of a `CREATE TABLE` statement.
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@ -1873,25 +2078,7 @@ pub enum Statement {
///
/// Note: this is a PostgreSQL-specific statement,
/// but may also compatible with other SQL.
Declare {
/// Cursor name
name: Ident,
/// Causes the cursor to return data in binary rather than in text format.
binary: bool,
/// None = Not specified
/// Some(true) = INSENSITIVE
/// Some(false) = ASENSITIVE
sensitive: Option<bool>,
/// None = Not specified
/// Some(true) = SCROLL
/// Some(false) = NO SCROLL
scroll: Option<bool>,
/// None = Not specified
/// Some(true) = WITH HOLD, specifies that the cursor can continue to be used after the transaction that created it successfully commits
/// Some(false) = WITHOUT HOLD, specifies that the cursor cannot be used outside of the transaction that created it
hold: Option<bool>,
query: Box<Query>,
},
Declare { stmts: Vec<Declare> },
/// ```sql
/// CREATE EXTENSION [ IF NOT EXISTS ] extension_name
/// [ WITH ] [ SCHEMA schema_name ]
@ -2447,47 +2634,9 @@ impl fmt::Display for Statement {
write!(f, "{statement}")
}
Statement::Query(s) => write!(f, "{s}"),
Statement::Declare {
name,
binary,
sensitive,
scroll,
hold,
query,
} => {
write!(f, "DECLARE {name} ")?;
if *binary {
write!(f, "BINARY ")?;
}
if let Some(sensitive) = sensitive {
if *sensitive {
write!(f, "INSENSITIVE ")?;
} else {
write!(f, "ASENSITIVE ")?;
}
}
if let Some(scroll) = scroll {
if *scroll {
write!(f, "SCROLL ")?;
} else {
write!(f, "NO SCROLL ")?;
}
}
write!(f, "CURSOR ")?;
if let Some(hold) = hold {
if *hold {
write!(f, "WITH HOLD ")?;
} else {
write!(f, "WITHOUT HOLD ")?;
}
}
write!(f, "FOR {query}")
Statement::Declare { stmts } => {
write!(f, "DECLARE ")?;
write!(f, "{}", display_separated(stmts, "; "))
}
Statement::Fetch {
name,

2
src/keywords.rs
View file

@ -261,6 +261,7 @@ define_keywords!(
EVENT,
EVERY,
EXCEPT,
EXCEPTION,
EXCLUDE,
EXCLUSIVE,
EXEC,
@ -562,6 +563,7 @@ define_keywords!(
RESTART,
RESTRICT,
RESULT,
RESULTSET,
RETAIN,
RETURN,
RETURNING,

218
src/parser/mod.rs
View file

@ -3915,14 +3915,26 @@ impl<'a> Parser<'a> {
Ok(DropFunctionDesc { name, args })
}
/// Parse a `DECLARE` statement.
///
/// ```sql
/// DECLARE name [ BINARY ] [ ASENSITIVE | INSENSITIVE ] [ [ NO ] SCROLL ]
/// CURSOR [ { WITH | WITHOUT } HOLD ] FOR query
/// ```
///
/// The syntax can vary significantly between warehouses. See the grammar
/// on the warehouse specific function in such cases.
pub fn parse_declare(&mut self) -> Result<Statement, ParserError> {
if dialect_of!(self is BigQueryDialect) {
return self.parse_big_query_declare();
}
if dialect_of!(self is SnowflakeDialect) {
return self.parse_snowflake_declare();
}
let name = self.parse_identifier(false)?;
let binary = self.parse_keyword(Keyword::BINARY);
let binary = Some(self.parse_keyword(Keyword::BINARY));
let sensitive = if self.parse_keyword(Keyword::INSENSITIVE) {
Some(true)
} else if self.parse_keyword(Keyword::ASENSITIVE) {
@ -3939,6 +3951,7 @@ impl<'a> Parser<'a> {
};
self.expect_keyword(Keyword::CURSOR)?;
let declare_type = Some(DeclareType::Cursor);
let hold = match self.parse_one_of_keywords(&[Keyword::WITH, Keyword::WITHOUT]) {
Some(keyword) => {
@ -3955,15 +3968,204 @@ impl<'a> Parser<'a> {
self.expect_keyword(Keyword::FOR)?;
let query = self.parse_query()?;
let query = Some(Box::new(self.parse_query()?));
Ok(Statement::Declare {
name,
binary,
sensitive,
scroll,
hold,
query: Box::new(query),
stmts: vec![Declare {
names: vec![name],
data_type: None,
assignment: None,
declare_type,
binary,
sensitive,
scroll,
hold,
for_query: query,
}],
})
}
/// Parse a [BigQuery] `DECLARE` statement.
///
/// Syntax:
/// ```text
/// DECLARE variable_name[, ...] [{ <variable_type> | <DEFAULT expression> }];
/// ```
/// [BigQuery]: https://cloud.google.com/bigquery/docs/reference/standard-sql/procedural-language#declare
pub fn parse_big_query_declare(&mut self) -> Result<Statement, ParserError> {
let names = self.parse_comma_separated(|parser| Parser::parse_identifier(parser, false))?;
let data_type = match self.peek_token().token {
Token::Word(w) if w.keyword == Keyword::DEFAULT => None,
_ => Some(self.parse_data_type()?),
};
let expr = if data_type.is_some() {
if self.parse_keyword(Keyword::DEFAULT) {
Some(self.parse_expr()?)
} else {
None
}
} else {
// If no variable type - default expression must be specified, per BQ docs.
// i.e `DECLARE foo;` is invalid.
self.expect_keyword(Keyword::DEFAULT)?;
Some(self.parse_expr()?)
};
Ok(Statement::Declare {
stmts: vec![Declare {
names,
data_type,
assignment: expr.map(|expr| DeclareAssignment::Default(Box::new(expr))),
declare_type: None,
binary: None,
sensitive: None,
scroll: None,
hold: None,
for_query: None,
}],
})
}
/// Parse a [Snowflake] `DECLARE` statement.
///
/// Syntax:
/// ```text
/// DECLARE
/// [{ <variable_declaration>
/// | <cursor_declaration>
/// | <resultset_declaration>
/// | <exception_declaration> }; ... ]
///
/// <variable_declaration>
/// <variable_name> [<type>] [ { DEFAULT | := } <expression>]
///
/// <cursor_declaration>
/// <cursor_name> CURSOR FOR <query>
///
/// <resultset_declaration>
/// <resultset_name> RESULTSET [ { DEFAULT | := } ( <query> ) ] ;
///
/// <exception_declaration>
/// <exception_name> EXCEPTION [ ( <exception_number> , '<exception_message>' ) ] ;
/// ```
///
/// [Snowflake]: https://docs.snowflake.com/en/sql-reference/snowflake-scripting/declare
pub fn parse_snowflake_declare(&mut self) -> Result<Statement, ParserError> {
let mut stmts = vec![];
loop {
let name = self.parse_identifier(false)?;
let (declare_type, for_query, assigned_expr, data_type) =
if self.parse_keyword(Keyword::CURSOR) {
self.expect_keyword(Keyword::FOR)?;
match self.peek_token().token {
Token::Word(w) if w.keyword == Keyword::SELECT => (
Some(DeclareType::Cursor),
Some(Box::new(self.parse_query()?)),
None,
None,
),
_ => (
Some(DeclareType::Cursor),
None,
Some(DeclareAssignment::For(Box::new(self.parse_expr()?))),
None,
),
}
} else if self.parse_keyword(Keyword::RESULTSET) {
let assigned_expr = if self.peek_token().token != Token::SemiColon {
self.parse_snowflake_variable_declaration_expression()?
} else {
// Nothing more to do. The statement has no further parameters.
None
};
(Some(DeclareType::ResultSet), None, assigned_expr, None)
} else if self.parse_keyword(Keyword::EXCEPTION) {
let assigned_expr = if self.peek_token().token == Token::LParen {
Some(DeclareAssignment::Expr(Box::new(self.parse_expr()?)))
} else {
// Nothing more to do. The statement has no further parameters.
None
};
(Some(DeclareType::Exception), None, assigned_expr, None)
} else {
// Without an explicit keyword, the only valid option is variable declaration.
let (assigned_expr, data_type) = if let Some(assigned_expr) =
self.parse_snowflake_variable_declaration_expression()?
{
(Some(assigned_expr), None)
} else if let Token::Word(_) = self.peek_token().token {
let data_type = self.parse_data_type()?;
(
self.parse_snowflake_variable_declaration_expression()?,
Some(data_type),
)
} else {
(None, None)
};
(None, None, assigned_expr, data_type)
};
let stmt = Declare {
names: vec![name],
data_type,
assignment: assigned_expr,
declare_type,
binary: None,
sensitive: None,
scroll: None,
hold: None,
for_query,
};
stmts.push(stmt);
if self.consume_token(&Token::SemiColon) {
match self.peek_token().token {
Token::Word(w)
if ALL_KEYWORDS
.binary_search(&w.value.to_uppercase().as_str())
.is_err() =>
{
// Not a keyword - start of a new declaration.
continue;
}
_ => {
// Put back the semi-colon, this is the end of the DECLARE statement.
self.prev_token();
}
}
}
break;
}
Ok(Statement::Declare { stmts })
}
/// Parses the assigned expression in a variable declaration.
///
/// Syntax:
/// ```text
/// [ { DEFAULT | := } <expression>]
/// ```
/// <https://docs.snowflake.com/en/sql-reference/snowflake-scripting/declare#variable-declaration-syntax>
pub fn parse_snowflake_variable_declaration_expression(
&mut self,
) -> Result<Option<DeclareAssignment>, ParserError> {
Ok(match self.peek_token().token {
Token::Word(w) if w.keyword == Keyword::DEFAULT => {
self.next_token(); // Skip `DEFAULT`
Some(DeclareAssignment::Default(Box::new(self.parse_expr()?)))
}
Token::DuckAssignment => {
self.next_token(); // Skip `:=`
Some(DeclareAssignment::DuckAssignment(Box::new(
self.parse_expr()?,
)))
}
_ => None,
})
}