From 3cdd74564e828bb429e4e7d919f9638bb231fe3f Mon Sep 17 00:00:00 2001
From: Yoav Cohen <yoav@satoricyber.com>
Date: Wed, 5 Feb 2025 21:46:21 +0100
Subject: [PATCH] Fix incorrect parsing of JsonAccess bracket notation after
 cast in Snowflake dialect

---
 src/dialect/duckdb.rs        |  5 +++++
 src/dialect/generic.rs       |  4 ++++
 src/dialect/mod.rs           |  6 ++++++
 src/dialect/postgresql.rs    |  5 +++++
 src/parser/mod.rs            | 17 +++++++----------
 tests/sqlparser_snowflake.rs | 26 ++++++++++++++++++++++++++
 6 files changed, 53 insertions(+), 10 deletions(-)

diff --git a/src/dialect/duckdb.rs b/src/dialect/duckdb.rs
index c41aec81d..e8dcf853c 100644
--- a/src/dialect/duckdb.rs
+++ b/src/dialect/duckdb.rs
@@ -80,4 +80,9 @@ impl Dialect for DuckDbDialect {
     fn supports_load_extension(&self) -> bool {
         true
     }
+
+    // See DuckDB <https://duckdb.org/docs/sql/data_types/array.html#defining-an-array-field>
+    fn supports_array_typedef_size(&self) -> bool {
+        true
+    }
 }
diff --git a/src/dialect/generic.rs b/src/dialect/generic.rs
index 4021b5753..7aaf00742 100644
--- a/src/dialect/generic.rs
+++ b/src/dialect/generic.rs
@@ -143,4 +143,8 @@ impl Dialect for GenericDialect {
     fn supports_string_escape_constant(&self) -> bool {
         true
     }
+
+    fn supports_array_typedef_size(&self) -> bool {
+        true
+    }
 }
diff --git a/src/dialect/mod.rs b/src/dialect/mod.rs
index 965e6c77f..6b04bacc1 100644
--- a/src/dialect/mod.rs
+++ b/src/dialect/mod.rs
@@ -890,6 +890,12 @@ pub trait Dialect: Debug + Any {
     fn requires_single_line_comment_whitespace(&self) -> bool {
         false
     }
+
+    /// Returns true if the dialect supports size definition for array types.
+    /// For example: ```CREATE TABLE my_table (my_array INT[3])```.
+    fn supports_array_typedef_size(&self) -> bool {
+        false
+    }
 }
 
 /// This represents the operators for which precedence must be defined
diff --git a/src/dialect/postgresql.rs b/src/dialect/postgresql.rs
index 5ce4250fb..74b963e82 100644
--- a/src/dialect/postgresql.rs
+++ b/src/dialect/postgresql.rs
@@ -253,6 +253,11 @@ impl Dialect for PostgreSqlDialect {
     fn supports_numeric_literal_underscores(&self) -> bool {
         true
     }
+
+    /// See: <https://www.postgresql.org/docs/current/arrays.html#ARRAYS-DECLARATION>
+    fn supports_array_typedef_size(&self) -> bool {
+        true
+    }
 }
 
 pub fn parse_create(parser: &mut Parser) -> Option<Result<Statement, ParserError>> {
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index 6d84ff843..9cc3bd612 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -8943,16 +8943,13 @@ impl<'a> Parser<'a> {
             _ => self.expected_at("a data type name", next_token_index),
         }?;
 
-        // Parse array data types. Note: this is postgresql-specific and different from
-        // Keyword::ARRAY syntax from above
-        while self.consume_token(&Token::LBracket) {
-            let size = if dialect_of!(self is GenericDialect | DuckDbDialect | PostgreSqlDialect) {
-                self.maybe_parse(|p| p.parse_literal_uint())?
-            } else {
-                None
-            };
-            self.expect_token(&Token::RBracket)?;
-            data = DataType::Array(ArrayElemTypeDef::SquareBracket(Box::new(data), size))
+        if self.dialect.supports_array_typedef_size() {
+            // Parse array data type size
+            while self.consume_token(&Token::LBracket) {
+                let size = self.maybe_parse(|p| p.parse_literal_uint())?;
+                self.expect_token(&Token::RBracket)?;
+                data = DataType::Array(ArrayElemTypeDef::SquareBracket(Box::new(data), size))
+            }
         }
         Ok((data, trailing_bracket))
     }
diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs
index a18f1a4d8..f0fc34069 100644
--- a/tests/sqlparser_snowflake.rs
+++ b/tests/sqlparser_snowflake.rs
@@ -1256,6 +1256,32 @@ fn parse_semi_structured_data_traversal() {
             .to_string(),
         "sql parser error: Expected: variant object key name, found: 42"
     );
+
+    // casting a json access and accessing an array element
+    assert_eq!(
+        snowflake().verified_expr("a:b::ARRAY[1]"),
+        Expr::JsonAccess {
+            value: Box::new(Expr::Cast {
+                kind: CastKind::DoubleColon,
+                data_type: DataType::Array(ArrayElemTypeDef::None),
+                format: None,
+                expr: Box::new(Expr::JsonAccess {
+                    value: Box::new(Expr::Identifier(Ident::new("a"))),
+                    path: JsonPath {
+                        path: vec![JsonPathElem::Dot {
+                            key: "b".to_string(),
+                            quoted: false
+                        }]
+                    }
+                })
+            }),
+            path: JsonPath {
+                path: vec![JsonPathElem::Bracket {
+                    key: Expr::Value(number("1"))
+                }]
+            }
+        }
+    );
 }
 
 #[test]