feat: add custom socket transport support for Postgres and MySQL

Add `connect_socket()` methods to `PgConnection` and `MySqlConnection`
that accept any pre-connected socket implementing the `Socket` trait.
This enables using custom transport layers (e.g., vsock for AWS Nitro
Enclaves, QUIC, or other non-TCP/UDS transports) without forking sqlx.

Re-export `Socket` and `ReadBuf` traits from `sqlx::net` so users can
implement custom socket types.

Author: jowparks

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_socket, connect_tcp, connect_uds, BufferedSocket, Socket, SocketIntoBox, WithSocket,
+    WriteBuffer,
 };

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

@@ -202,6 +202,20 @@ pub async fn connect_tcp<Ws: WithSocket>(
     }
 }
 
+/// Connect using a pre-connected socket that implements [`Socket`].
+///
+/// This allows using custom transport layers (e.g., vsock, QUIC, or any
+/// `AsyncRead + AsyncWrite` type) with SQLx database connections.
+///
+/// The socket will be passed through the `with_socket` handler, which
+/// typically performs TLS upgrade negotiation.
+pub async fn connect_socket<S: Socket, Ws: WithSocket>(
+    socket: S,
+    with_socket: Ws,
+) -> crate::Result<Ws::Output> {
+    Ok(with_socket.with_socket(socket).await)
+}
+
 /// Open a TCP socket to `host` and `port`.
 ///
 /// If `host` is a hostname, attempt to connect to each address it resolves to.

sqlx-mysql/src/connection/establish.rs

@@ -22,15 +22,29 @@ impl MySqlConnection {
 
         let stream = handshake?;
 
-        Ok(Self {
+        Ok(Self::establish_with_stream(stream, options))
+    }
+
+    pub(crate) async fn establish_with_socket<S: Socket>(
+        socket: S,
+        options: &MySqlConnectOptions,
+    ) -> Result<Self, Error> {
+        let do_handshake = DoHandshake::new(options)?;
+        let stream = do_handshake.with_socket(socket).await?;
+
+        Ok(Self::establish_with_stream(stream, options))
+    }
+
+    fn establish_with_stream(stream: MySqlStream, options: &MySqlConnectOptions) -> Self {
+        Self {
             inner: Box::new(MySqlConnectionInner {
                 stream,
                 transaction_depth: 0,
                 status_flags: Default::default(),
                 cache_statement: StatementCache::new(options.statement_cache_capacity),
                 log_settings: options.log_settings.clone(),
             }),
-        })
+        }
     }
 }
 

sqlx-mysql/src/connection/mod.rs

@@ -52,6 +52,38 @@ pub(crate) struct MySqlConnectionInner {
 }
 
 impl MySqlConnection {
+    /// Connect to a MySQL database using a pre-connected socket.
+    ///
+    /// This allows using custom transport layers such as vsock, QUIC,
+    /// or any type that implements [`sqlx_core::net::Socket`].
+    ///
+    /// The provided socket will go through TLS upgrade negotiation based on the
+    /// SSL mode configured in `options`.
+    ///
+    /// # Example
+    ///
+    /// ```rust,ignore
+    /// use sqlx::mysql::{MySqlConnectOptions, MySqlConnection};
+    ///
+    /// # async fn example() -> sqlx::Result<()> {
+    /// let socket: tokio::net::TcpStream = todo!();
+    /// let options = MySqlConnectOptions::new()
+    ///     .username("root")
+    ///     .database("mydb");
+    ///
+    /// let _conn = MySqlConnection::connect_socket(socket, &options).await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub async fn connect_socket<S: sqlx_core::net::Socket>(
+        socket: S,
+        options: &MySqlConnectOptions,
+    ) -> Result<Self, Error> {
+        let mut conn = Self::establish_with_socket(socket, options).await?;
+        crate::options::apply_connect_options(&mut conn, options).await?;
+        Ok(conn)
+    }
+
     pub(crate) fn in_transaction(&self) -> bool {
         self.inner
             .status_flags

sqlx-mysql/src/options/connect.rs

@@ -24,80 +24,92 @@ impl ConnectOptions for MySqlConnectOptions {
     {
         let mut conn = MySqlConnection::establish(self).await?;
 
-        // After the connection is established, we initialize by configuring a few
-        // connection parameters
+        apply_connect_options(&mut conn, self).await?;
 
-        // https://mariadb.com/kb/en/sql-mode/
+        Ok(conn)
+    }
 
-        // 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").
+    fn log_statements(mut self, level: LevelFilter) -> Self {
+        self.log_settings.log_statements(level);
+        self
+    }
 
-        // 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.
+    fn log_slow_statements(mut self, level: LevelFilter, duration: Duration) -> Self {
+        self.log_settings.log_slow_statements(level, duration);
+        self
+    }
+}
 
-        // NO_ZERO_DATE - Don't allow '0000-00-00'. This is invalid in Rust.
+pub(crate) async fn apply_connect_options(
+    conn: &mut MySqlConnection,
+    options: &MySqlConnectOptions,
+) -> Result<(), Error> {
+    // After the connection is established, we initialize by configuring a few
+    // connection parameters
 
-        // NO_ZERO_IN_DATE - Don't allow 'YYYY-00-00'. This is invalid in Rust.
+    // 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").
 
-        // Setting the time zone allows us to assume that the output
-        // from a TIMESTAMP field is UTC
+    // 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.
 
-        // https://mathiasbynens.be/notes/mysql-utf8mb4
+    // NO_ZERO_IN_DATE - Don't allow 'YYYY-00-00'. This is invalid in Rust.
 
-        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(",")
-            ));
-        }
+    // Setting the time zone allows us to assume that the output
+    // from a TIMESTAMP field is UTC
 
-        if let Some(timezone) = &self.timezone {
-            options.push(format!(r#"time_zone='{}'"#, timezone));
-        }
+    // --
 
-        if self.set_names {
-            // As it turns out, we don't _have_ to set a collation if we don't want to.
-            // We can let the server choose the default collation for the charset.
-            let set_names = if let Some(collation) = &self.collation {
-                format!(r#"NAMES {} COLLATE {collation}"#, self.charset,)
-            } else {
-                // Leaves the default collation up to the server,
-                // but ensures statements and results are encoded using the proper charset.
-                format!("NAMES {}", self.charset)
-            };
+    // https://mathiasbynens.be/notes/mysql-utf8mb4
 
-            options.push(set_names);
-        }
+    let mut sql_mode = Vec::new();
+    if options.pipes_as_concat {
+        sql_mode.push(r#"PIPES_AS_CONCAT"#);
+    }
+    if options.no_engine_substitution {
+        sql_mode.push(r#"NO_ENGINE_SUBSTITUTION"#);
+    }
 
-        if !options.is_empty() {
-            conn.execute(AssertSqlSafe(format!(r#"SET {};"#, options.join(","))))
-                .await?;
-        }
+    let mut session_options = Vec::new();
+    if !sql_mode.is_empty() {
+        session_options.push(format!(
+            r#"sql_mode=(SELECT CONCAT(@@sql_mode, ',{}'))"#,
+            sql_mode.join(",")
+        ));
+    }
 
-        Ok(conn)
+    if let Some(timezone) = &options.timezone {
+        session_options.push(format!(r#"time_zone='{}'"#, timezone));
     }
 
-    fn log_statements(mut self, level: LevelFilter) -> Self {
-        self.log_settings.log_statements(level);
-        self
+    if options.set_names {
+        // As it turns out, we don't _have_ to set a collation if we don't want to.
+        // We can let the server choose the default collation for the charset.
+        let set_names = if let Some(collation) = &options.collation {
+            format!(r#"NAMES {} COLLATE {collation}"#, options.charset,)
+        } else {
+            // Leaves the default collation up to the server,
+            // but ensures statements and results are encoded using the proper charset.
+            format!("NAMES {}", options.charset)
+        };
+
+        session_options.push(set_names);
     }
 
-    fn log_slow_statements(mut self, level: LevelFilter, duration: Duration) -> Self {
-        self.log_settings.log_slow_statements(level, duration);
-        self
+    if !session_options.is_empty() {
+        conn.execute(AssertSqlSafe(format!(
+            r#"SET {};"#,
+            session_options.join(",")
+        )))
+        .await?;
     }
+
+    Ok(())
 }

sqlx-mysql/src/options/mod.rs

@@ -5,6 +5,7 @@ mod parse;
 mod ssl_mode;
 
 use crate::{connection::LogSettings, net::tls::CertificateInput};
+pub(crate) use connect::apply_connect_options;
 pub use ssl_mode::MySqlSslMode;
 
 /// Options and flags which can be used to configure a MySQL connection.

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,22 @@ 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(stream, options).await
+    }
+
+    pub(crate) async fn establish_with_socket<S: Socket>(
+        socket: S,
+        options: &PgConnectOptions,
+    ) -> Result<Self, Error> {
+        let stream = PgStream::connect_socket(socket, options).await?;
+        Self::establish_with_stream(stream, options).await
+    }
 
+    async fn establish_with_stream(
+        mut stream: PgStream,
+        options: &PgConnectOptions,
+    ) -> Result<Self, Error> {
         // To begin a session, a frontend opens a connection to the server
         // and sends a startup message.
 

sqlx-postgres/src/connection/mod.rs

@@ -85,6 +85,36 @@ pub(crate) struct TableColumns {
 }
 
 impl PgConnection {
+    /// Connect to a PostgreSQL database using a pre-connected socket.
+    ///
+    /// This allows using custom transport layers such as vsock, QUIC,
+    /// or any type that implements [`sqlx_core::net::Socket`].
+    ///
+    /// The provided socket will go through TLS upgrade negotiation based on the
+    /// SSL mode configured in `options`.
+    ///
+    /// # Example
+    ///
+    /// ```rust,ignore
+    /// use sqlx::postgres::{PgConnectOptions, PgConnection};
+    ///
+    /// # async fn example() -> sqlx::Result<()> {
+    /// let socket: tokio::net::TcpStream = todo!();
+    /// let options = PgConnectOptions::new()
+    ///     .username("postgres")
+    ///     .database("mydb");
+    ///
+    /// let _conn = PgConnection::connect_socket(socket, &options).await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub async fn connect_socket<S: sqlx_core::net::Socket>(
+        socket: S,
+        options: &PgConnectOptions,
+    ) -> Result<Self, Error> {
+        Self::establish_with_socket(socket, options).await
+    }
+
     /// the version number of the server in `libpq` format
     pub fn server_version_num(&self) -> Option<u32> {
         self.inner.stream.server_version_num

sqlx-postgres/src/connection/stream.rs

@@ -57,6 +57,22 @@ impl PgStream {
         })
     }
 
+    pub(super) async fn connect_socket<S: Socket>(
+        socket: S,
+        options: &PgConnectOptions,
+    ) -> Result<Self, Error> {
+        let socket = net::connect_socket(socket, MaybeUpgradeTls(options)).await?;
+
+        let socket = socket?;
+
+        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))

src/lib.rs

@@ -153,6 +153,12 @@ pub mod decode {
 
 pub use self::decode::Decode;
 
+/// Networking traits for custom transport implementations.
+pub mod net {
+    pub use sqlx_core::io::ReadBuf;
+    pub use sqlx_core::net::Socket;
+}
+
 /// Types and traits for the `query` family of functions and macros.
 pub mod query {
     pub use sqlx_core::query::{Map, Query};