feat: add connect_with_socket() for custom transports

Add the ability to establish MySQL and PostgreSQL connections over
pre-connected sockets, enabling custom transports such as in-memory
pipes, simulation frameworks (e.g. turmoil for deterministic simulation
testing), SSH tunnels, or SOCKS proxies.

Changes:
- Add `net::connect_with()` in sqlx-core as the analog of `connect_tcp`
  and `connect_uds` for pre-established sockets
- Add `MySqlConnectOptions::connect_with_socket()` and
  `MySqlConnection::connect_with_socket()` for MySQL
- Add `PgConnectOptions::connect_with_socket()` and
  `PgConnection::connect_with_socket()` for PostgreSQL
- Re-export `Socket` trait and `connect_with` from top-level `sqlx::net`
- Refactor MySQL post-connect init into shared `after_connect()` to
  avoid duplication between `connect()` and `connect_with_socket()`

The Socket trait was already public; these changes simply expose a way
to pass a pre-connected Socket into the connection establishment flow
that was previously only reachable through connect_tcp/connect_uds.

Author: airhorns

sqlx-core/src/net/mod.rs

@@ -2,5 +2,6 @@ mod socket;
 pub mod tls;
 
 pub use socket::{
-    connect_tcp, connect_uds, BufferedSocket, Socket, SocketIntoBox, WithSocket, WriteBuffer,
+    connect_tcp, connect_uds, connect_with, BufferedSocket, Socket, SocketIntoBox, WithSocket,
+    WriteBuffer,
 };

sqlx-core/src/net/socket/mod.rs

@@ -242,6 +242,16 @@ pub async fn connect_tcp<Ws: WithSocket>(
     }
 }
 
+/// Use a pre-connected socket, passing it to the given [`WithSocket`] handler.
+///
+/// This is the analog of [`connect_tcp`] and [`connect_uds`] for sockets
+/// that have already been established by the caller. This enables custom
+/// transports such as in-memory pipes, simulation frameworks (e.g. turmoil),
+/// SSH tunnels, or SOCKS proxies.
+pub async fn connect_with<S: Socket, Ws: WithSocket>(socket: S, with_socket: Ws) -> Ws::Output {
+    with_socket.with_socket(socket).await
+}
+
 /// Connect a Unix Domain Socket at the given path.
 ///
 /// Returns an error if Unix Domain Sockets are not supported on this platform.

sqlx-mysql/src/connection/establish.rs

@@ -21,8 +21,31 @@ impl MySqlConnection {
             None => crate::net::connect_tcp(&options.host, options.port, do_handshake).await?,
         };
 
-        let stream = handshake?;
+        Self::from_stream(options, handshake?)
+    }
+
+    /// Establish a connection over a pre-connected socket.
+    ///
+    /// The provided socket must already be connected to a MySQL-compatible
+    /// server. The MySQL handshake and authentication will be performed
+    /// using the credentials from `options`.
+    ///
+    /// This enables custom transports such as in-memory pipes, simulation
+    /// frameworks (e.g. turmoil), SSH tunnels, or SOCKS proxies.
+    ///
+    /// Note: this only performs the low-level handshake and authentication.
+    /// Use [`MySqlConnectOptions::connect_with_socket()`] for a fully
+    /// initialized connection (with `SET NAMES`, `sql_mode`, etc.).
+    pub async fn connect_with_socket<S: Socket>(
+        options: &MySqlConnectOptions,
+        socket: S,
+    ) -> Result<Self, Error> {
+        let do_handshake = DoHandshake::new(options)?;
+        let stream = crate::net::connect_with(socket, do_handshake).await?;
+        Self::from_stream(options, stream)
+    }
 
+    fn from_stream(options: &MySqlConnectOptions, stream: MySqlStream) -> Result<Self, Error> {
         Ok(Self {
             inner: Box::new(MySqlConnectionInner {
                 stream,

sqlx-mysql/src/options/connect.rs

@@ -1,12 +1,106 @@
 use crate::connection::ConnectOptions;
 use crate::error::Error;
 use crate::executor::Executor;
+use crate::net::Socket;
 use crate::{MySqlConnectOptions, MySqlConnection};
 use futures_core::future::BoxFuture;
 use log::LevelFilter;
 use sqlx_core::Url;
 use std::time::Duration;
 
+impl MySqlConnectOptions {
+    /// Establish a fully initialized connection over a pre-connected socket.
+    ///
+    /// This performs the MySQL handshake, authentication, and post-connect
+    /// initialization (`SET NAMES`, `sql_mode`, `time_zone`) over the
+    /// provided socket.
+    ///
+    /// The socket must already be connected to a MySQL-compatible server.
+    /// This enables custom transports such as in-memory pipes, simulation
+    /// frameworks (e.g. turmoil), SSH tunnels, or SOCKS proxies.
+    ///
+    /// # Example
+    ///
+    /// ```rust,no_run
+    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
+    /// use sqlx::mysql::MySqlConnectOptions;
+    ///
+    /// let options = MySqlConnectOptions::new()
+    ///     .username("root")
+    ///     .database("mydb");
+    ///
+    /// let stream = tokio::net::TcpStream::connect("127.0.0.1:3306").await?;
+    /// let conn = options.connect_with_socket(stream).await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub async fn connect_with_socket<S: Socket>(
+        &self,
+        socket: S,
+    ) -> Result<MySqlConnection, Error> {
+        let mut conn = MySqlConnection::connect_with_socket(self, socket).await?;
+        self.after_connect(&mut conn).await?;
+        Ok(conn)
+    }
+
+    /// Post-connection initialization shared between `connect()` and
+    /// `connect_with_socket()`.
+    ///
+    /// After the connection is established, we initialize by configuring a few
+    /// connection parameters:
+    ///
+    /// - <https://mariadb.com/kb/en/sql-mode/>
+    ///
+    /// - `PIPES_AS_CONCAT` - Allows using the pipe character (ASCII 124) as string concatenation
+    ///   operator. This means that "A" || "B" can be used in place of CONCAT("A", "B").
+    ///
+    /// - `NO_ENGINE_SUBSTITUTION` - If not set, if the available storage engine specified by a
+    ///   CREATE TABLE is not available, a warning is given and the default storage engine is used
+    ///   instead.
+    ///
+    /// - `NO_ZERO_DATE` - Don't allow '0000-00-00'. This is invalid in Rust.
+    ///
+    /// - `NO_ZERO_IN_DATE` - Don't allow 'YYYY-00-00'. This is invalid in Rust.
+    ///
+    /// Setting the time zone allows us to assume that the output from a TIMESTAMP field is UTC.
+    ///
+    /// - <https://mathiasbynens.be/notes/mysql-utf8mb4>
+    async fn after_connect(&self, conn: &mut MySqlConnection) -> Result<(), Error> {
+        let mut sql_mode = Vec::new();
+        if self.pipes_as_concat {
+            sql_mode.push(r#"PIPES_AS_CONCAT"#);
+        }
+        if self.no_engine_substitution {
+            sql_mode.push(r#"NO_ENGINE_SUBSTITUTION"#);
+        }
+
+        let mut options = Vec::new();
+        if !sql_mode.is_empty() {
+            options.push(format!(
+                r#"sql_mode=(SELECT CONCAT(@@sql_mode, ',{}'))"#,
+                sql_mode.join(",")
+            ));
+        }
+        if let Some(timezone) = &self.timezone {
+            options.push(format!(r#"time_zone='{}'"#, timezone));
+        }
+        if self.set_names {
+            options.push(format!(
+                r#"NAMES {} COLLATE {}"#,
+                conn.inner.stream.charset.as_str(),
+                conn.inner.stream.collation.as_str()
+            ))
+        }
+
+        if !options.is_empty() {
+            conn.execute(&*format!(r#"SET {};"#, options.join(",")))
+                .await?;
+        }
+
+        Ok(())
+    }
+}
+
 impl ConnectOptions for MySqlConnectOptions {
     type Connection = MySqlConnection;
 
@@ -24,63 +118,7 @@ impl ConnectOptions for MySqlConnectOptions {
     {
         Box::pin(async move {
             let mut conn = MySqlConnection::establish(self).await?;
-
-            // After the connection is established, we initialize by configuring a few
-            // connection parameters
-
-            // https://mariadb.com/kb/en/sql-mode/
-
-            // PIPES_AS_CONCAT - Allows using the pipe character (ASCII 124) as string concatenation operator.
-            //                   This means that "A" || "B" can be used in place of CONCAT("A", "B").
-
-            // NO_ENGINE_SUBSTITUTION - If not set, if the available storage engine specified by a CREATE TABLE is
-            //                          not available, a warning is given and the default storage
-            //                          engine is used instead.
-
-            // NO_ZERO_DATE - Don't allow '0000-00-00'. This is invalid in Rust.
-
-            // NO_ZERO_IN_DATE - Don't allow 'YYYY-00-00'. This is invalid in Rust.
-
-            // --
-
-            // Setting the time zone allows us to assume that the output
-            // from a TIMESTAMP field is UTC
-
-            // --
-
-            // https://mathiasbynens.be/notes/mysql-utf8mb4
-
-            let mut sql_mode = Vec::new();
-            if self.pipes_as_concat {
-                sql_mode.push(r#"PIPES_AS_CONCAT"#);
-            }
-            if self.no_engine_substitution {
-                sql_mode.push(r#"NO_ENGINE_SUBSTITUTION"#);
-            }
-
-            let mut options = Vec::new();
-            if !sql_mode.is_empty() {
-                options.push(format!(
-                    r#"sql_mode=(SELECT CONCAT(@@sql_mode, ',{}'))"#,
-                    sql_mode.join(",")
-                ));
-            }
-            if let Some(timezone) = &self.timezone {
-                options.push(format!(r#"time_zone='{}'"#, timezone));
-            }
-            if self.set_names {
-                options.push(format!(
-                    r#"NAMES {} COLLATE {}"#,
-                    conn.inner.stream.charset.as_str(),
-                    conn.inner.stream.collation.as_str()
-                ))
-            }
-
-            if !options.is_empty() {
-                conn.execute(&*format!(r#"SET {};"#, options.join(",")))
-                    .await?;
-            }
-
+            self.after_connect(&mut conn).await?;
             Ok(conn)
         })
     }

sqlx-postgres/src/connection/establish.rs

@@ -7,6 +7,7 @@ use crate::io::StatementId;
 use crate::message::{
     Authentication, BackendKeyData, BackendMessageFormat, Password, ReadyForQuery, Startup,
 };
+use crate::net::Socket;
 use crate::{PgConnectOptions, PgConnection};
 
 use super::PgConnectionInner;
@@ -16,9 +17,35 @@ use super::PgConnectionInner;
 
 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 mut stream = PgStream::connect(options).await?;
+        let stream = PgStream::connect(options).await?;
+        Self::establish_with_stream(options, stream).await
+    }
+
+    /// Establish a connection over a pre-connected socket.
+    ///
+    /// The provided socket must already be connected to a
+    /// PostgreSQL-compatible server. The startup handshake, TLS upgrade
+    /// (if configured), and authentication will be performed over this
+    /// socket.
+    ///
+    /// This enables custom transports such as in-memory pipes, simulation
+    /// frameworks (e.g. turmoil), SSH tunnels, or SOCKS proxies.
+    ///
+    /// Note: this only performs the low-level handshake and authentication.
+    /// Use [`PgConnectOptions::connect_with_socket()`] for a fully
+    /// initialized connection.
+    pub async fn connect_with_socket<S: Socket>(
+        options: &PgConnectOptions,
+        socket: S,
+    ) -> Result<Self, Error> {
+        let stream = PgStream::connect_with_socket(options, socket).await?;
+        Self::establish_with_stream(options, stream).await
+    }
 
+    async fn establish_with_stream(
+        options: &PgConnectOptions,
+        mut stream: PgStream,
+    ) -> Result<Self, Error> {
         // To begin a session, a frontend opens a connection to the server
         // and sends a startup message.
 

sqlx-postgres/src/connection/stream.rs

@@ -57,6 +57,24 @@ impl PgStream {
         })
     }
 
+    /// Create a stream from a pre-connected socket.
+    ///
+    /// The socket must already be connected to a PostgreSQL server.
+    /// TLS upgrade will be attempted if configured in `options`.
+    pub(super) async fn connect_with_socket<S: Socket>(
+        options: &PgConnectOptions,
+        socket: S,
+    ) -> Result<Self, Error> {
+        let socket = net::connect_with(socket, MaybeUpgradeTls(options)).await?;
+
+        Ok(Self {
+            inner: BufferedSocket::new(socket),
+            notifications: None,
+            parameter_statuses: BTreeMap::default(),
+            server_version_num: None,
+        })
+    }
+
     #[inline(always)]
     pub(crate) fn write_msg(&mut self, message: impl FrontendMessage) -> Result<(), Error> {
         self.write(EncodeMessage(message))

sqlx-postgres/src/options/connect.rs

@@ -1,11 +1,42 @@
 use crate::connection::ConnectOptions;
 use crate::error::Error;
+use crate::net::Socket;
 use crate::{PgConnectOptions, PgConnection};
 use futures_core::future::BoxFuture;
 use log::LevelFilter;
 use sqlx_core::Url;
 use std::time::Duration;
 
+impl PgConnectOptions {
+    /// Establish a connection over a pre-connected socket.
+    ///
+    /// This performs the PostgreSQL startup handshake, TLS upgrade
+    /// (if configured), and authentication over the provided socket.
+    ///
+    /// The socket must already be connected to a PostgreSQL-compatible server.
+    /// This enables custom transports such as in-memory pipes, simulation
+    /// frameworks (e.g. turmoil), SSH tunnels, or SOCKS proxies.
+    ///
+    /// # Example
+    ///
+    /// ```rust,no_run
+    /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
+    /// use sqlx::postgres::PgConnectOptions;
+    ///
+    /// let options = PgConnectOptions::new()
+    ///     .username("postgres")
+    ///     .database("mydb");
+    ///
+    /// let stream = tokio::net::TcpStream::connect("127.0.0.1:5432").await?;
+    /// let conn = options.connect_with_socket(stream).await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub async fn connect_with_socket<S: Socket>(&self, socket: S) -> Result<PgConnection, Error> {
+        PgConnection::connect_with_socket(self, socket).await
+    }
+}
+
 impl ConnectOptions for PgConnectOptions {
     type Connection = PgConnection;
 

src/lib.rs

@@ -36,6 +36,18 @@ pub use sqlx_core::types::Type;
 pub use sqlx_core::value::{Value, ValueRef};
 pub use sqlx_core::Either;
 
+/// Networking primitives used by SQLx database drivers.
+///
+/// The [`Socket`][net::Socket] trait allows implementing custom transports
+/// for database connections (e.g. in-memory pipes, simulation frameworks,
+/// SSH tunnels, SOCKS proxies). Use with
+/// [`MySqlConnectOptions::connect_with_socket()`][crate::mysql::MySqlConnectOptions::connect_with_socket]
+/// or [`PgConnectOptions::connect_with_socket()`][crate::postgres::PgConnectOptions::connect_with_socket].
+pub mod net {
+    pub use sqlx_core::io::ReadBuf;
+    pub use sqlx_core::net::{connect_with, Socket};
+}
+
 #[doc(inline)]
 pub use sqlx_core::error::{self, Error, Result};