Fix SUBSTRING from/to argument construction for mssql (#947)

This commit is contained in:
Jeremy Maness 2023-08-17 06:17:57 -04:00 committed by GitHub
parent 173a6db818
commit 8bbb85356c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 153 additions and 24 deletions

View file

@ -477,6 +477,10 @@ pub enum Expr {
expr: Box<Expr>,
substring_from: Option<Box<Expr>>,
substring_for: Option<Box<Expr>>,
// Some dialects use `SUBSTRING(expr [FROM start] [FOR len])` syntax while others omit FROM,
// FOR keywords (e.g. Microsoft SQL Server). This flags is used for formatting.
special: bool,
},
/// ```sql
/// TRIM([BOTH | LEADING | TRAILING] [<expr> FROM] <expr>)
@ -830,13 +834,22 @@ impl fmt::Display for Expr {
expr,
substring_from,
substring_for,
special,
} => {
write!(f, "SUBSTRING({expr}")?;
if let Some(from_part) = substring_from {
write!(f, " FROM {from_part}")?;
if *special {
write!(f, ", {from_part}")?;
} else {
write!(f, " FROM {from_part}")?;
}
}
if let Some(for_part) = substring_for {
write!(f, " FOR {for_part}")?;
if *special {
write!(f, ", {for_part}")?;
} else {
write!(f, " FOR {for_part}")?;
}
}
write!(f, ")")

View file

@ -117,6 +117,10 @@ pub trait Dialect: Debug + Any {
fn supports_group_by_expr(&self) -> bool {
false
}
/// Returns true if the dialect supports `SUBSTRING(expr [FROM start] [FOR len])` expressions
fn supports_substring_from_for_expr(&self) -> bool {
true
}
/// Dialect-specific prefix parser override
fn parse_prefix(&self, _parser: &mut Parser) -> Option<Result<Expr, ParserError>> {
// return None to fall back to the default behavior

View file

@ -34,4 +34,8 @@ impl Dialect for MsSqlDialect {
|| ch == '#'
|| ch == '_'
}
fn supports_substring_from_for_expr(&self) -> bool {
false
}
}

View file

@ -1223,25 +1223,47 @@ impl<'a> Parser<'a> {
}
pub fn parse_substring_expr(&mut self) -> Result<Expr, ParserError> {
// PARSE SUBSTRING (EXPR [FROM 1] [FOR 3])
self.expect_token(&Token::LParen)?;
let expr = self.parse_expr()?;
let mut from_expr = None;
if self.parse_keyword(Keyword::FROM) || self.consume_token(&Token::Comma) {
from_expr = Some(self.parse_expr()?);
}
if self.dialect.supports_substring_from_for_expr() {
// PARSE SUBSTRING (EXPR [FROM 1] [FOR 3])
self.expect_token(&Token::LParen)?;
let expr = self.parse_expr()?;
let mut from_expr = None;
if self.parse_keyword(Keyword::FROM) || self.consume_token(&Token::Comma) {
from_expr = Some(self.parse_expr()?);
}
let mut to_expr = None;
if self.parse_keyword(Keyword::FOR) || self.consume_token(&Token::Comma) {
to_expr = Some(self.parse_expr()?);
}
self.expect_token(&Token::RParen)?;
let mut to_expr = None;
if self.parse_keyword(Keyword::FOR) || self.consume_token(&Token::Comma) {
to_expr = Some(self.parse_expr()?);
}
self.expect_token(&Token::RParen)?;
Ok(Expr::Substring {
expr: Box::new(expr),
substring_from: from_expr.map(Box::new),
substring_for: to_expr.map(Box::new),
})
Ok(Expr::Substring {
expr: Box::new(expr),
substring_from: from_expr.map(Box::new),
substring_for: to_expr.map(Box::new),
special: !self.dialect.supports_substring_from_for_expr(),
})
} else {
// PARSE SUBSTRING(EXPR, start, length)
self.expect_token(&Token::LParen)?;
let expr = self.parse_expr()?;
self.expect_token(&Token::Comma)?;
let from_expr = Some(self.parse_expr()?);
self.expect_token(&Token::Comma)?;
let to_expr = Some(self.parse_expr()?);
self.expect_token(&Token::RParen)?;
Ok(Expr::Substring {
expr: Box::new(expr),
substring_from: from_expr.map(Box::new),
substring_for: to_expr.map(Box::new),
special: !self.dialect.supports_substring_from_for_expr(),
})
}
}
pub fn parse_overlay_expr(&mut self) -> Result<Expr, ParserError> {

View file

@ -5084,19 +5084,45 @@ fn parse_scalar_subqueries() {
#[test]
fn parse_substring() {
one_statement_parses_to("SELECT SUBSTRING('1')", "SELECT SUBSTRING('1')");
let from_for_supported_dialects = TestedDialects {
dialects: vec![
Box::new(GenericDialect {}),
Box::new(PostgreSqlDialect {}),
Box::new(AnsiDialect {}),
Box::new(SnowflakeDialect {}),
Box::new(HiveDialect {}),
Box::new(RedshiftSqlDialect {}),
Box::new(MySqlDialect {}),
Box::new(BigQueryDialect {}),
Box::new(SQLiteDialect {}),
Box::new(DuckDbDialect {}),
],
options: None,
};
one_statement_parses_to(
let from_for_unsupported_dialects = TestedDialects {
dialects: vec![Box::new(MsSqlDialect {})],
options: None,
};
from_for_supported_dialects
.one_statement_parses_to("SELECT SUBSTRING('1')", "SELECT SUBSTRING('1')");
from_for_supported_dialects.one_statement_parses_to(
"SELECT SUBSTRING('1' FROM 1)",
"SELECT SUBSTRING('1' FROM 1)",
);
one_statement_parses_to(
from_for_supported_dialects.one_statement_parses_to(
"SELECT SUBSTRING('1' FROM 1 FOR 3)",
"SELECT SUBSTRING('1' FROM 1 FOR 3)",
);
one_statement_parses_to("SELECT SUBSTRING('1' FOR 3)", "SELECT SUBSTRING('1' FOR 3)");
from_for_unsupported_dialects
.one_statement_parses_to("SELECT SUBSTRING('1', 1, 3)", "SELECT SUBSTRING('1', 1, 3)");
from_for_supported_dialects
.one_statement_parses_to("SELECT SUBSTRING('1' FOR 3)", "SELECT SUBSTRING('1' FOR 3)");
}
#[test]

View file

@ -379,6 +379,65 @@ fn parse_similar_to() {
chk(true);
}
#[test]
fn parse_substring_in_select() {
let sql = "SELECT DISTINCT SUBSTRING(description, 0, 1) FROM test";
match ms().one_statement_parses_to(
sql,
"SELECT DISTINCT SUBSTRING(description, 0, 1) FROM test",
) {
Statement::Query(query) => {
assert_eq!(
Box::new(Query {
with: None,
body: Box::new(SetExpr::Select(Box::new(Select {
distinct: Some(Distinct::Distinct),
top: None,
projection: vec![SelectItem::UnnamedExpr(Expr::Substring {
expr: Box::new(Expr::Identifier(Ident {
value: "description".to_string(),
quote_style: None
})),
substring_from: Some(Box::new(Expr::Value(number("0")))),
substring_for: Some(Box::new(Expr::Value(number("1")))),
special: true,
})],
into: None,
from: vec![TableWithJoins {
relation: TableFactor::Table {
name: ObjectName(vec![Ident {
value: "test".to_string(),
quote_style: None
}]),
alias: None,
args: None,
with_hints: vec![]
},
joins: vec![]
}],
lateral_views: vec![],
selection: None,
group_by: vec![],
cluster_by: vec![],
distribute_by: vec![],
sort_by: vec![],
having: None,
named_window: vec![],
qualify: None
}))),
order_by: vec![],
limit: None,
offset: None,
fetch: None,
locks: vec![],
}),
query
);
}
_ => unreachable!(),
}
}
fn ms() -> TestedDialects {
TestedDialects {
dialects: vec![Box::new(MsSqlDialect {})],

View file

@ -1271,7 +1271,8 @@ fn parse_substring_in_select() {
quote_style: None
})),
substring_from: Some(Box::new(Expr::Value(number("0")))),
substring_for: Some(Box::new(Expr::Value(number("1"))))
substring_for: Some(Box::new(Expr::Value(number("1")))),
special: false,
})],
into: None,
from: vec![TableWithJoins {