Support arbitrary composite access expressions (#1600)

Author: ayman-sigma

src/ast/mod.rs

@@ -640,7 +640,7 @@ pub enum Expr {
         /// The path to the data to extract.
         path: JsonPath,
     },
-    /// CompositeAccess (postgres) eg: SELECT (information_schema._pg_expandarray(array['i','i'])).n
+    /// CompositeAccess eg: SELECT foo(bar).z, (information_schema._pg_expandarray(array['i','i'])).n
     CompositeAccess {
         expr: Box<Expr>,
         key: Ident,

src/parser/mod.rs

@@ -969,6 +969,14 @@ impl<'a> Parser<'a> {
         let _guard = self.recursion_counter.try_decrease()?;
         debug!("parsing expr");
         let mut expr = self.parse_prefix()?;
+        // Attempt to parse composite access. Example `SELECT f(x).a`
+        while self.consume_token(&Token::Period) {
+            expr = Expr::CompositeAccess {
+                expr: Box::new(expr),
+                key: self.parse_identifier(false)?,
+            }
+        }
+
         debug!("prefix: {:?}", expr);
         loop {
             let next_precedence = self.get_next_precedence()?;
@@ -1393,25 +1401,7 @@ impl<'a> Parser<'a> {
                     }
                 };
                 self.expect_token(&Token::RParen)?;
-                let expr = self.try_parse_method(expr)?;
-                if !self.consume_token(&Token::Period) {
-                    Ok(expr)
-                } else {
-                    let tok = self.next_token();
-                    let key = match tok.token {
-                        Token::Word(word) => word.to_ident(tok.span),
-                        _ => {
-                            return parser_err!(
-                                format!("Expected identifier, found: {tok}"),
-                                tok.span.start
-                            )
-                        }
-                    };
-                    Ok(Expr::CompositeAccess {
-                        expr: Box::new(expr),
-                        key,
-                    })
-                }
+                self.try_parse_method(expr)
             }
             Token::Placeholder(_) | Token::Colon | Token::AtSign => {
                 self.prev_token();

tests/sqlparser_common.rs

@@ -12341,6 +12341,92 @@ fn parse_create_table_with_bit_types() {
     }
 }
 
+#[test]
+fn parse_composite_access_expr() {
+    assert_eq!(
+        verified_expr("f(a).b"),
+        Expr::CompositeAccess {
+            expr: Box::new(Expr::Function(Function {
+                name: ObjectName(vec![Ident::new("f")]),
+                uses_odbc_syntax: false,
+                parameters: FunctionArguments::None,
+                args: FunctionArguments::List(FunctionArgumentList {
+                    duplicate_treatment: None,
+                    args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
+                        Expr::Identifier(Ident::new("a"))
+                    ))],
+                    clauses: vec![],
+                }),
+                null_treatment: None,
+                filter: None,
+                over: None,
+                within_group: vec![]
+            })),
+            key: Ident::new("b")
+        }
+    );
+
+    // Nested Composite Access
+    assert_eq!(
+        verified_expr("f(a).b.c"),
+        Expr::CompositeAccess {
+            expr: Box::new(Expr::CompositeAccess {
+                expr: Box::new(Expr::Function(Function {
+                    name: ObjectName(vec![Ident::new("f")]),
+                    uses_odbc_syntax: false,
+                    parameters: FunctionArguments::None,
+                    args: FunctionArguments::List(FunctionArgumentList {
+                        duplicate_treatment: None,
+                        args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
+                            Expr::Identifier(Ident::new("a"))
+                        ))],
+                        clauses: vec![],
+                    }),
+                    null_treatment: None,
+                    filter: None,
+                    over: None,
+                    within_group: vec![]
+                })),
+                key: Ident::new("b")
+            }),
+            key: Ident::new("c")
+        }
+    );
+
+    // Composite Access in Select and Where Clauses
+    let stmt = verified_only_select("SELECT f(a).b FROM t WHERE f(a).b IS NOT NULL");
+    let expr = Expr::CompositeAccess {
+        expr: Box::new(Expr::Function(Function {
+            name: ObjectName(vec![Ident::new("f")]),
+            uses_odbc_syntax: false,
+            parameters: FunctionArguments::None,
+            args: FunctionArguments::List(FunctionArgumentList {
+                duplicate_treatment: None,
+                args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
+                    Expr::Identifier(Ident::new("a")),
+                ))],
+                clauses: vec![],
+            }),
+            null_treatment: None,
+            filter: None,
+            over: None,
+            within_group: vec![],
+        })),
+        key: Ident::new("b"),
+    };
+
+    assert_eq!(stmt.projection[0], SelectItem::UnnamedExpr(expr.clone()));
+    assert_eq!(stmt.selection.unwrap(), Expr::IsNotNull(Box::new(expr)));
+
+    // Composite Access with quoted identifier
+    verified_only_select("SELECT f(a).\"an id\"");
+
+    // Composite Access in struct literal
+    all_dialects_where(|d| d.supports_struct_literal()).verified_stmt(
+        "SELECT * FROM t WHERE STRUCT(STRUCT(1 AS a, NULL AS b) AS c, NULL AS d).c.a IS NOT NULL",
+    );
+}
+
 #[test]
 fn parse_create_table_with_enum_types() {
     let sql = "CREATE TABLE t0 (foo ENUM8('a' = 1, 'b' = 2), bar ENUM16('a' = 1, 'b' = 2), baz ENUM('a', 'b'))";