fix(mongodb): resolve per-database pool for dump restore and export

crates/dbx-core/src/data/mongodb_dump.rs

@@ -406,10 +406,12 @@ fn command_cursor_id(cursor: &Document) -> Result<i64, String> {
 
 async fn dump_client(state: &AppState, id: &str, database: &str) -> Result<DumpClient, String> {
     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.

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 `<connection id>:<database>`, 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<PoolKind, String> {
+    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<Document>,
 ) -> Result<MongoInsertOutcome, MongoImportIssue> {
-    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<Box<dyn std::future::Future<Output = bool> + 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 {

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 `<connection id>:<database>`. 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::<u16>().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;
+}