From ab61864df1db0c714ab3ad0d688a2eac226bdc6f Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conradludgate@gmail.com>
Date: Wed, 21 May 2025 21:53:57 +0100
Subject: [PATCH] proxy(tokio-postgres): move statement cleanup to client drop

---
 libs/proxy/tokio-postgres2/src/client.rs    | 15 ++++++++++
 libs/proxy/tokio-postgres2/src/prepare.rs   |  2 +-
 libs/proxy/tokio-postgres2/src/statement.rs | 32 ++-------------------
 3 files changed, 19 insertions(+), 30 deletions(-)

diff --git a/libs/proxy/tokio-postgres2/src/client.rs b/libs/proxy/tokio-postgres2/src/client.rs
index 186eb07000..d1d87306a6 100644
--- a/libs/proxy/tokio-postgres2/src/client.rs
+++ b/libs/proxy/tokio-postgres2/src/client.rs
@@ -118,6 +118,21 @@ pub struct Client {
     secret_key: i32,
 }
 
+impl Drop for Client {
+    fn drop(&mut self) {
+        if let Some(stmt) = self.cached_typeinfo.typeinfo.take() {
+            let buf = self.inner.with_buf(|buf| {
+                frontend::close(b'S', stmt.name(), buf).unwrap();
+                frontend::sync(buf);
+                buf.split().freeze()
+            });
+            let _ = self
+                .inner
+                .send(RequestMessages::Single(FrontendMessage::Raw(buf)));
+        }
+    }
+}
+
 impl Client {
     pub(crate) fn new(
         sender: mpsc::UnboundedSender<Request>,
diff --git a/libs/proxy/tokio-postgres2/src/prepare.rs b/libs/proxy/tokio-postgres2/src/prepare.rs
index b27eabcb0e..8b6b363889 100644
--- a/libs/proxy/tokio-postgres2/src/prepare.rs
+++ b/libs/proxy/tokio-postgres2/src/prepare.rs
@@ -65,7 +65,7 @@ async fn prepare_typecheck(
         }
     }
 
-    Ok(Statement::new(client, name, parameters, columns))
+    Ok(Statement::new(name, parameters, columns))
 }
 
 fn encode(client: &InnerClient, name: &str, query: &str, types: &[Type]) -> Result<Bytes, Error> {
diff --git a/libs/proxy/tokio-postgres2/src/statement.rs b/libs/proxy/tokio-postgres2/src/statement.rs
index e4828db712..32c8317cc5 100644
--- a/libs/proxy/tokio-postgres2/src/statement.rs
+++ b/libs/proxy/tokio-postgres2/src/statement.rs
@@ -1,35 +1,16 @@
 use std::fmt;
-use std::sync::{Arc, Weak};
+use std::sync::Arc;
 
+use crate::types::Type;
 use postgres_protocol2::Oid;
 use postgres_protocol2::message::backend::Field;
-use postgres_protocol2::message::frontend;
-
-use crate::client::InnerClient;
-use crate::codec::FrontendMessage;
-use crate::connection::RequestMessages;
-use crate::types::Type;
 
 struct StatementInner {
-    client: Weak<InnerClient>,
     name: &'static str,
     params: Vec<Type>,
     columns: Vec<Column>,
 }
 
-impl Drop for StatementInner {
-    fn drop(&mut self) {
-        if let Some(client) = self.client.upgrade() {
-            let buf = client.with_buf(|buf| {
-                frontend::close(b'S', self.name, buf).unwrap();
-                frontend::sync(buf);
-                buf.split().freeze()
-            });
-            let _ = client.send(RequestMessages::Single(FrontendMessage::Raw(buf)));
-        }
-    }
-}
-
 /// A prepared statement.
 ///
 /// Prepared statements can only be used with the connection that created them.
@@ -37,14 +18,8 @@ impl Drop for StatementInner {
 pub struct Statement(Arc<StatementInner>);
 
 impl Statement {
-    pub(crate) fn new(
-        inner: &Arc<InnerClient>,
-        name: &'static str,
-        params: Vec<Type>,
-        columns: Vec<Column>,
-    ) -> Statement {
+    pub(crate) fn new(name: &'static str, params: Vec<Type>, columns: Vec<Column>) -> Statement {
         Statement(Arc::new(StatementInner {
-            client: Arc::downgrade(inner),
             name,
             params,
             columns,
@@ -53,7 +28,6 @@ impl Statement {
 
     pub(crate) fn new_anonymous(params: Vec<Type>, columns: Vec<Column>) -> Statement {
         Statement(Arc::new(StatementInner {
-            client: Weak::new(),
             name: "<anonymous>",
             params,
             columns,