feat(postgres, core): PgConnection::connect_with_socket and a wasm-safe Instant

sqlx-core/Cargo.toml

@@ -105,6 +105,9 @@ hashbrown = "0.16.0"
 
 thiserror.workspace = true
 
+[target.'cfg(all(target_arch = "wasm32", target_os = "unknown"))'.dependencies]
+web-time = "1.1.0"
+
 [dev-dependencies]
 tempfile = "3.10.1"
 tokio = { version = "1.25.0", features = ["rt"] }

sqlx-core/src/logger.rs

@@ -1,5 +1,5 @@
+use crate::rt::Instant;
 use crate::{connection::LogSettings, sql_str::SqlStr};
-use std::time::Instant;
 
 // Yes these look silly. `tracing` doesn't currently support dynamic levels
 // https://github.com/tokio-rs/tracing/issues/372

sqlx-core/src/rt/mod.rs

@@ -12,6 +12,16 @@ pub mod rt_async_io;
 #[cfg(feature = "_rt-tokio")]
 pub mod rt_tokio;
 
+/// A monotonic clock reading, for measuring elapsed time.
+///
+/// This is [`std::time::Instant`], except on `wasm32-unknown-unknown`: that target has no clock of
+/// its own, so `std::time::Instant::now()` panics there, and `web-time` reads the host's clock
+/// (`performance.now()`) instead.
+#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
+pub use std::time::Instant;
+#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
+pub use web_time::Instant;
+
 #[derive(Debug, thiserror::Error)]
 #[error("operation timed out")]
 pub struct TimeoutError;

sqlx-mysql/src/migrate.rs

@@ -1,6 +1,6 @@
+use sqlx_core::rt::Instant;
 use std::str::FromStr;
 use std::time::Duration;
-use std::time::Instant;
 
 use futures_core::future::BoxFuture;
 pub(crate) use sqlx_core::migrate::*;

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;
@@ -17,8 +18,32 @@ 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
+    }
+
+    /// Open a connection over `socket`, a transport the caller has already connected to a
+    /// Postgres server, instead of one sqlx opens itself.
+    ///
+    /// This is for environments where sqlx cannot open sockets on its own: no supported async
+    /// runtime, or a host that hands out its own socket type (e.g. `wasm32-unknown-unknown`
+    /// in Cloudflare Workers). The host, port and socket path in `options` are ignored;
+    /// everything else, including `ssl_mode`, applies as it does for
+    /// [`connect_with`][sqlx_core::connection::Connection::connect_with].
+    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(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/stream.rs

@@ -13,7 +13,7 @@ use crate::message::{
     BackendMessage, BackendMessageFormat, EncodeMessage, FrontendMessage, Notice, Notification,
     ParameterStatus, ReceivedMessage,
 };
-use crate::net::{self, BufferedSocket, Socket};
+use crate::net::{self, BufferedSocket, Socket, WithSocket};
 use crate::{PgConnectOptions, PgDatabaseError, PgSeverity};
 
 // the stream is a separate type from the connection to uphold the invariant where an instantiated
@@ -47,14 +47,27 @@ impl PgStream {
             None => net::connect_tcp(&options.host, options.port, MaybeUpgradeTls(options)).await?,
         };
 
-        let socket = socket_result?;
+        Ok(Self::from_socket(socket_result?))
+    }
+
+    // Like `connect`, but over a socket the caller already opened; TLS is still negotiated
+    // according to `options`.
+    pub(super) async fn connect_with_socket<S: Socket>(
+        options: &PgConnectOptions,
+        socket: S,
+    ) -> Result<Self, Error> {
+        let socket = MaybeUpgradeTls(options).with_socket(socket).await?;
 
-        Ok(Self {
+        Ok(Self::from_socket(socket))
+    }
+
+    fn from_socket(socket: Box<dyn Socket>) -> Self {
+        Self {
             inner: BufferedSocket::new(socket),
             notifications: None,
             parameter_statuses: BTreeMap::default(),
             server_version_num: None,
-        })
+        }
     }
 
     #[inline(always)]

sqlx-postgres/src/migrate.rs

@@ -1,6 +1,6 @@
+use sqlx_core::rt::Instant;
 use std::str::FromStr;
 use std::time::Duration;
-use std::time::Instant;
 
 use futures_core::future::BoxFuture;
 

sqlx-sqlite/src/migrate.rs

@@ -9,11 +9,11 @@ use crate::query::query;
 use crate::query_as::query_as;
 use crate::{Sqlite, SqliteConnectOptions, SqliteConnection, SqliteJournalMode};
 use futures_core::future::BoxFuture;
+use sqlx_core::rt::Instant;
 use sqlx_core::sql_str::AssertSqlSafe;
 use std::str::FromStr;
 use std::sync::atomic::Ordering;
 use std::time::Duration;
-use std::time::Instant;
 
 pub(crate) use sqlx_core::migrate::*;
 use sqlx_core::query_scalar::query_scalar;

tests/postgres/postgres.rs

@@ -28,6 +28,26 @@ async fn it_connects() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[cfg(feature = "_rt-tokio")]
+#[sqlx_macros::test]
+async fn it_connects_with_socket() -> anyhow::Result<()> {
+    let options: PgConnectOptions = env::var("DATABASE_URL")?.parse()?;
+    let socket = tokio::net::TcpStream::connect((options.get_host(), options.get_port())).await?;
+
+    let mut conn = PgConnection::connect_with_socket(&options, socket).await?;
+
+    let value = sqlx::query("select 1 + 1")
+        .try_map(|row: PgRow| row.try_get::<i32, _>(0))
+        .fetch_one(&mut conn)
+        .await?;
+
+    assert_eq!(2i32, value);
+
+    conn.close().await?;
+
+    Ok(())
+}
+
 #[sqlx_macros::test]
 async fn it_can_select_void() -> anyhow::Result<()> {
     let mut conn = new::<Postgres>().await?;