Handle derived tables with set operations

This commit adds support for derived tables (i.e., subqueries) that
incorporate set operations, like:

    SELECT * FROM (((SELECT 1) UNION (SELECT 2)) t1 AS NATURAL JOIN t2)

This introduces a bit of complexity around determining whether a left
paren starts a subquery, starts a nested join, or belongs to an
already-started subquery. The details are explained in a comment within
the patch.

Author: benesch

src/sqlparser.rs

@@ -42,6 +42,12 @@ pub enum IsOptional {
 }
 use IsOptional::*;
 
+pub enum IsLateral {
+    Lateral,
+    NotLateral,
+}
+use IsLateral::*;
+
 impl From<TokenizerError> for ParserError {
     fn from(e: TokenizerError) -> Self {
         ParserError::TokenizerError(format!("{:?}", e))
@@ -1670,30 +1676,42 @@ impl Parser {
 
     /// A table name or a parenthesized subquery, followed by optional `[AS] alias`
     pub fn parse_table_factor(&mut self) -> Result<TableFactor, ParserError> {
-        let lateral = self.parse_keyword("LATERAL");
+        if self.parse_keyword("LATERAL") {
+            // LATERAL must always be followed by a subquery.
+            if !self.consume_token(&Token::LParen) {
+                self.expected("subquery after LATERAL", self.peek_token())?;
+            }
+            return self.parse_derived_table_factor(Lateral);
+        }
+
         if self.consume_token(&Token::LParen) {
-            if self.parse_keyword("SELECT")
-                || self.parse_keyword("WITH")
-                || self.parse_keyword("VALUES")
-            {
-                self.prev_token();
-                let subquery = Box::new(self.parse_query()?);
-                self.expect_token(&Token::RParen)?;
-                let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;
-                Ok(TableFactor::Derived {
-                    lateral,
-                    subquery,
-                    alias,
-                })
-            } else if lateral {
-                parser_err!("Expected subquery after LATERAL, found nested join".to_string())
-            } else {
-                let table_reference = self.parse_table_and_joins()?;
-                self.expect_token(&Token::RParen)?;
-                Ok(TableFactor::NestedJoin(Box::new(table_reference)))
+            let index = self.index;
+            // A left paren introduces either a derived table (i.e., a subquery)
+            // or a nested join. It's nearly impossible to determine ahead of
+            // time which it is... so we just try to parse both.
+            //
+            // Here's an example that demonstrates the complexity:
+            //
+            //     SELECT * FROM ( ( (SELECT 1) UNION (SELECT 2) ) AS t1 NATURAL JOIN t2 )
+            //                   ^ ^ ^
+            //                   | | |
+            //                   | | belongs to the subquery
+            //                   | starts a derived table (subquery)
+            //                   starts a nested join
+            //
+            match self.parse_derived_table_factor(NotLateral) {
+                Ok(table_factor) => Ok(table_factor),
+                Err(_) => {
+                    // It wasn't a derived table. Ignore the error and back up
+                    // to where we were before. Either we'll be able to parse a
+                    // valid nested join, or we won't, and we'll return that
+                    // error instead.
+                    self.index = index;
+                    let table_and_joins = self.parse_table_and_joins()?;
+                    self.expect_token(&Token::RParen)?;
+                    Ok(TableFactor::NestedJoin(Box::new(table_and_joins)))
+                }
             }
-        } else if lateral {
-            self.expected("subquery after LATERAL", self.peek_token())
         } else {
             let name = self.parse_object_name()?;
             // Postgres, MSSQL: table-valued functions:
@@ -1723,6 +1741,23 @@ impl Parser {
         }
     }
 
+    pub fn parse_derived_table_factor(
+        &mut self,
+        lateral: IsLateral,
+    ) -> Result<TableFactor, ParserError> {
+        let subquery = Box::new(self.parse_query()?);
+        self.expect_token(&Token::RParen)?;
+        let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;
+        Ok(TableFactor::Derived {
+            lateral: match lateral {
+                Lateral => true,
+                NotLateral => false,
+            },
+            subquery,
+            alias,
+        })
+    }
+
     fn parse_join_constraint(&mut self, natural: bool) -> Result<JoinConstraint, ParserError> {
         if natural {
             Ok(JoinConstraint::Natural)

tests/sqlparser_common.rs

@@ -1837,6 +1837,32 @@ fn parse_derived_tables() {
     let sql = "SELECT * FROM t NATURAL JOIN (((SELECT 1)))";
     let _ = verified_only_select(sql);
     // TODO: add assertions
+
+    let sql = "SELECT * FROM (((SELECT 1) UNION (SELECT 2)) AS t1 NATURAL JOIN t2)";
+    let select = verified_only_select(sql);
+    let from = only(select.from);
+    assert_eq!(
+        from.relation,
+        TableFactor::NestedJoin(Box::new(TableWithJoins {
+            relation: TableFactor::Derived {
+                lateral: false,
+                subquery: Box::new(verified_query("(SELECT 1) UNION (SELECT 2)")),
+                alias: Some(TableAlias {
+                    name: "t1".into(),
+                    columns: vec![],
+                })
+            },
+            joins: vec![Join {
+                relation: TableFactor::Table {
+                    name: SQLObjectName(vec!["t2".into()]),
+                    alias: None,
+                    args: vec![],
+                    with_hints: vec![],
+                },
+                join_operator: JoinOperator::Inner(JoinConstraint::Natural),
+            }],
+        }))
+    )
 }
 
 #[test]
@@ -2349,7 +2375,9 @@ fn lateral_derived() {
     let sql = "SELECT * FROM a LEFT JOIN LATERAL (b CROSS JOIN c)";
     let res = parse_sql_statements(sql);
     assert_eq!(
-        ParserError::ParserError("Expected subquery after LATERAL, found nested join".to_string()),
+        ParserError::ParserError(
+            "Expected SELECT or a subquery in the query body, found: b".to_string()
+        ),
         res.unwrap_err()
     );
 }