From f53336feb359499cf38d54aee3496282916352e7 Mon Sep 17 00:00:00 2001
From: togami2864 <yoshiaki.togami@plaid.co.jp>
Date: Tue, 2 Aug 2022 16:28:43 +0900
Subject: [PATCH] feat: support SAFE_CAST for bigquery

---
 src/ast/mod.rs              |  8 ++++++++
 src/keywords.rs             |  1 +
 src/parser.rs               | 14 ++++++++++++++
 tests/sqlparser_bigquery.rs |  6 ++++++
 4 files changed, 29 insertions(+)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index c3dd65782..312a86813 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -285,6 +285,13 @@ pub enum Expr {
         expr: Box<Expr>,
         data_type: DataType,
     },
+    /// SAFE_CAST an expression to a different data type e.g. `SAFE_CAST(foo AS FLOAT64)`
+    //  only available for BigQuery: https://cloud.google.com/bigquery/docs/reference/standard-sql/functions-and-operators#safe_casting
+    //  this works the same as `TRY_CAST`
+    SafeCast {
+        expr: Box<Expr>,
+        data_type: DataType,
+    },
     /// AT a timestamp to a different timezone e.g. `FROM_UNIXTIME(0) AT TIME ZONE 'UTC-06:00'`
     AtTimeZone {
         timestamp: Box<Expr>,
@@ -442,6 +449,7 @@ impl fmt::Display for Expr {
             }
             Expr::Cast { expr, data_type } => write!(f, "CAST({} AS {})", expr, data_type),
             Expr::TryCast { expr, data_type } => write!(f, "TRY_CAST({} AS {})", expr, data_type),
+            Expr::SafeCast { expr, data_type } => write!(f, "SAFE_CAST({} AS {})", expr, data_type),
             Expr::Extract { field, expr } => write!(f, "EXTRACT({} FROM {})", field, expr),
             Expr::Position { expr, r#in } => write!(f, "POSITION({} IN {})", expr, r#in),
             Expr::Collate { expr, collation } => write!(f, "{} COLLATE {}", expr, collation),
diff --git a/src/keywords.rs b/src/keywords.rs
index dbcb67e71..5a12533ea 100644
--- a/src/keywords.rs
+++ b/src/keywords.rs
@@ -439,6 +439,7 @@ define_keywords!(
     ROWID,
     ROWS,
     ROW_NUMBER,
+    SAFE_CAST,
     SAVEPOINT,
     SCHEMA,
     SCOPE,
diff --git a/src/parser.rs b/src/parser.rs
index d98fddc83..1fd74e716 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -426,6 +426,7 @@ impl<'a> Parser<'a> {
                 Keyword::CASE => self.parse_case_expr(),
                 Keyword::CAST => self.parse_cast_expr(),
                 Keyword::TRY_CAST => self.parse_try_cast_expr(),
+                Keyword::SAFE_CAST => self.parse_safe_cast_expr(),
                 Keyword::EXISTS => self.parse_exists_expr(false),
                 Keyword::EXTRACT => self.parse_extract_expr(),
                 Keyword::POSITION => self.parse_position_expr(),
@@ -780,6 +781,19 @@ impl<'a> Parser<'a> {
         })
     }
 
+    /// Parse a BigQuery SAFE_CAST function e.g. `SAFE_CAST(expr AS FLOAT64)`
+    pub fn parse_safe_cast_expr(&mut self) -> Result<Expr, ParserError> {
+        self.expect_token(&Token::LParen)?;
+        let expr = self.parse_expr()?;
+        self.expect_keyword(Keyword::AS)?;
+        let data_type = self.parse_data_type()?;
+        self.expect_token(&Token::RParen)?;
+        Ok(Expr::SafeCast {
+            expr: Box::new(expr),
+            data_type,
+        })
+    }
+
     /// Parse a SQL EXISTS expression e.g. `WHERE EXISTS(SELECT ...)`.
     pub fn parse_exists_expr(&mut self, negated: bool) -> Result<Expr, ParserError> {
         self.expect_token(&Token::LParen)?;
diff --git a/tests/sqlparser_bigquery.rs b/tests/sqlparser_bigquery.rs
index 17240d22d..d866fc2f5 100644
--- a/tests/sqlparser_bigquery.rs
+++ b/tests/sqlparser_bigquery.rs
@@ -94,6 +94,12 @@ fn parse_table_identifiers() {
     test_table_ident("abc5.GROUP", vec![Ident::new("abc5"), Ident::new("GROUP")]);
 }
 
+#[test]
+fn parse_cast_type() {
+    let sql = r#"SELECT SAFE_CAST(1 AS INT64)"#;
+    bigquery().verified_only_select(sql);
+}
+
 fn bigquery() -> TestedDialects {
     TestedDialects {
         dialects: vec![Box::new(BigQueryDialect {})],