From 0bacb5effdd527b3be409e7e5c3d05941733112e Mon Sep 17 00:00:00 2001 From: zipg Date: Mon, 7 Sep 2026 23:11:01 +0800 Subject: [PATCH] feat(export): apply connection query timeout to database backup --- apps/desktop/src/App.vue | 4 + .../export/DatabaseExportDialog.vue | 28 ++++- .../src/components/layout/AppDialogs.vue | 2 + crates/dbx-core/src/database_export.rs | 95 +++++++++++--- crates/dbx-core/src/query.rs | 117 +++++++----------- 5 files changed, 155 insertions(+), 91 deletions(-) diff --git a/apps/desktop/src/App.vue b/apps/desktop/src/App.vue index c8df98c56d..a6629d5fa9 100644 --- a/apps/desktop/src/App.vue +++ b/apps/desktop/src/App.vue @@ -3746,6 +3746,10 @@ onUnmounted(() => { setConnectionDialogOpen(false); openSettings('tunnels'); " + @open-connection-settings=" + setConnectionDialogOpen(false); + openConnectionSettings($event, 'advanced'); + " @open-lineage-target="openLineageTarget" @open-database-search-target="openDatabaseSearchTarget" @open-diagram-target="openDiagramTarget" diff --git a/apps/desktop/src/components/export/DatabaseExportDialog.vue b/apps/desktop/src/components/export/DatabaseExportDialog.vue index 03fe062683..c9bcd863f4 100644 --- a/apps/desktop/src/components/export/DatabaseExportDialog.vue +++ b/apps/desktop/src/components/export/DatabaseExportDialog.vue @@ -17,8 +17,9 @@ import { buildSelectedTablesPayload, isDatabaseExportTableSelectionValid } from import { isTauriRuntime } from "@/lib/backend/tauriRuntime"; import { useToast } from "@/composables/useToast"; import { Input } from "@/components/ui/input"; -import { Download, Square, CheckSquare, Search, X, Loader2 } from "@lucide/vue"; +import { Download, Square, CheckSquare, Search, X, Loader2, Wrench } from "@lucide/vue"; import { formatDataTransferDuration, useExportTracker } from "@/composables/useExportTracker"; +import { isQueryTimeoutErrorMessage } from "@/lib/sql/queryError"; const { t } = useI18n(); const { toast } = useToast(); @@ -35,6 +36,10 @@ const props = defineProps<{ prefillAllDatabases?: boolean; }>(); +const emit = defineEmits<{ + openConnectionSettings: [connectionId: string]; +}>(); + // Connection / Database / Schema selectors const connectionId = ref(""); const database = ref(""); @@ -72,6 +77,7 @@ const exportProgress = ref(null); const exportId = ref(""); const exportDone = ref(false); const exportError = ref(null); +const exportWarning = ref(null); const exportCancelled = ref(false); const exportStartedAt = ref(null); const exportFinishedAt = ref(null); @@ -163,6 +169,8 @@ function toastDatabaseExportCompletion(errorCount: number, errorSummary: string toast(t("databaseExport.exportSuccess"), 3000); } +const canChangeQueryTimeout = computed(() => !!connectionId.value && !!exportWarning.value && isQueryTimeoutErrorMessage(exportWarning.value)); + async function loadDatabases(connId: string) { if (!connId) return; loadingMeta.value = true; @@ -334,6 +342,7 @@ async function startExport() { exportFinishedAt.value = null; exportDone.value = false; exportError.value = null; + exportWarning.value = null; exportCancelled.value = false; exportProgress.value = { exportId: exportId.value, @@ -380,11 +389,13 @@ async function startExport() { if (progress.status === "Done") { finishExportTiming(); exportDone.value = true; + exportWarning.value = progress.errorSummary ?? null; isExporting.value = false; toastDatabaseExportCompletion(progress.errorCount ?? 0, progress.errorSummary ?? null); } else if (progress.status === "Error") { finishExportTiming(); exportError.value = progress.error; + exportWarning.value = null; isExporting.value = false; } else if (progress.status === "Cancelled") { finishExportTiming(); @@ -397,6 +408,7 @@ async function startExport() { ); } catch (e: any) { exportError.value = e?.message || String(e); + exportWarning.value = null; const lastProgress = exportProgress.value as api.ExportProgress | null; const fallbackProgress: api.ExportProgress = { exportId: exportId.value, @@ -440,6 +452,7 @@ async function startAllDatabasesExport() { exportFinishedAt.value = null; exportDone.value = false; exportError.value = null; + exportWarning.value = null; exportCancelled.value = false; batchDatabaseIndex.value = 0; batchRowsExported.value = 0; @@ -532,6 +545,7 @@ async function startAllDatabasesExport() { if (progress.status === "Error") { finishExportTiming(); exportError.value = progress.error; + exportWarning.value = null; isExporting.value = false; } else if (progress.status === "Cancelled") { finishExportTiming(); @@ -551,6 +565,7 @@ async function startAllDatabasesExport() { if (!exportError.value && !exportCancelled.value) { exportDone.value = true; + exportWarning.value = batchFirstErrorSummary; finishExportTiming(); isExporting.value = false; const finalProgress: api.ExportProgress = { @@ -575,6 +590,7 @@ async function startAllDatabasesExport() { } } catch (e: any) { exportError.value = e?.message || String(e); + exportWarning.value = null; updateDatabaseExportTask(batchId, { exportId: batchId, currentObject: t("databaseExport.allDatabasesTask", { count: dbs.length }), @@ -629,6 +645,7 @@ function resetState() { exportProgress.value = null; exportDone.value = false; exportError.value = null; + exportWarning.value = null; exportCancelled.value = false; exportStartedAt.value = null; exportFinishedAt.value = null; @@ -940,6 +957,15 @@ watch(
{{ t("databaseExport.exportSuccess") }}
+
+
+ {{ exportWarning }} +
+ +
{{ t("databaseExport.exportError", { error: exportError }) }}
diff --git a/apps/desktop/src/components/layout/AppDialogs.vue b/apps/desktop/src/components/layout/AppDialogs.vue index 1143576e5f..010cc15e2a 100644 --- a/apps/desktop/src/components/layout/AppDialogs.vue +++ b/apps/desktop/src/components/layout/AppDialogs.vue @@ -59,6 +59,7 @@ const emit = defineEmits<{ connectFailed: [message: string]; openDriverStore: [focus?: DriverStoreFocus]; openTunnelProfileSettings: []; + openConnectionSettings: [connectionId: string, initialTab: "advanced"]; openLineageTarget: [ target: { connectionId: string; @@ -327,6 +328,7 @@ watch( :prefill-table="dialogs.databaseExportPrefillTable.value" :prefill-tables="dialogs.databaseExportPrefillTables.value" :prefill-all-databases="dialogs.databaseExportAllDatabases.value" + @open-connection-settings="emit('openConnectionSettings', $event, 'advanced')" /> String { task_client_session_id("database-export", export_id) } +async fn database_export_query_options( + state: &crate::connection::AppState, + connection_id: &str, + client_session_id: &str, + max_rows: Option, +) -> crate::query::QueryExecutionOptions { + let timeout_secs = + state.configs.read().await.get(connection_id).map(|config| config.effective_query_timeout_secs()).unwrap_or(0); + database_export_query_options_for_timeout(timeout_secs, client_session_id, max_rows) +} + +fn database_export_query_options_for_timeout( + timeout_secs: u64, + client_session_id: &str, + max_rows: Option, +) -> crate::query::QueryExecutionOptions { + crate::query::QueryExecutionOptions { + max_rows, + timeout_secs: Some(timeout_secs), + client_session_id: Some(client_session_id.to_string()), + ..Default::default() + } +} + #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub enum DatabaseExportOutputCompression { @@ -196,7 +220,9 @@ async fn list_mysql_export_view_dependencies( state: &crate::connection::AppState, connection_id: &str, database: &str, + client_session_id: &str, ) -> Result, String> { + let options = database_export_query_options(state, connection_id, client_session_id, Some(usize::MAX)).await; let result = crate::query::execute_sql_statement_with_options( state, connection_id, @@ -204,7 +230,7 @@ async fn list_mysql_export_view_dependencies( &mysql_view_dependencies_sql(database), None, None, - crate::query::QueryExecutionOptions { max_rows: Some(usize::MAX), ..Default::default() }, + options, ) .await?; Ok(mysql_view_dependencies_from_rows(&result.rows)) @@ -237,18 +263,21 @@ fn mysql_database_export_preamble(database: &str, charset: Option<&str>, collati async fn mysql_database_export_preamble_for_request( state: &crate::connection::AppState, request: &DatabaseExportRequest, + client_session_id: &str, ) -> String { let metadata_sql = format!( "SELECT DEFAULT_CHARACTER_SET_NAME, DEFAULT_COLLATION_NAME FROM information_schema.SCHEMATA WHERE SCHEMA_NAME = {}", mysql_sql_string_literal(&request.database) ); - let metadata = crate::query::execute_sql_statement( + let options = database_export_query_options(state, &request.connection_id, client_session_id, Some(1)).await; + let metadata = crate::query::execute_sql_statement_with_options( state, &request.connection_id, &request.database, &metadata_sql, None, None, + options, ) .await .ok() @@ -2663,7 +2692,7 @@ async fn export_database_sql_core_inner( }; let create_database_preamble = if request.include_create_database && matches!(db_type, DatabaseType::Mysql) { - Some(mysql_database_export_preamble_for_request(state, request).await) + Some(mysql_database_export_preamble_for_request(state, request, &client_session_id).await) } else { None }; @@ -2739,7 +2768,9 @@ async fn export_database_sql_core_inner( let mut tables: Vec<_> = all_tables.iter().filter(|t| !t.table_type.contains("VIEW")).collect(); let mut views: Vec<_> = all_tables.iter().filter(|t| t.table_type.contains("VIEW")).collect(); if request.include_objects && db_type == DatabaseType::Mysql && views.len() > 1 { - match list_mysql_export_view_dependencies(state, &request.connection_id, &request.database).await { + match list_mysql_export_view_dependencies(state, &request.connection_id, &request.database, &client_session_id) + .await + { Ok(dependencies) => views = sort_export_views_by_dependencies(&views, &dependencies), Err(error) => { log::debug!( @@ -3270,7 +3301,24 @@ async fn export_database_sql_core_inner( ); replace_database_export_select_list(sql, &col_names, &col_types, &db_type) }; - let result = match crate::transfer::execute_read_on_pool(state, &pool_key, &sql).await { + let options = database_export_query_options( + state, + &request.connection_id, + &client_session_id, + Some(batch_size), + ) + .await; + let result = match crate::query::execute_sql_statement_with_options( + state, + &request.connection_id, + &request.database, + &sql, + Some(&request.schema), + None, + options, + ) + .await + { Ok(result) => result, Err(error) => { record_export_error( @@ -3605,18 +3653,18 @@ mod tests { }; use super::{ build_database_export_object_source_sql, build_database_sql_export, build_export_insert_statements, - database_export_select_sql, database_export_total_objects, drop_table_if_exists_sql, - ensure_export_destination_dir, export_destination_identity_mismatch, filter_export_table_infos, - format_export_sql_literal, format_export_table_ddl, format_mysql_spatial_export_literal, - format_xugu_spatial_export_literal, generate_postgres_extension_ddl, generate_postgres_sequence_create_ddl, - generate_postgres_sequence_owner_ddl, generate_postgres_sequence_setval_sql, - is_postgres_extension_member_routine, mysql_database_export_preamble, mysql_view_dependencies_from_rows, - mysql_view_dependencies_sql, normalize_export_table_ddl, record_export_destination_identity, - record_export_error, replace_database_export_select_list, sort_export_views_by_dependencies, - split_postgres_export_table_triggers, write_database_export_rows, BuildDatabaseSqlExportOptions, - BuildExportInsertStatementsOptions, DatabaseExportObjectCounts, DatabaseExportRequest, DatabaseExportWriter, - DdlNormalizeOptions, ExportedTableSql, PostgresExportExtension, PostgresExportSequence, - PostgresExtensionMembers, DATABASE_EXPORT_INSERT_BATCH_SIZE, DATABASE_EXPORT_ROW_LIMIT, + database_export_query_options_for_timeout, database_export_select_sql, database_export_total_objects, + drop_table_if_exists_sql, ensure_export_destination_dir, export_destination_identity_mismatch, + filter_export_table_infos, format_export_sql_literal, format_export_table_ddl, + format_mysql_spatial_export_literal, format_xugu_spatial_export_literal, generate_postgres_extension_ddl, + generate_postgres_sequence_create_ddl, generate_postgres_sequence_owner_ddl, + generate_postgres_sequence_setval_sql, is_postgres_extension_member_routine, mysql_database_export_preamble, + mysql_view_dependencies_from_rows, mysql_view_dependencies_sql, normalize_export_table_ddl, + record_export_destination_identity, record_export_error, replace_database_export_select_list, + sort_export_views_by_dependencies, split_postgres_export_table_triggers, write_database_export_rows, + BuildDatabaseSqlExportOptions, BuildExportInsertStatementsOptions, DatabaseExportObjectCounts, + DatabaseExportRequest, DatabaseExportWriter, DdlNormalizeOptions, ExportedTableSql, PostgresExportExtension, + PostgresExportSequence, PostgresExtensionMembers, DATABASE_EXPORT_INSERT_BATCH_SIZE, DATABASE_EXPORT_ROW_LIMIT, }; use super::{ExportProgress, LenientExportErrors}; use crate::connection::AppState; @@ -3632,6 +3680,19 @@ mod tests { }; use tokio_util::sync::CancellationToken; + #[test] + fn database_export_query_options_preserve_connection_timeout() { + let options = database_export_query_options_for_timeout(7, "database-export-session", Some(500)); + + assert_eq!(options.timeout_secs, Some(7)); + assert_eq!(options.max_rows, Some(500)); + assert_eq!(options.client_session_id.as_deref(), Some("database-export-session")); + + let unlimited = database_export_query_options_for_timeout(0, "database-export-session", None); + assert_eq!(unlimited.timeout_secs, Some(0)); + assert_eq!(unlimited.max_rows, None); + } + #[tokio::test] async fn await_export_operation_drops_pending_metadata_after_cancel() { let export_id = format!("cancel-pending-{}", uuid::Uuid::new_v4()); diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index b7dfd1fe78..4d4ce66eba 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -5446,82 +5446,53 @@ where } } TxnConnection::Mysql(Some(conn)) => { - let query_result = match cancel_token.as_ref() { - Some(cancel_token) => { - tokio::select! { - _ = cancel_token.cancelled() => Err(QUERY_CANCELED.to_string()), - result = conn.query_iter(sql) => result.map_err(|error| format!("Query failed: {error}")), + // The query timeout is an inactivity budget reset by every received row, + // not a cap on the total duration of a long backup/export stream. + let progress_clock = Arc::new(StreamProgressClock::new()); + let progress_clock_for_rows = progress_clock.clone(); + let timeout_error = format!( + "Query timed out after {} seconds", + operation_budget.query_timeout.map_or(0, |timeout| timeout.as_secs()) + ); + let stream_future = async { + let mut result = conn.query_iter(sql).await.map_err(|error| format!("Query failed: {error}"))?; + let Some(mut stream) = + result.stream::().await.map_err(|error| format!("Query failed: {error}"))? + else { + return Err("Empty result set stream".to_string()); + }; + + let mut batch = Vec::with_capacity(batch_size); + let mut total_rows = 0_u64; + while let Some(row_result) = stream.next().await { + match row_result { + Ok(row) => { + batch.push( + (0..row.len()).map(|index| db::mysql::mysql_value_to_json(&row, index)).collect(), + ); + total_rows += 1; + if batch.len() >= batch_size { + on_batch(std::mem::take(&mut batch))?; + batch = Vec::with_capacity(batch_size); + } + } + Err(err) => return Err(format!("Query failed: {err}")), } + progress_clock_for_rows.mark(); } - None => conn.query_iter(sql).await.map_err(|error| format!("Query failed: {error}")), + if !batch.is_empty() { + on_batch(batch)?; + } + Ok(total_rows) }; - match query_result { - Ok(mut result) => { - let stream_result = match cancel_token.as_ref() { - Some(cancel_token) => { - tokio::select! { - _ = cancel_token.cancelled() => Err(QUERY_CANCELED.to_string()), - stream = result.stream::() => stream.map_err(|error| format!("Query failed: {error}")), - } - } - None => { - result.stream::().await.map_err(|error| format!("Query failed: {error}")) - } - }; - match stream_result { - Ok(Some(mut stream)) => { - let mut batch = Vec::with_capacity(batch_size); - let mut total_rows = 0_u64; - let mut error = None; - loop { - let next_row = match cancel_token.as_ref() { - Some(cancel_token) => { - tokio::select! { - _ = cancel_token.cancelled() => { - error = Some(QUERY_CANCELED.to_string()); - break; - } - row = stream.next() => row, - } - } - None => stream.next().await, - }; - let Some(row_result) = next_row else { break }; - match row_result { - Ok(row) => { - batch.push( - (0..row.len()) - .map(|index| db::mysql::mysql_value_to_json(&row, index)) - .collect(), - ); - total_rows += 1; - if batch.len() >= batch_size { - if let Err(err) = on_batch(std::mem::take(&mut batch)) { - error = Some(err); - break; - } - batch = Vec::with_capacity(batch_size); - } - } - Err(err) => { - error = Some(format!("Query failed: {err}")); - break; - } - } - } - if error.is_none() && !batch.is_empty() { - if let Err(err) = on_batch(batch) { - error = Some(err); - } - } - error.map_or(Ok(total_rows), Err) - } - Ok(None) => Err("Empty result set stream".to_string()), - Err(error) => Err(error), - } - } - Err(error) => Err(error), - } + await_stream_with_progress_timeout( + stream_future, + operation_budget.query_timeout, + progress_clock, + cancel_token.as_ref(), + timeout_error, + ) + .await } TxnConnection::Mysql(None) => Err(MANUAL_TRANSACTION_SESSION_NOT_FOUND_ERROR.to_string()), TxnConnection::Agent { .. } => {