fix(mysql): stream query result exports

crates/dbx-core/src/db/mysql.rs

@@ -44,6 +44,11 @@ impl MySqlQueryDialect {
     }
 }
 
+pub enum MySqlQueryStreamItem {
+    Columns { columns: Vec<String>, column_types: Vec<String> },
+    Row(Vec<serde_json::Value>),
+}
+
 fn quote_value(s: &str) -> String {
     format!("'{}'", s.replace('\\', "\\\\").replace('\'', "\\'"))
 }
@@ -2481,29 +2486,55 @@ pub async fn stream_query_rows(
     mut on_row: impl FnMut(&[serde_json::Value]) -> Result<(), String>,
 ) -> Result<u64, String> {
     let mut conn = get_conn_with_health_check(pool).await?;
+    stream_query_result_on_conn(&mut conn, sql, bare, max_rows, dialect, cancelled, |item| {
+        if let MySqlQueryStreamItem::Row(row) = item {
+            on_row(&row)?;
+        }
+        Ok(())
+    })
+    .await
+}
+
+pub async fn stream_query_result_on_conn(
+    conn: &mut mysql_async::Conn,
+    sql: &str,
+    bare: bool,
+    max_rows: Option<usize>,
+    dialect: MySqlQueryDialect,
+    cancelled: &AtomicBool,
+    mut on_item: impl FnMut(MySqlQueryStreamItem) -> Result<(), String>,
+) -> Result<u64, String> {
     let row_limit = max_rows.unwrap_or(usize::MAX);
 
     if bare || prefers_text_protocol_query(sql, dialect) {
-        stream_query_rows_text(&mut conn, sql, row_limit, cancelled, &mut on_row).await
+        stream_query_result_text(conn, sql, row_limit, cancelled, &mut on_item).await
     } else {
-        match stream_query_rows_prepared(&mut conn, sql, row_limit, cancelled, &mut on_row).await {
+        match stream_query_result_prepared(conn, sql, row_limit, cancelled, &mut on_item).await {
             Ok(rows) => Ok(rows),
             Err(err) if mysql_error_should_retry_with_text_protocol(&err) => {
-                stream_query_rows_text(&mut conn, sql, row_limit, cancelled, &mut on_row).await
+                stream_query_result_text(conn, sql, row_limit, cancelled, &mut on_item).await
             }
             Err(err) => Err(err),
         }
     }
 }
 
-async fn stream_query_rows_text(
+async fn stream_query_result_text(
     conn: &mut mysql_async::Conn,
     sql: &str,
     row_limit: usize,
     cancelled: &AtomicBool,
-    on_row: &mut impl FnMut(&[serde_json::Value]) -> Result<(), String>,
+    on_item: &mut impl FnMut(MySqlQueryStreamItem) -> Result<(), String>,
 ) -> Result<u64, String> {
     let mut result = conn.query_iter(sql).await.map_err(|e| e.to_string())?;
+    if !advance_to_result_set_with_columns(&mut result).await? {
+        return Ok(0);
+    }
+    let columns: Vec<String> = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect();
+    let column_types: Vec<String> =
+        result.columns_ref().iter().map(|c| mysql_column_type_name(c.column_type())).collect();
+    on_item(MySqlQueryStreamItem::Columns { columns, column_types })?;
+
     let mut stream = result
         .stream::<mysql_async::Row>()
         .await
@@ -2520,21 +2551,29 @@ async fn stream_query_rows_text(
         }
         let row = row.map_err(|e| e.to_string())?;
         let values: Vec<serde_json::Value> = (0..row.len()).map(|i| mysql_value_to_json(&row, i)).collect();
-        on_row(&values)?;
+        on_item(MySqlQueryStreamItem::Row(values))?;
         rows_exported += 1;
     }
 
     Ok(rows_exported)
 }
 
-async fn stream_query_rows_prepared(
+async fn stream_query_result_prepared(
     conn: &mut mysql_async::Conn,
     sql: &str,
     row_limit: usize,
     cancelled: &AtomicBool,
-    on_row: &mut impl FnMut(&[serde_json::Value]) -> Result<(), String>,
+    on_item: &mut impl FnMut(MySqlQueryStreamItem) -> Result<(), String>,
 ) -> Result<u64, String> {
     let mut result = conn.exec_iter(sql, ()).await.map_err(|e| e.to_string())?;
+    let columns: Vec<String> = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect();
+    if columns.is_empty() {
+        return Ok(0);
+    }
+    let column_types: Vec<String> =
+        result.columns_ref().iter().map(|c| mysql_column_type_name(c.column_type())).collect();
+    on_item(MySqlQueryStreamItem::Columns { columns, column_types })?;
+
     let mut stream = result
         .stream::<mysql_async::Row>()
         .await
@@ -2551,7 +2590,7 @@ async fn stream_query_rows_prepared(
         }
         let row = row.map_err(|e| e.to_string())?;
         let values: Vec<serde_json::Value> = (0..row.len()).map(|i| mysql_value_to_json(&row, i)).collect();
-        on_row(&values)?;
+        on_item(MySqlQueryStreamItem::Row(values))?;
         rows_exported += 1;
     }
 

crates/dbx-core/src/query_result_export.rs

@@ -1,6 +1,10 @@
 use serde::{Deserialize, Serialize};
 use std::fs::File;
 use std::io::{BufWriter, Write};
+use std::sync::{
+    atomic::{AtomicBool, Ordering},
+    Arc,
+};
 use std::time::{Duration, Instant};
 
 use crate::connection::{AppState, PoolKind};
@@ -323,6 +327,12 @@ async fn export_query_result_core_inner(
         return Ok(());
     }
 
+    // MySQL does not guarantee a stable row order for independent LIMIT/OFFSET
+    // executions without ORDER BY, so query-result export must stream one run.
+    if try_export_mysql_query_result_stream(state, request, &format, cancel_token.clone(), on_progress).await? {
+        return Ok(());
+    }
+
     let mut csv_file = if format == "csv" {
         Some(BufWriter::new(File::create(&request.file_path).map_err(|e| format!("Failed to create file: {e}"))?))
     } else {
@@ -705,6 +715,241 @@ async fn try_export_postgres_query_result_stream(
     Ok(true)
 }
 
+async fn try_export_mysql_query_result_stream(
+    state: &AppState,
+    request: &QueryResultExportRequest,
+    format: &str,
+    cancel_token: Option<CancellationToken>,
+    on_progress: &impl Fn(TableExportProgress),
+) -> Result<bool, String> {
+    if request.use_agent_cursor {
+        return Ok(false);
+    }
+
+    let pool_key = if request.database.trim().is_empty() {
+        state.get_or_create_pool_for_session(&request.connection_id, None, request.client_session_id.as_deref()).await?
+    } else {
+        state
+            .get_or_create_pool_for_session(
+                &request.connection_id,
+                Some(request.database.as_str()),
+                request.client_session_id.as_deref(),
+            )
+            .await?
+    };
+    let connections = state.connections.read().await;
+    let Some((pool, bare)) = connections.get(&pool_key).and_then(|pool| match pool {
+        PoolKind::Mysql(pool, mode) => Some((pool.clone(), *mode == crate::connection::MysqlMode::Bare)),
+        _ => None,
+    }) else {
+        return Ok(false);
+    };
+    drop(connections);
+
+    if let Some(execution_id) = request.execution_id.as_deref() {
+        state.running_queries.set_pool_key(execution_id, pool_key.clone());
+    }
+    state.touch_pool_activity(&pool_key).await;
+    let _activity_touch = state.pool_activity_touch(&pool_key);
+
+    let (mysql_dialect, read_only_connection_name) = {
+        let configs = state.configs.read().await;
+        let config = configs.get(&request.connection_id);
+        (
+            config
+                .map(|config| {
+                    crate::db::mysql::MySqlQueryDialect::for_connection(
+                        config.db_type,
+                        config.driver_profile.as_deref(),
+                    )
+                })
+                .unwrap_or_default(),
+            config.filter(|config| config.read_only).map(|config| config.name.clone()),
+        )
+    };
+    if let Some(name) = read_only_connection_name {
+        crate::query_execution_sql::check_read_only(&request.sql, &name)?;
+    }
+
+    let xlsx_hard_limit_active = xlsx_hard_limit_active(format, request);
+    let row_limit = effective_row_limit(format, request);
+    let stream_row_limit =
+        if xlsx_hard_limit_active { row_limit.map(|limit| limit.saturating_add(1)) } else { row_limit };
+    let progress_row_interval = request.page_size.max(1) as u64;
+    let mut columns: Vec<String> = Vec::new();
+    let mut rows_exported = 0_u64;
+    let mut last_progress_rows = 0_u64;
+    let mut last_progress_at = Instant::now();
+    let mut csv_file = if format == "csv" {
+        let mut file =
+            BufWriter::new(File::create(&request.file_path).map_err(|e| format!("Failed to create file: {e}"))?);
+        file.write_all(b"\xEF\xBB\xBF").map_err(|e| format!("Failed to write BOM: {e}"))?;
+        Some(file)
+    } else {
+        None
+    };
+    let mut xlsx = None;
+    let query_timeout = query_export_timeout(request.timeout_secs);
+    let operation_budget = operation_budget_for_pool_key(state, &pool_key, query_timeout).await;
+    let mut conn = crate::db::mysql::get_conn_with_health_check_with_cancel(
+        &pool,
+        operation_budget.checkout_timeout,
+        operation_budget.cleanup_timeout,
+        cancel_token.as_ref(),
+    )
+    .await?;
+    let mysql_connection_id = conn.id();
+    let kill_opts = conn.opts().clone();
+    if let Some(execution_id) = request.execution_id.clone() {
+        let interrupt_kill_opts = kill_opts.clone();
+        state.running_queries.register_interrupt(&execution_id, move || {
+            let kill_opts = interrupt_kill_opts.clone();
+            tokio::spawn(async move {
+                if let Err(error) = crate::db::mysql::kill_query_with_opts(kill_opts, mysql_connection_id).await {
+                    log::warn!("Failed to cancel MySQL export query {mysql_connection_id}: {error}");
+                }
+            });
+        });
+    }
+
+    let export_cancelled = Arc::new(AtomicBool::new(false));
+    let watcher_done = CancellationToken::new();
+    let watcher_done_task = watcher_done.clone();
+    let export_cancelled_task = export_cancelled.clone();
+    let export_id = request.export_id.clone();
+    let cancel_for_watcher = cancel_token.clone();
+    // Normal UI cancellation uses running query cancellation and KILL QUERY.
+    // This covers callers that only set the export-cancelled flag.
+    tokio::spawn(async move {
+        loop {
+            tokio::select! {
+                biased;
+                _ = watcher_done_task.cancelled() => break,
+                _ = tokio::time::sleep(Duration::from_millis(100)) => {}
+            }
+            if is_export_cancelled(&export_id).await {
+                export_cancelled_task.store(true, Ordering::SeqCst);
+                if let Some(token) = cancel_for_watcher.as_ref() {
+                    token.cancel();
+                }
+                break;
+            }
+        }
+    });
+
+    let stream_future = crate::db::mysql::stream_query_result_on_conn(
+        &mut conn,
+        &request.sql,
+        bare,
+        stream_row_limit,
+        mysql_dialect,
+        &export_cancelled,
+        |item| {
+            if export_cancelled.load(Ordering::SeqCst)
+                || cancel_token.as_ref().is_some_and(|token| token.is_cancelled())
+            {
+                return Err(canceled_error());
+            }
+            match item {
+                crate::db::mysql::MySqlQueryStreamItem::Columns { columns: stream_columns, .. } => {
+                    columns = stream_columns;
+                    if let Some(file) = csv_file.as_mut() {
+                        let csv = format_query_result_csv(&columns, &[]);
+                        let header = csv.strip_suffix('\n').unwrap_or(&csv);
+                        file.write_all(header.as_bytes()).map_err(|e| format!("Failed to write CSV: {e}"))?;
+                    } else {
+                        let xlsx_file =
+                            File::create(&request.file_path).map_err(|e| format!("Failed to create XLSX file: {e}"))?;
+                        xlsx =
+                            Some(start_streaming_xlsx_workbook(BufWriter::new(xlsx_file), Some("Result"), &columns)?);
+                    }
+                }
+                crate::db::mysql::MySqlQueryStreamItem::Row(row) => {
+                    if xlsx_hard_limit_active && rows_exported as usize >= XLSX_MAX_DATA_ROWS {
+                        return Err(XLSX_ROW_LIMIT_ERROR.to_string());
+                    }
+                    if let Some(file) = csv_file.as_mut() {
+                        let rows_csv = format_query_result_csv_rows(std::slice::from_ref(&row));
+                        write!(file, "\n{rows_csv}").map_err(|e| format!("Failed to write CSV rows: {e}"))?;
+                    } else if let Some(writer) = xlsx.as_mut() {
+                        writer.write_row(&row).map_err(|e| format!("Failed to write XLSX row: {e}"))?;
+                    } else {
+                        let xlsx_file =
+                            File::create(&request.file_path).map_err(|e| format!("Failed to create XLSX file: {e}"))?;
+                        xlsx =
+                            Some(start_streaming_xlsx_workbook(BufWriter::new(xlsx_file), Some("Result"), &columns)?);
+                        if let Some(writer) = xlsx.as_mut() {
+                            writer.write_row(&row).map_err(|e| format!("Failed to write XLSX row: {e}"))?;
+                        }
+                    }
+                    rows_exported += 1;
+                    let now = Instant::now();
+                    if should_emit_stream_progress(
+                        rows_exported,
+                        last_progress_rows,
+                        progress_row_interval,
+                        now.duration_since(last_progress_at),
+                    ) {
+                        on_progress(progress(request, rows_exported, ExportStatus::Running, None));
+                        last_progress_rows = rows_exported;
+                        last_progress_at = now;
+                    }
+                }
+            }
+            Ok(())
+        },
+    );
+    let stream_result = match query_timeout {
+        Some(timeout) => match tokio::time::timeout(timeout, stream_future).await {
+            Ok(result) => result,
+            Err(_) => {
+                let _ = crate::db::mysql::kill_query_with_opts(kill_opts, mysql_connection_id).await;
+                Err(format!("Query timed out after {} seconds", timeout.as_secs()))
+            }
+        },
+        None => stream_future.await,
+    };
+    watcher_done.cancel();
+
+    if let Err(error) = stream_result {
+        if error == QUERY_CANCELED
+            || export_cancelled.load(Ordering::SeqCst)
+            || cancel_token.as_ref().is_some_and(|token| token.is_cancelled())
+            || is_export_cancelled(&request.export_id).await
+        {
+            on_progress(progress(
+                request,
+                rows_exported,
+                ExportStatus::Cancelled,
+                Some("Export cancelled".to_string()),
+            ));
+            return Ok(true);
+        }
+        return Err(error);
+    }
+
+    if rows_exported != last_progress_rows {
+        on_progress(progress(request, rows_exported, ExportStatus::Running, None));
+    }
+    on_progress(progress(request, rows_exported, ExportStatus::Writing, None));
+    if let Some(file) = csv_file.as_mut() {
+        file.flush().map_err(|e| format!("Failed to flush CSV file: {e}"))?;
+    }
+    if let Some(writer) = xlsx {
+        let mut buf =
+            finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?;
+        buf.flush().map_err(|e| format!("Failed to flush XLSX file: {e}"))?;
+    } else if format == "xlsx" {
+        let xlsx_file = File::create(&request.file_path).map_err(|e| format!("Failed to create XLSX file: {e}"))?;
+        let writer = start_streaming_xlsx_workbook(BufWriter::new(xlsx_file), Some("Result"), &columns)?;
+        let mut buf =
+            finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?;
+        buf.flush().map_err(|e| format!("Failed to flush XLSX file: {e}"))?;
+    }
+    on_progress(progress(request, rows_exported, ExportStatus::Done, None));
+    Ok(true)
+}
+
 async fn try_export_sqlserver_query_result_stream(
     state: &AppState,
     request: &QueryResultExportRequest,