WIP refactor: new #[sqlx::test] architecture

Author: abonander

sqlx-core/Cargo.toml

@@ -21,11 +21,11 @@ json = ["serde", "serde_json"]
 
 # for conditional compilation
 _rt-async-global-executor = ["async-global-executor", "_rt-async-io", "_rt-async-task"]
-_rt-async-io = ["async-io", "async-fs"] # see note at async-fs declaration
+_rt-async-io = ["async-io", "async-fs", "async-lock"] # see note at async-fs declaration
 _rt-async-std = ["async-std", "_rt-async-io"]
 _rt-async-task = ["async-task"]
 _rt-smol = ["smol", "_rt-async-io", "_rt-async-task"]
-_rt-tokio = ["tokio", "tokio-stream"]
+_rt-tokio = ["tokio", "tokio-stream", "tokio/rt"] # `rt` feature is almost always going to be enabled anyway
 
 _tls-native-tls = ["native-tls"]
 _tls-rustls-aws-lc-rs = ["_tls-rustls", "rustls/aws-lc-rs", "webpki-roots"]
@@ -72,6 +72,7 @@ uuid = { workspace = true, optional = true }
 
 # work around bug in async-fs 2.0.0, which references futures-lite dependency wrongly, see https://github.com/launchbadge/sqlx/pull/3791#issuecomment-3043363281
 async-fs = { version = "2.1", optional = true }
+async-lock = { version = "3.4.2", optional = true }
 async-io = { version = "2.4.1", optional = true }
 async-task = { version = "4.7.1", optional = true }
 
@@ -83,7 +84,6 @@ crossbeam-queue = "0.3.2"
 either = "1.6.1"
 futures-core = { version = "0.3.32", default-features = false }
 futures-io = "0.3.32"
-futures-intrusive = "0.5.0"
 futures-util = { version = "0.3.32", default-features = false, features = ["alloc", "sink", "io"] }
 log = { version = "0.4.18", default-features = false }
 memchr = { version = "2.5.0", default-features = false }
@@ -103,10 +103,9 @@ indexmap = "2.0"
 event-listener = "5.2.0"
 hashbrown = "0.16.0"
 
-thiserror.workspace = true
+futures-intrusive = "0.5.0"
 
-[dev-dependencies]
-tokio = { version = "1.25.0", features = ["rt"] }
+thiserror.workspace = true
 
 [dev-dependencies.sqlx]
 # FIXME: https://github.com/rust-lang/cargo/issues/15622

sqlx-core/src/sync.rs

@@ -1,4 +1,5 @@
 use cfg_if::cfg_if;
+use std::future::Future;
 
 // For types with identical signatures that don't require runtime support,
 // we can just arbitrarily pick one to use based on what's enabled.
@@ -204,3 +205,60 @@ impl AsyncSemaphoreReleaser<'_> {
         }
     }
 }
+
+pub struct AsyncOnceCell<T> {
+    #[cfg(feature = "_rt-tokio")]
+    inner: tokio::sync::OnceCell<T>,
+
+    #[cfg(all(feature = "_rt-async-io", not(feature = "_rt-tokio")))]
+    inner: async_lock::OnceCell<T>,
+
+    #[cfg(not(any(feature = "_rt-async-std", feature = "_rt-tokio")))]
+    phantom: std::marker::PhantomData<T>,
+}
+
+impl<T> AsyncOnceCell<T> {
+    pub fn new() -> Self {
+        cfg_if! {
+            if #[cfg(feature = "_rt-tokio")] {
+                Self { inner: tokio::sync::OnceCell::new() }
+            } else if #[cfg(feature = "_rt-async-io")] {
+                Self { inner: async_lock::OnceCell::new() }
+            } else {
+                crate::rt::missing_rt(());
+            }
+        }
+    }
+
+    pub const fn const_new() -> Self {
+        cfg_if! {
+            if #[cfg(feature = "_rt-tokio")] {
+                Self { inner: tokio::sync::OnceCell::const_new() }
+            } else if #[cfg(feature = "_rt-async-io")] {
+                Self { inner: async_lock::OnceCell::new() }
+            } else {
+                crate::rt::missing_rt(());
+            }
+        }
+    }
+
+    pub async fn get_or_try_init<F, Fut, E>(&self, f: F) -> Result<&T, E>
+    where
+        F: FnOnce() -> Fut,
+        Fut: Future<Output = Result<T, E>>,
+    {
+        cfg_if! {
+            if #[cfg(any(feature = "_rt-tokio", feature = "_rt-async-io"))] {
+                self.inner.get_or_try_init(f).await
+            } else {
+                crate::rt::missing_rt(f)
+            }
+        }
+    }
+}
+
+impl<T> Default for AsyncOnceCell<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}

sqlx-core/src/testing/mod.rs

@@ -13,6 +13,7 @@ use crate::migrate::{Migrate, Migrator};
 use crate::pool::{Pool, PoolConnection, PoolOptions};
 
 mod fixtures;
+mod pool;
 
 pub trait TestSupport: Database {
     /// Get parameters to construct a `Pool` suitable for testing.
@@ -66,6 +67,7 @@ pub struct TestArgs {
     pub test_path: &'static str,
     pub migrator: Option<&'static Migrator>,
     pub fixtures: &'static [TestFixture],
+    pub max_connections: usize,
 }
 
 pub trait TestFn {
@@ -158,6 +160,7 @@ impl TestArgs {
             test_path,
             migrator: None,
             fixtures: &[],
+            max_connections: 5,
         }
     }
 
@@ -168,6 +171,10 @@ impl TestArgs {
     pub fn fixtures(&mut self, fixtures: &'static [TestFixture]) {
         self.fixtures = fixtures;
     }
+
+    pub fn max_connections(&mut self, max_connections: usize) {
+        self.max_connections = max_connections;
+    }
 }
 
 impl TestTermination for () {

sqlx-core/src/testing/pool.rs

@@ -0,0 +1,20 @@
+use std::rc::Weak;
+use crate::database::Database;
+use crate::pool::Pool;
+use crate::sync::AsyncOnceCell;
+
+pub struct TestMasterPool<DB: Database> {
+    inner: AsyncOnceCell<Inner<DB>>,
+}
+
+struct Inner<DB: Database> {
+    pool: Pool<DB>,
+
+    #[cfg(feature = "_rt-tokio")]
+
+}
+
+
+impl<DB: Database> TestMasterPool<DB> {
+
+}

sqlx-postgres/src/testing/mod.rs

@@ -146,7 +146,51 @@ async fn test_context(args: &TestArgs) -> Result<TestContext<Postgres>, Error> {
         create index if not exists databases_created_at 
             on _sqlx_test.databases(created_at);
 
-        create sequence if not exists _sqlx_test.database_ids;
+        create table if not exists _sqlx_test.tests (
+            test_id int8 primary key generated always as identity,
+            -- Automatically cleans up leaked test runs as well
+            db_name text not null references _sqlx_test.databases(db_name) on delete cascade,
+            required_connections int4 not null
+                check (required_connections > 0 and required_connections <= max_connections),
+            -- Each test's `SQLX_TEST_MAX_CONNECTIONS`, ideally all the same
+            max_connections int4 not null check (max_connections > 0),
+            started_at timestamptz not null default now()
+        );
+
+        create or replace function _sqlx_test.tests_check_max_connections()
+            returns trigger as
+        $$
+            declare
+                    used_connections int4;
+                    max_required_connections int4;
+                    max_connections int4;
+            begin
+                select
+                    sum(required_connections),
+                    max(required_connections),
+                    -- Abide by the highest `SQLX_TEST_MAX_CONNECTIONS`
+                    max(max_connections)
+                into
+                    used_connections,
+                    max_required_connections,
+                    max_connections
+                from _sqlx_test.tests;
+
+                if max_required_connections > max_connections then
+                    raise
+                        'max(required_connections) exceeds min(max_connections) of any process'
+                    using constraint = 'required_connections_exceeds_max';
+                elsif max_connections > max_connections then
+                    raise 'not enough spare connections available; used: %i, total: %i',
+                        used_connections, max_connections
+                    using constraint = 'insufficient_connections_available';
+                end if;
+            end;
+        $$
+        language plpgsql;
+
+        create or replace constraint trigger check_max_connections after insert on _sqlx_test.tests
+        for each statement execute function _sqlx_test.check_max_connections();
     "#,
     )
     .await?;