xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(mysql): stream query result exports

This commit is contained in:
t8y2 2026-07-07 00:50:14 +08:00
parent 9cbe857a2a
commit 3609641fc8
3 changed files with 412 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

57
crates/dbx-core/src/db/mysql.rs
View File

@ -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;
}

245
crates/dbx-core/src/query_result_export.rs
View File

@ -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,

119
crates/dbx-core/tests/live_mysql57.rs
View File

@ -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() {