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,

crates/dbx-core/tests/live_mysql57.rs

@@ -1,9 +1,14 @@
+use std::collections::BTreeSet;
+use std::sync::atomic::{AtomicBool, Ordering};
+
 use dbx_core::connection::AppState;
 use dbx_core::models::connection::{ConnectionConfig, DatabaseType};
 use dbx_core::query::execute_sql_statement;
+use dbx_core::query_result_export::{export_query_result_core, ExportStatus, QueryResultExportRequest};
 use dbx_core::sql::{split_sql_statements_for_database, SqlFileRequest};
 use dbx_core::sql_file_import::execute_sql_file_content;
 use dbx_core::storage::Storage;
+use dbx_core::table_import::parse_xlsx_file;
 use tokio_util::sync::CancellationToken;
 
 fn live_mysql_sql_file_config(id: &str) -> ConnectionConfig {
@@ -38,6 +43,41 @@ async fn app_state_with_config(config: ConnectionConfig) -> (AppState, std::path
     (state, db_path)
 }
 
+fn live_mysql_query_export_config(
+    id: &str,
+    host: &str,
+    port: u16,
+    user: &str,
+    password: &str,
+    database: &str,
+) -> ConnectionConfig {
+    serde_json::from_value(serde_json::json!({
+        "id": id,
+        "name": id,
+        "db_type": DatabaseType::Mysql,
+        "host": host,
+        "port": port,
+        "username": user,
+        "password": password,
+        "database": database,
+        "connect_timeout_secs": 10,
+        "query_timeout_secs": 30,
+        "idle_timeout_secs": 60,
+        "keepalive_interval_secs": 0
+    }))
+    .expect("live MySQL query export config should deserialize")
+}
+
+fn json_cell_text(value: &serde_json::Value) -> String {
+    match value {
+        serde_json::Value::String(value) => value.clone(),
+        serde_json::Value::Number(value) => value.to_string(),
+        serde_json::Value::Bool(value) => value.to_string(),
+        serde_json::Value::Null => String::new(),
+        other => other.to_string(),
+    }
+}
+
 #[tokio::test]
 #[ignore = "requires the remote DBX MySQL 5.7 smoke-test container"]
 async fn live_mysql57_text_protocol_select_succeeds() {
@@ -74,6 +114,85 @@ async fn live_mysql_compatible_limited_text_protocol_query_succeeds() {
     assert!(result.rows.len() <= 100);
 }
 
+#[tokio::test]
+#[ignore = "requires a writable MySQL endpoint for query-result XLSX export"]
+async fn live_mysql_query_result_export_xlsx_streams_single_query_without_duplicate_batches() {
+    let host = std::env::var("DBX_LIVE_MYSQL_EXPORT_HOST").expect("DBX_LIVE_MYSQL_EXPORT_HOST");
+    let port = std::env::var("DBX_LIVE_MYSQL_EXPORT_PORT").expect("DBX_LIVE_MYSQL_EXPORT_PORT").parse::<u16>().unwrap();
+    let user = std::env::var("DBX_LIVE_MYSQL_EXPORT_USER").expect("DBX_LIVE_MYSQL_EXPORT_USER");
+    let password = std::env::var("DBX_LIVE_MYSQL_EXPORT_PASSWORD").expect("DBX_LIVE_MYSQL_EXPORT_PASSWORD");
+    let database = std::env::var("DBX_LIVE_MYSQL_EXPORT_DATABASE").expect("DBX_LIVE_MYSQL_EXPORT_DATABASE");
+    let suffix = uuid::Uuid::new_v4().simple().to_string();
+    let table = format!("dbx_query_export_{}", &suffix[..8]);
+    let connection_id = format!("live-mysql-query-export-{suffix}");
+    let config = live_mysql_query_export_config(&connection_id, &host, port, &user, &password, &database);
+    let dir = std::env::temp_dir().join(format!("dbx-live-mysql-query-export-{suffix}"));
+    std::fs::create_dir_all(&dir).unwrap();
+    let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
+    let state = AppState::new(storage);
+    state.configs.write().await.insert(config.id.clone(), config);
+
+    let values = (1..=250).map(|id| format!("({id}, 'row-{id}')")).collect::<Vec<_>>().join(", ");
+    let cleanup_sql = format!("DROP TABLE IF EXISTS `{table}`");
+    let create_sql = format!("CREATE TABLE `{table}` (id INT PRIMARY KEY, label VARCHAR(32) NOT NULL)");
+    let insert_sql = format!("INSERT INTO `{table}` (id, label) VALUES {values}");
+    let _ = execute_sql_statement(&state, &connection_id, &database, &cleanup_sql, None, None).await;
+    execute_sql_statement(&state, &connection_id, &database, &create_sql, None, None)
+        .await
+        .expect("create live export table");
+    execute_sql_statement(&state, &connection_id, &database, &insert_sql, None, None)
+        .await
+        .expect("insert live export rows");
+
+    let file_path = dir.join("result.xlsx");
+    let sql = format!("SELECT id, label FROM `{table}`");
+    let request = QueryResultExportRequest {
+        export_id: format!("live-mysql-query-export-{suffix}"),
+        connection_id: connection_id.clone(),
+        database: database.clone(),
+        schema: None,
+        sql: sql.clone(),
+        query_base_sql: sql,
+        database_type: DatabaseType::Mysql,
+        use_agent_cursor: false,
+        file_path: file_path.to_string_lossy().to_string(),
+        format: "xlsx".to_string(),
+        page_size: 50,
+        row_limit: None,
+        total_rows: Some(250),
+        timeout_secs: Some(30),
+        keyset_optimization_enabled: false,
+        client_session_id: None,
+        execution_id: Some(format!("live-mysql-query-export-{suffix}")),
+    };
+    let done_seen = AtomicBool::new(false);
+    let result = export_query_result_core(&state, &request, None, |progress| {
+        if matches!(progress.status, ExportStatus::Done) {
+            done_seen.store(true, Ordering::Relaxed);
+        }
+    })
+    .await;
+
+    let cleanup_result = execute_sql_statement(&state, &connection_id, &database, &cleanup_sql, None, None).await;
+    result.expect("export MySQL query result to XLSX");
+    cleanup_result.expect("cleanup live export table");
+    assert!(done_seen.load(Ordering::Relaxed));
+
+    let parsed = parse_xlsx_file(&file_path.to_string_lossy(), 300).expect("parse exported XLSX");
+    let _ = std::fs::remove_dir_all(&dir);
+    assert_eq!(parsed.columns, vec!["id", "label"]);
+    assert_eq!(parsed.total_rows, 250);
+    assert_eq!(parsed.rows.len(), 250);
+
+    let exported_ids = parsed
+        .rows
+        .iter()
+        .map(|row| json_cell_text(&row[0]).parse::<i64>().expect("numeric id"))
+        .collect::<BTreeSet<_>>();
+    let expected_ids = (1..=250).collect::<BTreeSet<_>>();
+    assert_eq!(exported_ids, expected_ids);
+}
+
 #[tokio::test]
 #[ignore = "requires a remote MySQL endpoint"]
 async fn live_mysql_call_procedure_returns_select_result_set() {