feat: update pgwire to 0.37 (#265)

* feat: update pgwire to 0.37

* chore: update msrv

Author: sunng87

.github/workflows/ci.yml

@@ -91,6 +91,6 @@ jobs:
     - uses: actions/checkout@v4
     - uses: actions-rs/toolchain@v1
       with:
-        toolchain: "1.88.0"
+        toolchain: "1.89"
         override: true
     - run: cargo build --all-features

Cargo.lock

@@ -5,7 +5,7 @@ members = ["datafusion-postgres", "datafusion-postgres-cli", "arrow-pg", "datafu
 [workspace.package]
 edition = "2021"
 license = "Apache-2.0"
-rust-version = "1.88.0"
+rust-version = "1.89"
 authors = ["Ning Sun <[email protected]>"]
 keywords = ["database", "postgresql", "datafusion"]
 homepage = "https://github.com/datafusion-contrib/datafusion-postgres/"
@@ -18,7 +18,7 @@ bytes = "1.11.0"
 chrono = { version = "0.4", features = ["std"] }
 datafusion = { version = "51", default-features = false }
 futures = "0.3"
-pgwire = { version = "0.36.3", default-features = false }
+pgwire = { version = "0.37", default-features = false }
 postgres-types = "0.2"
 rust_decimal = { version = "1.39", features = ["db-postgres"] }
 tokio = { version = "1", default-features = false }

Cargo.toml

@@ -17,30 +17,31 @@ pub struct RowEncoder {
     rb: RecordBatch,
     curr_idx: usize,
     fields: Arc<Vec<FieldInfo>>,
+    row_encoder: DataRowEncoder,
 }
 
 impl RowEncoder {
     pub fn new(rb: RecordBatch, fields: Arc<Vec<FieldInfo>>) -> Self {
         assert_eq!(rb.num_columns(), fields.len());
         Self {
             rb,
-            fields,
+            fields: fields.clone(),
             curr_idx: 0,
+            row_encoder: DataRowEncoder::new(fields),
         }
     }
 
     pub fn next_row(&mut self) -> Option<PgWireResult<DataRow>> {
         if self.curr_idx == self.rb.num_rows() {
             return None;
         }
-        let mut encoder = DataRowEncoder::new(self.fields.clone());
         for col in 0..self.rb.num_columns() {
             let array = self.rb.column(col);
             let field = &self.fields[col];
 
-            encode_value(&mut encoder, array, self.curr_idx, field).unwrap();
+            encode_value(&mut self.row_encoder, array, self.curr_idx, field).unwrap();
         }
         self.curr_idx += 1;
-        Some(encoder.finish())
+        Some(Ok(self.row_encoder.take_row()))
     }
 }

arrow-pg/src/row_encoder.rs

@@ -13,11 +13,8 @@ use pgwire::api::auth::noop::NoopStartupHandler;
 use pgwire::api::auth::StartupHandler;
 use pgwire::api::portal::{Format, Portal};
 use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler};
-use pgwire::api::results::{
-    DescribePortalResponse, DescribeResponse, DescribeStatementResponse, Response, Tag,
-};
+use pgwire::api::results::{FieldInfo, Response, Tag};
 use pgwire::api::stmt::QueryParser;
-use pgwire::api::stmt::StoredStatement;
 use pgwire::api::{ClientInfo, ErrorHandler, PgWireServerHandlers, Type};
 use pgwire::error::{PgWireError, PgWireResult};
 use pgwire::types::format::FormatOptions;
@@ -202,58 +199,6 @@ impl ExtendedQueryHandler for DfSessionService {
         self.parser.clone()
     }
 
-    async fn do_describe_statement<C>(
-        &self,
-        _client: &mut C,
-        target: &StoredStatement<Self::Statement>,
-    ) -> PgWireResult<DescribeStatementResponse>
-    where
-        C: ClientInfo + Unpin + Send + Sync,
-    {
-        if let (_, Some((_, plan))) = &target.statement {
-            let schema = plan.schema();
-            let fields =
-                arrow_schema_to_pg_fields(schema.as_arrow(), &Format::UnifiedBinary, None)?;
-            let params = plan
-                .get_parameter_types()
-                .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
-
-            let mut param_types = Vec::with_capacity(params.len());
-            for param_type in ordered_param_types(&params).iter() {
-                // Fixed: Use &params
-                if let Some(datatype) = param_type {
-                    let pgtype = into_pg_type(datatype)?;
-                    param_types.push(pgtype);
-                } else {
-                    param_types.push(Type::UNKNOWN);
-                }
-            }
-
-            Ok(DescribeStatementResponse::new(param_types, fields))
-        } else {
-            Ok(DescribeStatementResponse::no_data())
-        }
-    }
-
-    async fn do_describe_portal<C>(
-        &self,
-        _client: &mut C,
-        target: &Portal<Self::Statement>,
-    ) -> PgWireResult<DescribePortalResponse>
-    where
-        C: ClientInfo + Unpin + Send + Sync,
-    {
-        if let (_, Some((_, plan))) = &target.statement.statement {
-            let format = &target.result_column_format;
-            let schema = plan.schema();
-            let fields = arrow_schema_to_pg_fields(schema.as_arrow(), format, None)?;
-
-            Ok(DescribePortalResponse::new(fields))
-        } else {
-            Ok(DescribePortalResponse::no_data())
-        }
-    }
-
     async fn do_query<C>(
         &self,
         client: &mut C,
@@ -433,6 +378,48 @@ impl QueryParser for Parser {
             .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
         Ok((query, Some((statement, logical_plan))))
     }
+
+    fn get_parameter_types(&self, stmt: &Self::Statement) -> PgWireResult<Vec<Type>> {
+        if let (_, Some((_, plan))) = stmt {
+            let params = plan
+                .get_parameter_types()
+                .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
+
+            let mut param_types = Vec::with_capacity(params.len());
+            for param_type in ordered_param_types(&params).iter() {
+                // Fixed: Use &params
+                if let Some(datatype) = param_type {
+                    let pgtype = into_pg_type(datatype)?;
+                    param_types.push(pgtype);
+                } else {
+                    param_types.push(Type::UNKNOWN);
+                }
+            }
+
+            Ok(param_types)
+        } else {
+            Ok(vec![])
+        }
+    }
+
+    fn get_result_schema(
+        &self,
+        stmt: &Self::Statement,
+        column_format: Option<&Format>,
+    ) -> PgWireResult<Vec<FieldInfo>> {
+        if let (_, Some((_, plan))) = stmt {
+            let schema = plan.schema();
+            let fields = arrow_schema_to_pg_fields(
+                schema.as_arrow(),
+                column_format.unwrap_or(&Format::UnifiedBinary),
+                None,
+            )?;
+
+            Ok(fields)
+        } else {
+            Ok(vec![])
+        }
+    }
 }
 
 fn ordered_param_types(types: &HashMap<String, Option<DataType>>) -> Vec<Option<&DataType>> {

datafusion-postgres/src/handlers.rs

@@ -111,7 +111,7 @@ fn mock_show_response(name: &str, value: &str) -> PgWireResult<QueryResponse> {
     let row = {
         let mut encoder = DataRowEncoder::new(Arc::new(fields.clone()));
         encoder.encode_field(&Some(value))?;
-        encoder.finish()
+        Ok(encoder.take_row())
     };
 
     let row_stream = futures::stream::once(async move { row });