Skip to content

Commit d282c13

Browse files
committed
Fix connection hang in database_is_empty() for SerenDB/Neon targets
Reuse the existing target_client connection instead of creating a new connection when checking if a database is empty. This prevents connection pool exhaustion on serverless PostgreSQL providers (SerenDB, Neon) that have strict concurrent connection limits. The previous implementation would create a second connection while the first was still open, causing indefinite hangs when the connection pool was exhausted. Since tokio_postgres::connect has no built-in timeout, the operation would hang forever. Fixes #23
1 parent 69f37c6 commit d282c13

1 file changed

Lines changed: 12 additions & 8 deletions

File tree

src/commands/init.rs

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -405,8 +405,8 @@ pub async fn init(
405405
db_info.name
406406
);
407407

408-
// Check if empty
409-
if database_is_empty(target_url, &db_info.name).await? {
408+
// Check if empty (reuse existing connection to avoid pool exhaustion)
409+
if database_is_empty(&target_client).await? {
410410
tracing::info!(
411411
" Database '{}' is empty, proceeding with restore",
412412
db_info.name
@@ -677,11 +677,10 @@ fn confirm_replication(sizes: &[migration::DatabaseSizeInfo]) -> Result<bool> {
677677
}
678678

679679
/// Checks if a database is empty (no user tables)
680-
async fn database_is_empty(target_url: &str, db_name: &str) -> Result<bool> {
681-
// Need to connect to the specific database to check tables
682-
let db_url = replace_database_in_url(target_url, db_name)?;
683-
let client = postgres::connect_with_retry(&db_url).await?;
684-
680+
///
681+
/// Uses the existing connection to avoid connection pool exhaustion on
682+
/// serverless PostgreSQL providers (SerenDB, Neon) that have strict limits.
683+
async fn database_is_empty(client: &tokio_postgres::Client) -> Result<bool> {
685684
let query = "
686685
SELECT COUNT(*)
687686
FROM information_schema.tables
@@ -1223,9 +1222,14 @@ mod tests {
12231222
async fn test_database_is_empty() {
12241223
let url = std::env::var("TEST_TARGET_URL").expect("TEST_TARGET_URL not set");
12251224

1225+
// Connect to the database first
1226+
let client = crate::postgres::connect_with_retry(&url)
1227+
.await
1228+
.expect("Failed to connect");
1229+
12261230
// postgres database might be empty of user tables
12271231
// This test just verifies the function doesn't crash
1228-
let result = database_is_empty(&url, "postgres").await;
1232+
let result = database_is_empty(&client).await;
12291233
assert!(result.is_ok());
12301234
}
12311235
}

0 commit comments

Comments
 (0)