feat: Support FETCH (cursors) (#510)

Author: ovr

src/ast/mod.rs

@@ -895,6 +895,17 @@ pub enum Statement {
         /// deleted along with the dropped table
         purge: bool,
     },
+    /// FETCH - retrieve rows from a query using a cursor
+    ///
+    /// Note: this is a PostgreSQL-specific statement,
+    /// but may also compatible with other SQL.
+    Fetch {
+        /// Cursor name
+        name: Ident,
+        direction: FetchDirection,
+        /// Optional, It's possible to fetch rows form cursor to the table
+        into: Option<ObjectName>,
+    },
     /// DISCARD [ ALL | PLANS | SEQUENCES | TEMPORARY | TEMP ]
     ///
     /// Note: this is a PostgreSQL-specific statement,
@@ -1114,6 +1125,21 @@ impl fmt::Display for Statement {
                 write!(f, "{}", statement)
             }
             Statement::Query(s) => write!(f, "{}", s),
+            Statement::Fetch {
+                name,
+                direction,
+                into,
+            } => {
+                write!(f, "FETCH {} ", direction)?;
+
+                write!(f, "IN {}", name)?;
+
+                if let Some(into) = into {
+                    write!(f, " INTO {}", into)?;
+                }
+
+                Ok(())
+            }
             Statement::Directory {
                 overwrite,
                 local,
@@ -1859,6 +1885,69 @@ impl fmt::Display for Privileges {
     }
 }
 
+/// Specific direction for FETCH statement
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum FetchDirection {
+    Count { limit: Value },
+    Next,
+    Prior,
+    First,
+    Last,
+    Absolute { limit: Value },
+    Relative { limit: Value },
+    All,
+    // FORWARD
+    // FORWARD count
+    Forward { limit: Option<Value> },
+    ForwardAll,
+    // BACKWARD
+    // BACKWARD count
+    Backward { limit: Option<Value> },
+    BackwardAll,
+}
+
+impl fmt::Display for FetchDirection {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            FetchDirection::Count { limit } => f.write_str(&limit.to_string())?,
+            FetchDirection::Next => f.write_str("NEXT")?,
+            FetchDirection::Prior => f.write_str("PRIOR")?,
+            FetchDirection::First => f.write_str("FIRST")?,
+            FetchDirection::Last => f.write_str("LAST")?,
+            FetchDirection::Absolute { limit } => {
+                f.write_str("ABSOLUTE ")?;
+                f.write_str(&limit.to_string())?;
+            }
+            FetchDirection::Relative { limit } => {
+                f.write_str("RELATIVE ")?;
+                f.write_str(&limit.to_string())?;
+            }
+            FetchDirection::All => f.write_str("ALL")?,
+            FetchDirection::Forward { limit } => {
+                f.write_str("FORWARD")?;
+
+                if let Some(l) = limit {
+                    f.write_str(" ")?;
+                    f.write_str(&l.to_string())?;
+                }
+            }
+            FetchDirection::ForwardAll => f.write_str("FORWARD ALL")?,
+            FetchDirection::Backward { limit } => {
+                f.write_str("BACKWARD")?;
+
+                if let Some(l) = limit {
+                    f.write_str(" ")?;
+                    f.write_str(&l.to_string())?;
+                }
+            }
+            FetchDirection::BackwardAll => f.write_str("BACKWARD ALL")?,
+        };
+
+        Ok(())
+    }
+}
+
 /// A privilege on a database object (table, sequence, etc.).
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]

src/keywords.rs

@@ -67,6 +67,7 @@ macro_rules! define_keywords {
 define_keywords!(
     ABORT,
     ABS,
+    ABSOLUTE,
     ACTION,
     ADD,
     ALL,
@@ -93,6 +94,7 @@ define_keywords!(
     AUTO_INCREMENT,
     AVG,
     AVRO,
+    BACKWARD,
     BEGIN,
     BEGIN_FRAME,
     BEGIN_PARTITION,
@@ -239,6 +241,7 @@ define_keywords!(
     FORCE_QUOTE,
     FOREIGN,
     FORMAT,
+    FORWARD,
     FRAME_ROW,
     FREE,
     FREEZE,
@@ -387,6 +390,7 @@ define_keywords!(
     PREPARE,
     PRESERVE,
     PRIMARY,
+    PRIOR,
     PRIVILEGES,
     PROCEDURE,
     PROGRAM,
@@ -415,6 +419,7 @@ define_keywords!(
     REGR_SXX,
     REGR_SXY,
     REGR_SYY,
+    RELATIVE,
     RELEASE,
     RENAME,
     REPAIR,

src/parser.rs

@@ -167,6 +167,7 @@ impl<'a> Parser<'a> {
                 Keyword::CREATE => Ok(self.parse_create()?),
                 Keyword::DROP => Ok(self.parse_drop()?),
                 Keyword::DISCARD => Ok(self.parse_discard()?),
+                Keyword::FETCH => Ok(self.parse_fetch_statement()?),
                 Keyword::DELETE => Ok(self.parse_delete()?),
                 Keyword::INSERT => Ok(self.parse_insert()?),
                 Keyword::UPDATE => Ok(self.parse_update()?),
@@ -1824,6 +1825,67 @@ impl<'a> Parser<'a> {
         })
     }
 
+    // FETCH [ direction { FROM | IN } ] cursor INTO target;
+    pub fn parse_fetch_statement(&mut self) -> Result<Statement, ParserError> {
+        let direction = if self.parse_keyword(Keyword::NEXT) {
+            FetchDirection::Next
+        } else if self.parse_keyword(Keyword::PRIOR) {
+            FetchDirection::Prior
+        } else if self.parse_keyword(Keyword::FIRST) {
+            FetchDirection::First
+        } else if self.parse_keyword(Keyword::LAST) {
+            FetchDirection::Last
+        } else if self.parse_keyword(Keyword::ABSOLUTE) {
+            FetchDirection::Absolute {
+                limit: self.parse_number_value()?,
+            }
+        } else if self.parse_keyword(Keyword::RELATIVE) {
+            FetchDirection::Relative {
+                limit: self.parse_number_value()?,
+            }
+        } else if self.parse_keyword(Keyword::FORWARD) {
+            if self.parse_keyword(Keyword::ALL) {
+                FetchDirection::ForwardAll
+            } else {
+                FetchDirection::Forward {
+                    // TODO: Support optional
+                    limit: Some(self.parse_number_value()?),
+                }
+            }
+        } else if self.parse_keyword(Keyword::BACKWARD) {
+            if self.parse_keyword(Keyword::ALL) {
+                FetchDirection::BackwardAll
+            } else {
+                FetchDirection::Backward {
+                    // TODO: Support optional
+                    limit: Some(self.parse_number_value()?),
+                }
+            }
+        } else if self.parse_keyword(Keyword::ALL) {
+            FetchDirection::All
+        } else {
+            FetchDirection::Count {
+                limit: self.parse_number_value()?,
+            }
+        };
+
+        self.expect_one_of_keywords(&[Keyword::FROM, Keyword::IN])?;
+
+        let name = self.parse_identifier()?;
+
+        let into = if self.parse_keyword(Keyword::INTO) {
+            Some(self.parse_object_name()?)
+        } else {
+            None
+        };
+
+        Ok(Statement::Fetch {
+            name,
+            direction,
+            into,
+        })
+    }
+
     pub fn parse_discard(&mut self) -> Result<Statement, ParserError> {
         let object_type = if self.parse_keyword(Keyword::ALL) {
             DiscardObject::ALL

tests/sqlparser_postgres.rs

@@ -1513,3 +1513,27 @@ fn parse_escaped_literal_string() {
         "sql parser error: Unterminated encoded string literal at Line: 1, Column 8"
     );
 }
+
+#[test]
+fn parse_fetch() {
+    pg_and_generic().verified_stmt("FETCH 2048 IN \"SQL_CUR0x7fa44801bc00\"");
+    pg_and_generic().verified_stmt("FETCH 2048 IN \"SQL_CUR0x7fa44801bc00\" INTO \"new_table\"");
+    pg_and_generic().verified_stmt("FETCH NEXT IN \"SQL_CUR0x7fa44801bc00\" INTO \"new_table\"");
+    pg_and_generic().verified_stmt("FETCH PRIOR IN \"SQL_CUR0x7fa44801bc00\" INTO \"new_table\"");
+    pg_and_generic().verified_stmt("FETCH FIRST IN \"SQL_CUR0x7fa44801bc00\" INTO \"new_table\"");
+    pg_and_generic().verified_stmt("FETCH LAST IN \"SQL_CUR0x7fa44801bc00\" INTO \"new_table\"");
+    pg_and_generic()
+        .verified_stmt("FETCH ABSOLUTE 2048 IN \"SQL_CUR0x7fa44801bc00\" INTO \"new_table\"");
+    pg_and_generic()
+        .verified_stmt("FETCH RELATIVE 2048 IN \"SQL_CUR0x7fa44801bc00\" INTO \"new_table\"");
+    pg_and_generic().verified_stmt("FETCH ALL IN \"SQL_CUR0x7fa44801bc00\" INTO \"new_table\"");
+    pg_and_generic().verified_stmt("FETCH ALL IN \"SQL_CUR0x7fa44801bc00\" INTO \"new_table\"");
+    pg_and_generic()
+        .verified_stmt("FETCH FORWARD 2048 IN \"SQL_CUR0x7fa44801bc00\" INTO \"new_table\"");
+    pg_and_generic()
+        .verified_stmt("FETCH FORWARD ALL IN \"SQL_CUR0x7fa44801bc00\" INTO \"new_table\"");
+    pg_and_generic()
+        .verified_stmt("FETCH BACKWARD 2048 IN \"SQL_CUR0x7fa44801bc00\" INTO \"new_table\"");
+    pg_and_generic()
+        .verified_stmt("FETCH BACKWARD ALL IN \"SQL_CUR0x7fa44801bc00\" INTO \"new_table\"");
+}