Skip to content

Commit

Permalink
schema-engine: try force drop in shadow databases (#4722)
Browse files Browse the repository at this point in the history
  • Loading branch information
skyzh committed Feb 15, 2024
1 parent 58fef65 commit 3d78035
Showing 1 changed file with 38 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,10 @@ impl PostgresFlavour {
.unwrap_or(false)
}

pub(crate) fn is_postgres(&self) -> bool {
self.provider == PostgresProvider::PostgreSql && !self.is_cockroachdb()
}

pub(crate) fn schema_name(&self) -> &str {
self.state.params().map(|p| p.url.schema()).unwrap_or("public")
}
Expand Down Expand Up @@ -430,6 +434,7 @@ impl SqlFlavour for PostgresFlavour {
shadow_db::sql_schema_from_migrations_history(migrations, shadow_database, namespaces).await
}),
None => {
let is_postgres = self.is_postgres();
with_connection(self, move |params, _circumstances, main_connection| async move {
let shadow_database_name = crate::new_shadow_database_name();

Expand Down Expand Up @@ -462,8 +467,12 @@ impl SqlFlavour for PostgresFlavour {
let ret =
shadow_db::sql_schema_from_migrations_history(migrations, shadow_database, namespaces).await;

let drop_database = format!("DROP DATABASE IF EXISTS \"{shadow_database_name}\"");
main_connection.raw_cmd(&drop_database, &params.url).await?;
if is_postgres {
drop_db_try_force(main_connection, &params.url, &shadow_database_name).await?;
} else {
let drop_database = format!("DROP DATABASE IF EXISTS \"{shadow_database_name}\"");
main_connection.raw_cmd(&drop_database, &params.url).await?;
}

ret
})
Expand All @@ -482,6 +491,33 @@ impl SqlFlavour for PostgresFlavour {
}
}

/// Drop a database using `WITH (FORCE)` syntax.
///
/// When drop database is routed through pgbouncer, the database may still be used in other pooled connections.
/// In this case, given that we (as a user) know the database will not be used any more, we can forcefully drop
/// the database. Note that `with (force)` is added in Postgres 13, and therefore we will need to
/// fallback to the normal drop if it errors with syntax error.
///
/// TL;DR,
/// 1. pg >= 13 -> it works.
/// 2. pg < 13 -> syntax error on WITH (FORCE), and then fail with db in use if pgbouncer is used.
async fn drop_db_try_force(conn: &mut Connection, url: &PostgresUrl, database_name: &str) -> ConnectorResult<()> {
let drop_database = format!("DROP DATABASE IF EXISTS \"{database_name}\" WITH (FORCE)");
if let Err(err) = conn.raw_cmd(&drop_database, url).await {
if let Some(msg) = err.message() {
if msg.contains("syntax error") {
let drop_database_alt = format!("DROP DATABASE IF EXISTS \"{database_name}\"");
conn.raw_cmd(&drop_database_alt, url).await?;
} else {
return Err(err);
}
} else {
return Err(err);
}
}
Ok(())
}

fn strip_schema_param_from_url(url: &mut Url) {
let mut params: HashMap<String, String> = url.query_pairs().into_owned().collect();
params.remove("schema");
Expand Down

0 comments on commit 3d78035

Please sign in to comment.