@@ -5,7 +5,6 @@ use crate::migration::dump::remove_restricted_role_grants;
55use crate :: { checkpoint, migration, postgres} ;
66use anyhow:: { bail, Context , Result } ;
77use std:: io:: { self , Write } ;
8- use tokio_postgres:: Client ;
98
109/// Initial replication command for snapshot schema and data copy
1110///
@@ -479,7 +478,7 @@ pub async fn init(
479478 } ;
480479
481480 if should_drop {
482- drop_database_if_exists ( & target_client , & db_info. name ) . await ?;
481+ drop_database_if_exists ( target_url , & db_info. name ) . await ?;
483482
484483 // Recreate the database
485484 let create_query = format ! (
@@ -768,13 +767,24 @@ fn prompt_drop_database(db_name: &str) -> Result<bool> {
768767}
769768
770769/// Drops a database if it exists
771- async fn drop_database_if_exists ( target_conn : & Client , db_name : & str ) -> Result < ( ) > {
770+ ///
771+ /// NOTE: This function connects to the `postgres` database to issue the DROP command,
772+ /// because PostgreSQL does not allow dropping the currently connected database.
773+ async fn drop_database_if_exists ( target_url : & str , db_name : & str ) -> Result < ( ) > {
772774 // Validate database name to prevent SQL injection
773775 crate :: utils:: validate_postgres_identifier ( db_name)
774776 . with_context ( || format ! ( "Invalid database name: '{}'" , db_name) ) ?;
775777
776778 tracing:: info!( " Dropping existing database '{}'..." , db_name) ;
777779
780+ // Connect to 'postgres' database to issue DROP command
781+ // (PostgreSQL doesn't allow dropping the currently connected database)
782+ let admin_url = replace_database_in_url ( target_url, "postgres" )
783+ . context ( "Failed to construct admin connection URL" ) ?;
784+ let admin_conn = postgres:: connection:: connect_with_retry ( & admin_url)
785+ . await
786+ . context ( "Failed to connect to postgres database for DROP operation" ) ?;
787+
778788 // Terminate existing connections to the database
779789 // Skip connections owned by SUPERUSER roles (we can't terminate those on managed PostgreSQL)
780790 let terminate_query = "
@@ -785,7 +795,7 @@ async fn drop_database_if_exists(target_conn: &Client, db_name: &str) -> Result<
785795 AND sa.pid <> pg_backend_pid()
786796 AND NOT r.rolsuper
787797 " ;
788- target_conn . execute ( terminate_query, & [ & db_name] ) . await ?;
798+ admin_conn . execute ( terminate_query, & [ & db_name] ) . await ?;
789799
790800 // Check if any connections remain (including SUPERUSER connections we couldn't terminate)
791801 let remaining_query = "
@@ -794,7 +804,7 @@ async fn drop_database_if_exists(target_conn: &Client, db_name: &str) -> Result<
794804 WHERE sa.datname = $1
795805 AND sa.pid <> pg_backend_pid()
796806 " ;
797- let row = target_conn
807+ let row = admin_conn
798808 . query_one ( remaining_query, & [ & db_name] )
799809 . await
800810 . context ( "Failed to check remaining connections" ) ?;
@@ -825,7 +835,7 @@ async fn drop_database_if_exists(target_conn: &Client, db_name: &str) -> Result<
825835 "DROP DATABASE IF EXISTS {}" ,
826836 crate :: utils:: quote_ident( db_name)
827837 ) ;
828- target_conn
838+ admin_conn
829839 . execute ( & drop_query, & [ ] )
830840 . await
831841 . with_context ( || format ! ( "Failed to drop database '{}'" , db_name) ) ?;
0 commit comments