Skip to content

feat: add custom socket transport support for Postgres and MySQL#4187

Open
jowparks wants to merge 1 commit intolaunchbadge:mainfrom
jowparks:feat/custom-socket-transport
Open

feat: add custom socket transport support for Postgres and MySQL#4187
jowparks wants to merge 1 commit intolaunchbadge:mainfrom
jowparks:feat/custom-socket-transport

Conversation

@jowparks
Copy link

@jowparks jowparks commented Mar 12, 2026

Summary

Add connect_socket() methods to PgConnection and MySqlConnection that accept any pre-connected socket implementing the Socket trait. This enables using custom transport layers without forking sqlx.

Motivation: In environments like AWS Nitro Enclaves, vsock is the only available networking transport. Currently, sqlx hardcodes its connection establishment to TCP and Unix domain sockets, making it impossible to use with non-standard transports without forking. This PR makes the Socket trait usable as a public extension point.

Before / After

Before

There is no way to provide a custom socket to sqlx. The only connection paths are TCP and Unix domain sockets, hardcoded in PgStream::connect() and MySqlConnection::establish():

// The only options: TCP or Unix domain socket via connection URL
let conn = PgConnection::connect("postgres://user:pass@host/db").await?;
let conn = MySqlConnection::connect("mysql://user:pass@host/db").await?;

// For vsock or other transports, you had to:
// 1. Fork sqlx entirely, OR
// 2. Run a TCP-to-vsock proxy (adds latency + operational complexity)

After

Users can connect over any transport by providing a pre-connected socket:

use sqlx::net::Socket;

// 1. Implement Socket for your transport type
impl Socket for tokio_vsock::VsockStream {
    fn try_read(&mut self, buf: &mut dyn ReadBuf) -> io::Result<usize> { /* ... */ }
    fn try_write(&mut self, buf: &[u8]) -> io::Result<usize> { /* ... */ }
    fn poll_read_ready(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> { /* ... */ }
    fn poll_write_ready(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> { /* ... */ }
    fn poll_shutdown(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> { /* ... */ }
}

// 2. Connect with your custom socket — TLS upgrade works automatically
let socket = VsockStream::connect(VsockAddr::new(3, 5432)).await?;
let options = PgConnectOptions::new()
    .username("postgres")
    .database("mydb");
let conn = PgConnection::connect_socket(socket, &options).await?;

// MySQL works the same way
let socket = VsockStream::connect(VsockAddr::new(3, 3306)).await?;
let options = MySqlConnectOptions::new()
    .username("root")
    .database("mydb");
let conn = MySqlConnection::connect_socket(socket, &options).await?;

Changes

New public API

  • PgConnection::connect_socket(socket, &options) — connect to PostgreSQL over any Socket impl
  • MySqlConnection::connect_socket(socket, &options) — connect to MySQL over any Socket impl
  • sqlx::net::Socket — re-exported trait for implementing custom transports
  • sqlx::net::ReadBuf — re-exported trait needed by Socket implementations

Files changed

  • sqlx-core/src/net/socket/mod.rs: Add connect_socket() helper function
  • sqlx-core/src/net/mod.rs: Re-export connect_socket
  • sqlx-postgres/src/connection/stream.rs: Add PgStream::connect_socket()
  • sqlx-postgres/src/connection/establish.rs: Refactor establish() to extract establish_with_stream(), add establish_with_socket()
  • sqlx-postgres/src/connection/mod.rs: Add public PgConnection::connect_socket()
  • sqlx-mysql/src/connection/establish.rs: Refactor to extract establish_with_stream(), add establish_with_socket()
  • sqlx-mysql/src/connection/mod.rs: Add public MySqlConnection::connect_socket()
  • sqlx-mysql/src/options/connect.rs: Extract configure_session() from ConnectOptions::connect() for reuse by connect_socket
  • src/lib.rs: Add sqlx::net module re-exporting Socket and ReadBuf

Design decisions

  • Separate methods, not modified options: PgConnectOptions/MySqlConnectOptions derive Clone + Debug, which prevents storing Box<dyn Socket>. A separate connect_socket() method is cleaner and more ergonomic.
  • TLS works automatically: Custom sockets go through the same MaybeUpgradeTls / DoHandshake flow, so TLS upgrade negotiation works out of the box.
  • Minimal refactoring: The establish logic was factored into establish_with_stream() to share code between the TCP/UDS path and the custom socket path. No behavioral changes to existing code.
  • No new dependencies or feature flags: Custom socket support is always available since it only depends on the existing Socket trait.

Use cases

  • AWS Nitro Enclaves: Connect to databases over vsock (via tokio-vsock)
  • QUIC-based transports: Use QUIC streams as database connections
  • Custom proxied connections: Pre-connected sockets through SSH tunnels, SOCKS proxies, etc.
  • Testing: Inject mock sockets for testing database connection logic

Checklist

  • cargo check --all-features passes
  • No changes to existing tests
  • Doc comments on all new public APIs
  • No new dependencies added
  • No feature flags added

@jowparks jowparks force-pushed the feat/custom-socket-transport branch 3 times, most recently from cbedb56 to 9655878 Compare March 12, 2026 20:24
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.
@jowparks jowparks force-pushed the feat/custom-socket-transport branch from 9655878 to f7dcf2d Compare March 12, 2026 20:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant