bigquery: CAST AS x FORMAT support

apache · Sep 25, 2023 · 0d42a8a · 0d42a8a
1 parent 7723ea5
commit 0d42a8a
Show file tree

Hide file tree

Showing 6 changed files with 109 additions and 4 deletions.
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -437,19 +437,22 @@ pub enum Expr {
     Cast {
         expr: Box<Expr>,
         data_type: DataType,
+        format: Option<Value>,
     },
     /// TRY_CAST an expression to a different data type e.g. `TRY_CAST(foo AS VARCHAR(123))`
     //  this differs from CAST in the choice of how to implement invalid conversions
     TryCast {
         expr: Box<Expr>,
         data_type: DataType,
+        format: Option<Value>,
     },
     /// 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,
+        format: Option<Value>,
     },
     /// AT a timestamp to a different timezone e.g. `FROM_UNIXTIME(0) AT TIME ZONE 'UTC-06:00'`
     AtTimeZone {
@@ -751,9 +754,39 @@ impl fmt::Display for Expr {
                     write!(f, "{op}{expr}")
                 }
             }
-            Expr::Cast { expr, data_type } => write!(f, "CAST({expr} AS {data_type})"),
-            Expr::TryCast { expr, data_type } => write!(f, "TRY_CAST({expr} AS {data_type})"),
-            Expr::SafeCast { expr, data_type } => write!(f, "SAFE_CAST({expr} AS {data_type})"),
+            Expr::Cast {
+                expr,
+                data_type,
+                format,
+            } => {
+                if let Some(format) = format {
+                    write!(f, "CAST({expr} AS {data_type} FORMAT {format})")
+                } else {
+                    write!(f, "CAST({expr} AS {data_type})")
+                }
+            }
+            Expr::TryCast {
+                expr,
+                data_type,
+                format,
+            } => {
+                if let Some(format) = format {
+                    write!(f, "TRY_CAST({expr} AS {data_type} FORMAT {format})")
+                } else {
+                    write!(f, "TRY_CAST({expr} AS {data_type})")
+                }
+            }
+            Expr::SafeCast {
+                expr,
+                data_type,
+                format,
+            } => {
+                if let Some(format) = format {
+                    write!(f, "SAFE_CAST({expr} AS {data_type} FORMAT {format})")
+                } else {
+                    write!(f, "SAFE_CAST({expr} AS {data_type})")
+                }
+            }
             Expr::Extract { field, expr } => write!(f, "EXTRACT({field} FROM {expr})"),
             Expr::Ceil { expr, field } => {
                 if field == &DateTimeField::NoDateTime {

diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -1132,10 +1132,19 @@ impl<'a> Parser<'a> {
         let expr = self.parse_expr()?;
         self.expect_keyword(Keyword::AS)?;
         let data_type = self.parse_data_type()?;
+
+        let format = if dialect_of!(self is BigQueryDialect) && self.parse_keyword(Keyword::FORMAT)
+        {
+            Some(self.parse_value()?)
+        } else {
+            None
+        };
+
         self.expect_token(&Token::RParen)?;
         Ok(Expr::Cast {
             expr: Box::new(expr),
             data_type,
+            format,
         })
     }
 
@@ -1145,10 +1154,19 @@ impl<'a> Parser<'a> {
         let expr = self.parse_expr()?;
         self.expect_keyword(Keyword::AS)?;
         let data_type = self.parse_data_type()?;
+
+        let format = if dialect_of!(self is BigQueryDialect) && self.parse_keyword(Keyword::FORMAT)
+        {
+            Some(self.parse_value()?)
+        } else {
+            None
+        };
+
         self.expect_token(&Token::RParen)?;
         Ok(Expr::TryCast {
             expr: Box::new(expr),
             data_type,
+            format,
         })
     }
 
@@ -1158,10 +1176,19 @@ impl<'a> Parser<'a> {
         let expr = self.parse_expr()?;
         self.expect_keyword(Keyword::AS)?;
         let data_type = self.parse_data_type()?;
+
+        let format = if dialect_of!(self is BigQueryDialect) && self.parse_keyword(Keyword::FORMAT)
+        {
+            Some(self.parse_value()?)
+        } else {
+            None
+        };
+
         self.expect_token(&Token::RParen)?;
         Ok(Expr::SafeCast {
             expr: Box::new(expr),
             data_type,
+            format,
         })
     }
 
@@ -2070,6 +2097,7 @@ impl<'a> Parser<'a> {
         Ok(Expr::Cast {
             expr: Box::new(expr),
             data_type: self.parse_data_type()?,
+            format: None,
         })
     }
 

diff --git a/tests/sqlparser_bigquery.rs b/tests/sqlparser_bigquery.rs
@@ -222,6 +222,38 @@ fn parse_cast_type() {
     bigquery().verified_only_select(sql);
 }
 
+#[test]
+fn parse_cast_date_format() {
+    let sql =
+        r#"SELECT CAST(date_valid_from AS DATE FORMAT 'YYYY-MM-DD') AS date_valid_from FROM foo"#;
+    bigquery().verified_only_select(sql);
+}
+
+#[test]
+fn parse_cast_time_format() {
+    let sql = r#"SELECT CAST(TIME '21:30:00' AS STRING FORMAT 'PM') AS date_time_to_string"#;
+    bigquery().verified_only_select(sql);
+}
+
+#[test]
+#[ignore]
+fn parse_cast_timestamp_format_tz() {
+    let sql = r#"SELECT CAST(TIMESTAMP '2008-12-25 00:00:00+00:00' AS STRING FORMAT 'TZH' AT TIME ZONE 'Asia/Kolkata') AS date_time_to_string"#;
+    bigquery().verified_only_select(sql);
+}
+
+#[test]
+fn parse_cast_string_to_bytes_format() {
+    let sql = r#"SELECT CAST('Hello' AS BYTES FORMAT 'ASCII') AS string_to_bytes"#;
+    bigquery().verified_only_select(sql);
+}
+
+#[test]
+fn parse_cast_bytes_to_string_format() {
+    let sql = r#"SELECT CAST(B'\x48\x65\x6c\x6c\x6f' AS STRING FORMAT 'ASCII') AS bytes_to_string"#;
+    bigquery().verified_only_select(sql);
+}
+
 #[test]
 fn parse_like() {
     fn chk(negated: bool) {

diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
@@ -1931,6 +1931,7 @@ fn parse_cast() {
         &Expr::Cast {
             expr: Box::new(Expr::Identifier(Ident::new("id"))),
             data_type: DataType::BigInt(None),
+            format: None,
         },
         expr_from_projection(only(&select.projection))
     );
@@ -1941,6 +1942,7 @@ fn parse_cast() {
         &Expr::Cast {
             expr: Box::new(Expr::Identifier(Ident::new("id"))),
             data_type: DataType::TinyInt(None),
+            format: None,
         },
         expr_from_projection(only(&select.projection))
     );
@@ -1967,6 +1969,7 @@ fn parse_cast() {
         &Expr::Cast {
             expr: Box::new(Expr::Identifier(Ident::new("id"))),
             data_type: DataType::Nvarchar(Some(50)),
+            format: None,
         },
         expr_from_projection(only(&select.projection))
     );
@@ -1977,6 +1980,7 @@ fn parse_cast() {
         &Expr::Cast {
             expr: Box::new(Expr::Identifier(Ident::new("id"))),
             data_type: DataType::Clob(None),
+            format: None,
         },
         expr_from_projection(only(&select.projection))
     );
@@ -1987,6 +1991,7 @@ fn parse_cast() {
         &Expr::Cast {
             expr: Box::new(Expr::Identifier(Ident::new("id"))),
             data_type: DataType::Clob(Some(50)),
+            format: None,
         },
         expr_from_projection(only(&select.projection))
     );
@@ -1997,6 +2002,7 @@ fn parse_cast() {
         &Expr::Cast {
             expr: Box::new(Expr::Identifier(Ident::new("id"))),
             data_type: DataType::Binary(Some(50)),
+            format: None,
         },
         expr_from_projection(only(&select.projection))
     );
@@ -2007,6 +2013,7 @@ fn parse_cast() {
         &Expr::Cast {
             expr: Box::new(Expr::Identifier(Ident::new("id"))),
             data_type: DataType::Varbinary(Some(50)),
+            format: None,
         },
         expr_from_projection(only(&select.projection))
     );
@@ -2017,6 +2024,7 @@ fn parse_cast() {
         &Expr::Cast {
             expr: Box::new(Expr::Identifier(Ident::new("id"))),
             data_type: DataType::Blob(None),
+            format: None,
         },
         expr_from_projection(only(&select.projection))
     );
@@ -2027,6 +2035,7 @@ fn parse_cast() {
         &Expr::Cast {
             expr: Box::new(Expr::Identifier(Ident::new("id"))),
             data_type: DataType::Blob(Some(50)),
+            format: None,
         },
         expr_from_projection(only(&select.projection))
     );
@@ -2040,6 +2049,7 @@ fn parse_try_cast() {
         &Expr::TryCast {
             expr: Box::new(Expr::Identifier(Ident::new("id"))),
             data_type: DataType::BigInt(None),
+            format: None,
         },
         expr_from_projection(only(&select.projection))
     );

diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
@@ -1781,7 +1781,8 @@ fn parse_array_index_expr() {
                 })),
                 data_type: DataType::Array(Some(Box::new(DataType::Array(Some(Box::new(
                     DataType::Int(None)
-                ))))))
+                )))))),
+                format: None,
             }))),
             indexes: vec![num[1].clone(), num[2].clone()],
         },

diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs
@@ -167,6 +167,7 @@ fn parse_array() {
         &Expr::Cast {
             expr: Box::new(Expr::Identifier(Ident::new("a"))),
             data_type: DataType::Array(None),
+            format: None,
         },
         expr_from_projection(only(&select.projection))
     );