Sqlx Cli: Added force flag to drop database for postgres (#2873)

* Updated ahash so it can compile on mac

* Updated MigrateDatabase Trait + related functions

* Postgres force drop database flag impl

* Update migrate.rs

* Reverted MigrateDatabase Trait

* Update migrate.rs

* Update migrate.rs

* Added force drop database fn impl

* Add Migrate Error

* Fixed changed function name

Author: Vrajs16

Cargo.lock

@@ -23,7 +23,7 @@ pub async fn create(connect_opts: &ConnectOpts) -> anyhow::Result<()> {
     Ok(())
 }
 
-pub async fn drop(connect_opts: &ConnectOpts, confirm: bool) -> anyhow::Result<()> {
+pub async fn drop(connect_opts: &ConnectOpts, confirm: bool, force: bool) -> anyhow::Result<()> {
     if confirm && !ask_to_continue_drop(connect_opts.required_db_url()?) {
         return Ok(());
     }
@@ -33,7 +33,11 @@ pub async fn drop(connect_opts: &ConnectOpts, confirm: bool) -> anyhow::Result<(
     let exists = crate::retry_connect_errors(connect_opts, Any::database_exists).await?;
 
     if exists {
-        Any::drop_database(connect_opts.required_db_url()?).await?;
+        if force {
+            Any::force_drop_database(connect_opts.required_db_url()?).await?;
+        } else {
+            Any::drop_database(connect_opts.required_db_url()?).await?;
+        }
     }
 
     Ok(())
@@ -43,8 +47,9 @@ pub async fn reset(
     migration_source: &str,
     connect_opts: &ConnectOpts,
     confirm: bool,
+    force: bool,
 ) -> anyhow::Result<()> {
-    drop(connect_opts, confirm).await?;
+    drop(connect_opts, confirm, force).await?;
     setup(migration_source, connect_opts).await
 }
 

sqlx-cli/src/database.rs

@@ -74,12 +74,14 @@ pub async fn run(opt: Opt) -> Result<()> {
             DatabaseCommand::Drop {
                 confirmation,
                 connect_opts,
-            } => database::drop(&connect_opts, !confirmation.yes).await?,
+                force,
+            } => database::drop(&connect_opts, !confirmation.yes, force).await?,
             DatabaseCommand::Reset {
                 confirmation,
                 source,
                 connect_opts,
-            } => database::reset(&source, &connect_opts, !confirmation.yes).await?,
+                force,
+            } => database::reset(&source, &connect_opts, !confirmation.yes, force).await?,
             DatabaseCommand::Setup {
                 source,
                 connect_opts,

sqlx-cli/src/lib.rs

@@ -76,6 +76,10 @@ pub enum DatabaseCommand {
 
         #[clap(flatten)]
         connect_opts: ConnectOpts,
+
+        /// PostgreSQL only: force drops the database.
+        #[clap(long, short, default_value = "false")]
+        force: bool,
     },
 
     /// Drops the database specified in your DATABASE_URL, re-creates it, and runs any pending migrations.
@@ -88,6 +92,10 @@ pub enum DatabaseCommand {
 
         #[clap(flatten)]
         connect_opts: ConnectOpts,
+
+        /// PostgreSQL only: force drops the database.
+        #[clap(long, short, default_value = "false")]
+        force: bool,
     },
 
     /// Creates the database specified in your DATABASE_URL and runs any pending migrations.

sqlx-cli/src/opt.rs

@@ -51,7 +51,7 @@ uuid = { workspace = true, optional = true }
 
 async-io = { version = "1.9.0", optional = true }
 paste = "1.0.6"
-ahash = "0.8"
+ahash = "0.8.6"
 atoi = "2.0"
 
 bytes = "1.1.0"

sqlx-core/Cargo.toml

@@ -70,6 +70,7 @@ impl AnyDriver {
                 create_database: DebugFn(DB::create_database),
                 database_exists: DebugFn(DB::database_exists),
                 drop_database: DebugFn(DB::drop_database),
+                force_drop_database: DebugFn(DB::force_drop_database),
             }),
             ..Self::without_migrate::<DB>()
         }
@@ -94,6 +95,7 @@ pub struct AnyMigrateDatabase {
     create_database: DebugFn<fn(&str) -> BoxFuture<'_, crate::Result<()>>>,
     database_exists: DebugFn<fn(&str) -> BoxFuture<'_, crate::Result<bool>>>,
     drop_database: DebugFn<fn(&str) -> BoxFuture<'_, crate::Result<()>>>,
+    force_drop_database: DebugFn<fn(&str) -> BoxFuture<'_, crate::Result<()>>>,
 }
 
 impl AnyMigrateDatabase {
@@ -108,6 +110,10 @@ impl AnyMigrateDatabase {
     pub fn drop_database<'a>(&self, url: &'a str) -> BoxFuture<'a, crate::Result<()>> {
         (self.drop_database)(url)
     }
+
+    pub fn force_drop_database<'a>(&self, url: &'a str) -> BoxFuture<'a, crate::Result<()>> {
+        (self.force_drop_database)(url)
+    }
 }
 
 /// Install the list of drivers for [`AnyConnection`] to use.

sqlx-core/src/any/driver.rs

@@ -32,6 +32,15 @@ impl MigrateDatabase for Any {
                 .await
         })
     }
+
+    fn force_drop_database(url: &str) -> BoxFuture<'_, Result<(), Error>> {
+        Box::pin(async {
+            driver::from_url_str(url)?
+                .get_migrate_database()?
+                .force_drop_database(url)
+                .await
+        })
+    }
 }
 
 impl Migrate for AnyConnection {

sqlx-core/src/any/migrate.rs

@@ -24,6 +24,9 @@ pub enum MigrateError {
     #[error("migration {0} is newer than the latest applied migration {1}")]
     VersionTooNew(i64, i64),
 
+    #[error("database driver does not support force-dropping a database (Only PostgreSQL)")]
+    ForceNotSupported,
+
     #[deprecated = "migration types are now inferred"]
     #[error("cannot mix reversible migrations with simple migrations. All migrations should be reversible or simple migrations")]
     InvalidMixReversibleAndSimple,

sqlx-core/src/migrate/error.rs

@@ -15,6 +15,12 @@ pub trait MigrateDatabase {
     // drop database in url
     // uses a maintenance database depending on driver
     fn drop_database(url: &str) -> BoxFuture<'_, Result<(), Error>>;
+
+    // force drop database in url
+    // uses a maintenance database depending on driver
+    fn force_drop_database(_url: &str) -> BoxFuture<'_, Result<(), Error>> {
+        Box::pin(async { Err(MigrateError::ForceNotSupported)? })
+    }
 }
 
 // 'e = Executor

sqlx-core/src/migrate/migrate.rs

@@ -85,6 +85,30 @@ impl MigrateDatabase for Postgres {
             Ok(())
         })
     }
+
+    fn force_drop_database(url: &str) -> BoxFuture<'_, Result<(), Error>> {
+        Box::pin(async move {
+            let (options, database) = parse_for_maintenance(url)?;
+            let mut conn = options.connect().await?;
+
+            let row: (String,) = query_as("SELECT current_setting('server_version_num')")
+                .fetch_one(&mut conn)
+                .await?;
+
+            let version = row.0.parse::<i32>().unwrap();
+
+            let pid_type = if version >= 90200 { "pid" } else { "procpid" };
+
+            conn.execute(&*format!(
+                "SELECT pg_terminate_backend(pg_stat_activity.{pid_type}) FROM pg_stat_activity \
+                 WHERE pg_stat_activity.datname = {} AND {pid_type} <> pg_backend_pid()",
+                database.replace('"', "\"\""),
+            ))
+            .await?;
+
+            Self::drop_database(url).await
+        })
+    }
 }
 
 impl Migrate for PgConnection {