Implement async db drop

blockscout · Oct 15, 2023 · 40206e8 · 40206e8
1 parent 6233072
commit 40206e8
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 2 deletions.
diff --git a/libs/blockscout-service-launcher/Cargo.toml b/libs/blockscout-service-launcher/Cargo.toml
@@ -41,6 +41,9 @@ sea-orm-migration-0_11 = { package = "sea-orm-migration", version = "0.11", opti
 sea-orm-0_12 = { package = "sea-orm", version = "0.12.2", optional = true }
 sea-orm-migration-0_12 = { package = "sea-orm-migration", version = "0.12.2", optional = true }
 
+async-dropper = { version = "0.2", features = [ "tokio" ], optional = true }
+async-trait = { version = "0.1", optional = true }
+
 [features]
 default = ["launcher", "tracing"]
 launcher = [
@@ -96,4 +99,6 @@ test-server = [
 test-database = [
     "database",
     "dep:keccak-hash",
+    "dep:async-dropper",
+    "dep:async-trait",
 ]
diff --git a/libs/blockscout-service-launcher/src/test_database.rs b/libs/blockscout-service-launcher/src/test_database.rs
@@ -1,14 +1,34 @@
 use crate::database::{
     ConnectionTrait, Database, DatabaseConnection, DbErr, MigratorTrait, Statement,
 };
-use std::{ops::Deref, sync::Arc};
+use async_dropper::derive::AsyncDrop;
+use async_trait::async_trait;
+use std::{ops::Deref, sync::Arc, time::Duration};
 
-#[derive(Clone, Debug)]
+#[derive(AsyncDrop, Clone, Debug, Default)]
 pub struct TestDbGuard {
     conn_with_db: Arc<DatabaseConnection>,
     db_url: String,
+    db_name: String,
+    should_drop: bool,
 }
 
+impl PartialEq<Self> for TestDbGuard {
+    fn eq(&self, other: &Self) -> bool {
+        self.db_url == other.db_url
+            && self.db_name == other.db_name
+            && matches!(
+                (self.conn_with_db.deref(), other.conn_with_db.deref()),
+                (
+                    DatabaseConnection::Disconnected,
+                    DatabaseConnection::Disconnected
+                )
+            )
+    }
+}
+
+impl Eq for TestDbGuard {}
+
 impl TestDbGuard {
     pub async fn new<Migrator: MigratorTrait>(db_name: &str) -> Self {
         let db_url = std::env::var("DATABASE_URL")
@@ -20,9 +40,16 @@ impl TestDbGuard {
         TestDbGuard {
             conn_with_db: Arc::new(conn_with_db),
             db_url,
+            db_name,
+            should_drop: true,
         }
     }
 
+    pub fn without_drop(mut self) -> Self {
+        self.should_drop = false;
+        self
+    }
+
     pub fn client(&self) -> Arc<DatabaseConnection> {
         self.conn_with_db.clone()
     }
@@ -89,3 +116,21 @@ impl Deref for TestDbGuard {
         &self.conn_with_db
     }
 }
+
+#[async_trait]
+impl AsyncDrop for TestDbGuard {
+    async fn async_drop(&mut self) -> Result<(), AsyncDropError> {
+        if self.should_drop {
+            // Workaround for postgres error `cannot drop the currently open database`
+            let conn_without_db = Database::connect(self.db_url())
+                .await
+                .expect("Connection to postgres (without database) failed");
+            self.conn_with_db = Arc::new(conn_without_db);
+
+            TestDbGuard::drop_database(self, self.db_name.as_str())
+                .await
+                .expect("Failed to drop database");
+        }
+        Ok(())
+    }
+}