From bbe48e9e8be3e4faffeabf8fc65dc3a261073d72 Mon Sep 17 00:00:00 2001
From: Ning Sun <sunng@protonmail.com>
Date: Thu, 7 Aug 2025 14:17:52 +0800
Subject: [PATCH] feat: update pgwire to 0.32 (#6674)

* feat: update pgwire api

* feat: update pgwire and override on_query/on_execute

* feat: update pgwire to 0.32

* chore: remove code example

Signed-off-by: Ning Sun <sunning@greptime.com>

---------

Signed-off-by: Ning Sun <sunning@greptime.com>
---
 Cargo.lock                               | 15 +++++++++++----
 src/servers/Cargo.toml                   |  5 +----
 src/servers/src/postgres.rs              | 24 +++++++-----------------
 src/servers/src/postgres/auth_handler.rs |  5 +++--
 4 files changed, 22 insertions(+), 27 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 73a28f1eed..dcb3480d46 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -7287,6 +7287,12 @@ version = "0.7.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "490cc448043f947bae3cbee9c203358d62dbee0db12107a74be5c30ccfd09771"
 
+[[package]]
+name = "md5"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ae960838283323069879657ca3de837e9f7bbb4c7bf6ea7f1b290d5e9476d2e0"
+
 [[package]]
 name = "measure_time"
 version = "0.9.0"
@@ -8349,7 +8355,7 @@ dependencies = [
  "futures",
  "humantime-serde",
  "lazy_static",
- "md5",
+ "md5 0.7.0",
  "moka",
  "opendal",
  "prometheus",
@@ -9182,8 +9188,9 @@ dependencies = [
 
 [[package]]
 name = "pgwire"
-version = "0.30.2"
-source = "git+https://github.com/sunng87/pgwire?rev=127573d997228cfb70c7699881c568eae8131270#127573d997228cfb70c7699881c568eae8131270"
+version = "0.32.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "017b8b74f9e8c7aff0087d4ef2b91676a5509e8928100d6a3510fd472210feb5"
 dependencies = [
  "async-trait",
  "bytes",
@@ -9192,7 +9199,7 @@ dependencies = [
  "futures",
  "hex",
  "lazy-regex",
- "md5",
+ "md5 0.8.0",
  "postgres-types",
  "rand 0.9.0",
  "ring",
diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml
index ad9fd0988c..b1d8c42996 100644
--- a/src/servers/Cargo.toml
+++ b/src/servers/Cargo.toml
@@ -89,10 +89,7 @@ opensrv-mysql = { git = "https://github.com/datafuselabs/opensrv", rev = "a1fb4d
 opentelemetry-proto.workspace = true
 otel-arrow-rust.workspace = true
 parking_lot.workspace = true
-#pgwire = { version = "0.30", default-features = false, features = ["server-api-ring"] }
-pgwire = { git = "https://github.com/sunng87/pgwire", rev = "127573d997228cfb70c7699881c568eae8131270", default-features = false, features = [
-    "server-api-ring",
-] }
+pgwire = { version = "0.32", default-features = false, features = ["server-api-ring"] }
 pin-project = "1.0"
 pipeline.workspace = true
 postgres-types = { version = "0.2", features = ["with-chrono-0_4", "with-serde_json-1"] }
diff --git a/src/servers/src/postgres.rs b/src/servers/src/postgres.rs
index 112d6b39dd..ca396dc6d2 100644
--- a/src/servers/src/postgres.rs
+++ b/src/servers/src/postgres.rs
@@ -32,9 +32,9 @@ use std::sync::Arc;
 
 use ::auth::UserProviderRef;
 use derive_builder::Builder;
-use pgwire::api::auth::ServerParameterProvider;
-use pgwire::api::copy::NoopCopyHandler;
-use pgwire::api::{ClientInfo, PgWireServerHandlers};
+use pgwire::api::auth::{ServerParameterProvider, StartupHandler};
+use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler};
+use pgwire::api::{ClientInfo, ErrorHandler, PgWireServerHandlers};
 pub use server::PostgresServer;
 use session::context::Channel;
 use session::Session;
@@ -92,29 +92,19 @@ pub(crate) struct MakePostgresServerHandler {
 pub(crate) struct PostgresServerHandler(Arc<PostgresServerHandlerInner>);
 
 impl PgWireServerHandlers for PostgresServerHandler {
-    type StartupHandler = PostgresServerHandlerInner;
-    type SimpleQueryHandler = PostgresServerHandlerInner;
-    type ExtendedQueryHandler = PostgresServerHandlerInner;
-    type CopyHandler = NoopCopyHandler;
-    type ErrorHandler = PostgresServerHandlerInner;
-
-    fn simple_query_handler(&self) -> Arc<Self::SimpleQueryHandler> {
+    fn simple_query_handler(&self) -> Arc<impl SimpleQueryHandler> {
         self.0.clone()
     }
 
-    fn extended_query_handler(&self) -> Arc<Self::ExtendedQueryHandler> {
+    fn extended_query_handler(&self) -> Arc<impl ExtendedQueryHandler> {
         self.0.clone()
     }
 
-    fn startup_handler(&self) -> Arc<Self::StartupHandler> {
+    fn startup_handler(&self) -> Arc<impl StartupHandler> {
         self.0.clone()
     }
 
-    fn copy_handler(&self) -> Arc<Self::CopyHandler> {
-        Arc::new(NoopCopyHandler)
-    }
-
-    fn error_handler(&self) -> Arc<Self::ErrorHandler> {
+    fn error_handler(&self) -> Arc<impl ErrorHandler> {
         self.0.clone()
     }
 }
diff --git a/src/servers/src/postgres/auth_handler.rs b/src/servers/src/postgres/auth_handler.rs
index 134ee54b19..d6388b7012 100644
--- a/src/servers/src/postgres/auth_handler.rs
+++ b/src/servers/src/postgres/auth_handler.rs
@@ -24,7 +24,7 @@ use pgwire::api::auth::StartupHandler;
 use pgwire::api::{auth, ClientInfo, PgWireConnectionState};
 use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
 use pgwire::messages::response::ErrorResponse;
-use pgwire::messages::startup::Authentication;
+use pgwire::messages::startup::{Authentication, SecretKey};
 use pgwire::messages::{PgWireBackendMessage, PgWireFrontendMessage};
 use session::Session;
 use snafu::IntoError;
@@ -127,7 +127,8 @@ where
 
     // pass generated process id and secret key to client, this information will
     // be sent to postgres client for query cancellation.
-    client.set_pid_and_secret_key(session.process_id() as i32, rand::random::<i32>());
+    // use all 0 before we actually supported query cancellation
+    client.set_pid_and_secret_key(0, SecretKey::I32(0));
     // set userinfo outside
 }