feat: implement cancel handler for postgres

Author: sunng87

src/servers/src/postgres.rs

@@ -30,9 +30,11 @@ use std::net::SocketAddr;
 use std::sync::Arc;
 
 use ::auth::UserProviderRef;
+use catalog::process_manager::ProcessManagerRef;
 use derive_builder::Builder;
 use pgwire::api::auth::ServerParameterProvider;
 use pgwire::api::auth::StartupHandler;
+use pgwire::api::cancel::CancelHandler;
 use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler};
 use pgwire::api::ErrorHandler;
 use pgwire::api::{ClientInfo, PgWireServerHandlers};
@@ -73,6 +75,7 @@ impl ServerParameterProvider for GreptimeDBStartupParameters {
 
 pub struct PostgresServerHandlerInner {
     query_handler: ServerSqlQueryHandlerRef,
+    process_manager: ProcessManagerRef,
     login_verifier: PgLoginVerifier,
     force_tls: bool,
     param_provider: Arc<GreptimeDBStartupParameters>,
@@ -84,6 +87,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>,
@@ -108,10 +112,15 @@ impl PgWireServerHandlers for PostgresServerHandler {
     fn error_handler(&self) -> Arc<impl ErrorHandler> {
         self.0.clone()
     }
+
+    fn cancel_handler(&self) -> Arc<impl CancelHandler> {
+        self.0.clone()
+    }
 }
 
 impl MakePostgresServerHandler {
-    fn make(&self, addr: Option<SocketAddr>, process_id: u32) -> PostgresServerHandler {
+    fn make(&self, addr: Option<SocketAddr>) -> PostgresServerHandler {
+        let process_id = self.process_manager.next_id();
         let session = Arc::new(Session::new(
             addr,
             Channel::Postgres,
@@ -120,6 +129,7 @@ impl MakePostgresServerHandler {
         ));
         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(),

src/servers/src/postgres/handler.rs

@@ -25,6 +25,7 @@ use datafusion_common::ParamValues;
 use datatypes::prelude::ConcreteDataType;
 use datatypes::schema::SchemaRef;
 use futures::{future, stream, Sink, SinkExt, Stream, StreamExt};
+use pgwire::api::cancel::CancelHandler;
 use pgwire::api::portal::{Format, Portal};
 use pgwire::api::query::{ExtendedQueryHandler, SimpleQueryHandler};
 use pgwire::api::results::{
@@ -33,6 +34,7 @@ use pgwire::api::results::{
 use pgwire::api::stmt::{QueryParser, StoredStatement};
 use pgwire::api::{ClientInfo, ErrorHandler, Type};
 use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
+use pgwire::messages::cancel::CancelRequest;
 use pgwire::messages::PgWireBackendMessage;
 use query::query_engine::DescribeResult;
 use session::context::QueryContextRef;
@@ -424,3 +426,17 @@ impl ErrorHandler for PostgresServerHandlerInner {
         debug!("Postgres interface error {}", error)
     }
 }
+
+#[async_trait]
+impl CancelHandler for PostgresServerHandlerInner {
+    async fn on_cancel_request(&self, cancel_request: CancelRequest) {
+        let pid = cancel_request.pid as u32;
+        let _secret_key = cancel_request.secret_key;
+
+        // FIXME:
+        let _ = self
+            .process_manager
+            .kill_process("todo".to_string(), "todo".to_string(), pid)
+            .await;
+    }
+}

src/servers/src/postgres/server.rs

@@ -38,7 +38,6 @@ pub struct PostgresServer {
     tls_server_config: Arc<ReloadableTlsServerConfig>,
     keep_alive_secs: u64,
     bind_addr: Option<SocketAddr>,
-    process_manager: Option<ProcessManagerRef>,
 }
 
 impl PostgresServer {
@@ -50,11 +49,12 @@ impl PostgresServer {
         keep_alive_secs: u64,
         io_runtime: Runtime,
         user_provider: Option<UserProviderRef>,
-        process_manager: Option<ProcessManagerRef>,
+        process_manager: ProcessManagerRef,
     ) -> PostgresServer {
         let make_handler = Arc::new(
             MakePostgresServerHandlerBuilder::default()
                 .query_handler(query_handler.clone())
+                .process_manager(process_manager.clone())
                 .user_provider(user_provider.clone())
                 .force_tls(force_tls)
                 .build()
@@ -66,7 +66,6 @@ impl PostgresServer {
             tls_server_config,
             keep_alive_secs,
             bind_addr: None,
-            process_manager,
         }
     }
 
@@ -77,12 +76,10 @@ impl PostgresServer {
     ) -> impl Future<Output = ()> {
         let handler_maker = self.make_handler.clone();
         let tls_server_config = self.tls_server_config.clone();
-        let process_manager = self.process_manager.clone();
         accepting_stream.for_each(move |tcp_stream| {
             let io_runtime = io_runtime.clone();
             let tls_acceptor = tls_server_config.get_server_config().map(TlsAcceptor::from);
             let handler_maker = handler_maker.clone();
-            let process_id = process_manager.as_ref().map(|p| p.next_id()).unwrap_or(0);
 
             async move {
                 match tcp_stream {
@@ -101,7 +98,7 @@ impl PostgresServer {
 
                         let _handle = io_runtime.spawn(async move {
                             crate::metrics::METRIC_POSTGRES_CONNECTIONS.inc();
-                            let pg_handler = Arc::new(handler_maker.make(addr, process_id));
+                            let pg_handler = Arc::new(handler_maker.make(addr));
                             let r =
                                 process_socket(io_stream, tls_acceptor.clone(), pg_handler).await;
                             crate::metrics::METRIC_POSTGRES_CONNECTIONS.dec();