feat: add cursor statements (#5094)

* feat: add sql parsers for cursor operations

* feat: cursor operator

* feat: implement RecordBatchStreamCursor

* feat: implement cursor storage and execution

* test: add tests

* chore: update docstring

* feat: add a temporary sql rewrite for cast in limit

this issue is described in #5097

* test: add more sql for cursor integration test

* feat: reject non-select query for cursor statement

* refactor: address review issues

* test: add empty result case

* feat: address review comments

src/sql/src/parser.rs

@@ -167,6 +167,12 @@ impl ParserContext<'_> {
                         self.parse_tql()
                     }
 
+                    Keyword::DECLARE => self.parse_declare_cursor(),
+
+                    Keyword::FETCH => self.parse_fetch_cursor(),
+
+                    Keyword::CLOSE => self.parse_close_cursor(),
+
                     Keyword::USE => {
                         let _ = self.parser.next_token();
 

src/sql/src/parsers.rs

@@ -16,6 +16,7 @@ pub(crate) mod admin_parser;
 mod alter_parser;
 pub(crate) mod copy_parser;
 pub(crate) mod create_parser;
+pub(crate) mod cursor_parser;
 pub(crate) mod deallocate_parser;
 pub(crate) mod delete_parser;
 pub(crate) mod describe_parser;

src/sql/src/parsers/cursor_parser.rs

@@ -0,0 +1,157 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use snafu::{ensure, ResultExt};
+use sqlparser::keywords::Keyword;
+use sqlparser::tokenizer::Token;
+
+use crate::error::{self, Result};
+use crate::parser::ParserContext;
+use crate::statements::cursor::{CloseCursor, DeclareCursor, FetchCursor};
+use crate::statements::statement::Statement;
+
+impl ParserContext<'_> {
+    pub(crate) fn parse_declare_cursor(&mut self) -> Result<Statement> {
+        let _ = self.parser.expect_keyword(Keyword::DECLARE);
+        let cursor_name = self
+            .parser
+            .parse_object_name(false)
+            .context(error::SyntaxSnafu)?;
+        let _ = self
+            .parser
+            .expect_keywords(&[Keyword::CURSOR, Keyword::FOR]);
+
+        let mut is_select = false;
+        if let Token::Word(w) = self.parser.peek_token().token {
+            match w.keyword {
+                Keyword::SELECT | Keyword::WITH => {
+                    is_select = true;
+                }
+                _ => {}
+            }
+        };
+        ensure!(
+            is_select,
+            error::InvalidSqlSnafu {
+                msg: "Expect select query in cursor statement".to_string(),
+            }
+        );
+
+        let query_stmt = self.parse_query()?;
+        match query_stmt {
+            Statement::Query(query) => Ok(Statement::DeclareCursor(DeclareCursor {
+                cursor_name: ParserContext::canonicalize_object_name(cursor_name),
+                query,
+            })),
+            _ => error::InvalidSqlSnafu {
+                msg: format!("Expect query, found {}", query_stmt),
+            }
+            .fail(),
+        }
+    }
+
+    pub(crate) fn parse_fetch_cursor(&mut self) -> Result<Statement> {
+        let _ = self.parser.expect_keyword(Keyword::FETCH);
+
+        let fetch_size = self
+            .parser
+            .parse_literal_uint()
+            .context(error::SyntaxSnafu)?;
+        let _ = self.parser.parse_keyword(Keyword::FROM);
+
+        let cursor_name = self
+            .parser
+            .parse_object_name(false)
+            .context(error::SyntaxSnafu)?;
+
+        Ok(Statement::FetchCursor(FetchCursor {
+            cursor_name: ParserContext::canonicalize_object_name(cursor_name),
+            fetch_size,
+        }))
+    }
+
+    pub(crate) fn parse_close_cursor(&mut self) -> Result<Statement> {
+        let _ = self.parser.expect_keyword(Keyword::CLOSE);
+        let cursor_name = self
+            .parser
+            .parse_object_name(false)
+            .context(error::SyntaxSnafu)?;
+
+        Ok(Statement::CloseCursor(CloseCursor {
+            cursor_name: ParserContext::canonicalize_object_name(cursor_name),
+        }))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use super::*;
+    use crate::dialect::GreptimeDbDialect;
+    use crate::parser::ParseOptions;
+
+    #[test]
+    fn test_parse_declare_cursor() {
+        let sql = "DECLARE c1 CURSOR FOR\nSELECT * FROM numbers";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+
+        if let Statement::DeclareCursor(dc) = &result[0] {
+            assert_eq!("c1", dc.cursor_name.to_string());
+            assert_eq!(
+                "DECLARE c1 CURSOR FOR SELECT * FROM numbers",
+                dc.to_string()
+            );
+        } else {
+            panic!("Unexpected statement");
+        }
+
+        let sql = "DECLARE c1 CURSOR FOR\nINSERT INTO numbers VALUES (1);";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_parese_fetch_cursor() {
+        let sql = "FETCH 1000 FROM c1";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+
+        if let Statement::FetchCursor(fc) = &result[0] {
+            assert_eq!("c1", fc.cursor_name.to_string());
+            assert_eq!("1000", fc.fetch_size.to_string());
+            assert_eq!(sql, fc.to_string());
+        } else {
+            panic!("Unexpected statement")
+        }
+    }
+
+    #[test]
+    fn test_close_fetch_cursor() {
+        let sql = "CLOSE c1";
+        let result =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+
+        if let Statement::CloseCursor(cc) = &result[0] {
+            assert_eq!("c1", cc.cursor_name.to_string());
+            assert_eq!(sql, cc.to_string());
+        } else {
+            panic!("Unexpected statement")
+        }
+    }
+}

src/sql/src/statements.rs

@@ -16,6 +16,7 @@ pub mod admin;
 pub mod alter;
 pub mod copy;
 pub mod create;
+pub mod cursor;
 pub mod delete;
 pub mod describe;
 pub mod drop;
@@ -224,7 +225,7 @@ pub fn sql_number_to_value(data_type: &ConcreteDataType, n: &str) -> Result<Valu
     // TODO(hl): also Date/DateTime
 }
 
-fn parse_sql_number<R: FromStr + std::fmt::Debug>(n: &str) -> Result<R>
+pub(crate) fn parse_sql_number<R: FromStr + std::fmt::Debug>(n: &str) -> Result<R>
 where
     <R as FromStr>::Err: std::fmt::Debug,
 {

src/sql/src/statements/cursor.rs

@@ -0,0 +1,60 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::fmt::Display;
+
+use sqlparser::ast::ObjectName;
+use sqlparser_derive::{Visit, VisitMut};
+
+use super::query::Query;
+
+/// Represents a DECLARE CURSOR statement
+///
+/// This statement will carry a SQL query
+#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
+pub struct DeclareCursor {
+    pub cursor_name: ObjectName,
+    pub query: Box<Query>,
+}
+
+impl Display for DeclareCursor {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "DECLARE {} CURSOR FOR {}", self.cursor_name, self.query)
+    }
+}
+
+/// Represents a FETCH FROM cursor statement
+#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
+pub struct FetchCursor {
+    pub cursor_name: ObjectName,
+    pub fetch_size: u64,
+}
+
+impl Display for FetchCursor {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "FETCH {} FROM {}", self.fetch_size, self.cursor_name)
+    }
+}
+
+/// Represents a CLOSE cursor statement
+#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
+pub struct CloseCursor {
+    pub cursor_name: ObjectName,
+}
+
+impl Display for CloseCursor {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "CLOSE {}", self.cursor_name)
+    }
+}

src/sql/src/statements/statement.rs

@@ -24,6 +24,7 @@ use crate::statements::alter::{AlterDatabase, AlterTable};
 use crate::statements::create::{
     CreateDatabase, CreateExternalTable, CreateFlow, CreateTable, CreateTableLike, CreateView,
 };
+use crate::statements::cursor::{CloseCursor, DeclareCursor, FetchCursor};
 use crate::statements::delete::Delete;
 use crate::statements::describe::DescribeTable;
 use crate::statements::drop::{DropDatabase, DropFlow, DropTable, DropView};
@@ -118,6 +119,12 @@ pub enum Statement {
     Use(String),
     // Admin statement(extension)
     Admin(Admin),
+    // DECLARE ... CURSOR FOR ...
+    DeclareCursor(DeclareCursor),
+    // FETCH ... FROM ...
+    FetchCursor(FetchCursor),
+    // CLOSE
+    CloseCursor(CloseCursor),
 }
 
 impl Display for Statement {
@@ -165,6 +172,9 @@ impl Display for Statement {
             Statement::CreateView(s) => s.fmt(f),
             Statement::Use(s) => s.fmt(f),
             Statement::Admin(admin) => admin.fmt(f),
+            Statement::DeclareCursor(s) => s.fmt(f),
+            Statement::FetchCursor(s) => s.fmt(f),
+            Statement::CloseCursor(s) => s.fmt(f),
         }
     }
 }