feat: add custom socket transport support for Postgres and MySQL

Add methods to PgConnection and MySqlConnection that accept pre-connected
sockets implementing AsyncRead + AsyncWrite, enabling custom transport
layers (vsock, QUIC, turmoil, SSH tunnels, etc.) without forking sqlx.

Per maintainer feedback on transact-rs#4187, this uses AsyncRead + AsyncWrite traits
instead of exposing the internal Socket trait. Two separate methods are
provided for each runtime's trait set:

- connect_with_custom_tokio(): accepts tokio::io::{AsyncRead, AsyncWrite}
- connect_with_custom_futures(): accepts futures_io::{AsyncRead, AsyncWrite}

Also adds PoolOptions::connector() so pools can use custom transports:

    PgPoolOptions::new()
        .connector(|options| async move {
            let socket = VsockStream::connect(addr).await?;
            PgConnection::connect_with_custom_tokio(socket, &options).await
        })
        .connect_with(options)
        .await?

TLS upgrade is negotiated automatically based on the connection options.
No new public trait exposure. No behavioral changes to existing code.

Author: Benjscho

Cargo.toml

@@ -94,10 +94,10 @@ _unstable-docs = [
 ]
 
 # Base runtime features without TLS
-runtime-async-global-executor = ["_rt-async-global-executor", "sqlx-core/_rt-async-global-executor", "sqlx-macros?/_rt-async-global-executor"]
-runtime-async-std = ["_rt-async-std", "sqlx-core/_rt-async-std", "sqlx-macros?/_rt-async-std"]
-runtime-smol = ["_rt-smol", "sqlx-core/_rt-smol", "sqlx-macros?/_rt-smol"]
-runtime-tokio = ["_rt-tokio", "sqlx-core/_rt-tokio", "sqlx-macros?/_rt-tokio"]
+runtime-async-global-executor = ["_rt-async-global-executor", "sqlx-core/_rt-async-global-executor", "sqlx-macros?/_rt-async-global-executor", "sqlx-postgres?/_rt-async-io", "sqlx-mysql?/_rt-async-io"]
+runtime-async-std = ["_rt-async-std", "sqlx-core/_rt-async-std", "sqlx-macros?/_rt-async-std", "sqlx-postgres?/_rt-async-io", "sqlx-mysql?/_rt-async-io"]
+runtime-smol = ["_rt-smol", "sqlx-core/_rt-smol", "sqlx-macros?/_rt-smol", "sqlx-postgres?/_rt-async-io", "sqlx-mysql?/_rt-async-io"]
+runtime-tokio = ["_rt-tokio", "sqlx-core/_rt-tokio", "sqlx-macros?/_rt-tokio", "sqlx-postgres?/_rt-tokio", "sqlx-mysql?/_rt-tokio"]
 
 # TLS features
 tls-native-tls = ["sqlx-core/_tls-native-tls", "sqlx-macros?/_tls-native-tls"]

sqlx-core/src/net/mod.rs

@@ -4,3 +4,9 @@ pub mod tls;
 pub use socket::{
     connect_tcp, connect_uds, BufferedSocket, Socket, SocketIntoBox, WithSocket, WriteBuffer,
 };
+
+#[cfg(feature = "_rt-tokio")]
+pub use socket::async_rw_adapter::TokioStream;
+
+#[cfg(feature = "_rt-async-io")]
+pub use socket::async_rw_adapter::FuturesStream;

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

@@ -0,0 +1,346 @@
+use std::io;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use bytes::BufMut;
+
+use crate::io::ReadBuf;
+use crate::net::Socket;
+
+// Internal buffer size for the read-ahead used by `poll_read_ready`.
+const ADAPTER_BUF_SIZE: usize = 8192;
+
+/// Adapter that wraps a tokio [`AsyncRead`][tokio::io::AsyncRead] + [`AsyncWrite`][tokio::io::AsyncWrite]
+/// into a [`Socket`] implementation.
+#[cfg(feature = "_rt-tokio")]
+pub struct TokioStream<S> {
+    inner: S,
+    read_buf: Vec<u8>,
+    read_len: usize,
+    read_pos: usize,
+}
+
+#[cfg(feature = "_rt-tokio")]
+impl<S> TokioStream<S> {
+    pub fn new(inner: S) -> Self {
+        Self {
+            inner,
+            read_buf: vec![0u8; ADAPTER_BUF_SIZE],
+            read_len: 0,
+            read_pos: 0,
+        }
+    }
+
+    fn buffered(&self) -> &[u8] {
+        &self.read_buf[self.read_pos..self.read_len]
+    }
+}
+
+#[cfg(feature = "_rt-tokio")]
+impl<S> Socket for TokioStream<S>
+where
+    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Sync + Unpin + 'static,
+{
+    fn try_read(&mut self, buf: &mut dyn ReadBuf) -> io::Result<usize> {
+        let buffered = self.buffered();
+        if !buffered.is_empty() {
+            let to_copy = buffered.len().min(buf.remaining_mut());
+            buf.put_slice(&buffered[..to_copy]);
+            self.read_pos += to_copy;
+            if self.read_pos == self.read_len {
+                self.read_pos = 0;
+                self.read_len = 0;
+            }
+            return Ok(to_copy);
+        }
+        Err(io::Error::from(io::ErrorKind::WouldBlock))
+    }
+
+    fn try_write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        // Use a noop waker to attempt a non-blocking write.
+        let waker = futures_util::task::noop_waker();
+        let mut cx = Context::from_waker(&waker);
+        let pin = Pin::new(&mut self.inner);
+        match pin.poll_write(&mut cx, buf) {
+            Poll::Ready(result) => result,
+            Poll::Pending => Err(io::Error::from(io::ErrorKind::WouldBlock)),
+        }
+    }
+
+    fn poll_read_ready(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        // If we already have buffered data, we're ready.
+        if !self.buffered().is_empty() {
+            return Poll::Ready(Ok(()));
+        }
+
+        // Reset buffer positions.
+        self.read_pos = 0;
+        self.read_len = 0;
+
+        let mut read_buf = tokio::io::ReadBuf::new(&mut self.read_buf);
+        let pin = Pin::new(&mut self.inner);
+        match pin.poll_read(cx, &mut read_buf) {
+            Poll::Ready(Ok(())) => {
+                let n = read_buf.filled().len();
+                if n == 0 {
+                    return Poll::Ready(Err(io::Error::from(io::ErrorKind::UnexpectedEof)));
+                }
+                self.read_len = n;
+                Poll::Ready(Ok(()))
+            }
+            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+
+    fn poll_write_ready(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        // For tokio AsyncWrite, we just attempt a zero-length write to check readiness.
+        // Most implementations will return Ready(Ok(0)) if writable.
+        // A more robust approach: just return Ready since we rely on try_write to handle Pending.
+        let pin = Pin::new(&mut self.inner);
+        match pin.poll_write(cx, &[]) {
+            Poll::Ready(Ok(_)) => Poll::Ready(Ok(())),
+            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+
+    fn poll_flush(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.inner).poll_flush(cx)
+    }
+
+    fn poll_shutdown(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.inner).poll_shutdown(cx)
+    }
+}
+
+/// Adapter that wraps a futures-io [`AsyncRead`][futures_io::AsyncRead] + [`AsyncWrite`][futures_io::AsyncWrite]
+/// into a [`Socket`] implementation.
+#[cfg(feature = "_rt-async-io")]
+pub struct FuturesStream<S> {
+    inner: S,
+    read_buf: Vec<u8>,
+    read_len: usize,
+    read_pos: usize,
+}
+
+#[cfg(feature = "_rt-async-io")]
+impl<S> FuturesStream<S> {
+    pub fn new(inner: S) -> Self {
+        Self {
+            inner,
+            read_buf: vec![0u8; ADAPTER_BUF_SIZE],
+            read_len: 0,
+            read_pos: 0,
+        }
+    }
+
+    fn buffered(&self) -> &[u8] {
+        &self.read_buf[self.read_pos..self.read_len]
+    }
+}
+
+#[cfg(feature = "_rt-async-io")]
+impl<S> Socket for FuturesStream<S>
+where
+    S: futures_io::AsyncRead + futures_io::AsyncWrite + Send + Sync + Unpin + 'static,
+{
+    fn try_read(&mut self, buf: &mut dyn ReadBuf) -> io::Result<usize> {
+        let buffered = self.buffered();
+        if !buffered.is_empty() {
+            let to_copy = buffered.len().min(buf.remaining_mut());
+            buf.put_slice(&buffered[..to_copy]);
+            self.read_pos += to_copy;
+            if self.read_pos == self.read_len {
+                self.read_pos = 0;
+                self.read_len = 0;
+            }
+            return Ok(to_copy);
+        }
+        Err(io::Error::from(io::ErrorKind::WouldBlock))
+    }
+
+    fn try_write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        let waker = futures_util::task::noop_waker();
+        let mut cx = Context::from_waker(&waker);
+        let pin = Pin::new(&mut self.inner);
+        match pin.poll_write(&mut cx, buf) {
+            Poll::Ready(result) => result,
+            Poll::Pending => Err(io::Error::from(io::ErrorKind::WouldBlock)),
+        }
+    }
+
+    fn poll_read_ready(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        if !self.buffered().is_empty() {
+            return Poll::Ready(Ok(()));
+        }
+
+        self.read_pos = 0;
+        self.read_len = 0;
+
+        let pin = Pin::new(&mut self.inner);
+        match pin.poll_read(cx, &mut self.read_buf) {
+            Poll::Ready(Ok(n)) => {
+                if n == 0 {
+                    return Poll::Ready(Err(io::Error::from(io::ErrorKind::UnexpectedEof)));
+                }
+                self.read_len = n;
+                Poll::Ready(Ok(()))
+            }
+            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+
+    fn poll_write_ready(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        let pin = Pin::new(&mut self.inner);
+        match pin.poll_write(cx, &[]) {
+            Poll::Ready(Ok(_)) => Poll::Ready(Ok(())),
+            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+
+    fn poll_flush(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.inner).poll_flush(cx)
+    }
+
+    fn poll_shutdown(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        Pin::new(&mut self.inner).poll_close(cx)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::io::Cursor;
+
+    // Cursor<Vec<u8>> implements both tokio and futures-io AsyncRead/AsyncWrite
+
+    #[cfg(feature = "_rt-tokio")]
+    mod tokio_adapter {
+        use super::*;
+        use crate::net::Socket;
+        use bytes::BytesMut;
+        use std::task::{Context, Poll};
+
+        fn noop_cx() -> Context<'static> {
+            Context::from_waker(futures_util::task::noop_waker_ref())
+        }
+
+        #[test]
+        fn try_read_returns_would_block_when_empty() {
+            let stream = tokio::io::duplex(64).0;
+            let mut adapter = TokioStream::new(stream);
+            let mut buf = BytesMut::with_capacity(32);
+            let err = adapter.try_read(&mut buf).unwrap_err();
+            assert_eq!(err.kind(), io::ErrorKind::WouldBlock);
+        }
+
+        #[test]
+        fn poll_read_ready_fills_buffer_then_try_read_drains() {
+            // Use a duplex stream with pre-written data
+            let (client, mut server) = tokio::io::duplex(64);
+
+            // Write data into the server side so client can read it
+            let rt = tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()
+                .unwrap();
+
+            rt.block_on(async {
+                use tokio::io::AsyncWriteExt;
+                server.write_all(b"hello world").await.unwrap();
+
+                let mut adapter = TokioStream::new(client);
+                let mut buf = BytesMut::with_capacity(32);
+
+                // poll_read_ready should fill internal buffer
+                let poll =
+                    std::future::poll_fn(|cx| adapter.poll_read_ready(cx)).await;
+                assert!(poll.is_ok());
+
+                // try_read should drain from internal buffer
+                let n = adapter.try_read(&mut buf).unwrap();
+                assert_eq!(&buf[..n], b"hello world");
+
+                // After draining, try_read should return WouldBlock again
+                let mut buf2 = BytesMut::with_capacity(32);
+                let err = adapter.try_read(&mut buf2).unwrap_err();
+                assert_eq!(err.kind(), io::ErrorKind::WouldBlock);
+            });
+        }
+
+        #[test]
+        fn try_write_writes_data() {
+            let (client, mut server) = tokio::io::duplex(64);
+
+            let rt = tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()
+                .unwrap();
+
+            rt.block_on(async {
+                use tokio::io::AsyncReadExt;
+                let mut adapter = TokioStream::new(client);
+
+                // Write via the adapter
+                let n = std::future::poll_fn(|cx| {
+                    match adapter.try_write(b"test data") {
+                        Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
+                            match adapter.poll_write_ready(cx) {
+                                Poll::Ready(Ok(())) => {
+                                    Poll::Ready(adapter.try_write(b"test data"))
+                                }
+                                Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
+                                Poll::Pending => Poll::Pending,
+                            }
+                        }
+                        other => Poll::Ready(other),
+                    }
+                })
+                .await
+                .unwrap();
+
+                assert_eq!(n, 9);
+
+                // Read it back from the other end
+                let mut read_buf = vec![0u8; 32];
+                let n = server.read(&mut read_buf).await.unwrap();
+                assert_eq!(&read_buf[..n], b"test data");
+            });
+        }
+
+        #[test]
+        fn partial_drain_preserves_remaining() {
+            let (client, mut server) = tokio::io::duplex(64);
+
+            let rt = tokio::runtime::Builder::new_current_thread()
+                .enable_all()
+                .build()
+                .unwrap();
+
+            rt.block_on(async {
+                use tokio::io::AsyncWriteExt;
+                server.write_all(b"abcdefghij").await.unwrap();
+
+                let mut adapter = TokioStream::new(client);
+
+                // Fill internal buffer
+                std::future::poll_fn(|cx| adapter.poll_read_ready(cx)).await.unwrap();
+
+                // Read only 4 bytes using a fixed-size slice
+                let mut buf = [0u8; 4];
+                let n = adapter.try_read(&mut buf.as_mut_slice()).unwrap();
+                assert_eq!(n, 4);
+                assert_eq!(&buf, b"abcd");
+
+                // Remaining 6 bytes should still be available
+                let mut buf2 = [0u8; 32];
+                let n = adapter.try_read(&mut buf2.as_mut_slice()).unwrap();
+                assert_eq!(n, 6);
+                assert_eq!(&buf2[..6], b"efghij");
+            });
+        }
+    }
+}

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

@@ -10,6 +10,8 @@ use cfg_if::cfg_if;
 
 use crate::io::ReadBuf;
 
+#[cfg(any(feature = "_rt-tokio", feature = "_rt-async-io"))]
+pub(crate) mod async_rw_adapter;
 mod buffered;
 
 pub trait Socket: Send + Sync + Unpin + 'static {