mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-12-23 11:12:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Add CASE and IF statement support (#1741)

Browse source
This commit is contained in:
Ifeanyi Ubah 2025-03-14 07:49:25 +01:00 • committed by GitHub
parent cf4ab7f9ab
commit 862e887a66
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 473 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

198
src/ast/mod.rs
View file

@ -151,6 +151,15 @@ where
DisplaySeparated { slice, sep: ", " }
}
/// Writes the given statements to the formatter, each ending with
/// a semicolon and space separated.
fn format_statement_list(f: &mut fmt::Formatter, statements: &[Statement]) -> fmt::Result {
write!(f, "{}", display_separated(statements, "; "))?;
// We manually insert semicolon for the last statement,
// since display_separated doesn't handle that case.
write!(f, ";")
}
/// An identifier, decomposed into its value or character data and the quote style.
#[derive(Debug, Clone, PartialOrd, Ord)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@ -2080,6 +2089,173 @@ pub enum Password {
NullPassword,
}
/// A `CASE` statement.
///
/// Examples:
/// ```sql
/// CASE
/// WHEN EXISTS(SELECT 1)
/// THEN SELECT 1 FROM T;
/// WHEN EXISTS(SELECT 2)
/// THEN SELECT 1 FROM U;
/// ELSE
/// SELECT 1 FROM V;
/// END CASE;
/// ```
///
/// [BigQuery](https://cloud.google.com/bigquery/docs/reference/standard-sql/procedural-language#case_search_expression)
/// [Snowflake](https://docs.snowflake.com/en/sql-reference/snowflake-scripting/case)
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct CaseStatement {
pub match_expr: Option<Expr>,
pub when_blocks: Vec<ConditionalStatements>,
pub else_block: Option<Vec<Statement>>,
/// TRUE if the statement ends with `END CASE` (vs `END`).
pub has_end_case: bool,
}
impl fmt::Display for CaseStatement {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let CaseStatement {
match_expr,
when_blocks,
else_block,
has_end_case,
} = self;
write!(f, "CASE")?;
if let Some(expr) = match_expr {
write!(f, " {expr}")?;
}
if !when_blocks.is_empty() {
write!(f, " {}", display_separated(when_blocks, " "))?;
}
if let Some(else_block) = else_block {
write!(f, " ELSE ")?;
format_statement_list(f, else_block)?;
}
write!(f, " END")?;
if *has_end_case {
write!(f, " CASE")?;
}
Ok(())
}
}
/// An `IF` statement.
///
/// Examples:
/// ```sql
/// IF TRUE THEN
/// SELECT 1;
/// SELECT 2;
/// ELSEIF TRUE THEN
/// SELECT 3;
/// ELSE
/// SELECT 4;
/// END IF
/// ```
///
/// [BigQuery](https://cloud.google.com/bigquery/docs/reference/standard-sql/procedural-language#if)
/// [Snowflake](https://docs.snowflake.com/en/sql-reference/snowflake-scripting/if)
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct IfStatement {
pub if_block: ConditionalStatements,
pub elseif_blocks: Vec<ConditionalStatements>,
pub else_block: Option<Vec<Statement>>,
}
impl fmt::Display for IfStatement {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let IfStatement {
if_block,
elseif_blocks,
else_block,
} = self;
write!(f, "{if_block}")?;
if !elseif_blocks.is_empty() {
write!(f, " {}", display_separated(elseif_blocks, " "))?;
}
if let Some(else_block) = else_block {
write!(f, " ELSE ")?;
format_statement_list(f, else_block)?;
}
write!(f, " END IF")?;
Ok(())
}
}
/// Represents a type of [ConditionalStatements]
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum ConditionalStatementKind {
/// `WHEN <condition> THEN <statements>`
When,
/// `IF <condition> THEN <statements>`
If,
/// `ELSEIF <condition> THEN <statements>`
ElseIf,
}
/// A block within a [Statement::Case] or [Statement::If]-like statement
///
/// Examples:
/// ```sql
/// WHEN EXISTS(SELECT 1) THEN SELECT 1;
///
/// IF TRUE THEN SELECT 1; SELECT 2;
/// ```
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct ConditionalStatements {
/// The condition expression.
pub condition: Expr,
/// Statement list of the `THEN` clause.
pub statements: Vec<Statement>,
pub kind: ConditionalStatementKind,
}
impl fmt::Display for ConditionalStatements {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let ConditionalStatements {
condition: expr,
statements,
kind,
} = self;
let kind = match kind {
ConditionalStatementKind::When => "WHEN",
ConditionalStatementKind::If => "IF",
ConditionalStatementKind::ElseIf => "ELSEIF",
};
write!(f, "{kind} {expr} THEN")?;
if !statements.is_empty() {
write!(f, " ")?;
format_statement_list(f, statements)?;
}
Ok(())
}
}
/// Represents an expression assignment within a variable `DECLARE` statement.
///
/// Examples:
@ -2647,6 +2823,10 @@ pub enum Statement {
file_format: Option<FileFormat>,
source: Box<Query>,
},
/// A `CASE` statement.
Case(CaseStatement),
/// An `IF` statement.
If(IfStatement),
/// ```sql
/// CALL <function>
/// ```
@ -3940,6 +4120,12 @@ impl fmt::Display for Statement {
}
Ok(())
}
Statement::Case(stmt) => {
write!(f, "{stmt}")
}
Statement::If(stmt) => {
write!(f, "{stmt}")
}
Statement::AttachDatabase {
schema_name,
database_file_name,
@ -4942,18 +5128,14 @@ impl fmt::Display for Statement {
write!(f, " {}", display_comma_separated(modes))?;
}
if !statements.is_empty() {
write!(f, " {}", display_separated(statements, "; "))?;
// We manually insert semicolon for the last statement,
// since display_separated doesn't handle that case.
write!(f, ";")?;
write!(f, " ")?;
format_statement_list(f, statements)?;
}
if let Some(exception_statements) = exception_statements {
write!(f, " EXCEPTION WHEN ERROR THEN")?;
if !exception_statements.is_empty() {
write!(f, " {}", display_separated(exception_statements, "; "))?;
// We manually insert semicolon for the last statement,
// since display_separated doesn't handle that case.
write!(f, ";")?;
write!(f, " ")?;
format_statement_list(f, exception_statements)?;
}
}
if *has_end_keyword {

78
src/ast/spans.rs
View file

@ -22,20 +22,21 @@ use crate::tokenizer::Span;
use super::{
dcl::SecondaryRoles, value::ValueWithSpan, AccessExpr, AlterColumnOperation,
AlterIndexOperation, AlterTableOperation, Array, Assignment, AssignmentTarget, CloseCursor,
ClusteredIndex, ColumnDef, ColumnOption, ColumnOptionDef, ConflictTarget, ConnectBy,
ConstraintCharacteristics, CopySource, CreateIndex, CreateTable, CreateTableOptions, Cte,
Delete, DoUpdate, ExceptSelectItem, ExcludeSelectItem, Expr, ExprWithAlias, Fetch, FromTable,
Function, FunctionArg, FunctionArgExpr, FunctionArgumentClause, FunctionArgumentList,
FunctionArguments, GroupByExpr, HavingBound, IlikeSelectItem, Insert, Interpolate,
InterpolateExpr, Join, JoinConstraint, JoinOperator, JsonPath, JsonPathElem, LateralView,
LimitClause, MatchRecognizePattern, Measure, NamedWindowDefinition, ObjectName, ObjectNamePart,
Offset, OnConflict, OnConflictAction, OnInsert, OrderBy, OrderByExpr, OrderByKind, Partition,
PivotValueSource, ProjectionSelect, Query, ReferentialAction, RenameSelectItem,
ReplaceSelectElement, ReplaceSelectItem, Select, SelectInto, SelectItem, SetExpr, SqlOption,
Statement, Subscript, SymbolDefinition, TableAlias, TableAliasColumnDef, TableConstraint,
TableFactor, TableObject, TableOptionsClustered, TableWithJoins, UpdateTableFromKind, Use,
Value, Values, ViewColumnDef, WildcardAdditionalOptions, With, WithFill,
AlterIndexOperation, AlterTableOperation, Array, Assignment, AssignmentTarget, CaseStatement,
CloseCursor, ClusteredIndex, ColumnDef, ColumnOption, ColumnOptionDef, ConditionalStatements,
ConflictTarget, ConnectBy, ConstraintCharacteristics, CopySource, CreateIndex, CreateTable,
CreateTableOptions, Cte, Delete, DoUpdate, ExceptSelectItem, ExcludeSelectItem, Expr,
ExprWithAlias, Fetch, FromTable, Function, FunctionArg, FunctionArgExpr,
FunctionArgumentClause, FunctionArgumentList, FunctionArguments, GroupByExpr, HavingBound,
IfStatement, IlikeSelectItem, Insert, Interpolate, InterpolateExpr, Join, JoinConstraint,
JoinOperator, JsonPath, JsonPathElem, LateralView, LimitClause, MatchRecognizePattern, Measure,
NamedWindowDefinition, ObjectName, ObjectNamePart, Offset, OnConflict, OnConflictAction,
OnInsert, OrderBy, OrderByExpr, OrderByKind, Partition, PivotValueSource, ProjectionSelect,
Query, ReferentialAction, RenameSelectItem, ReplaceSelectElement, ReplaceSelectItem, Select,
SelectInto, SelectItem, SetExpr, SqlOption, Statement, Subscript, SymbolDefinition, TableAlias,
TableAliasColumnDef, TableConstraint, TableFactor, TableObject, TableOptionsClustered,
TableWithJoins, UpdateTableFromKind, Use, Value, Values, ViewColumnDef,
WildcardAdditionalOptions, With, WithFill,
};
/// Given an iterator of spans, return the [Span::union] of all spans.
@ -334,6 +335,8 @@ impl Spanned for Statement {
file_format: _,
source,
} => source.span(),
Statement::Case(stmt) => stmt.span(),
Statement::If(stmt) => stmt.span(),
Statement::Call(function) => function.span(),
Statement::Copy {
source,
@ -732,6 +735,53 @@ impl Spanned for CreateIndex {
}
}
impl Spanned for CaseStatement {
fn span(&self) -> Span {
let CaseStatement {
match_expr,
when_blocks,
else_block,
has_end_case: _,
} = self;
union_spans(
match_expr
.iter()
.map(|e| e.span())
.chain(when_blocks.iter().map(|b| b.span()))
.chain(else_block.iter().flat_map(|e| e.iter().map(|s| s.span()))),
)
}
}
impl Spanned for IfStatement {
fn span(&self) -> Span {
let IfStatement {
if_block,
elseif_blocks,
else_block,
} = self;
union_spans(
iter::once(if_block.span())
.chain(elseif_blocks.iter().map(|b| b.span()))
.chain(else_block.iter().flat_map(|e| e.iter().map(|s| s.span()))),
)
}
}
impl Spanned for ConditionalStatements {
fn span(&self) -> Span {
let ConditionalStatements {
condition,
statements,
kind: _,
} = self;
union_spans(iter::once(condition.span()).chain(statements.iter().map(|s| s.span())))
}
}
/// # partial span
///
/// Missing spans:

1
src/keywords.rs
View file

@ -297,6 +297,7 @@ define_keywords!(
ELEMENT,
ELEMENTS,
ELSE,
ELSEIF,
EMPTY,
ENABLE,
ENABLE_SCHEMA_EVOLUTION,

104
src/parser/mod.rs
View file

@ -528,6 +528,14 @@ impl<'a> Parser<'a> {
Keyword::DESCRIBE => self.parse_explain(DescribeAlias::Describe),
Keyword::EXPLAIN => self.parse_explain(DescribeAlias::Explain),
Keyword::ANALYZE => self.parse_analyze(),
Keyword::CASE => {
self.prev_token();
self.parse_case_stmt()
}
Keyword::IF => {
self.prev_token();
self.parse_if_stmt()
}
Keyword::SELECT | Keyword::WITH | Keyword::VALUES | Keyword::FROM => {
self.prev_token();
self.parse_query().map(Statement::Query)
@ -615,6 +623,102 @@ impl<'a> Parser<'a> {
}
}
/// Parse a `CASE` statement.
///
/// See [Statement::Case]
pub fn parse_case_stmt(&mut self) -> Result<Statement, ParserError> {
self.expect_keyword_is(Keyword::CASE)?;
let match_expr = if self.peek_keyword(Keyword::WHEN) {
None
} else {
Some(self.parse_expr()?)
};
self.expect_keyword_is(Keyword::WHEN)?;
let when_blocks = self.parse_keyword_separated(Keyword::WHEN, |parser| {
parser.parse_conditional_statements(
ConditionalStatementKind::When,
&[Keyword::WHEN, Keyword::ELSE, Keyword::END],
)
})?;
let else_block = if self.parse_keyword(Keyword::ELSE) {
Some(self.parse_statement_list(&[Keyword::END])?)
} else {
None
};
self.expect_keyword_is(Keyword::END)?;
let has_end_case = self.parse_keyword(Keyword::CASE);
Ok(Statement::Case(CaseStatement {
match_expr,
when_blocks,
else_block,
has_end_case,
}))
}
/// Parse an `IF` statement.
///
/// See [Statement::If]
pub fn parse_if_stmt(&mut self) -> Result<Statement, ParserError> {
self.expect_keyword_is(Keyword::IF)?;
let if_block = self.parse_conditional_statements(
ConditionalStatementKind::If,
&[Keyword::ELSE, Keyword::ELSEIF, Keyword::END],
)?;
let elseif_blocks = if self.parse_keyword(Keyword::ELSEIF) {
self.parse_keyword_separated(Keyword::ELSEIF, |parser| {
parser.parse_conditional_statements(
ConditionalStatementKind::ElseIf,
&[Keyword::ELSEIF, Keyword::ELSE, Keyword::END],
)
})?
} else {
vec![]
};
let else_block = if self.parse_keyword(Keyword::ELSE) {
Some(self.parse_statement_list(&[Keyword::END])?)
} else {
None
};
self.expect_keywords(&[Keyword::END, Keyword::IF])?;
Ok(Statement::If(IfStatement {
if_block,
elseif_blocks,
else_block,
}))
}
/// Parses an expression and associated list of statements
/// belonging to a conditional statement like `IF` or `WHEN`.
///
/// Example:
/// ```sql
/// IF condition THEN statement1; statement2;
/// ```
fn parse_conditional_statements(
&mut self,
kind: ConditionalStatementKind,
terminal_keywords: &[Keyword],
) -> Result<ConditionalStatements, ParserError> {
let condition = self.parse_expr()?;
self.expect_keyword_is(Keyword::THEN)?;
let statements = self.parse_statement_list(terminal_keywords)?;
Ok(ConditionalStatements {
condition,
statements,
kind,
})
}
pub fn parse_comment(&mut self) -> Result<Statement, ParserError> {
let if_exists = self.parse_keywords(&[Keyword::IF, Keyword::EXISTS]);