fix: tests

Author: abonander

sqlx-core/src/pool/connect.rs

@@ -33,18 +33,19 @@ use std::io;
 /// use sqlx::PgConnection;
 /// use sqlx::postgres::PgPoolOptions;
 /// use sqlx::Connection;
+/// use sqlx::pool::PoolConnectMetadata;
 ///
-/// # async fn _example() -> sqlx::Result<()> {
-/// // `PoolConnector` is implemented for closures but has restrictions on returning borrows
-/// // due to current language limitations.
+/// async fn _example() -> sqlx::Result<()> {
+/// // `PoolConnector` is implemented for closures but this has restrictions on returning borrows
+/// // due to current language limitations. Custom implementations are not subject to this.
 /// //
 /// // This example shows how to get around this using `Arc`.
 /// let database_url: Arc<str> = "postgres://...".into();
 ///
 /// let pool = PgPoolOptions::new()
 ///     .min_connections(5)
 ///     .max_connections(30)
-///     .connect_with_connector(move |meta| {
+///     .connect_with_connector(move |meta: PoolConnectMetadata| {
 ///         let database_url = database_url.clone();
 ///         async move {
 ///             println!(
@@ -57,7 +58,9 @@ use std::io;
 ///             let mut conn = PgConnection::connect(&database_url).await?;
 ///
 ///             // Override the time zone of the connection.
-///             sqlx::raw_sql("SET TIME ZONE 'Europe/Berlin'").await?;
+///             sqlx::raw_sql("SET TIME ZONE 'Europe/Berlin'")
+///                 .execute(&mut conn)
+///                 .await?;
 ///
 ///             Ok(conn)
 ///         }
@@ -76,7 +79,7 @@ use std::io;
 ///
 /// ```rust,no_run
 /// use std::sync::Arc;
-/// use tokio::sync::{Mutex, RwLock};
+/// use tokio::sync::RwLock;
 /// use sqlx::PgConnection;
 /// use sqlx::postgres::PgConnectOptions;
 /// use sqlx::postgres::PgPoolOptions;

tests/any/pool.rs

@@ -5,6 +5,8 @@ use std::sync::{
     Arc, Mutex,
 };
 use std::time::Duration;
+use sqlx_core::connection::ConnectOptions;
+use sqlx_core::pool::PoolConnectMetadata;
 
 #[sqlx_macros::test]
 async fn pool_should_invoke_after_connect() -> anyhow::Result<()> {
@@ -83,38 +85,13 @@ async fn test_pool_callbacks() -> anyhow::Result<()> {
 
     sqlx_test::setup_if_needed();
 
-    let conn_options: AnyConnectOptions = std::env::var("DATABASE_URL")?.parse()?;
+    let conn_options: Arc<AnyConnectOptions> = Arc::new(std::env::var("DATABASE_URL")?.parse()?);
 
     let current_id = AtomicI32::new(0);
 
     let pool = AnyPoolOptions::new()
         .max_connections(1)
         .acquire_timeout(Duration::from_secs(5))
-        .after_connect(move |conn, meta| {
-            assert_eq!(meta.age, Duration::ZERO);
-            assert_eq!(meta.idle_for, Duration::ZERO);
-
-            let id = current_id.fetch_add(1, Ordering::AcqRel);
-
-            Box::pin(async move {
-                let statement = format!(
-                    // language=SQL
-                    r#"
-                    CREATE TEMPORARY TABLE conn_stats(
-                        id int primary key,
-                        before_acquire_calls int default 0,
-                        after_release_calls int default 0
-                    );
-                    INSERT INTO conn_stats(id) VALUES ({});
-                    "#,
-                    // Until we have generalized bind parameters
-                    id
-                );
-
-                conn.execute(&statement[..]).await?;
-                Ok(())
-            })
-        })
         .before_acquire(|conn, meta| {
             // `age` and `idle_for` should both be nonzero
             assert_ne!(meta.age, Duration::ZERO);
@@ -165,7 +142,31 @@ async fn test_pool_callbacks() -> anyhow::Result<()> {
             })
         })
         // Don't establish a connection yet.
-        .connect_lazy_with(conn_options);
+        .connect_lazy_with_connector(move |meta: PoolConnectMetadata| {
+            let connect_opts = Arc::clone(&conn_options);
+            let id = current_id.fetch_add(1, Ordering::AcqRel);
+            
+            async move {
+                let mut conn = connect_opts.connect().await?;
+                
+                let statement = format!(
+                    // language=SQL
+                    r#"
+                    CREATE TEMPORARY TABLE conn_stats(
+                        id int primary key,
+                        before_acquire_calls int default 0,
+                        after_release_calls int default 0
+                    );
+                    INSERT INTO conn_stats(id) VALUES ({});
+                    "#,
+                    // Until we have generalized bind parameters
+                    id
+                );
+
+                conn.execute(&statement[..]).await?;
+                Ok(conn)
+            }
+        });
 
     // Expected pattern of (id, before_acquire_calls, after_release_calls)
     let pattern = [