Drop database on drop Using Sqlx and Rust

Question

I'm following step by step the zero2prod Rust book. So far, I'm in the last parts of Chapter 3 in which we setup some integration tests and create a database per test run.

I'd like to extend this functionality and add the capability to drop the database after each test has run. This lead me to the Drop trait, in which I would execute the DROP DATABASE command and call it a day (I'd store the db name and its connection/pool so I could access it later).

The problem, though, is that the app hangs at executing the DROP DATABASE query and and times out after 60 seconds. I don't know what's causing this as I am unable to print or debug the connection.

Here's the code that I have:

use futures::executor;
use sqlx::{Connection, Executor, PgConnection, PgPool};
use std::net::TcpListener;
use uuid::Uuid;
use zero2prod::configuration::{get_configuration, DatabaseSettings};

const BASE_URL: &str = "127.0.0.1";
pub struct TestApp {
    db_name: String,
    connection_string: String,
    pub address: String,
    pub db_pool: PgPool,
    pub connection: PgConnection,
}

/**
 * We need to refactor our project into a library and a binary: all our logic will live in the library crate
while the binary itself will be just an entrypoint with a very slim main function
 */

pub async fn init(url: &str) -> TestApp {
    let mut app = spawn_app().await;
    app.address = format!("{}{}", app.address, url);
    return app;
}

// Launch our application in the background ~somehow~
async fn spawn_app() -> TestApp {
    // We take the BASE_URL const and assign it a port 0. We then
    // pass the listener to the server
    let base_url = format!("{}:0", BASE_URL);
    let listener = TcpListener::bind(base_url).expect("Failed to bind random port");

    // We retrieve the port assigned by the OS
    let port = listener.local_addr().unwrap().port();

    let (connection, db_connection, db_name, connection_string) = init_db().await;

    // We pass the port now to our server
    let server = zero2prod::run(listener, db_connection.clone()).expect("Failed to bind address");
    let _ = actix_web::rt::spawn(server);
    let address = format!("http://{}:{}", BASE_URL, port);

    TestApp {
        db_name: String::from(db_name),
        address,
        db_pool: db_connection,
        connection,
        connection_string,
    }
}

async fn init_db() -> (PgConnection, PgPool, String, String) {
    let mut configuration = get_configuration().expect("Failed to read configuration");

    // We change the db name in each run as we need to run the test multiple times
    configuration.database.database_name = Uuid::new_v4().to_string();

    let (connection, pool) = configure_database(&configuration.database).await;

    return (
        connection,
        pool,
        String::from(&configuration.database.database_name),
        configuration.database.connection_string_without_db(),
    );
}

async fn configure_database(config: &DatabaseSettings) -> (PgConnection, PgPool) {

   // The following returns:
   //   format!(
   //       "postgres://{}:{}@{}:{}",
   //       self.username, self.password, self.host, self.port
   //   )
    let mut connection = PgConnection::connect(&config.connection_string_without_db())
        .await
        .expect("Failed to connect to Postgres.");

    connection
        .execute(format!(r#"CREATE DATABASE "{}""#, config.database_name).as_str())
        .await
        .expect("Failed to create the db");

    // Migrate the database

    let connection_pool = PgPool::connect(&config.connection_string())
        .await
        .expect("Failed to connect to Postgres");

    sqlx::migrate!("./migrations")
        .run(&connection_pool)
        .await
        .expect("Failed to migrate db");

    return (connection, connection_pool);

The entry point is the init() function which basically returns a TestApp struct (which originally contained the db_pool and address fields). Everything on the code above is working.

The problem is down below. Here's everything that I've tried:

1. Using Smol's runtime to run async in drop - Tried initializing a new connection to the Postgres database

impl Drop for TestApp {
    fn drop(&mut self) {
        smol::block_on(async {
            let mut connection = PgConnection::connect(&self.connection_string)
                .await
                .expect("Failed to connect to Postgres.");

            let result = connection
                .execute(format!(r#"DROP DATABASE "{}""#, self.db_name).as_str())
                .await
                .expect("Error while querying the drop database");
            println!("{:?}", result);
        });
    }
}

2. Using Smol's runtime to run async in drop - tried using the exiting db_pool

    fn drop(&mut self) {
        smol::block_on(async {
            let result = self
                .db_pool
                .execute(format!(r#"DROP DATABASE "{}""#, self.db_name).as_str())
                .await.expect("Error while querying");
            println!("{:?}", result);
        });
    }

3. Using Future's crate executor - using the existing db_pool

     let result = executor::block_on(
            self.db_pool
                .execute(format!(r#"DROP DATABASE "{}""#, self.db_name).as_str()),
        )
        .expect("Failed to drop database");

        println!("{:?}", result);

4. Using Future's crate executor - Running db_pool.acquire() and then the pool (This hangs at db_pool.acquire.

     executor::block_on(self.db_pool.acquire()).expect("Failed to acquire pool");
     let result = executor::block_on(
            self.db_pool
                .execute(format!(r#"DROP DATABASE "{}""#, self.db_name).as_str()),
        )
        .expect("Failed to drop database");

        println!("{:?}", result);

5. Using Future's crate executor - Running the existing connection.

        let result = executor::block_on(
            self.connection
                .execute(format!(r#"DROP DATABASE "{}""#, self.db_name).as_str()),
        )
        .expect("Failed to drop database");

        println!("{:?}", result);

Note that the code isn't the prettiest, as I'm trying to find a working solution first.

Unfortunately I don't know what the problem is as there are no errors thrown.

Any ideas?

Answer 1

sqlx >0.6.1 update

The solution by @Ziliang Lin is indeed much cleaner now that sqlx::test is supported in sqlx. My implementation looks like this:

// Start the app
async fn spawn_app(pool: Pool<Postgres>) -> TestApp {
    let listener = TcpListener::bind("127.0.0.1:0").expect("Failed to bind random port");
    let port = listener.local_addr().unwrap().port();

    let address = format!("http://127.0.0.1:{}", port);

    let server = run(listener, pool.clone()).expect("Failed to bind address");
    let _ = tokio::spawn(server);
    TestApp {
        address,
        db_pool: pool,
    }
}

// sqlx::test will automatically create the database, run migrations,
// and delete the database when the test is complete
#[sqlx::test]
async fn health_check_works(pool: Pool<Postgres>) {
    let app = spawn_app(pool).await;
    let client = reqwest::Client::new();

    // Act
    let response = client
        // Use the returned application address
        .get(&format!("{}/health_check", &app.address))
        .send()
        .await
        .expect("Failed to execute request.");

    // Assert
    assert!(response.status().is_success());
    assert_eq!(Some(0), response.content_length());
}

Just update all test routines from tokio::test -> sqlx::test, delete the configure_database functions, and add the pool: Pool<Postgres> argument to the test functions.

Original post

I ran into this same issue. Thanks to some of the answers here, I ended up with the following solution (Note that I added an assertion for the db name starting with "test_db-", so you'll have to prepend your database names or delete the assertion):

pub struct TestApp {
    pub address: String,
    pub db_pool: PgPool,
    pub db_name: String,
    pub db_connection_string: String,
}

impl TestApp {
    async fn terminate(&mut self) {
        assert!(self.db_name.starts_with("test_db-"));
        println!("Cleaning up database: {}", self.db_name);
        self.db_pool.close().await;
        let mut connection = PgConnection::connect(&self.db_connection_string)
            .await
            .expect("Failed to connect to Postgres");

        // Force drop all active connections to database
        // TODO: see if there is a softer way to handle this (i.e. close connection when DB access is complete)
        connection
            .execute(
                format!(
                    r#"
                    SELECT pg_terminate_backend(pg_stat_activity.pid)
                    FROM pg_stat_activity
                    WHERE pg_stat_activity.datname = '{}'
                    AND pid <> pg_backend_pid()
                    "#,
                    self.db_name
                )
                .as_str(),
            )
            .await
            .expect("Failed to terminate current connections to test db");

        connection
            .execute(format!(r#"DROP DATABASE "{}";"#, self.db_name).as_str())
            .await
            .expect("Failed to drop database.");
        println!("Database cleaned up successfully.")
    }
}

impl Drop for TestApp {
    fn drop(&mut self) {
        std::thread::scope(|s| {
            s.spawn(|| {
                let runtime = tokio::runtime::Builder::new_multi_thread()
                    .enable_all()
                    .build()
                    .unwrap();
                runtime.block_on(self.terminate());
            });
        });
    }
}

Answer 2

For people later meeting this question, I would suggest having a look at sqlx's sqlx::test

This procedural macro would achieve what the OP wants to do automatically:

1. create a database for each integration test annotated by [sqlx::test]
2. drop the database after the test run

here's my code sample:

#[sqlx::test]
async fn confirmations_without_token_rejected_with_400(pool: Pool<Postgres>) -> Result<()> {
    // here I am using sea_orm, so I would do the migration inside `get_test_ap`
    let app = get_test_app(pool).await?;
    let cli = &app.cli;
    let resp = cli.get("/subscriptions/confirm").send().await;
    resp.assert_status(StatusCode::BAD_REQUEST);
    Ok(())
}

if you are using pure sqlx, you could just follow the link above and find the following example:

use sqlx::PgPool;

#[sqlx::test(migrations = "foo_migrations")]
async fn basic_test(pool: PgPool) -> sqlx::Result<()> {
    let mut conn = pool.acquire().await?;

    sqlx::query("SELECT * FROM foo")
        .fetch_one(&mut conn)
        .await?;
        
    assert_eq!(foo.get::<String>("bar"), "foobar!");
    
    Ok(())
}