Add FETCH and OFFSET support

Author: thomas-jeepe

src/dialect/keywords.rs

@@ -155,6 +155,7 @@ define_keywords!(
     EXTRACT,
     FALSE,
     FETCH,
+    FIRST,
     FILTER,
     FIRST_VALUE,
     FLOAT,
@@ -229,6 +230,7 @@ define_keywords!(
     NATURAL,
     NCHAR,
     NCLOB,
+    NEXT,
     NEW,
     NO,
     NONE,
@@ -341,6 +343,7 @@ define_keywords!(
     TABLESAMPLE,
     TEXT,
     THEN,
+    TIES,
     TIME,
     TIMESTAMP,
     TIMEZONE_HOUR,
@@ -396,7 +399,7 @@ pub const RESERVED_FOR_TABLE_ALIAS: &[&str] = &[
     // Reserved as both a table and a column alias:
     WITH, SELECT, WHERE, GROUP, ORDER, UNION, EXCEPT, INTERSECT,
     // Reserved only as a table alias in the `FROM`/`JOIN` clauses:
-    ON, JOIN, INNER, CROSS, FULL, LEFT, RIGHT, NATURAL, USING, LIMIT,
+    ON, JOIN, INNER, CROSS, FULL, LEFT, RIGHT, NATURAL, USING, LIMIT, OFFSET, FETCH,
 ];
 
 /// Can't be used as a column alias, so that `SELECT <expr> alias`

src/sqlast/mod.rs

@@ -21,8 +21,8 @@ mod table_key;
 mod value;
 
 pub use self::query::{
-    Cte, Join, JoinConstraint, JoinOperator, SQLOrderByExpr, SQLQuery, SQLSelect, SQLSelectItem,
-    SQLSetExpr, SQLSetOperator, TableFactor,
+    Cte, Fetch, Join, JoinConstraint, JoinOperator, SQLOrderByExpr, SQLQuery, SQLSelect,
+    SQLSelectItem, SQLSetExpr, SQLSetOperator, TableFactor,
 };
 pub use self::sqltype::SQLType;
 pub use self::table_key::{AlterOperation, Key, TableKey};

src/sqlast/query.rs

@@ -10,8 +10,12 @@ pub struct SQLQuery {
     pub body: SQLSetExpr,
     /// ORDER BY
     pub order_by: Vec<SQLOrderByExpr>,
-    /// LIMIT
+    /// LIMIT { <N> | ALL }
     pub limit: Option<ASTNode>,
+    /// OFFSET <N> { ROW | ROWS }
+    pub offset: Option<ASTNode>,
+    /// FETCH { FIRST | NEXT } <N> [ PERCENT ] { ROW | ROWS } | { ONLY | WITH TIES }
+    pub fetch: Option<Fetch>,
 }
 
 impl ToString for SQLQuery {
@@ -27,6 +31,13 @@ impl ToString for SQLQuery {
         if let Some(ref limit) = self.limit {
             s += &format!(" LIMIT {}", limit.to_string());
         }
+        if let Some(ref offset) = self.offset {
+            s += &format!(" OFFSET {} ROWS", offset.to_string());
+        }
+        if let Some(ref fetch) = self.fetch {
+            s.push(' ');
+            s += &fetch.to_string();
+        }
         s
     }
 }
@@ -320,3 +331,27 @@ impl ToString for SQLOrderByExpr {
         }
     }
 }
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct Fetch {
+    pub with_ties: bool,
+    pub percent: bool,
+    pub quantity: Option<ASTNode>,
+}
+
+impl ToString for Fetch {
+    fn to_string(&self) -> String {
+        let extension = if self.with_ties { "WITH TIES" } else { "ONLY" };
+        if let Some(ref quantity) = self.quantity {
+            let percent = if self.percent { " PERCENT" } else { "" };
+            format!(
+                "FETCH FIRST {}{} ROWS {}",
+                quantity.to_string(),
+                percent,
+                extension
+            )
+        } else {
+            format!("FETCH FIRST ROWS {}", extension)
+        }
+    }
+}

src/sqlparser.rs

@@ -684,6 +684,40 @@ impl Parser {
         true
     }
 
+    /// Look for one of the given keywords and return the one that matches.
+    #[must_use]
+    pub fn parse_one_of_keywords(&mut self, keywords: &[&'static str]) -> Option<&'static str> {
+        for keyword in keywords {
+            assert!(keywords::ALL_KEYWORDS.contains(keyword));
+        }
+        match self.peek_token() {
+            Some(Token::SQLWord(ref k)) => keywords
+                .iter()
+                .find(|keyword| keyword.eq_ignore_ascii_case(&k.keyword))
+                .map(|keyword| {
+                    self.next_token();
+                    *keyword
+                }),
+            _ => None,
+        }
+    }
+
+    /// Bail out if the current token is not one of the expected keywords, or consume it if it is
+    #[must_use]
+    pub fn expect_one_of_keywords(
+        &mut self,
+        keywords: &[&'static str],
+    ) -> Result<&'static str, ParserError> {
+        if let Some(keyword) = self.parse_one_of_keywords(keywords) {
+            Ok(keyword)
+        } else {
+            self.expected(
+                &format!("one of {}", keywords.join(" or ")),
+                self.peek_token(),
+            )
+        }
+    }
+
     /// Bail out if the current token is not an expected keyword, or consume it if it is
     pub fn expect_keyword(&mut self, expected: &'static str) -> Result<(), ParserError> {
         if self.parse_keyword(expected) {
@@ -1279,11 +1313,25 @@ impl Parser {
             None
         };
 
+        let offset = if self.parse_keyword("OFFSET") {
+            Some(self.parse_offset()?)
+        } else {
+            None
+        };
+
+        let fetch = if self.parse_keyword("FETCH") {
+            Some(self.parse_fetch()?)
+        } else {
+            None
+        };
+
         Ok(SQLQuery {
             ctes,
             body,
             limit,
             order_by,
+            offset,
+            fetch,
         })
     }
 
@@ -1655,6 +1703,40 @@ impl Parser {
                 .map(|n| Some(ASTNode::SQLValue(Value::Long(n))))
         }
     }
+
+    /// Parse an OFFSET clause
+    pub fn parse_offset(&mut self) -> Result<ASTNode, ParserError> {
+        let value = self
+            .parse_literal_int()
+            .map(|n| ASTNode::SQLValue(Value::Long(n)))?;
+        self.expect_one_of_keywords(&["ROW", "ROWS"])?;
+        Ok(value)
+    }
+
+    /// Parse a FETCH clause
+    pub fn parse_fetch(&mut self) -> Result<Fetch, ParserError> {
+        self.expect_one_of_keywords(&["FIRST", "NEXT"])?;
+        let (quantity, percent) = if self.parse_one_of_keywords(&["ROW", "ROWS"]).is_some() {
+            (None, false)
+        } else {
+            let quantity = self.parse_sql_value()?;
+            let percent = self.parse_keyword("PERCENT");
+            self.expect_one_of_keywords(&["ROW", "ROWS"])?;
+            (Some(quantity), percent)
+        };
+        let with_ties = if self.parse_keyword("ONLY") {
+            false
+        } else if self.parse_keywords(vec!["WITH", "TIES"]) {
+            true
+        } else {
+            return self.expected("one of ONLY or WITH TIES", self.peek_token());
+        };
+        Ok(Fetch {
+            with_ties,
+            percent,
+            quantity,
+        })
+    }
 }
 
 impl SQLWord {

tests/sqlparser_common.rs

@@ -1388,6 +1388,187 @@ fn parse_invalid_subquery_without_parens() {
     );
 }
 
+#[test]
+fn parse_offset() {
+    let ast = verified_query("SELECT foo FROM bar OFFSET 2 ROWS");
+    assert_eq!(ast.offset, Some(ASTNode::SQLValue(Value::Long(2))));
+    let ast = verified_query("SELECT foo FROM bar WHERE foo = 4 OFFSET 2 ROWS");
+    assert_eq!(ast.offset, Some(ASTNode::SQLValue(Value::Long(2))));
+    let ast = verified_query("SELECT foo FROM bar ORDER BY baz OFFSET 2 ROWS");
+    assert_eq!(ast.offset, Some(ASTNode::SQLValue(Value::Long(2))));
+    let ast = verified_query("SELECT foo FROM bar WHERE foo = 4 ORDER BY baz OFFSET 2 ROWS");
+    assert_eq!(ast.offset, Some(ASTNode::SQLValue(Value::Long(2))));
+    let ast = verified_query("SELECT foo FROM (SELECT * FROM bar OFFSET 2 ROWS) OFFSET 2 ROWS");
+    assert_eq!(ast.offset, Some(ASTNode::SQLValue(Value::Long(2))));
+    match ast.body {
+        SQLSetExpr::Select(s) => match s.relation {
+            Some(TableFactor::Derived { subquery, .. }) => {
+                assert_eq!(subquery.offset, Some(ASTNode::SQLValue(Value::Long(2))));
+            }
+            _ => panic!("Test broke"),
+        },
+        _ => panic!("Test broke"),
+    }
+}
+
+#[test]
+fn parse_singular_row_offset() {
+    one_statement_parses_to(
+        "SELECT foo FROM bar OFFSET 1 ROW",
+        "SELECT foo FROM bar OFFSET 1 ROWS",
+    );
+}
+
+#[test]
+fn parse_fetch() {
+    let ast = verified_query("SELECT foo FROM bar FETCH FIRST 2 ROWS ONLY");
+    assert_eq!(
+        ast.fetch,
+        Some(Fetch {
+            with_ties: false,
+            percent: false,
+            quantity: Some(ASTNode::SQLValue(Value::Long(2))),
+        })
+    );
+    let ast = verified_query("SELECT foo FROM bar FETCH FIRST ROWS ONLY");
+    assert_eq!(
+        ast.fetch,
+        Some(Fetch {
+            with_ties: false,
+            percent: false,
+            quantity: None,
+        })
+    );
+    let ast = verified_query("SELECT foo FROM bar WHERE foo = 4 FETCH FIRST 2 ROWS ONLY");
+    assert_eq!(
+        ast.fetch,
+        Some(Fetch {
+            with_ties: false,
+            percent: false,
+            quantity: Some(ASTNode::SQLValue(Value::Long(2))),
+        })
+    );
+    let ast = verified_query("SELECT foo FROM bar ORDER BY baz FETCH FIRST 2 ROWS ONLY");
+    assert_eq!(
+        ast.fetch,
+        Some(Fetch {
+            with_ties: false,
+            percent: false,
+            quantity: Some(ASTNode::SQLValue(Value::Long(2))),
+        })
+    );
+    let ast = verified_query(
+        "SELECT foo FROM bar WHERE foo = 4 ORDER BY baz FETCH FIRST 2 ROWS WITH TIES",
+    );
+    assert_eq!(
+        ast.fetch,
+        Some(Fetch {
+            with_ties: true,
+            percent: false,
+            quantity: Some(ASTNode::SQLValue(Value::Long(2))),
+        })
+    );
+    let ast = verified_query("SELECT foo FROM bar FETCH FIRST 50 PERCENT ROWS ONLY");
+    assert_eq!(
+        ast.fetch,
+        Some(Fetch {
+            with_ties: false,
+            percent: true,
+            quantity: Some(ASTNode::SQLValue(Value::Long(50))),
+        })
+    );
+    let ast = verified_query(
+        "SELECT foo FROM bar WHERE foo = 4 ORDER BY baz OFFSET 2 ROWS FETCH FIRST 2 ROWS ONLY",
+    );
+    assert_eq!(ast.offset, Some(ASTNode::SQLValue(Value::Long(2))));
+    assert_eq!(
+        ast.fetch,
+        Some(Fetch {
+            with_ties: false,
+            percent: false,
+            quantity: Some(ASTNode::SQLValue(Value::Long(2))),
+        })
+    );
+    let ast = verified_query(
+        "SELECT foo FROM (SELECT * FROM bar FETCH FIRST 2 ROWS ONLY) FETCH FIRST 2 ROWS ONLY",
+    );
+    assert_eq!(
+        ast.fetch,
+        Some(Fetch {
+            with_ties: false,
+            percent: false,
+            quantity: Some(ASTNode::SQLValue(Value::Long(2))),
+        })
+    );
+    match ast.body {
+        SQLSetExpr::Select(s) => match s.relation {
+            Some(TableFactor::Derived { subquery, .. }) => {
+                assert_eq!(
+                    subquery.fetch,
+                    Some(Fetch {
+                        with_ties: false,
+                        percent: false,
+                        quantity: Some(ASTNode::SQLValue(Value::Long(2))),
+                    })
+                );
+            }
+            _ => panic!("Test broke"),
+        },
+        _ => panic!("Test broke"),
+    }
+    let ast = verified_query("SELECT foo FROM (SELECT * FROM bar OFFSET 2 ROWS FETCH FIRST 2 ROWS ONLY) OFFSET 2 ROWS FETCH FIRST 2 ROWS ONLY");
+    assert_eq!(ast.offset, Some(ASTNode::SQLValue(Value::Long(2))));
+    assert_eq!(
+        ast.fetch,
+        Some(Fetch {
+            with_ties: false,
+            percent: false,
+            quantity: Some(ASTNode::SQLValue(Value::Long(2))),
+        })
+    );
+    match ast.body {
+        SQLSetExpr::Select(s) => match s.relation {
+            Some(TableFactor::Derived { subquery, .. }) => {
+                assert_eq!(subquery.offset, Some(ASTNode::SQLValue(Value::Long(2))));
+                assert_eq!(
+                    subquery.fetch,
+                    Some(Fetch {
+                        with_ties: false,
+                        percent: false,
+                        quantity: Some(ASTNode::SQLValue(Value::Long(2))),
+                    })
+                );
+            }
+            _ => panic!("Test broke"),
+        },
+        _ => panic!("Test broke"),
+    }
+}
+
+#[test]
+fn parse_fetch_variations() {
+    one_statement_parses_to(
+        "SELECT foo FROM bar FETCH FIRST 10 ROW ONLY",
+        "SELECT foo FROM bar FETCH FIRST 10 ROWS ONLY",
+    );
+    one_statement_parses_to(
+        "SELECT foo FROM bar FETCH NEXT 10 ROW ONLY",
+        "SELECT foo FROM bar FETCH FIRST 10 ROWS ONLY",
+    );
+    one_statement_parses_to(
+        "SELECT foo FROM bar FETCH NEXT 10 ROWS WITH TIES",
+        "SELECT foo FROM bar FETCH FIRST 10 ROWS WITH TIES",
+    );
+    one_statement_parses_to(
+        "SELECT foo FROM bar FETCH NEXT ROWS WITH TIES",
+        "SELECT foo FROM bar FETCH FIRST ROWS WITH TIES",
+    );
+    one_statement_parses_to(
+        "SELECT foo FROM bar FETCH FIRST ROWS ONLY",
+        "SELECT foo FROM bar FETCH FIRST ROWS ONLY",
+    );
+}
+
 #[test]
 #[should_panic(
     expected = "Parse results with GenericSqlDialect are different from PostgreSqlDialect"