db-pool

db-pool is a thread-safe database pool for running database-tied integration tests in parallel with:

Easy setup
Proper isolation
Automatic creation, reuse, and cleanup
Async support

Description

Rather than simply providing a database connection pool that allows multiple connections to the same database, db-pool maintains a pool of separate isolated databases in order to allow running database-tied tests in parallel. It also handles the lifecycles of those databases: whenever you pick a database out of the pool, you can be sure that the database is clean and ready to be used, and that no other tests are connected to the database you are using in any one test.

Motivation

When running tests against a database-tied service, such as a web server, a test database is generally used. However, this comes with its own set of difficulties:

The database has to be either (a) dropped and re-created or (b) cleaned before every test.
Tests have to run serially in order to avoid cross-contamination.

This leads to several issues when running tests serially:

Test setup and teardown is now required.
Dropping and creating a database from scratch can be expensive.
Cleaning a database instead of dropping and re-creating one requires careful execution of dialect-specific statements.

When switching to parallel execution of tests, even more difficulties arise:

Creating and dropping a database for each test can be expensive.
Sharing temporary databases across tests requires:
- isolating databases in concurrent use
- cleaning each database before reuse by a subsequent test
- restricting user privileges to prevent schema modification by rogue tests
- dropping temporary databases before or after a test run to reduce clutter

db-pool takes care of all of these concerns while supporting multiple database types, backends, and connection pools.

Databases

MySQL (MariaDB)
PostgreSQL

Backends & Pools

Sync

Backend	Pool	Feature
diesel/mysql	r2d2	`diesel-mysql`
diesel/postgres	r2d2	`diesel-postgres`
mysql	r2d2	`mysql`
postgres	r2d2	`postgres`

Async

Backend	Pool	Features
diesel-async/mysql	bb8	`diesel-async-mysql`, `diesel-async-bb8`
diesel-async/mysql	mobc	`diesel-async-mysql`, `diesel-async-mobc`
diesel-async/postgres	bb8	`diesel-async-postgres`, `diesel-async-bb8`
diesel-async/postgres	mobc	`diesel-async-postgres`, `diesel-async-mobc`
sea-orm/sqlx-mysql	sqlx	`sea-orm-mysql`
sea-orm/sqlx-postgres	sqlx	`sea-orm-postgres`
sqlx/mysql	sqlx	`sqlx-mysql`
sqlx/postgres	sqlx	`sqlx-postgres`
tokio-postgres	bb8	`tokio-postgres`, `tokio-postgres-bb8`
tokio-postgres	mobc	`tokio-postgres`, `tokio-postgres-mobc`

We will use the Diesel Postgres backend in this tutorial.

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

fn main() {}

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

    // import OnceLock
    use std::sync::OnceLock;

    fn get_connection_pool() {
        // add "lazy static"
        static POOL: OnceLock<()> = OnceLock::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.

Environment Variable	Default
POSTGRES_USERNAME	postgres
POSTGRES_PASSWORD	{blank}
POSTGRES_HOST	localhost
POSTGRES_PORT	3306

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)]

    use std::sync::OnceLock;

    // import privileged configuration
    use db_pool::PrivilegedPostgresConfig;
    // import dotenvy
    use dotenvy::dotenv;

    fn get_connection_pool() {
        static POOL: OnceLock<()> = OnceLock::new();

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

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

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)]

    use std::sync::OnceLock;

    use db_pool::{
        // import backend
        sync::DieselPostgresBackend,
        PrivilegedPostgresConfig,
    };
    // import diesel-specific constructs
    use diesel::{sql_query, RunQueryDsl};
    use dotenvy::dotenv;
    // import connection pool
    use r2d2::Pool;

    fn get_connection_pool() {
        static POOL: OnceLock<()> = OnceLock::new();

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

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

            // create backend
            let backend = DieselPostgresBackend::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),
                // create entities
                move |conn| {
                    sql_query("CREATE TABLE book(id SERIAL PRIMARY KEY, title TEXT NOT NULL)")
                        .execute(conn)
                        .unwrap();
                },
            )
            .unwrap();
        });
    }
}

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

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 std::sync::OnceLock;

    use db_pool::{
        sync::{
            // import database pool
            DatabasePool,
            // import database pool builder trait
            DatabasePoolBuilderTrait,
            DieselPostgresBackend,
        },
        PrivilegedPostgresConfig,
    };
    use diesel::{sql_query, RunQueryDsl};
    use dotenvy::dotenv;
    use r2d2::Pool;

    fn get_connection_pool() {
        // change OnceLock inner type
        static POOL: OnceLock<DatabasePool<DieselPostgresBackend>> = OnceLock::new();

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

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

            let backend = DieselPostgresBackend::new(
                config,
                || Pool::builder().max_size(10),
                || Pool::builder().max_size(2),
                move |conn| {
                    sql_query("CREATE TABLE book(id SERIAL PRIMARY KEY, title TEXT NOT NULL)")
                        .execute(conn)
                        .unwrap();
                },
            )
            .unwrap();

            // create database pool
            backend.create_database_pool().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 std::sync::OnceLock;

    use db_pool::{
        sync::{
            DatabasePool,
            DatabasePoolBuilderTrait,
            DieselPostgresBackend,
            // import reusable connection pool
            ReusableConnectionPool,
        },
        PrivilegedPostgresConfig,
    };
    use diesel::{sql_query, RunQueryDsl};
    use dotenvy::dotenv;
    use r2d2::Pool;

    // change return type
    fn get_connection_pool() -> ReusableConnectionPool<'static, DieselPostgresBackend> {
        static POOL: OnceLock<DatabasePool<DieselPostgresBackend>> = OnceLock::new();

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

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

            let backend = DieselPostgresBackend::new(
                config,
                || Pool::builder().max_size(10),
                || Pool::builder().max_size(2),
                move |conn| {
                    sql_query("CREATE TABLE book(id SERIAL PRIMARY KEY, title TEXT NOT NULL)")
                        .execute(conn)
                        .unwrap();
                },
            )
            .unwrap();

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

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

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 std::sync::OnceLock;

    use db_pool::{
        sync::{
            DatabasePool, DatabasePoolBuilderTrait, DieselPostgresBackend, ReusableConnectionPool,
        },
        PrivilegedPostgresConfig,
    };
    // import extra diesel-specific constructs
    use diesel::{insert_into, sql_query, table, Insertable, QueryDsl, RunQueryDsl};
    use dotenvy::dotenv;
    use r2d2::Pool;

    fn get_connection_pool() -> ReusableConnectionPool<'static, DieselPostgresBackend> {
        static POOL: OnceLock<DatabasePool<DieselPostgresBackend>> = OnceLock::new();

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

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

            let backend = DieselPostgresBackend::new(
                config,
                || Pool::builder().max_size(10),
                || Pool::builder().max_size(2),
                move |conn| {
                    sql_query("CREATE TABLE book(id SERIAL PRIMARY KEY, title TEXT NOT NULL)")
                        .execute(conn)
                        .unwrap();
                },
            )
            .unwrap();

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

        db_pool.pull_immutable()
    }

    // add test case
    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();
        let conn = &mut conn_pool.get().unwrap();

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

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

        let count = book::table.count().get_result::<i64>(conn).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 {
    use std::sync::OnceLock;

    use db_pool::{
        sync::{
            DatabasePool, DatabasePoolBuilderTrait, DieselPostgresBackend, ReusableConnectionPool,
        },
        PrivilegedPostgresConfig,
    };
    use diesel::{insert_into, sql_query, table, Insertable, QueryDsl, RunQueryDsl};
    use dotenvy::dotenv;
    use r2d2::Pool;

    fn get_connection_pool() -> ReusableConnectionPool<'static, DieselPostgresBackend> {
        static POOL: OnceLock<DatabasePool<DieselPostgresBackend>> = OnceLock::new();

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

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

            let backend = DieselPostgresBackend::new(
                config,
                || Pool::builder().max_size(10),
                || Pool::builder().max_size(2),
                move |conn| {
                    sql_query("CREATE TABLE book(id SERIAL PRIMARY KEY, title TEXT NOT NULL)")
                        .execute(conn)
                        .unwrap();
                },
            )
            .unwrap();

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

        db_pool.pull_immutable()
    }

    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();
        let conn = &mut conn_pool.get().unwrap();

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

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

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

    // add first test
    #[test]
    fn test1() {
        test();
    }

    // add second test
    #[test]
    fn test2() {
        test();
    }
}

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.25s

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();
    }
}