Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 11 commits into from
Nov 21, 2023
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(connect_opts) {
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.database_url).await?;
if force {
Any::force_drop_database(&connect_opts.database_url).await?;
} else {
Any::drop_database(&connect_opts.database_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
Loading