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

Cargo.lock

@@ -30,9 +30,9 @@ dependencies = [
 
 [[package]]
 name = "ahash"
-version = "0.8.4"
+version = "0.8.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "72832d73be48bac96a5d7944568f305d829ed55b0ce3b483647089dfaf6cf704"
+checksum = "91429305e9f0a25f6205c5b8e0d2db09e0708a7a6df0f42212bb56c32c8ac97a"
 dependencies = [
  "cfg-if",
  "getrandom",
@@ -1503,7 +1503,7 @@ version = "0.13.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "43a3c133739dddd0d2990f9a4bdf8eb4b21ef50e4851ca85ab661199821d510e"
 dependencies = [
- "ahash 0.8.4",
+ "ahash 0.8.6",
 ]
 
 [[package]]
@@ -1512,7 +1512,7 @@ version = "0.14.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f93e7192158dbcda357bdec5fb5788eebf8bbac027f3f33e719d29135ae84156"
 dependencies = [
- "ahash 0.8.4",
+ "ahash 0.8.6",
  "allocator-api2",
 ]
 
@@ -3220,7 +3220,7 @@ dependencies = [
 name = "sqlx-core"
 version = "0.7.2"
 dependencies = [
- "ahash 0.8.4",
+ "ahash 0.8.6",
  "async-io",
  "async-std",
  "atoi",
@@ -4410,18 +4410,18 @@ dependencies = [
 
 [[package]]
 name = "zerocopy"
-version = "0.7.11"
+version = "0.7.25"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4c19fae0c8a9efc6a8281f2e623db8af1db9e57852e04cde3e754dd2dc29340f"
+checksum = "8cd369a67c0edfef15010f980c3cbe45d7f651deac2cd67ce097cd801de16557"
 dependencies = [
  "zerocopy-derive",
 ]
 
 [[package]]
 name = "zerocopy-derive"
-version = "0.7.11"
+version = "0.7.25"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fc56589e9ddd1f1c28d4b4b5c773ce232910a6bb67a70133d61c9e347585efe9"
+checksum = "c2f140bda219a26ccc0cdb03dba58af72590c53b22642577d88a927bc5c87d6b"
 dependencies = [
  "proc-macro2",
  "quote",

sqlx-cli/src/database.rs

@@ -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/lib.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/opt.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-core/Cargo.toml

@@ -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/src/any/driver.rs

@@ -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/migrate.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/migrate/error.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/migrate.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-postgres/src/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 {