Skip to content

Commit

Permalink
fix: ensure fresh test db's aren't accidentally deleted by do_cleanup (
Browse files Browse the repository at this point in the history
…#2454)

If the first test thread is a bit slow after it acquires the
`DO_CLEANUP` permit, it can accidentally delete a fresh test db created
by another thread right before that other thread has a chance to open
its connection.

This fix simply records the current timestamp _before_ we acquire the
`DO_CLEANUP` permit and only cleans up test db's created before then.
  • Loading branch information
phlip9 authored May 4, 2023
1 parent cbf8dd3 commit 1227d5d
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 13 deletions.
24 changes: 19 additions & 5 deletions sqlx-mysql/src/testing/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use std::fmt::Write;
use std::ops::Deref;
use std::str::FromStr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use std::time::{Duration, SystemTime};

use futures_core::future::BoxFuture;

Expand Down Expand Up @@ -60,7 +60,12 @@ impl TestSupport for MySql {
let url = dotenvy::var("DATABASE_URL").expect("DATABASE_URL must be set");

let mut conn = MySqlConnection::connect(&url).await?;
let num_deleted = do_cleanup(&mut conn).await?;

let now = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap();

let num_deleted = do_cleanup(&mut conn, now).await?;
let _ = conn.close().await;
Ok(Some(num_deleted))
})
Expand Down Expand Up @@ -123,9 +128,16 @@ async fn test_context(args: &TestArgs) -> Result<TestContext<MySql>, Error> {
)
.await?;

// Record the current time _before_ we acquire the `DO_CLEANUP` permit. This
// prevents the first test thread from accidentally deleting new test dbs
// created by other test threads if we're a bit slow.
let now = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap();

// Only run cleanup if the test binary just started.
if DO_CLEANUP.swap(false, Ordering::SeqCst) {
do_cleanup(&mut conn).await?;
do_cleanup(&mut conn, now).await?;
}

query("insert into _sqlx_test_databases(test_path) values (?)")
Expand Down Expand Up @@ -163,10 +175,12 @@ async fn test_context(args: &TestArgs) -> Result<TestContext<MySql>, Error> {
})
}

async fn do_cleanup(conn: &mut MySqlConnection) -> Result<usize, Error> {
async fn do_cleanup(conn: &mut MySqlConnection, created_before: Duration) -> Result<usize, Error> {
let delete_db_ids: Vec<u64> = query_scalar(
"select db_id from _sqlx_test_databases where created_at < current_timestamp()",
"select db_id from _sqlx_test_databases \
where created_at < (cast(from_unixtime($1) as timestamp))",
)
.bind(&created_before.as_secs())
.fetch_all(&mut *conn)
.await?;

Expand Down
33 changes: 25 additions & 8 deletions sqlx-postgres/src/testing/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use std::fmt::Write;
use std::ops::Deref;
use std::str::FromStr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use std::time::{Duration, SystemTime};

use futures_core::future::BoxFuture;

Expand Down Expand Up @@ -57,7 +57,12 @@ impl TestSupport for Postgres {
let url = dotenvy::var("DATABASE_URL").expect("DATABASE_URL must be set");

let mut conn = PgConnection::connect(&url).await?;
let num_deleted = do_cleanup(&mut conn).await?;

let now = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap();

let num_deleted = do_cleanup(&mut conn, now).await?;
let _ = conn.close().await;
Ok(Some(num_deleted))
})
Expand Down Expand Up @@ -133,9 +138,16 @@ async fn test_context(args: &TestArgs) -> Result<TestContext<Postgres>, Error> {
)
.await?;

// Record the current time _before_ we acquire the `DO_CLEANUP` permit. This
// prevents the first test thread from accidentally deleting new test dbs
// created by other test threads if we're a bit slow.
let now = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap();

// Only run cleanup if the test binary just started.
if DO_CLEANUP.swap(false, Ordering::SeqCst) {
do_cleanup(&mut conn).await?;
do_cleanup(&mut conn, now).await?;
}

let new_db_name: String = query_scalar(
Expand Down Expand Up @@ -170,11 +182,16 @@ async fn test_context(args: &TestArgs) -> Result<TestContext<Postgres>, Error> {
})
}

async fn do_cleanup(conn: &mut PgConnection) -> Result<usize, Error> {
let delete_db_names: Vec<String> =
query_scalar("select db_name from _sqlx_test.databases where created_at < now()")
.fetch_all(&mut *conn)
.await?;
async fn do_cleanup(conn: &mut PgConnection, created_before: Duration) -> Result<usize, Error> {
let created_before = i64::try_from(created_before.as_secs()).unwrap();

let delete_db_names: Vec<String> = query_scalar(
"select db_name from _sqlx_test.databases \
where created_at < (to_timestamp($1) at time zone 'UTC')",
)
.bind(&created_before)
.fetch_all(&mut *conn)
.await?;

if delete_db_names.is_empty() {
return Ok(0);
Expand Down

0 comments on commit 1227d5d

Please sign in to comment.