diff --git a/crates/dbx-core/src/data/mongodb_dump.rs b/crates/dbx-core/src/data/mongodb_dump.rs index a55e1d93cc..e8943d6a04 100644 --- a/crates/dbx-core/src/data/mongodb_dump.rs +++ b/crates/dbx-core/src/data/mongodb_dump.rs @@ -406,10 +406,12 @@ fn command_cursor_id(cursor: &Document) -> Result { async fn dump_client(state: &AppState, id: &str, database: &str) -> Result { metadata::validate_database(database)?; - state.get_or_create_pool(id, Some(database)).await?; - match state.pool_handle(id).await { - Some(PoolKind::MongoDb(client)) => Ok(DumpClient::Native(client.clone())), - Some(PoolKind::Agent(client)) => { + // Propagate pool resolution errors (e.g. connection establishment failures) instead of + // swallowing them into the generic message below, so users see the real cause. + let pool = crate::mongodb_import_export::mongo_pool_for_database(state, id, database).await?; + match pool { + PoolKind::MongoDb(client) => Ok(DumpClient::Native(client.clone())), + PoolKind::Agent(client) => { // Check what every dump, restore, or preview shares up front, so an outdated // agent is refused before any work starts rather than mid-operation. The find // cursor only the dump's BSON export relies on is gated on that path instead. diff --git a/crates/dbx-core/src/data/mongodb_import_export.rs b/crates/dbx-core/src/data/mongodb_import_export.rs index 68a266e84a..bfeabb9105 100644 --- a/crates/dbx-core/src/data/mongodb_import_export.rs +++ b/crates/dbx-core/src/data/mongodb_import_export.rs @@ -1669,6 +1669,22 @@ fn progress( } } +/// Resolve the pool that serves `database` on this connection. +/// +/// MongoDB pools are keyed `:`, so looking one up by the bare +/// connection id misses the pool `get_or_create_pool` just returned. Depending on what else the +/// app has opened, the bare key is either absent — "Not a MongoDB connection" on a perfectly good +/// connection — or holds a *different*, possibly already closed, pool, which surfaces later as +/// "Agent stdin not available". Always resolve through the key the pool was created under. +pub(crate) async fn mongo_pool_for_database( + state: &AppState, + connection_id: &str, + database: &str, +) -> Result { + let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; + state.pool_handle(&pool_key).await.ok_or_else(|| "Not found".to_string()) +} + async fn insert_documents_batch( state: &AppState, connection_id: &str, @@ -1676,8 +1692,9 @@ async fn insert_documents_batch( collection: &str, documents: Vec, ) -> Result { - let pool = - state.pool_handle(connection_id).await.ok_or_else(|| MongoImportIssue::new("CONNECTION", "Not found"))?; + let pool = mongo_pool_for_database(state, connection_id, database) + .await + .map_err(|error| MongoImportIssue::new("CONNECTION", error))?; match &pool { PoolKind::MongoDb(client) => { insert_bson_documents(client, database, collection, documents).await.map_err(bulk_write_issue) @@ -2339,8 +2356,7 @@ where return Err("Export cancelled".to_string()); } - state.get_or_create_pool(&request.connection_id, Some(&request.database)).await?; - let pool = state.pool_handle(&request.connection_id).await.ok_or_else(|| "Not found".to_string())?; + let pool = mongo_pool_for_database(state, &request.connection_id, &request.database).await?; match &pool { PoolKind::MongoDb(_) | PoolKind::Agent(_) => {} _ => return Err("Not a MongoDB connection".to_string()), @@ -2420,7 +2436,7 @@ async fn count_export_documents(state: &AppState, request: &MongoExportRequest) if request.filter.as_deref().is_some_and(|filter| !filter.trim().is_empty() && filter.trim() != "{}") { return None; } - let pool = state.pool_handle(&request.connection_id).await?; + let pool = mongo_pool_for_database(state, &request.connection_id, &request.database).await.ok()?; match pool { PoolKind::MongoDb(client) => mongo_driver::count_documents( &client, @@ -2457,7 +2473,7 @@ where C: FnMut(&str) -> std::pin::Pin + Send>>, F: FnMut(serde_json::Value) -> Result<(), String>, { - let pool = state.pool_handle(&request.connection_id).await.ok_or_else(|| "Not found".to_string())?; + let pool = mongo_pool_for_database(state, &request.connection_id, &request.database).await?; match pool { PoolKind::MongoDb(client) => { for_each_find_document( @@ -2849,7 +2865,7 @@ where if is_cancelled(&request.export_id).await { return Err("Export cancelled".to_string()); } - let pool = state.pool_handle(&request.connection_id).await.ok_or_else(|| "Not found".to_string())?; + let pool = mongo_pool_for_database(state, &request.connection_id, &request.database).await?; let file = File::create(temp).map_err(|error| error.to_string())?; let buffered = BufWriter::new(file); let mut writer = if request.gzip { diff --git a/crates/dbx-core/tests/live_mongodb_legacy_agent_commands.rs b/crates/dbx-core/tests/live_mongodb_legacy_agent_commands.rs index 5bc6084b90..4a74b0b627 100644 --- a/crates/dbx-core/tests/live_mongodb_legacy_agent_commands.rs +++ b/crates/dbx-core/tests/live_mongodb_legacy_agent_commands.rs @@ -285,3 +285,87 @@ async fn collection_stats_and_create_database_run_over_the_legacy_agent() { command(&state, id, &created, doc! { "dropDatabase": 1 }).await; command(&state, id, &database, doc! { "dropDatabase": 1 }).await; } + +/// MongoDB pools are keyed `:`. Dump, restore and collection +/// import/export used to resolve the pool by the bare connection id instead, which only worked +/// while some earlier call happened to have left a pool under that key. +#[tokio::test] +#[ignore = "opt-in: DBX_MONGO_LEGACY_DUMP_TEST_HOST (host:port, MongoDB 3.6+ without auth) and an installed MongoDB Legacy Agent; creates a temporary database"] +async fn dump_and_export_resolve_the_per_database_pool() { + use dbx_core::mongodb_dump::*; + + let endpoint = std::env::var("DBX_MONGO_LEGACY_DUMP_TEST_HOST").expect("DBX_MONGO_LEGACY_DUMP_TEST_HOST"); + let (host, port) = endpoint.split_once(':').expect("host:port"); + let files = tempfile::tempdir().unwrap(); + let database = format!("dbx_legacy_pool_{}", uuid::Uuid::new_v4().simple()); + let state = AppState::new(Storage::open(&files.path().join("storage.db")).await.unwrap()); + let id = "legacy-pool-key-test"; + let config: ConnectionConfig = serde_json::from_value(serde_json::json!({ "id": id, "name": "Legacy pool key test", "db_type": "mongodb", "host": host, "port": port.parse::().unwrap(), "username": "", "password": "", "database": database, "driver_profile": "mongodb-legacy" })).unwrap(); + state.configs.write().await.insert(id.into(), config); + + command( + &state, + id, + &database, + doc! { "insert": "orders", "documents": [ { "_id": 1, "n": 1 }, { "_id": 2, "n": 2 } ] }, + ) + .await; + + // Seeding went through the bare-connection pool. Drop it, leaving only the per-database pool + // the dump itself creates — the state the app is normally in. + state.remove_pool_by_key(id).await; + let key = state.get_or_create_pool(id, Some(&database)).await.unwrap(); + assert_eq!(key, format!("{id}:{database}"), "MongoDB pools are keyed per database"); + assert!( + state.pool_handle(id).await.is_none(), + "the bare-connection pool must be gone for this test to mean anything" + ); + + let catalog = inspect_mongodb_database_dump(&state, id, &database).await.unwrap(); + assert_eq!(catalog.collections.len(), 1, "{catalog:?}"); + + let output = files.path().join("pool-key.archive"); + let dump = MongoDatabaseDumpRequest { + task_id: uuid::Uuid::new_v4().to_string(), + connection_id: id.into(), + database: database.clone(), + file_path: output.to_str().unwrap().into(), + format: MongoDumpFormat::Archive, + gzip: false, + collections: None, + }; + let dumped = dump_mongodb_database(&state, &dump, |_| Box::pin(async { false }), |_| {}).await.unwrap(); + assert_eq!(dumped.collections_done, 1); + assert_eq!(dumped.documents_read, 2); + + // Restore writes through the same lookup. + let preview = prepare_mongodb_restore_source(MongoRestoreSourceRequest { + path: output.to_str().unwrap().into(), + format: MongoDumpFormat::Archive, + gzip: false, + }) + .await + .unwrap(); + let restored = format!("{database}_restored"); + let restore = MongoDatabaseRestoreRequest { + task_id: uuid::Uuid::new_v4().to_string(), + connection_id: id.into(), + database: restored.clone(), + source_database: database.clone(), + source_ref: preview.source_ref.clone(), + collections: None, + drop_existing: false, + restore_options: true, + restore_indexes: true, + stop_on_error: true, + objcheck: false, + batch_size: 500, + execution_id: None, + }; + let result = restore_mongodb_database(&state, &restore, |_| Box::pin(async { false }), |_| {}).await.unwrap(); + assert_eq!(result.documents_written, 2); + assert!(release_mongodb_restore_source(&preview.source_ref)); + + command(&state, id, &restored, doc! { "dropDatabase": 1 }).await; + command(&state, id, &database, doc! { "dropDatabase": 1 }).await; +}