PartiQL queries in Redshift (#1534)

Author: yoavcloud

src/ast/query.rs

@@ -974,6 +974,8 @@ pub enum TableFactor {
         with_ordinality: bool,
         /// [Partition selection](https://dev.mysql.com/doc/refman/8.0/en/partitioning-selection.html), supported by MySQL.
         partitions: Vec<Ident>,
+        /// Optional PartiQL JsonPath: <https://partiql.org/dql/from.html>
+        json_path: Option<JsonPath>,
     },
     Derived {
         lateral: bool,
@@ -1375,8 +1377,12 @@ impl fmt::Display for TableFactor {
                 version,
                 partitions,
                 with_ordinality,
+                json_path,
             } => {
                 write!(f, "{name}")?;
+                if let Some(json_path) = json_path {
+                    write!(f, "{json_path}")?;
+                }
                 if !partitions.is_empty() {
                     write!(f, "PARTITION ({})", display_comma_separated(partitions))?;
                 }

src/dialect/mod.rs

@@ -675,6 +675,12 @@ pub trait Dialect: Debug + Any {
     fn supports_create_table_select(&self) -> bool {
         false
     }
+
+    /// Returns true if the dialect supports PartiQL for querying semi-structured data
+    /// <https://partiql.org/index.html>
+    fn supports_partiql(&self) -> bool {
+        false
+    }
 }
 
 /// This represents the operators for which precedence must be defined

src/dialect/redshift.rs

@@ -74,4 +74,9 @@ impl Dialect for RedshiftSqlDialect {
     fn supports_top_before_distinct(&self) -> bool {
         true
     }
+
+    /// Redshift supports PartiQL: <https://docs.aws.amazon.com/redshift/latest/dg/super-overview.html>
+    fn supports_partiql(&self) -> bool {
+        true
+    }
 }

src/parser/mod.rs

@@ -2936,7 +2936,7 @@ impl<'a> Parser<'a> {
         } else if Token::LBracket == tok {
             if dialect_of!(self is PostgreSqlDialect | DuckDbDialect | GenericDialect) {
                 self.parse_subscript(expr)
-            } else if dialect_of!(self is SnowflakeDialect) {
+            } else if dialect_of!(self is SnowflakeDialect) || self.dialect.supports_partiql() {
                 self.prev_token();
                 self.parse_json_access(expr)
             } else {
@@ -3072,6 +3072,14 @@ impl<'a> Parser<'a> {
     }
 
     fn parse_json_access(&mut self, expr: Expr) -> Result<Expr, ParserError> {
+        let path = self.parse_json_path()?;
+        Ok(Expr::JsonAccess {
+            value: Box::new(expr),
+            path,
+        })
+    }
+
+    fn parse_json_path(&mut self) -> Result<JsonPath, ParserError> {
         let mut path = Vec::new();
         loop {
             match self.next_token().token {
@@ -3095,10 +3103,7 @@ impl<'a> Parser<'a> {
         }
 
         debug_assert!(!path.is_empty());
-        Ok(Expr::JsonAccess {
-            value: Box::new(expr),
-            path: JsonPath { path },
-        })
+        Ok(JsonPath { path })
     }
 
     pub fn parse_map_access(&mut self, expr: Expr) -> Result<Expr, ParserError> {
@@ -10338,6 +10343,11 @@ impl<'a> Parser<'a> {
         } else {
             let name = self.parse_object_name(true)?;
 
+            let json_path = match self.peek_token().token {
+                Token::LBracket if self.dialect.supports_partiql() => Some(self.parse_json_path()?),
+                _ => None,
+            };
+
             let partitions: Vec<Ident> = if dialect_of!(self is MySqlDialect | GenericDialect)
                 && self.parse_keyword(Keyword::PARTITION)
             {
@@ -10380,6 +10390,7 @@ impl<'a> Parser<'a> {
                 version,
                 partitions,
                 with_ordinality,
+                json_path,
             };
 
             while let Some(kw) = self.parse_one_of_keywords(&[Keyword::PIVOT, Keyword::UNPIVOT]) {

src/test_utils.rs

@@ -345,6 +345,7 @@ pub fn table(name: impl Into<String>) -> TableFactor {
         version: None,
         partitions: vec![],
         with_ordinality: false,
+        json_path: None,
     }
 }
 
@@ -360,6 +361,7 @@ pub fn table_with_alias(name: impl Into<String>, alias: impl Into<String>) -> Ta
         version: None,
         partitions: vec![],
         with_ordinality: false,
+        json_path: None,
     }
 }
 

tests/sqlparser_bigquery.rs

@@ -229,6 +229,7 @@ fn parse_delete_statement() {
                     version: None,
                     partitions: vec![],
                     with_ordinality: false,
+                    json_path: None,
                 },
                 from[0].relation
             );
@@ -1373,6 +1374,7 @@ fn parse_table_identifiers() {
                     version: None,
                     partitions: vec![],
                     with_ordinality: false,
+                    json_path: None,
                 },
                 joins: vec![]
             },]
@@ -1546,6 +1548,7 @@ fn parse_table_time_travel() {
                 ))),
                 partitions: vec![],
                 with_ordinality: false,
+                json_path: None,
             },
             joins: vec![]
         },]
@@ -1644,6 +1647,7 @@ fn parse_merge() {
                     version: Default::default(),
                     partitions: Default::default(),
                     with_ordinality: false,
+                    json_path: None,
                 },
                 table
             );
@@ -1659,6 +1663,7 @@ fn parse_merge() {
                     version: Default::default(),
                     partitions: Default::default(),
                     with_ordinality: false,
+                    json_path: None,
                 },
                 source
             );

tests/sqlparser_clickhouse.rs

@@ -67,6 +67,7 @@ fn parse_map_access_expr() {
                     version: None,
                     partitions: vec![],
                     with_ordinality: false,
+                    json_path: None,
                 },
                 joins: vec![],
             }],
@@ -172,6 +173,7 @@ fn parse_delimited_identifiers() {
             version,
             with_ordinality: _,
             partitions: _,
+            json_path: _,
         } => {
             assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);
             assert_eq!(Ident::with_quote('"', "alias"), alias.unwrap().name);