Skip to content

Commit

Permalink
Support Snowflake/BigQuery TRIM. (apache#975)
Browse files Browse the repository at this point in the history
  • Loading branch information
zdenal authored and serprex committed Nov 6, 2023
1 parent 85a6229 commit 9f45634
Show file tree
Hide file tree
Showing 5 changed files with 95 additions and 0 deletions.
6 changes: 6 additions & 0 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -496,12 +496,14 @@ pub enum Expr {
/// ```sql
/// TRIM([BOTH | LEADING | TRAILING] [<expr> FROM] <expr>)
/// TRIM(<expr>)
/// TRIM(<expr>, [, characters]) -- only Snowflake or Bigquery
/// ```
Trim {
expr: Box<Expr>,
// ([BOTH | LEADING | TRAILING]
trim_where: Option<TrimWhereField>,
trim_what: Option<Box<Expr>>,
trim_characters: Option<Vec<Expr>>,
},
/// ```sql
/// OVERLAY(<expr> PLACING <expr> FROM <expr>[ FOR <expr> ]
Expand Down Expand Up @@ -895,6 +897,7 @@ impl fmt::Display for Expr {
expr,
trim_where,
trim_what,
trim_characters,
} => {
write!(f, "TRIM(")?;
if let Some(ident) = trim_where {
Expand All @@ -905,6 +908,9 @@ impl fmt::Display for Expr {
} else {
write!(f, "{expr}")?;
}
if let Some(characters) = trim_characters {
write!(f, ", {}", display_comma_separated(characters))?;
}

write!(f, ")")
}
Expand Down
14 changes: 14 additions & 0 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1319,6 +1319,7 @@ impl<'a> Parser<'a> {
/// ```sql
/// TRIM ([WHERE] ['text' FROM] 'text')
/// TRIM ('text')
/// TRIM(<expr>, [, characters]) -- only Snowflake or BigQuery
/// ```
pub fn parse_trim_expr(&mut self) -> Result<Expr, ParserError> {
self.expect_token(&Token::LParen)?;
Expand All @@ -1340,13 +1341,26 @@ impl<'a> Parser<'a> {
expr: Box::new(expr),
trim_where,
trim_what: Some(trim_what),
trim_characters: None,
})
} else if self.consume_token(&Token::Comma)
&& dialect_of!(self is SnowflakeDialect | BigQueryDialect | GenericDialect)
{
let characters = self.parse_comma_separated(Parser::parse_expr)?;
self.expect_token(&Token::RParen)?;
Ok(Expr::Trim {
expr: Box::new(expr),
trim_where: None,
trim_what: None,
trim_characters: Some(characters),
})
} else {
self.expect_token(&Token::RParen)?;
Ok(Expr::Trim {
expr: Box::new(expr),
trim_where,
trim_what: None,
trim_characters: None,
})
}
}
Expand Down
26 changes: 26 additions & 0 deletions tests/sqlparser_bigquery.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ use std::ops::Deref;

use sqlparser::ast::*;
use sqlparser::dialect::{BigQueryDialect, GenericDialect};
use sqlparser::parser::ParserError;
use test_utils::*;

#[test]
Expand Down Expand Up @@ -549,3 +550,28 @@ fn parse_map_access_offset() {
bigquery().verified_only_select(sql);
}
}

#[test]
fn test_bigquery_trim() {
let real_sql = r#"SELECT customer_id, TRIM(item_price_id, '"', "a") AS item_price_id FROM models_staging.subscriptions"#;
assert_eq!(bigquery().verified_stmt(real_sql).to_string(), real_sql);

let sql_only_select = "SELECT TRIM('xyz', 'a')";
let select = bigquery().verified_only_select(sql_only_select);
assert_eq!(
&Expr::Trim {
expr: Box::new(Expr::Value(Value::SingleQuotedString("xyz".to_owned()))),
trim_where: None,
trim_what: None,
trim_characters: Some(vec![Expr::Value(Value::SingleQuotedString("a".to_owned()))]),
},
expr_from_projection(only(&select.projection))
);

// missing comma separation
let error_sql = "SELECT TRIM('xyz' 'a')";
assert_eq!(
ParserError::ParserError("Expected ), found: 'a'".to_owned()),
bigquery().parse_sql_statements(error_sql).unwrap_err()
);
}
24 changes: 24 additions & 0 deletions tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5225,6 +5225,30 @@ fn parse_trim() {
ParserError::ParserError("Expected ), found: 'xyz'".to_owned()),
parse_sql_statements("SELECT TRIM(FOO 'xyz' FROM 'xyzfooxyz')").unwrap_err()
);

//keep Snowflake/BigQuery TRIM syntax failing
let all_expected_snowflake = TestedDialects {
dialects: vec![
//Box::new(GenericDialect {}),
Box::new(PostgreSqlDialect {}),
Box::new(MsSqlDialect {}),
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,
};
assert_eq!(
ParserError::ParserError("Expected ), found: 'a'".to_owned()),
all_expected_snowflake
.parse_sql_statements("SELECT TRIM('xyz', 'a')")
.unwrap_err()
);
}

#[test]
Expand Down
25 changes: 25 additions & 0 deletions tests/sqlparser_snowflake.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1039,3 +1039,28 @@ fn test_snowflake_stage_object_names() {
}
}
}

#[test]
fn test_snowflake_trim() {
let real_sql = r#"SELECT customer_id, TRIM(sub_items.value:item_price_id, '"', "a") AS item_price_id FROM models_staging.subscriptions"#;
assert_eq!(snowflake().verified_stmt(real_sql).to_string(), real_sql);

let sql_only_select = "SELECT TRIM('xyz', 'a')";
let select = snowflake().verified_only_select(sql_only_select);
assert_eq!(
&Expr::Trim {
expr: Box::new(Expr::Value(Value::SingleQuotedString("xyz".to_owned()))),
trim_where: None,
trim_what: None,
trim_characters: Some(vec![Expr::Value(Value::SingleQuotedString("a".to_owned()))]),
},
expr_from_projection(only(&select.projection))
);

// missing comma separation
let error_sql = "SELECT TRIM('xyz' 'a')";
assert_eq!(
ParserError::ParserError("Expected ), found: 'a'".to_owned()),
snowflake().parse_sql_statements(error_sql).unwrap_err()
);
}

0 comments on commit 9f45634

Please sign in to comment.