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/operator/src/error.rs

@@ -786,6 +786,12 @@ pub enum Error {
         #[snafu(source)]
         error: Elapsed,
     },
+
+    #[snafu(display("Cursor {name} is not found"))]
+    CursorNotFound { name: String },
+
+    #[snafu(display("A cursor named {name} already exists"))]
+    CursorExists { name: String },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -825,7 +831,9 @@ impl ErrorExt for Error {
             | Error::FunctionArityMismatch { .. }
             | Error::InvalidPartition { .. }
             | Error::PhysicalExpr { .. }
-            | Error::InvalidJsonFormat { .. } => StatusCode::InvalidArguments,
+            | Error::InvalidJsonFormat { .. }
+            | Error::CursorNotFound { .. }
+            | Error::CursorExists { .. } => StatusCode::InvalidArguments,
 
             Error::TableAlreadyExists { .. } | Error::ViewAlreadyExists { .. } => {
                 StatusCode::TableAlreadyExists

src/operator/src/statement.rs

@@ -16,6 +16,7 @@ mod admin;
 mod copy_database;
 mod copy_table_from;
 mod copy_table_to;
+mod cursor;
 mod ddl;
 mod describe;
 mod dml;
@@ -133,6 +134,16 @@ impl StatementExecutor {
                 self.plan_exec(QueryStatement::Sql(stmt), query_ctx).await
             }
 
+            Statement::DeclareCursor(declare_cursor) => {
+                self.declare_cursor(declare_cursor, query_ctx).await
+            }
+            Statement::FetchCursor(fetch_cursor) => {
+                self.fetch_cursor(fetch_cursor, query_ctx).await
+            }
+            Statement::CloseCursor(close_cursor) => {
+                self.close_cursor(close_cursor, query_ctx).await
+            }
+
             Statement::Insert(insert) => self.insert(insert, query_ctx).await,
 
             Statement::Tql(tql) => self.execute_tql(tql, query_ctx).await,

src/operator/src/statement/cursor.rs

@@ -0,0 +1,98 @@
+// 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 common_query::{Output, OutputData};
+use common_recordbatch::cursor::RecordBatchStreamCursor;
+use common_recordbatch::RecordBatches;
+use common_telemetry::tracing;
+use query::parser::QueryStatement;
+use session::context::QueryContextRef;
+use snafu::ResultExt;
+use sql::statements::cursor::{CloseCursor, DeclareCursor, FetchCursor};
+use sql::statements::statement::Statement;
+
+use crate::error::{self, Result};
+use crate::statement::StatementExecutor;
+
+impl StatementExecutor {
+    #[tracing::instrument(skip_all)]
+    pub(super) async fn declare_cursor(
+        &self,
+        declare_cursor: DeclareCursor,
+        query_ctx: QueryContextRef,
+    ) -> Result<Output> {
+        let cursor_name = declare_cursor.cursor_name.to_string();
+
+        if query_ctx.get_cursor(&cursor_name).is_some() {
+            error::CursorExistsSnafu {
+                name: cursor_name.to_string(),
+            }
+            .fail()?;
+        }
+
+        let query_stmt = Statement::Query(declare_cursor.query);
+
+        let output = self
+            .plan_exec(QueryStatement::Sql(query_stmt), query_ctx.clone())
+            .await?;
+        match output.data {
+            OutputData::RecordBatches(rb) => {
+                let rbs = rb.as_stream();
+                query_ctx.insert_cursor(cursor_name, RecordBatchStreamCursor::new(rbs));
+            }
+            OutputData::Stream(rbs) => {
+                query_ctx.insert_cursor(cursor_name, RecordBatchStreamCursor::new(rbs));
+            }
+            // Should not happen because we have query type ensured from parser.
+            OutputData::AffectedRows(_) => error::NotSupportedSnafu {
+                feat: "Non-query statement on cursor",
+            }
+            .fail()?,
+        }
+
+        Ok(Output::new_with_affected_rows(0))
+    }
+
+    #[tracing::instrument(skip_all)]
+    pub(super) async fn fetch_cursor(
+        &self,
+        fetch_cursor: FetchCursor,
+        query_ctx: QueryContextRef,
+    ) -> Result<Output> {
+        let cursor_name = fetch_cursor.cursor_name.to_string();
+        let fetch_size = fetch_cursor.fetch_size;
+        if let Some(rb) = query_ctx.get_cursor(&cursor_name) {
+            let record_batch = rb
+                .take(fetch_size as usize)
+                .await
+                .context(error::BuildRecordBatchSnafu)?;
+            let record_batches =
+                RecordBatches::try_new(record_batch.schema.clone(), vec![record_batch])
+                    .context(error::BuildRecordBatchSnafu)?;
+            Ok(Output::new_with_record_batches(record_batches))
+        } else {
+            error::CursorNotFoundSnafu { name: cursor_name }.fail()
+        }
+    }
+
+    #[tracing::instrument(skip_all)]
+    pub(super) async fn close_cursor(
+        &self,
+        close_cursor: CloseCursor,
+        query_ctx: QueryContextRef,
+    ) -> Result<Output> {
+        query_ctx.remove_cursor(&close_cursor.cursor_name.to_string());
+        Ok(Output::new_with_affected_rows(0))
+    }
+}