From 3609641fc84fa4dfa13e444c2126cc3ebb62b63b Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Tue, 7 Jul 2026 00:50:14 +0800 Subject: [PATCH] fix(mysql): stream query result exports --- crates/dbx-core/src/db/mysql.rs | 57 ++++- crates/dbx-core/src/query_result_export.rs | 245 +++++++++++++++++++++ crates/dbx-core/tests/live_mysql57.rs | 119 ++++++++++ 3 files changed, 412 insertions(+), 9 deletions(-) diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 9f501cab6..5554ee8c2 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -44,6 +44,11 @@ impl MySqlQueryDialect { } } +pub enum MySqlQueryStreamItem { + Columns { columns: Vec, column_types: Vec }, + Row(Vec), +} + 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 { 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, + dialect: MySqlQueryDialect, + cancelled: &AtomicBool, + mut on_item: impl FnMut(MySqlQueryStreamItem) -> Result<(), String>, +) -> Result { 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 { 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 = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect(); + let column_types: Vec = + 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::() .await @@ -2520,21 +2551,29 @@ async fn stream_query_rows_text( } let row = row.map_err(|e| e.to_string())?; let values: Vec = (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 { let mut result = conn.exec_iter(sql, ()).await.map_err(|e| e.to_string())?; + let columns: Vec = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect(); + if columns.is_empty() { + return Ok(0); + } + let column_types: Vec = + 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::() .await @@ -2551,7 +2590,7 @@ async fn stream_query_rows_prepared( } let row = row.map_err(|e| e.to_string())?; let values: Vec = (0..row.len()).map(|i| mysql_value_to_json(&row, i)).collect(); - on_row(&values)?; + on_item(MySqlQueryStreamItem::Row(values))?; rows_exported += 1; } diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs index 6176ca687..3d889d222 100644 --- a/crates/dbx-core/src/query_result_export.rs +++ b/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, + on_progress: &impl Fn(TableExportProgress), +) -> Result { + 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 = 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, diff --git a/crates/dbx-core/tests/live_mysql57.rs b/crates/dbx-core/tests/live_mysql57.rs index 8fa91aa39..c643dc7bc 100644 --- a/crates/dbx-core/tests/live_mysql57.rs +++ b/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::().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::>().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::().expect("numeric id")) + .collect::>(); + let expected_ids = (1..=250).collect::>(); + assert_eq!(exported_ids, expected_ids); +} + #[tokio::test] #[ignore = "requires a remote MySQL endpoint"] async fn live_mysql_call_procedure_returns_select_result_set() {