sqlx-postgres: Remove PgStream from PgConnection

joeydewaal · joeydewaal · commit 014e12255f82 · 2025-07-04T15:59:39.000+02:00
diff --git a/sqlx-postgres/src/connection/establish.rs b/sqlx-postgres/src/connection/establish.rs
@@ -13,15 +13,13 @@ use super::worker::Worker;
 impl PgConnection {
     pub(crate) async fn establish(options: &PgConnectOptions) -> Result<Self, Error> {
         // Upgrade to TLS if we were asked to and the server supports it
-        let pg_stream = PgStream::connect(options).await?;
-
         let stream = PgStream::connect(options).await?;
 
         let (notif_tx, notif_rx) = unbounded();
 
-        let (x, shared) = Worker::spawn(stream.into_inner(), notif_tx);
+        let (channel, shared) = Worker::spawn(stream.into_inner(), notif_tx);
 
-        let mut conn = PgConnection::new(pg_stream, options, x, notif_rx, shared);
+        let mut conn = PgConnection::new(options, channel, notif_rx, shared);
 
         // To begin a session, a frontend opens a connection to the server
         // and sends a startup message.
diff --git a/sqlx-postgres/src/connection/mod.rs b/sqlx-postgres/src/connection/mod.rs
@@ -23,8 +23,6 @@ use crate::{PgConnectOptions, PgTypeInfo, Postgres};
 
 pub(crate) use sqlx_core::connection::*;
 
-pub use self::stream::PgStream;
-
 pub(crate) mod describe;
 mod establish;
 mod executor;
@@ -43,11 +41,6 @@ pub struct PgConnection {
 }
 
 pub struct PgConnectionInner {
-    // underlying TCP or UDS stream,
-    // wrapped in a potentially TLS stream,
-    // wrapped in a buffered stream
-    pub(crate) stream: PgStream,
-
     chan: UnboundedSender<IoRequest>,
 
     pub(crate) notifications: UnboundedReceiver<Notification>,
@@ -112,7 +105,6 @@ impl PgConnection {
     }
 
     fn new(
-        stream: PgStream,
         options: &PgConnectOptions,
         chan: UnboundedSender<IoRequest>,
         notifications: UnboundedReceiver<Notification>,
@@ -132,7 +124,6 @@ impl PgConnection {
                 cache_elem_type_to_array: HashMap::new(),
                 cache_table_to_column_names: HashMap::new(),
                 transaction_depth: 0,
-                stream,
                 server_version_num: None,
                 shared,
             }),
@@ -289,7 +280,7 @@ impl Connection for PgConnection {
     }
 
     fn shrink_buffers(&mut self) {
-        self.inner.stream.shrink_buffers();
+        // No-op
     }
 
     #[doc(hidden)]
@@ -299,7 +290,7 @@ impl Connection for PgConnection {
 
     #[doc(hidden)]
     fn should_flush(&self) -> bool {
-        !self.inner.stream.write_buffer().is_empty()
+        false
     }
 }
 
diff --git a/sqlx-postgres/src/connection/stream.rs b/sqlx-postgres/src/connection/stream.rs
@@ -1,16 +1,9 @@
-use std::ops::{ControlFlow, Deref, DerefMut};
 use std::str::FromStr;
 
-use log::Level;
-use sqlx_core::bytes::Buf;
-
 use crate::connection::tls::MaybeUpgradeTls;
 use crate::error::Error;
-use crate::message::{
-    BackendMessage, BackendMessageFormat, EncodeMessage, FrontendMessage, Notice, ReceivedMessage,
-};
 use crate::net::{self, BufferedSocket, Socket};
-use crate::{PgConnectOptions, PgDatabaseError, PgSeverity};
+use crate::PgConnectOptions;
 
 // the stream is a separate type from the connection to uphold the invariant where an instantiated
 // [PgConnection] is a **valid** connection to postgres
@@ -44,154 +37,6 @@ impl PgStream {
             inner: BufferedSocket::new(socket),
         })
     }
-
-    #[inline(always)]
-    pub(crate) fn write_msg(&mut self, message: impl FrontendMessage) -> Result<(), Error> {
-        self.write(EncodeMessage(message))
-    }
-
-    pub(crate) async fn send<T>(&mut self, message: T) -> Result<(), Error>
-    where
-        T: FrontendMessage,
-    {
-        self.write_msg(message)?;
-        self.flush().await?;
-        Ok(())
-    }
-
-    // Expect a specific type and format
-    pub(crate) async fn recv_expect<B: BackendMessage>(&mut self) -> Result<B, Error> {
-        self.recv().await?.decode()
-    }
-
-    pub(crate) async fn recv_unchecked(&mut self) -> Result<ReceivedMessage, Error> {
-        // NOTE: to not break everything, this should be cancel-safe;
-        // DO NOT modify `buf` unless a full message has been read
-        self.inner
-            .try_read(|buf| {
-                // all packets in postgres start with a 5-byte header
-                // this header contains the message type and the total length of the message
-                let Some(mut header) = buf.get(..5) else {
-                    return Ok(ControlFlow::Continue(5));
-                };
-
-                let format = BackendMessageFormat::try_from_u8(header.get_u8())?;
-
-                let message_len = header.get_u32() as usize;
-
-                let expected_len = message_len
-                    .checked_add(1)
-                    // this shouldn't really happen but is mostly a sanity check
-                    .ok_or_else(|| {
-                        err_protocol!("message_len + 1 overflows usize: {message_len}")
-                    })?;
-
-                if buf.len() < expected_len {
-                    return Ok(ControlFlow::Continue(expected_len));
-                }
-
-                // `buf` SHOULD NOT be modified ABOVE this line
-
-                // pop off the format code since it's not counted in `message_len`
-                buf.advance(1);
-
-                // consume the message, including the length prefix
-                let mut contents = buf.split_to(message_len).freeze();
-
-                // cut off the length prefix
-                contents.advance(4);
-
-                Ok(ControlFlow::Break(ReceivedMessage { format, contents }))
-            })
-            .await
-    }
-
-    // Get the next message from the server
-    // May wait for more data from the server
-    pub(crate) async fn recv(&mut self) -> Result<ReceivedMessage, Error> {
-        loop {
-            let message = self.recv_unchecked().await?;
-
-            match message.format {
-                BackendMessageFormat::ErrorResponse => {
-                    // An error returned from the database server.
-                    return Err(message.decode::<PgDatabaseError>()?.into());
-                }
-
-                // BackendMessageFormat::ParameterStatus => {
-                //     // informs the frontend about the current (initial)
-                //     // setting of backend parameters
-
-                //     let ParameterStatus { name, value } = message.decode()?;
-                //     // TODO: handle `client_encoding`, `DateStyle` change
-
-                //     match name.as_str() {
-                //         "server_version" => {
-                //             self.server_version_num = parse_server_version(&value);
-                //         }
-                //         _ => {
-                //             self.parameter_statuses.insert(name, value);
-                //         }
-                //     }
-
-                //     continue;
-                // }
-                BackendMessageFormat::NoticeResponse => {
-                    // do we need this to be more configurable?
-                    // if you are reading this comment and think so, open an issue
-
-                    let notice: Notice = message.decode()?;
-
-                    let (log_level, tracing_level) = match notice.severity() {
-                        PgSeverity::Fatal | PgSeverity::Panic | PgSeverity::Error => {
-                            (Level::Error, tracing::Level::ERROR)
-                        }
-                        PgSeverity::Warning => (Level::Warn, tracing::Level::WARN),
-                        PgSeverity::Notice => (Level::Info, tracing::Level::INFO),
-                        PgSeverity::Debug => (Level::Debug, tracing::Level::DEBUG),
-                        PgSeverity::Info | PgSeverity::Log => (Level::Trace, tracing::Level::TRACE),
-                    };
-
-                    let log_is_enabled = log::log_enabled!(
-                        target: "sqlx::postgres::notice",
-                        log_level
-                    ) || sqlx_core::private_tracing_dynamic_enabled!(
-                        target: "sqlx::postgres::notice",
-                        tracing_level
-                    );
-                    if log_is_enabled {
-                        sqlx_core::private_tracing_dynamic_event!(
-                            target: "sqlx::postgres::notice",
-                            tracing_level,
-                            message = notice.message()
-                        );
-                    }
-
-                    continue;
-                }
-
-                _ => {}
-            }
-
-            return Ok(message);
-        }
-    }
-}
-
-impl Deref for PgStream {
-    type Target = BufferedSocket<Box<dyn Socket>>;
-
-    #[inline]
-    fn deref(&self) -> &Self::Target {
-        &self.inner
-    }
-}
-
-impl DerefMut for PgStream {
-    #[inline]
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.inner
-    }
 }
 
 // reference:
diff --git a/sqlx-postgres/src/connection/worker.rs b/sqlx-postgres/src/connection/worker.rs
@@ -8,7 +8,7 @@ use std::{
 };
 
 use crate::message::{
-    BackendMessageFormat, FrontendMessage, Notification, ParameterStatus, ReadyForQuery,
+    BackendMessageFormat, FrontendMessage, Notice, Notification, ParameterStatus, ReadyForQuery,
     ReceivedMessage, Terminate, TransactionStatus,
 };
 use futures_channel::mpsc::{unbounded, UnboundedReceiver, UnboundedSender};
@@ -153,12 +153,16 @@ impl Worker {
                 }
                 BackendMessageFormat::ParameterStatus => {
                     // Asynchronous response
-                    //
                     let ParameterStatus { name, value } = response.decode()?;
                     self.shared.insert_parameter_status(name, value);
                 }
                 BackendMessageFormat::NoticeResponse => {
-                    // Asynchronous response - todo
+                    // do we need this to be more configurable?
+                    // if you are reading this comment and think so, open an issue
+
+                    let notice: Notice = response.decode()?;
+
+                    notice.emit_notice();
                 }
                 _ => self.send_back(response)?,
             }
diff --git a/sqlx-postgres/src/message/response.rs b/sqlx-postgres/src/message/response.rs
@@ -1,6 +1,7 @@
 use std::ops::Range;
 use std::str::from_utf8;
 
+use log::Level;
 use memchr::memchr;
 
 use sqlx_core::bytes::Bytes;
@@ -90,6 +91,33 @@ impl Notice {
             .map(|(_, range)| &self.storage[range])
             .next()
     }
+
+    pub(crate) fn emit_notice(&self) {
+        let (log_level, tracing_level) = match self.severity() {
+            PgSeverity::Fatal | PgSeverity::Panic | PgSeverity::Error => {
+                (Level::Error, tracing::Level::ERROR)
+            }
+            PgSeverity::Warning => (Level::Warn, tracing::Level::WARN),
+            PgSeverity::Notice => (Level::Info, tracing::Level::INFO),
+            PgSeverity::Debug => (Level::Debug, tracing::Level::DEBUG),
+            PgSeverity::Info | PgSeverity::Log => (Level::Trace, tracing::Level::TRACE),
+        };
+
+        let log_is_enabled = log::log_enabled!(
+            target: "sqlx::postgres::notice",
+            log_level
+        ) || sqlx_core::private_tracing_dynamic_enabled!(
+            target: "sqlx::postgres::notice",
+            tracing_level
+        );
+        if log_is_enabled {
+            sqlx_core::private_tracing_dynamic_event!(
+                target: "sqlx::postgres::notice",
+                tracing_level,
+                message = self.message()
+            );
+        }
+    }
 }
 
 impl Notice {
