db-pool

We will use the Diesel async Postgres backend with BB8 in this tutorial.

The database pool has to live for the duration of the test run. We use a OnceCell to store a "lazy static".

fn main() {}

#[cfg(test)]
mod tests {
    #![allow(dead_code)]

    // import OnceCell
    use tokio::sync::OnceCell;

    async fn get_connection_pool() {
        // add "lazy static"
        static POOL: OnceCell<()> = OnceCell::const_new();
    }
}

We then need to create a privileged database configuration loaded from environment variables.

We create a .env file.

export POSTGRES_USERNAME=postgres
export POSTGRES_PASSWORD=postgres
export POSTGRES_HOST=localhost
export POSTGRES_PORT=3306

The environment variables used are optional.

We load the environment variables from the .env file and create a privileged configuration.

fn main() {}

#[cfg(test)]
mod tests {
    #![allow(dead_code, unused_variables)]

    // import privileged configuration
    use db_pool::PrivilegedPostgresConfig;
    // import dotenvy
    use dotenvy::dotenv;
    use tokio::sync::OnceCell;

    async fn get_connection_pool() {
        static POOL: OnceCell<()> = OnceCell::const_new();

        let db_pool = POOL
            .get_or_init(|| async {
                // load environment variables from .env
                dotenv().ok();

                // create privileged configuration from environment variables
                let config = PrivilegedPostgresConfig::from_env().unwrap();
            })
            .await;
    }
}

PrivilegedPostgresConfig is the database connection configuration that includes username, password, host, and port. It is used for administration of the created databases - creation, cleaning, and dropping.

We then create the backend with the privileged configuration, among others.

fn main() {}

#[cfg(test)]
mod tests {
    #![allow(dead_code, unused_variables)]

    // import connection pool
    use bb8::Pool;
    use db_pool::{
        // import backend
        r#async::{DieselAsyncPostgresBackend, DieselBb8},
        PrivilegedPostgresConfig,
    };
    // import diesel-specific constructs
    use diesel::sql_query;
    use diesel_async::RunQueryDsl;
    use dotenvy::dotenv;
    use tokio::sync::OnceCell;

    async fn get_connection_pool() {
        static POOL: OnceCell<()> = OnceCell::const_new();
        let db_pool = POOL
            .get_or_init(|| async {
                dotenv().ok();

                let config = PrivilegedPostgresConfig::from_env().unwrap();

                // create backend for BB8 connection pools
                let backend = DieselAsyncPostgresBackend::<DieselBb8>::new(
                    config,
                    // create privileged connection pool with max 10 connections
                    || Pool::builder().max_size(10),
                    // create restricted connection pool with max 2 connections
                    || Pool::builder().max_size(2),
                    // no custom create connection
                    None,
                    // create entities
                    move |mut conn| {
                        Box::pin(async {
                            sql_query(
                                "CREATE TABLE book(id SERIAL PRIMARY KEY, title TEXT NOT NULL)",
                            )
                            .execute(&mut conn)
                            .await
                            .unwrap();

                            conn
                        })
                    },
                )
                .await
                .unwrap();
            })
            .await;
    }
}

DieselPostgresBackend is the backend that will communicate with the PostgreSQL instance using Diesel-specific constructs.

DieselBb8 is a construct that allows the backend to work with and issue BB8 connection pools.

PrivilegedPostgresConfig::from_env().unwrap() creates a privileged Postgres configuration from environment variables.

|| Pool::builder().max_size(10) creates a privileged connection pool with a max of 10 connections. This pool is used for administration of the created databases and relies on the privileged configuration to establish connections.

|| Pool::builder().max_size(2) creates a restricted connection pool meant to be used by a test. A restricted connection pool is created every time the database pool runs out of available connection pools to lend.

The last closure creates the database entities required to be available for tests. This happens every time a new restricted connection pool needs to be created - a new database is created, its entities are created, and the restricted connection pool is bound to it.

We create a database pool from the backend.

fn main() {}

#[cfg(test)]
mod tests {
    #![allow(dead_code, unused_variables)]

    use bb8::Pool;
    use db_pool::{
        r#async::{
            // import database pool
            DatabasePool,
            // import database pool builder trait
            DatabasePoolBuilderTrait,
            DieselAsyncPostgresBackend,
            DieselBb8,
        },
        PrivilegedPostgresConfig,
    };
    use diesel::sql_query;
    use diesel_async::RunQueryDsl;
    use dotenvy::dotenv;
    use tokio::sync::OnceCell;

    async fn get_connection_pool() {
        // change OnceCell inner type
        static POOL: OnceCell<DatabasePool<DieselAsyncPostgresBackend<DieselBb8>>> =
            OnceCell::const_new();

        let db_pool = POOL.get_or_init(|| async {
            dotenv().ok();

            let config = PrivilegedPostgresConfig::from_env().unwrap();

            // Diesel pool association type can be inferred now
            let backend = DieselAsyncPostgresBackend::new(
                config,
                || Pool::builder().max_size(10),
                || Pool::builder().max_size(2),
                None,
                move |mut conn| {
                    Box::pin(async {
                        sql_query("CREATE TABLE book(id SERIAL PRIMARY KEY, title TEXT NOT NULL)")
                            .execute(&mut conn)
                            .await
                            .unwrap();

                        conn
                    })
                },
            )
            .await
            .unwrap();

            // create database pool
            backend.create_database_pool().await.unwrap()
        });
    }
}

DatabasePool is the returned pool of connection pools that will be assigned to tests in isolation. The connection pools can be reused after a test has finished and no longer needs the connection pool assigned to it.

We pull a connection pool out of the database pool.

fn main() {}

#[cfg(test)]
mod tests {
    #![allow(dead_code)]

    use bb8::Pool;
    use db_pool::{
        r#async::{
            DatabasePool,
            DatabasePoolBuilderTrait,
            DieselAsyncPostgresBackend,
            DieselBb8,
            // import reusable connection pool
            ReusableConnectionPool,
        },
        PrivilegedPostgresConfig,
    };
    use diesel::sql_query;
    use diesel_async::RunQueryDsl;
    use dotenvy::dotenv;
    use tokio::sync::OnceCell;

    // change return type
    async fn get_connection_pool(
    ) -> ReusableConnectionPool<'static, DieselAsyncPostgresBackend<DieselBb8>> {
        static POOL: OnceCell<DatabasePool<DieselAsyncPostgresBackend<DieselBb8>>> =
            OnceCell::const_new();

        let db_pool = POOL
            .get_or_init(|| async {
                dotenv().ok();

                let config = PrivilegedPostgresConfig::from_env().unwrap();

                let backend = DieselAsyncPostgresBackend::new(
                    config,
                    || Pool::builder().max_size(10),
                    || Pool::builder().max_size(2),
                    None,
                    move |mut conn| {
                        Box::pin(async {
                            sql_query(
                                "CREATE TABLE book(id SERIAL PRIMARY KEY, title TEXT NOT NULL)",
                            )
                            .execute(&mut conn)
                            .await
                            .unwrap();

                            conn
                        })
                    },
                )
                .await
                .unwrap();

                backend.create_database_pool().await.unwrap()
            })
            .await;

        // pull connection pool
        db_pool.pull_immutable().await
    }
}

ConnectionPool is a connection pool assigned to a test.

Reusable is a wrapper around ConnectionPool that allows access to the connection pool by the assigned test and orchestrates its return to the database pool after a test has finished and no longer needs it.

We add a simple test case that inserts a row then counts the number of rows in a table.

fn main() {}

#[cfg(test)]
mod tests {
    #![allow(dead_code)]

    use bb8::Pool;
    use db_pool::{
        r#async::{
            DatabasePool, DatabasePoolBuilderTrait, DieselAsyncPostgresBackend, DieselBb8,
            ReusableConnectionPool,
        },
        PrivilegedPostgresConfig,
    };
    // import extra diesel-specific constructs
    use diesel::{insert_into, sql_query, table, Insertable, QueryDsl};
    use diesel_async::RunQueryDsl;
    use dotenvy::dotenv;
    use tokio::sync::OnceCell;

    async fn get_connection_pool(
    ) -> ReusableConnectionPool<'static, DieselAsyncPostgresBackend<DieselBb8>> {
        static POOL: OnceCell<DatabasePool<DieselAsyncPostgresBackend<DieselBb8>>> =
            OnceCell::const_new();

        let db_pool = POOL
            .get_or_init(|| async {
                dotenv().ok();

                let config = PrivilegedPostgresConfig::from_env().unwrap();

                let backend = DieselAsyncPostgresBackend::new(
                    config,
                    || Pool::builder().max_size(10),
                    || Pool::builder().max_size(2),
                    None,
                    move |mut conn| {
                        Box::pin(async {
                            sql_query(
                                "CREATE TABLE book(id SERIAL PRIMARY KEY, title TEXT NOT NULL)",
                            )
                            .execute(&mut conn)
                            .await
                            .unwrap();

                            conn
                        })
                    },
                )
                .await
                .unwrap();

                backend.create_database_pool().await.unwrap()
            })
            .await;

        db_pool.pull_immutable().await
    }

    // add test case
    async fn test() {
        table! {
            book (id) {
                id -> Int4,
                title -> Text
            }
        }

        #[derive(Insertable)]
        #[diesel(table_name = book)]
        struct NewBook<'a> {
            title: &'a str,
        }

        // get connection pool from database pool
        let conn_pool = get_connection_pool().await;
        let conn = &mut conn_pool.get().await.unwrap();

        let new_book = NewBook { title: "Title" };

        insert_into(book::table)
            .values(&new_book)
            .execute(conn)
            .await
            .unwrap();

        let count = book::table.count().get_result::<i64>(conn).await.unwrap();
        assert_eq!(count, 1);
    }
}

The test gets a connection pool from the database pool every time it runs.

Finally, we add a couple of separate tests that call the same test case.

fn main() {}

#[cfg(test)]
mod tests {
    // allow needless returns generated by test macro
    #![allow(clippy::needless_return)]

    use bb8::Pool;
    use db_pool::{
        r#async::{
            DatabasePool, DatabasePoolBuilderTrait, DieselAsyncPostgresBackend, DieselBb8,
            ReusableConnectionPool,
        },
        PrivilegedPostgresConfig,
    };
    use diesel::{insert_into, sql_query, table, Insertable, QueryDsl};
    use diesel_async::RunQueryDsl;
    use dotenvy::dotenv;
    use tokio::sync::OnceCell;
    // import test macro
    use tokio_shared_rt::test;

    async fn get_connection_pool(
    ) -> ReusableConnectionPool<'static, DieselAsyncPostgresBackend<DieselBb8>> {
        static POOL: OnceCell<DatabasePool<DieselAsyncPostgresBackend<DieselBb8>>> =
            OnceCell::const_new();

        let db_pool = POOL
            .get_or_init(|| async {
                dotenv().ok();

                let config = PrivilegedPostgresConfig::from_env().unwrap();

                let backend = DieselAsyncPostgresBackend::new(
                    config,
                    || Pool::builder().max_size(10),
                    || Pool::builder().max_size(2),
                    None,
                    move |mut conn| {
                        Box::pin(async {
                            sql_query(
                                "CREATE TABLE book(id SERIAL PRIMARY KEY, title TEXT NOT NULL)",
                            )
                            .execute(&mut conn)
                            .await
                            .unwrap();

                            conn
                        })
                    },
                )
                .await
                .unwrap();

                backend.create_database_pool().await.unwrap()
            })
            .await;

        db_pool.pull_immutable().await
    }

    async fn test() {
        table! {
            book (id) {
                id -> Int4,
                title -> Text
            }
        }

        #[derive(Insertable)]
        #[diesel(table_name = book)]
        struct NewBook<'a> {
            title: &'a str,
        }

        let conn_pool = get_connection_pool().await;
        let conn = &mut conn_pool.get().await.unwrap();

        let new_book = NewBook { title: "Title" };

        insert_into(book::table)
            .values(&new_book)
            .execute(conn)
            .await
            .unwrap();

        let count = book::table.count().get_result::<i64>(conn).await.unwrap();
        assert_eq!(count, 1);
    }

    // add first test
    #[test(shared)]
    async fn test1() {
        test().await;
    }

    // add second test
    #[test(shared)]
    async fn test2() {
        test().await;
    }
}

We use the test macro from tokio_shared_rt instead of tokio::test in order to use a shared Tokio runtime for all tests. This is essential so that the static database pool is created and used by the same runtime and therefore remains valid for all tests.

We run the test cases in parallel and both pass.

running 2 tests
test tests::test1 ... ok
test tests::test2 ... ok

test result: ok. 2 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.24s