From 1d78ba247666a0f0006cc4b50755b901c4467866 Mon Sep 17 00:00:00 2001
From: Philip Hayes <philiphayes9@gmail.com>
Date: Thu, 13 Apr 2023 11:27:36 -0700
Subject: [PATCH] fix: ensure fresh test db's aren't accidentally deleted by
 do_cleanup

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.
---
 sqlx-mysql/src/testing/mod.rs    | 24 ++++++++++++++++++-----
 sqlx-postgres/src/testing/mod.rs | 33 ++++++++++++++++++++++++--------
 2 files changed, 44 insertions(+), 13 deletions(-)

diff --git a/sqlx-mysql/src/testing/mod.rs b/sqlx-mysql/src/testing/mod.rs
index 73fc9d3fa1..9addd4c6d0 100644
--- a/sqlx-mysql/src/testing/mod.rs
+++ b/sqlx-mysql/src/testing/mod.rs
@@ -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;
 
@@ -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))
         })
@@ -123,12 +128,19 @@ 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) {
         // help trigger the test-db-race
         std::thread::sleep(Duration::from_secs(1));
 
-        do_cleanup(&mut conn).await?;
+        do_cleanup(&mut conn, now).await?;
     }
 
     query("insert into _sqlx_test_databases(test_path) values (?)")
@@ -166,10 +178,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?;
 
diff --git a/sqlx-postgres/src/testing/mod.rs b/sqlx-postgres/src/testing/mod.rs
index f882753901..3835dba36a 100644
--- a/sqlx-postgres/src/testing/mod.rs
+++ b/sqlx-postgres/src/testing/mod.rs
@@ -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;
 
@@ -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))
         })
@@ -133,12 +138,19 @@ 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) {
         // help trigger the test-db-race
         std::thread::sleep(Duration::from_secs(1));
 
-        do_cleanup(&mut conn).await?;
+        do_cleanup(&mut conn, now).await?;
     }
 
     let new_db_name: String = query_scalar(
@@ -173,11 +185,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);