feat(desktop): show batch drop progress

crates/dbx-core/src/query.rs

@@ -78,6 +78,8 @@ pub struct ExecuteMultiResult {
     pub statement_index: Option<usize>,
 }
 
+pub type ExecuteMultiProgressCallback = Arc<dyn Fn(usize, usize, bool) + Send + Sync>;
+
 impl ExecuteMultiResult {
     fn execution_error(result: db::QueryResult) -> Self {
         Self { result, execution_error: true, statement_index: None }
@@ -2019,6 +2021,30 @@ pub async fn execute_multi_core_with_options_for_client(
     schema: Option<&str>,
     cancel_token: Option<CancellationToken>,
     options: QueryExecutionOptions,
+) -> Result<Vec<ExecuteMultiResult>, String> {
+    execute_multi_core_with_options_for_client_and_progress(
+        state,
+        connection_id,
+        database,
+        sql,
+        schema,
+        cancel_token,
+        options,
+        None,
+    )
+    .await
+}
+
+/// Executes a SQL batch and reports each completed statement to the optional callback.
+pub async fn execute_multi_core_with_options_for_client_and_progress(
+    state: &AppState,
+    connection_id: &str,
+    database: &str,
+    sql: &str,
+    schema: Option<&str>,
+    cancel_token: Option<CancellationToken>,
+    options: QueryExecutionOptions,
+    progress: Option<ExecuteMultiProgressCallback>,
 ) -> Result<Vec<ExecuteMultiResult>, String> {
     // Reject MongoDB queries that fall through to the generic executor.
     if connection_is_mongodb(state, connection_id).await {
@@ -2122,6 +2148,7 @@ pub async fn execute_multi_core_with_options_for_client(
             &statements,
             cancel_token,
             options,
+            progress.as_ref(),
         )
         .await;
     }
@@ -2143,10 +2170,18 @@ pub async fn execute_multi_core_with_options_for_client(
         )
         .await
         {
-            Ok(r) => results.push(ExecuteMultiResult::success_with_index(r, statement_index)),
+            Ok(r) => {
+                results.push(ExecuteMultiResult::success_with_index(r, statement_index));
+                if let Some(progress) = progress.as_ref() {
+                    progress(statement_index + 1, statements.len(), true);
+                }
+            }
             Err(e) => {
                 let action = query_pool_error_action(db_type, stmt, &e);
                 results.push(ExecuteMultiResult::execution_error_with_index(error_query_result(e), statement_index));
+                if let Some(progress) = progress.as_ref() {
+                    progress(statement_index + 1, statements.len(), false);
+                }
                 if !should_continue_batch_after_error(options.continue_on_error, action) {
                     break;
                 }
@@ -2193,6 +2228,7 @@ async fn execute_mysql_batch_statements<E>(
     db_type: Option<DatabaseType>,
     cancel_token: Option<CancellationToken>,
     continue_on_error: bool,
+    progress: Option<&ExecuteMultiProgressCallback>,
 ) -> (Vec<ExecuteMultiResult>, Option<PoolErrorAction>)
 where
     E: MysqlBatchStatementExecutor,
@@ -2205,10 +2241,18 @@ where
         }
 
         match executor.execute_statement(statement).await {
-            Ok(result) => results.push(ExecuteMultiResult::success_with_index(result, statement_index)),
+            Ok(result) => {
+                results.push(ExecuteMultiResult::success_with_index(result, statement_index));
+                if let Some(progress) = progress {
+                    progress(statement_index + 1, statements.len(), true);
+                }
+            }
             Err(err) => {
                 let action = pool_error_action(db_type, &err);
                 results.push(ExecuteMultiResult::execution_error_with_index(error_query_result(err), statement_index));
+                if let Some(progress) = progress {
+                    progress(statement_index + 1, statements.len(), false);
+                }
                 // Statement errors are safe to collect, but connection-level failures leave
                 // the protocol state unusable and must still trigger pool cleanup.
                 if !should_continue_batch_after_error(continue_on_error, action) {
@@ -2231,6 +2275,7 @@ async fn execute_multi_mysql(
     statements: &[String],
     cancel_token: Option<CancellationToken>,
     options: QueryExecutionOptions,
+    progress: Option<&ExecuteMultiProgressCallback>,
 ) -> Result<Vec<ExecuteMultiResult>, String> {
     let query_timeout = resolve_query_timeout(options.timeout_secs);
     let operation_budget = operation_budget_for_pool_key(state, pool_key, query_timeout).await;
@@ -2264,9 +2309,15 @@ async fn execute_multi_mysql(
         max_rows,
         dialect,
     };
-    let (results, error_action) =
-        execute_mysql_batch_statements(&mut executor, statements, db_type, cancel_token, options.continue_on_error)
-            .await;
+    let (results, error_action) = execute_mysql_batch_statements(
+        &mut executor,
+        statements,
+        db_type,
+        cancel_token,
+        options.continue_on_error,
+        progress,
+    )
+    .await;
     drop(executor);
 
     if matches!(error_action, Some(PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry)) {
@@ -3769,7 +3820,8 @@ mod tests {
         };
 
         let (results, error_action) =
-            execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, false).await;
+            execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, false, None)
+                .await;
 
         assert_eq!(executor.executed, vec!["first", "fails"]);
         assert_eq!(results.len(), 2);
@@ -3779,6 +3831,35 @@ mod tests {
         assert_eq!(error_action, Some(PoolErrorAction::Keep));
     }
 
+    #[tokio::test]
+    async fn mysql_batch_reports_progress_for_each_completed_statement() {
+        let statements = vec!["first".to_string(), "fails".to_string(), "must-not-run".to_string()];
+        let mut executor = FakeMysqlBatchExecutor {
+            outcomes: std::collections::VecDeque::from([Ok(empty_query_result(0)), Err("Duplicate entry".to_string())]),
+            executed: Vec::new(),
+        };
+        let progress_events = Arc::new(std::sync::Mutex::new(Vec::new()));
+        let progress: ExecuteMultiProgressCallback = {
+            let progress_events = Arc::clone(&progress_events);
+            Arc::new(move |completed, total, success| progress_events.lock().unwrap().push((completed, total, success)))
+        };
+
+        let (results, error_action) = execute_mysql_batch_statements(
+            &mut executor,
+            &statements,
+            Some(DatabaseType::Mysql),
+            None,
+            false,
+            Some(&progress),
+        )
+        .await;
+
+        assert_eq!(executor.executed, vec!["first", "fails"]);
+        assert_eq!(results.len(), 2);
+        assert_eq!(*progress_events.lock().unwrap(), vec![(1, 3, true), (2, 3, false)]);
+        assert_eq!(error_action, Some(PoolErrorAction::Keep));
+    }
+
     #[tokio::test]
     async fn mysql_batch_stops_when_the_first_statement_fails() {
         let statements = vec!["fails".to_string(), "must-not-run".to_string()];
@@ -3788,7 +3869,8 @@ mod tests {
         };
 
         let (results, error_action) =
-            execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, false).await;
+            execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, false, None)
+                .await;
 
         assert_eq!(executor.executed, vec!["fails"]);
         assert_eq!(results.len(), 1);
@@ -3809,7 +3891,8 @@ mod tests {
         };
 
         let (results, error_action) =
-            execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, true).await;
+            execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, true, None)
+                .await;
 
         assert_eq!(executor.executed, statements);
         assert_eq!(results.len(), 3);
@@ -3830,7 +3913,8 @@ mod tests {
         };
 
         let (results, error_action) =
-            execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, true).await;
+            execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, true, None)
+                .await;
 
         assert_eq!(executor.executed, statements);
         assert_eq!(results.len(), 2);
@@ -3851,7 +3935,8 @@ mod tests {
         };
 
         let (results, error_action) =
-            execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, true).await;
+            execute_mysql_batch_statements(&mut executor, &statements, Some(DatabaseType::Mysql), None, true, None)
+                .await;
 
         assert_eq!(executor.executed, vec!["first", "disconnects"]);
         assert_eq!(results.len(), 2);