From 307d10b111fd89f6ffa7b0e9326477b057bc8727 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Arpad=20M=C3=BCller?= <arpad@neon.tech>
Date: Sun, 28 Apr 2024 01:01:21 +0200
Subject: [PATCH] Add HandlerSync trait

---
 libs/postgres_backend/src/lib.rs |  5 +++--
 pageserver/src/page_service.rs   | 10 +++++++---
 proxy/src/console/mgmt.rs        |  2 ++
 safekeeper/src/handler.rs        |  8 ++++++--
 4 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/libs/postgres_backend/src/lib.rs b/libs/postgres_backend/src/lib.rs
index 260018ad89..df6a95e326 100644
--- a/libs/postgres_backend/src/lib.rs
+++ b/libs/postgres_backend/src/lib.rs
@@ -79,7 +79,7 @@ pub fn is_expected_io_error(e: &io::Error) -> bool {
 }
 
 #[async_trait::async_trait]
-pub trait Handler<IO> {
+pub trait Handler<IO>: HandlerSync<IO> {
     /// Handle single query.
     /// postgres_backend will issue ReadyForQuery after calling this (this
     /// might be not what we want after CopyData streaming, but currently we don't
@@ -89,7 +89,8 @@ pub trait Handler<IO> {
         pgb: &mut PostgresBackend<IO>,
         query_string: &str,
     ) -> Result<(), QueryError>;
-
+}
+pub trait HandlerSync<IO> {
     /// Called on startup packet receival, allows to process params.
     ///
     /// If Ok(false) is returned postgres_backend will skip auth -- that is needed for new users
diff --git a/pageserver/src/page_service.rs b/pageserver/src/page_service.rs
index 7d01ad4157..f8b6df36cd 100644
--- a/pageserver/src/page_service.rs
+++ b/pageserver/src/page_service.rs
@@ -1368,8 +1368,7 @@ impl PageServerHandler {
     }
 }
 
-#[async_trait::async_trait]
-impl<IO> postgres_backend::Handler<IO> for PageServerHandler
+impl<IO> postgres_backend::HandlerSync<IO> for PageServerHandler
 where
     IO: AsyncRead + AsyncWrite + Send + Sync + Unpin,
 {
@@ -1409,7 +1408,12 @@ where
     ) -> Result<(), QueryError> {
         Ok(())
     }
-
+}
+#[async_trait::async_trait]
+impl<IO> postgres_backend::Handler<IO> for PageServerHandler
+where
+    IO: AsyncRead + AsyncWrite + Send + Sync + Unpin,
+{
     #[instrument(skip_all, fields(tenant_id, timeline_id))]
     async fn process_query(
         &mut self,
diff --git a/proxy/src/console/mgmt.rs b/proxy/src/console/mgmt.rs
index c7a2d467c0..633252275d 100644
--- a/proxy/src/console/mgmt.rs
+++ b/proxy/src/console/mgmt.rs
@@ -89,6 +89,8 @@ impl postgres_backend::Handler<tokio::net::TcpStream> for MgmtHandler {
     }
 }
 
+impl postgres_backend::HandlerSync<tokio::net::TcpStream> for MgmtHandler {}
+
 fn try_process_query(pgb: &mut PostgresBackendTCP, query: &str) -> Result<(), QueryError> {
     let resp: KickSession = serde_json::from_str(query).context("Failed to parse query as json")?;
 
diff --git a/safekeeper/src/handler.rs b/safekeeper/src/handler.rs
index d411aea41d..bb3619ae15 100644
--- a/safekeeper/src/handler.rs
+++ b/safekeeper/src/handler.rs
@@ -95,8 +95,7 @@ fn cmd_to_string(cmd: &SafekeeperPostgresCommand) -> &str {
     }
 }
 
-#[async_trait::async_trait]
-impl<IO: AsyncRead + AsyncWrite + Unpin + Send> postgres_backend::Handler<IO>
+impl<IO: AsyncRead + AsyncWrite + Unpin + Send> postgres_backend::HandlerSync<IO>
     for SafekeeperPostgresHandler
 {
     // tenant_id and timeline_id are passed in connection string params
@@ -191,7 +190,12 @@ impl<IO: AsyncRead + AsyncWrite + Unpin + Send> postgres_backend::Handler<IO>
         self.claims = Some(data.claims);
         Ok(())
     }
+}
 
+#[async_trait::async_trait]
+impl<IO: AsyncRead + AsyncWrite + Unpin + Send> postgres_backend::Handler<IO>
+    for SafekeeperPostgresHandler
+{
     async fn process_query(
         &mut self,
         pgb: &mut PostgresBackend<IO>,