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/session/Cargo.toml

@@ -17,9 +17,11 @@ auth.workspace = true
 common-catalog.workspace = true
 common-error.workspace = true
 common-macro.workspace = true
+common-recordbatch.workspace = true
 common-telemetry.workspace = true
 common-time.workspace = true
 derive_builder.workspace = true
+derive_more = { version = "1", default-features = false, features = ["debug"] }
 meter-core.workspace = true
 snafu.workspace = true
 sql.workspace = true

src/session/src/context.rs

@@ -23,6 +23,8 @@ use arc_swap::ArcSwap;
 use auth::UserInfoRef;
 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
 use common_catalog::{build_db_string, parse_catalog_and_schema_from_db_string};
+use common_recordbatch::cursor::RecordBatchStreamCursor;
+use common_telemetry::warn;
 use common_time::timezone::parse_timezone;
 use common_time::Timezone;
 use derive_builder::Builder;
@@ -34,6 +36,8 @@ use crate::MutableInner;
 pub type QueryContextRef = Arc<QueryContext>;
 pub type ConnInfoRef = Arc<ConnInfo>;
 
+const CURSOR_COUNT_WARNING_LIMIT: usize = 10;
+
 #[derive(Debug, Builder, Clone)]
 #[builder(pattern = "owned")]
 #[builder(build_fn(skip))]
@@ -299,6 +303,27 @@ impl QueryContext {
     pub fn set_query_timeout(&self, timeout: Duration) {
         self.mutable_session_data.write().unwrap().query_timeout = Some(timeout);
     }
+
+    pub fn insert_cursor(&self, name: String, rb: RecordBatchStreamCursor) {
+        let mut guard = self.mutable_session_data.write().unwrap();
+        guard.cursors.insert(name, Arc::new(rb));
+
+        let cursor_count = guard.cursors.len();
+        if cursor_count > CURSOR_COUNT_WARNING_LIMIT {
+            warn!("Current connection has {} open cursors", cursor_count);
+        }
+    }
+
+    pub fn remove_cursor(&self, name: &str) {
+        let mut guard = self.mutable_session_data.write().unwrap();
+        guard.cursors.remove(name);
+    }
+
+    pub fn get_cursor(&self, name: &str) -> Option<Arc<RecordBatchStreamCursor>> {
+        let guard = self.mutable_session_data.read().unwrap();
+        let rb = guard.cursors.get(name);
+        rb.cloned()
+    }
 }
 
 impl QueryContextBuilder {

src/session/src/lib.rs

@@ -16,6 +16,7 @@ pub mod context;
 pub mod session_config;
 pub mod table_name;
 
+use std::collections::HashMap;
 use std::net::SocketAddr;
 use std::sync::{Arc, RwLock};
 use std::time::Duration;
@@ -23,9 +24,11 @@ use std::time::Duration;
 use auth::UserInfoRef;
 use common_catalog::build_db_string;
 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
+use common_recordbatch::cursor::RecordBatchStreamCursor;
 use common_time::timezone::get_timezone;
 use common_time::Timezone;
 use context::{ConfigurationVariables, QueryContextBuilder};
+use derive_more::Debug;
 
 use crate::context::{Channel, ConnInfo, QueryContextRef};
 
@@ -47,6 +50,8 @@ pub(crate) struct MutableInner {
     user_info: UserInfoRef,
     timezone: Timezone,
     query_timeout: Option<Duration>,
+    #[debug(skip)]
+    pub(crate) cursors: HashMap<String, Arc<RecordBatchStreamCursor>>,
 }
 
 impl Default for MutableInner {
@@ -56,6 +61,7 @@ impl Default for MutableInner {
             user_info: auth::userinfo_by_name(None),
             timezone: get_timezone(None).clone(),
             query_timeout: None,
+            cursors: HashMap::with_capacity(0),
         }
     }
 }