feat: update cursor to support extended query

Author: sunng87

Cargo.lock

@@ -30,3 +30,6 @@ opt-level = "z"
 lto = true
 codegen-units = 1
 panic = "abort"
+
+[patch.crates-io]
+pgwire = { git = "https://github.com/sunng87/pgwire", branch = "feature/cursor-for-portal" }

Cargo.toml

@@ -27,7 +27,7 @@ datafusion = { workspace = true, optional = true }
 futures.workspace = true
 geoarrow = { version = "0.8", optional = true }
 geoarrow-schema = { version = "0.8", optional = true }
-pg_interval = { version = "0.5.1", package = "pg_interval_2" }
+pg_interval = { version = "0.5.0" }
 pgwire = { workspace = true, default-features = false, features = ["server-api", "pg-ext-types"] }
 postgres-types.workspace = true
 rust_decimal.workspace = true

arrow-pg/Cargo.toml

@@ -2,6 +2,7 @@ use std::sync::Arc;
 
 use datafusion::sql::sqlparser::ast::Statement;
 use datafusion::sql::sqlparser::dialect::PostgreSqlDialect;
+use datafusion::sql::sqlparser::keywords::Keyword;
 use datafusion::sql::sqlparser::parser::Parser;
 use datafusion::sql::sqlparser::parser::ParserError;
 use datafusion::sql::sqlparser::tokenizer::Token;
@@ -247,10 +248,19 @@ impl PostgresCompatibilityParser {
 
         // Get token values (without spans) and filter out only whitespace
         // Keep semicolons as they separate statements
+        // Also rewrite ABORT to ROLLBACK for postgres compatibility
+        // remove this when https://github.com/apache/datafusion-sqlparser-rs/pull/2332 is ready
         let filtered_tokens: Vec<Token> = tokens
             .iter()
             .map(|t| t.token.clone())
             .filter(|t| !matches!(t, Token::Whitespace(_)))
+            .map(|t| {
+                if matches!(&t, Token::Word(w) if w.keyword == Keyword::ABORT) {
+                    Token::make_keyword("ROLLBACK")
+                } else {
+                    t
+                }
+            })
             .collect();
 
         // Handle empty input

datafusion-pg-catalog/src/sql/parser.rs

@@ -1,4 +1,3 @@
-use std::ops::DerefMut;
 use std::sync::Arc;
 
 use async_trait::async_trait;
@@ -7,9 +6,8 @@ use datafusion::logical_expr::LogicalPlan;
 use datafusion::prelude::SessionContext;
 use datafusion::sql::sqlparser;
 use datafusion::sql::sqlparser::ast::{CloseCursor, DeclareType, FetchDirection};
-use futures::StreamExt;
 use pgwire::api::ClientInfo;
-use pgwire::api::portal::{Format, Portal, PortalExecutionState};
+use pgwire::api::portal::{Format, Portal};
 use pgwire::api::results::{QueryResponse, Response, Tag};
 use pgwire::api::stmt::StoredStatement;
 use pgwire::api::store::{MemPortalStore, PortalStore};
@@ -48,22 +46,43 @@ impl QueryHook for CursorStatementHook {
 
     async fn handle_extended_parse_query(
         &self,
-        _statement: &sqlparser::ast::Statement,
+        statement: &sqlparser::ast::Statement,
         _session_context: &SessionContext,
         _client: &(dyn ClientInfo + Send + Sync),
     ) -> Option<PgWireResult<LogicalPlan>> {
-        None
+        match statement {
+            sqlparser::ast::Statement::Declare { .. }
+            | sqlparser::ast::Statement::Fetch { .. }
+            | sqlparser::ast::Statement::Close { .. } => Some(Ok(LogicalPlan::EmptyRelation(
+                datafusion::logical_expr::EmptyRelation {
+                    produce_one_row: false,
+                    schema: Arc::new(datafusion::common::DFSchema::empty()),
+                },
+            ))),
+            _ => None,
+        }
     }
 
     async fn handle_extended_query(
         &self,
-        _statement: &sqlparser::ast::Statement,
+        statement: &sqlparser::ast::Statement,
         _logical_plan: &LogicalPlan,
         _params: &ParamValues,
-        _session_context: &SessionContext,
-        _client: &mut dyn HookClient,
+        session_context: &SessionContext,
+        client: &mut dyn HookClient,
     ) -> Option<PgWireResult<Response>> {
-        None
+        let store = client.portal_store();
+
+        match statement {
+            sqlparser::ast::Statement::Declare { stmts } => {
+                Some(handle_declare(store, stmts, session_context).await)
+            }
+            sqlparser::ast::Statement::Fetch {
+                name, direction, ..
+            } => Some(handle_fetch(store, name, direction).await),
+            sqlparser::ast::Statement::Close { cursor } => Some(handle_close(store, cursor)),
+            _ => None,
+        }
     }
 }
 
@@ -122,20 +141,9 @@ async fn handle_declare(
             vec![],
         ));
 
-        let bind = pgwire::messages::extendedquery::Bind::new(
-            Some(cursor_name.clone()),
-            None,
-            vec![],
-            vec![],
-            vec![],
-        );
-        let portal = Portal::try_new(&bind, stored_stmt)?;
+        let portal = Portal::new_cursor(cursor_name.clone(), stored_stmt);
 
-        let state = portal.state();
-        {
-            let mut portal_state = state.lock().await;
-            *portal_state = PortalExecutionState::Suspended(query_response);
-        }
+        portal.start(query_response).await;
 
         store.put_portal(Arc::new(portal));
     }
@@ -150,7 +158,7 @@ async fn handle_fetch(
 ) -> PgWireResult<Response> {
     let cursor_name = &name.value;
 
-    let limit = match direction {
+    let max_rows = match direction {
         FetchDirection::Next | FetchDirection::Forward { limit: None } => Some(1),
         FetchDirection::Forward { limit: Some(v) } | FetchDirection::Count { limit: v } => {
             parse_value_as_usize(v)
@@ -187,70 +195,19 @@ async fn handle_fetch(
         )))
     })?;
 
-    let state = portal.state();
-    let mut state = state.lock().await;
-
-    let query_response = match state.deref_mut() {
-        PortalExecutionState::Suspended(qr) => qr,
-        PortalExecutionState::Finished => {
-            return Ok(Response::Execution(Tag::new("FETCH").with_rows(0)));
-        }
-        PortalExecutionState::Initial => {
-            return Err(PgWireError::UserError(Box::new(
-                pgwire::error::ErrorInfo::new(
-                    "ERROR".to_string(),
-                    "24000".to_string(),
-                    "cursor is in invalid state".to_string(),
-                ),
-            )));
-        }
-    };
-
-    let schema = query_response.row_schema();
-    let mut fetched_rows: Vec<pgwire::messages::data::DataRow> = vec![];
-    let mut stream_exhausted = false;
-
-    if let Some(n) = limit {
-        for _ in 0..n {
-            match query_response.data_rows().next().await {
-                Some(Ok(row)) => fetched_rows.push(row),
-                Some(Err(e)) => return Err(e),
-                None => {
-                    stream_exhausted = true;
-                    break;
-                }
-            }
-        }
-    } else {
-        loop {
-            match query_response.data_rows().next().await {
-                Some(Ok(row)) => fetched_rows.push(row),
-                Some(Err(e)) => return Err(e),
-                None => {
-                    stream_exhausted = true;
-                    break;
-                }
-            }
-        }
-    }
-
-    if stream_exhausted {
-        *state = PortalExecutionState::Finished;
-    }
-
-    drop(state);
+    let fetch_result = portal.fetch(max_rows.unwrap_or(0)).await?;
 
-    if fetched_rows.is_empty() {
+    if fetch_result.rows.is_empty() {
         return Ok(Response::Execution(Tag::new("FETCH").with_rows(0)));
     }
 
-    let mut fetched_response = QueryResponse::new(
-        schema,
-        futures::stream::iter(fetched_rows.into_iter().map(Ok)),
+    let mut response = QueryResponse::new(
+        fetch_result.row_schema,
+        futures::stream::iter(fetch_result.rows.into_iter().map(Ok)),
     );
-    fetched_response.set_command_tag("FETCH");
+    response.set_command_tag("FETCH");
 
-    Ok(Response::Query(fetched_response))
+    Ok(Response::Query(response))
 }
 
 fn handle_close(