feat: implement cancel request on postgres protocol

src/catalog/Cargo.toml

@@ -43,6 +43,7 @@ moka = { workspace = true, features = ["future", "sync"] }
 partition.workspace = true
 paste.workspace = true
 prometheus.workspace = true
+rand.workspace = true
 rustc-hash.workspace = true
 serde_json.workspace = true
 session.workspace = true

src/catalog/src/process_manager.rs

@@ -210,6 +210,10 @@ impl ProcessManager {
             Ok(false)
         }
     }
+
+    pub fn server_addr(&self) -> &str {
+        &self.server_addr
+    }
 }
 
 pub struct Ticket {

src/frontend/src/server.rs

@@ -258,7 +258,7 @@ where
                 opts.keep_alive.as_secs(),
                 common_runtime::global_runtime(),
                 user_provider.clone(),
-                Some(self.instance.process_manager().clone()),
+                self.instance.process_manager().clone(),
             )) as Box<dyn Server>;
 
             handlers.insert((pg_server, pg_addr));

src/servers/Cargo.toml

@@ -88,10 +88,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.31", 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"] }
@@ -149,7 +146,7 @@ session = { workspace = true, features = ["testing"] }
 table.workspace = true
 tempfile = "3.0.0"
 tokio-postgres = "0.7"
-tokio-postgres-rustls = "0.12"
+tokio-postgres-rustls = "0.13"
 
 [target.'cfg(unix)'.dev-dependencies]
 pprof = { version = "0.14", features = ["criterion", "flamegraph"] }

src/servers/src/postgres.rs

@@ -29,11 +29,13 @@ use std::collections::HashMap;
 use std::net::SocketAddr;
 use std::sync::Arc;
 
-use ::auth::UserProviderRef;
+use auth::UserProviderRef;
+use catalog::process_manager::ProcessManagerRef;
 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::cancel::CancelHandler;
+use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler};
+use pgwire::api::{ClientInfo, ErrorHandler, PgWireServerHandlers};
 pub use server::PostgresServer;
 use session::context::Channel;
 use session::Session;
@@ -71,6 +73,7 @@ impl ServerParameterProvider for GreptimeDBStartupParameters {
 
 pub struct PostgresServerHandlerInner {
     query_handler: ServerSqlQueryHandlerRef,
+    process_manager: ProcessManagerRef,
     login_verifier: PgLoginVerifier,
     force_tls: bool,
     param_provider: Arc<GreptimeDBStartupParameters>,
@@ -82,6 +85,7 @@ pub struct PostgresServerHandlerInner {
 #[derive(Builder)]
 pub(crate) struct MakePostgresServerHandler {
     query_handler: ServerSqlQueryHandlerRef,
+    process_manager: ProcessManagerRef,
     user_provider: Option<UserProviderRef>,
     #[builder(default = "Arc::new(GreptimeDBStartupParameters::new())")]
     param_provider: Arc<GreptimeDBStartupParameters>,
@@ -91,43 +95,40 @@ 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<impl ErrorHandler> {
+        self.0.clone()
     }
 
-    fn error_handler(&self) -> Arc<Self::ErrorHandler> {
+    fn cancel_handler(&self) -> Arc<impl CancelHandler> {
         self.0.clone()
     }
 }
 
 impl MakePostgresServerHandler {
-    fn make(&self, addr: Option<SocketAddr>, process_id: u32) -> PostgresServerHandler {
-        let session = Arc::new(Session::new(
-            addr,
-            Channel::Postgres,
-            Default::default(),
-            process_id,
-        ));
+    fn make(&self, addr: Option<SocketAddr>) -> PostgresServerHandler {
+        let process_id = self.process_manager.next_id();
+        let secret_key = rand::random();
+
+        let session = Arc::new(
+            Session::new(addr, Channel::Postgres, Default::default(), process_id)
+                .with_secret_key(secret_key),
+        );
+
         let handler = PostgresServerHandlerInner {
             query_handler: self.query_handler.clone(),
+            process_manager: self.process_manager.clone(),
             login_verifier: PgLoginVerifier::new(self.user_provider.clone()),
             force_tls: self.force_tls,
             param_provider: self.param_provider.clone(),