Postgres force drop database flag impl

launchbadge · Nov 10, 2023 · 77c071d · 77c071d
1 parent 73ab65c
commit 77c071d
Showing 1 changed file with 31 additions and 6 deletions.
diff --git a/sqlx-postgres/src/migrate.rs b/sqlx-postgres/src/migrate.rs
@@ -75,12 +75,37 @@ impl MigrateDatabase for Postgres {
             let (options, database) = parse_for_maintenance(url)?;
             let mut conn = options.connect().await?;
 
-            let _ = conn
-                .execute(&*format!(
-                    "DROP DATABASE IF EXISTS \"{}\"",
-                    database.replace('"', "\"\"")
-                ))
-                .await?;
+            let with_force = if force {
+                // language=SQL
+                let row: (String,) = query_as("SELECT current_setting('server_version_num')")
+                    .fetch_one(&mut conn)
+                    .await?;
+
+                let version = row.0.parse::<i32>().unwrap();
+
+                if version >= 130000 {
+                    "WITH (FORCE)"
+                } else {
+                    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?;
+
+                    ""
+                }
+            } else {
+                ""
+            };
+
+            conn.execute(&*format!(
+                "DROP DATABASE IF EXISTS \"{}\" {with_force}",
+                database.replace('"', "\"\""),
+            ))
+            .await?;
 
             Ok(())
         })