Skip to content

Commit

Permalink
Sqlx Cli: Added force flag to drop database for postgres (#2873)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
Vrajs16 authored Nov 21, 2023
1 parent ed1b030 commit fda4159
Show file tree
Hide file tree
Showing 10 changed files with 78 additions and 15 deletions.
18 changes: 9 additions & 9 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

11 changes: 8 additions & 3 deletions sqlx-cli/src/database.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(());
}
Expand All @@ -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(())
Expand All @@ -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
}

Expand Down
6 changes: 4 additions & 2 deletions sqlx-cli/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
8 changes: 8 additions & 0 deletions sqlx-cli/src/opt.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.
Expand Down
2 changes: 1 addition & 1 deletion sqlx-core/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
6 changes: 6 additions & 0 deletions sqlx-core/src/any/driver.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>()
}
Expand All @@ -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 {
Expand All @@ -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.
Expand Down
9 changes: 9 additions & 0 deletions sqlx-core/src/any/migrate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
3 changes: 3 additions & 0 deletions sqlx-core/src/migrate/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
6 changes: 6 additions & 0 deletions sqlx-core/src/migrate/migrate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
24 changes: 24 additions & 0 deletions sqlx-postgres/src/migrate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down

0 comments on commit fda4159

Please sign in to comment.