diff --git a/apps/desktop/src/components/sql-file/SqlFileExecutionDialog.vue b/apps/desktop/src/components/sql-file/SqlFileExecutionDialog.vue index 929f91a2ef..ee9ff9507a 100644 --- a/apps/desktop/src/components/sql-file/SqlFileExecutionDialog.vue +++ b/apps/desktop/src/components/sql-file/SqlFileExecutionDialog.vue @@ -85,6 +85,7 @@ const database = ref(""); const databaseOptions = ref([]); const loadingDatabases = ref(false); const continueOnError = ref(false); +const skipRelationalConstraints = ref(false); const running = ref(false); const cancelling = ref(false); @@ -132,6 +133,15 @@ function resetPerFileState() { } const sqlConnections = computed(() => store.connections.filter((c) => !["redis", "mongodb", "elasticsearch", "easysearch", "meilisearch", "qdrant", "milvus", "weaviate", "chromadb", "etcd", "zookeeper", "consul", "mq", "nacos"].includes(c.db_type))); +// Mirrors the core executor gate (`supports_connection_level_database_bootstrap_target`): the +// MySQL-family types it runs for, so the constraint toggle appears wherever the backend honors it. +const MYSQL_BOOTSTRAP_IMPORT_TYPES = new Set(["mysql", "doris", "starrocks", "goldendb"]); +const MYSQL_BOOTSTRAP_IMPORT_PROFILES = new Set(["mariadb", "tidb", "oceanbase", "custom_mysql", "doris", "starrocks", "selectdb", "goldendb"]); +const isMysqlCompatibleTarget = computed(() => { + const config = store.getConfig(connectionId.value); + if (!config) return false; + return MYSQL_BOOTSTRAP_IMPORT_TYPES.has(config.db_type) || (!!config.driver_profile && MYSQL_BOOTSTRAP_IMPORT_PROFILES.has(config.driver_profile.toLowerCase())); +}); const selectedConnection = computed(() => sqlConnections.value.find((c) => c.id === connectionId.value)); @@ -309,6 +319,7 @@ function resetState() { databaseOptions.value = []; loadingDatabases.value = false; continueOnError.value = false; + skipRelationalConstraints.value = false; restoreSelectedTables.value = false; resetExecution(); } @@ -386,7 +397,7 @@ async function selectFile() { const { open } = await import("@tauri-apps/plugin-dialog"); const selected = await open({ multiple: true, - filters: [{ name: "SQL", extensions: ["sql", "gz"] }], + filters: [{ name: "SQL package", extensions: ["sql", "gz", "zip"] }], }); const paths = Array.isArray(selected) ? selected : selected ? [selected] : []; if (paths.length > 0) { @@ -539,16 +550,19 @@ async function startExecution() { try { executionStarted.value = true; + const executionPaths = previews.value.flatMap((item) => item.packageFilePaths ?? [item.filePath]); await executeSqlFiles( { executionId: batchId, connectionId: connectionId.value, database: database.value.trim(), - filePath: previews.value[0]!.filePath, + filePath: executionPaths[0]!, continueOnError: continueOnError.value, ...(restoreSelectedTables.value ? { selectedTables: selectedTables.value.map((table) => ({ ...table })) } : {}), + partCooldownMs: previews.value.some((item) => item.packageFilePaths) ? 500 : 0, + skipRelationalConstraints: skipRelationalConstraints.value, }, - previews.value.map((item) => item.filePath), + executionPaths, ); const terminal = await terminalProgress; if (terminal.status === "error") { @@ -656,7 +670,7 @@ watch(
- +
@@ -791,6 +806,11 @@ watch( {{ t("sqlFile.continueOnError") }} +
diff --git a/apps/desktop/src/i18n/locales/az.ts b/apps/desktop/src/i18n/locales/az.ts index cf53640831..4e147dc4b2 100644 --- a/apps/desktop/src/i18n/locales/az.ts +++ b/apps/desktop/src/i18n/locales/az.ts @@ -6270,6 +6270,8 @@ export default withEnglishFallback({ selectDatabase: "Verilənlər bazası seç", databasePlaceholder: "Verilənlər bazasının adı", options: "Seçimlər", + packageParts: "{count} SQL hissəsi", + skipRelationalConstraints: "İdxal zamanı əlaqəli məhdudiyyətləri müvəqqəti söndür", continueOnError: "Xəta olduqda davam et", runInBackground: "Arxa planda işlət", cancelling: "Ləğv edilir...", diff --git a/apps/desktop/src/i18n/locales/en.ts b/apps/desktop/src/i18n/locales/en.ts index 44ab6375f4..67656f9108 100644 --- a/apps/desktop/src/i18n/locales/en.ts +++ b/apps/desktop/src/i18n/locales/en.ts @@ -6424,6 +6424,8 @@ export default { databasePlaceholder: "Database name", options: "Options", continueOnError: "Continue on error", + skipRelationalConstraints: "Temporarily disable relational constraints during import", + packageParts: "{count} SQL parts", runInBackground: "Run in background", cancelling: "Cancelling...", cancel: "Cancel", diff --git a/apps/desktop/src/i18n/locales/es.ts b/apps/desktop/src/i18n/locales/es.ts index 101889ffa7..2eee542b7e 100644 --- a/apps/desktop/src/i18n/locales/es.ts +++ b/apps/desktop/src/i18n/locales/es.ts @@ -6040,6 +6040,8 @@ export default withEnglishFallback({ selectDatabase: "Seleccionar base de datos", databasePlaceholder: "Nombre de la base de datos", options: "Opciones", + packageParts: "{count} partes SQL", + skipRelationalConstraints: "Desactivar temporalmente las restricciones relacionales durante la importación", continueOnError: "Continuar en caso de error", runInBackground: "Ejecutar en segundo plano", cancelling: "Cancelando...", diff --git a/apps/desktop/src/i18n/locales/it.ts b/apps/desktop/src/i18n/locales/it.ts index e0fc35e270..cc40dab48e 100644 --- a/apps/desktop/src/i18n/locales/it.ts +++ b/apps/desktop/src/i18n/locales/it.ts @@ -6038,6 +6038,8 @@ export default withEnglishFallback({ selectDatabase: "Seleziona database", databasePlaceholder: "Nome database", options: "Opzioni", + packageParts: "{count} parti SQL", + skipRelationalConstraints: "Disattiva temporaneamente i vincoli relazionali durante l'importazione", continueOnError: "Continua in caso di errore", runInBackground: "Esegui in background", cancelling: "Annullamento in corso...", diff --git a/apps/desktop/src/i18n/locales/ja.ts b/apps/desktop/src/i18n/locales/ja.ts index 45e03987c7..cf52de62d1 100644 --- a/apps/desktop/src/i18n/locales/ja.ts +++ b/apps/desktop/src/i18n/locales/ja.ts @@ -6077,6 +6077,8 @@ export default withEnglishFallback({ selectDatabase: "データベースを選択", databasePlaceholder: "データベース名", options: "オプション", + packageParts: "{count} 個のSQLパート", + skipRelationalConstraints: "インポート中に関連性制約を一時的に無効化", continueOnError: "エラー時に続行", cancelling: "キャンセル中...", cancel: "キャンセル", diff --git a/apps/desktop/src/i18n/locales/ko.ts b/apps/desktop/src/i18n/locales/ko.ts index 675a671d34..4b9786e8b5 100644 --- a/apps/desktop/src/i18n/locales/ko.ts +++ b/apps/desktop/src/i18n/locales/ko.ts @@ -5755,6 +5755,8 @@ export default withEnglishFallback({ selectDatabase: "데이터베이스 선택", databasePlaceholder: "데이터베이스 이름", options: "옵션", + packageParts: "SQL 파트 {count}개", + skipRelationalConstraints: "가져오기 중 관계형 제약 조건을 일시적으로 해제", continueOnError: "오류 시 계속", runInBackground: "백그라운드에서 실행", cancelling: "취소하는 중...", diff --git a/apps/desktop/src/i18n/locales/pt-BR.ts b/apps/desktop/src/i18n/locales/pt-BR.ts index 9ac79fde97..a3261eddea 100644 --- a/apps/desktop/src/i18n/locales/pt-BR.ts +++ b/apps/desktop/src/i18n/locales/pt-BR.ts @@ -6040,6 +6040,8 @@ export default withEnglishFallback({ selectDatabase: "Selecionar banco de dados", databasePlaceholder: "Nome do banco de dados", options: "Opções", + packageParts: "{count} partes SQL", + skipRelationalConstraints: "Desabilitar temporariamente restrições relacionais durante a importação", continueOnError: "Continuar em caso de erro", runInBackground: "Executar em segundo plano", cancelling: "Cancelando...", diff --git a/apps/desktop/src/i18n/locales/tr.ts b/apps/desktop/src/i18n/locales/tr.ts index 0574c0fbcb..9d9614ad7f 100644 --- a/apps/desktop/src/i18n/locales/tr.ts +++ b/apps/desktop/src/i18n/locales/tr.ts @@ -6152,6 +6152,8 @@ export default withEnglishFallback({ selectDatabase: "Veritabanı seç", databasePlaceholder: "Veritabanı adı", options: "Seçenekler", + packageParts: "{count} SQL bölümü", + skipRelationalConstraints: "İçe aktarma sırasında ilişkisel kısıtlamaları geçici olarak devre dışı bırak", continueOnError: "Hatada devam et", runInBackground: "Arka planda çalıştır", cancelling: "İptal ediliyor...", diff --git a/apps/desktop/src/i18n/locales/zh-CN.ts b/apps/desktop/src/i18n/locales/zh-CN.ts index d31ba9be9f..a228ee05f3 100644 --- a/apps/desktop/src/i18n/locales/zh-CN.ts +++ b/apps/desktop/src/i18n/locales/zh-CN.ts @@ -6398,6 +6398,8 @@ export default withEnglishFallback({ selectDatabase: "选择数据库", databasePlaceholder: "数据库名称", options: "选项", + packageParts: "{count} 个 SQL 分片", + skipRelationalConstraints: "导入期间临时禁用关系约束", continueOnError: "出错后继续", runInBackground: "后台运行", cancelling: "正在取消...", diff --git a/apps/desktop/src/i18n/locales/zh-TW.ts b/apps/desktop/src/i18n/locales/zh-TW.ts index a37c6828d2..2a1767109e 100644 --- a/apps/desktop/src/i18n/locales/zh-TW.ts +++ b/apps/desktop/src/i18n/locales/zh-TW.ts @@ -5360,6 +5360,8 @@ export default withEnglishFallback({ selectDatabase: "選擇資料庫", databasePlaceholder: "資料庫名稱", options: "選項", + packageParts: "{count} 個 SQL 分片", + skipRelationalConstraints: "匯入期間暫時停用關聯式約束", continueOnError: "錯誤後繼續", cancelling: "正在取消……", cancel: "取消", diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 7f0cf7c61b..6bc3d716b9 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -4769,6 +4769,8 @@ export interface SqlFileRequest { filePath: string; continueOnError: boolean; selectedTables?: SqlFileTable[]; + partCooldownMs?: number; + skipRelationalConstraints?: boolean; } export interface SqlFileTable { @@ -4787,6 +4789,8 @@ export interface SqlFilePreview { preview: string; canExecuteWithoutSelectedDatabase: boolean; establishesDatabaseContext?: boolean; + packageFilePaths?: string[]; + packagePartCount?: number; } export interface SqlFileProgress { diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs index 7738bcf8ae..5af975ae09 100644 --- a/crates/dbx-core/src/lib.rs +++ b/crates/dbx-core/src/lib.rs @@ -95,6 +95,7 @@ pub mod sql_dialect; pub mod sql_editability; pub mod sql_error_position; pub mod sql_file_import; +pub mod sql_file_zip_package; pub mod sql_parser; pub mod sql_risk; pub mod sqlite_backup; diff --git a/crates/dbx-core/src/sql.rs b/crates/dbx-core/src/sql.rs index ccbdb85778..28f1181280 100644 --- a/crates/dbx-core/src/sql.rs +++ b/crates/dbx-core/src/sql.rs @@ -61,6 +61,13 @@ pub struct SqlFileRequest { pub continue_on_error: bool, #[serde(default)] pub selected_tables: Option>, + #[serde(default)] + pub part_cooldown_ms: u64, + /// Temporarily disable MySQL `FOREIGN_KEY_CHECKS` for this import and + /// restore them on completion, error, or cancellation. Only applies to + /// MySQL-compatible connections that reuse one pinned session. + #[serde(default)] + pub skip_relational_constraints: bool, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -73,6 +80,10 @@ pub struct SqlFilePreview { pub can_execute_without_selected_database: bool, #[serde(default)] pub establishes_database_context: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub package_file_paths: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub package_part_count: Option, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] diff --git a/crates/dbx-core/src/sql_file_import.rs b/crates/dbx-core/src/sql_file_import.rs index 31d0eac8a1..f89cdc24ed 100644 --- a/crates/dbx-core/src/sql_file_import.rs +++ b/crates/dbx-core/src/sql_file_import.rs @@ -3,6 +3,7 @@ use std::path::Path; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; +use mysql_async::prelude::Queryable; use tokio::io::{AsyncReadExt, BufReader}; use tokio_util::sync::CancellationToken; @@ -183,6 +184,7 @@ struct MySqlSqlFileExecutor { dialect: db::mysql::MySqlQueryDialect, budget: DbOperationBudget, conn: Option, + constraints_disabled: bool, } impl MySqlSqlFileExecutor { @@ -229,6 +231,7 @@ impl MySqlSqlFileExecutor { ), budget, conn: None, + constraints_disabled: false, })) } @@ -320,6 +323,11 @@ impl MySqlSqlFileExecutor { let database = self.database.trim(); let database = (!database.is_empty()).then_some(database); self.pool_key = state.reconnect_pool_for_session(&self.connection_id, database, None).await?; + // A fresh session restores the server default FOREIGN_KEY_CHECKS = 1; + // re-issue the bypass so the remaining parts match the caller's toggle. + if self.constraints_disabled { + self.set_foreign_key_checks(state, child_token, false).await?; + } continue; } // Cancelled, or the retry itself failed with another @@ -338,6 +346,19 @@ impl MySqlSqlFileExecutor { unreachable!("MySQL SQL file executor retry loop runs at most twice") } + async fn set_foreign_key_checks( + &mut self, + state: &AppState, + token: &CancellationToken, + enabled: bool, + ) -> Result<(), String> { + self.ensure_conn(state, token).await?; + let conn = self.conn.as_mut().ok_or("MySQL SQL file executor is missing a connection".to_string())?; + conn.query_drop(if enabled { "SET FOREIGN_KEY_CHECKS = 1" } else { "SET FOREIGN_KEY_CHECKS = 0" }) + .await + .map_err(|error| error.to_string()) + } + async fn ensure_conn(&mut self, state: &AppState, token: &CancellationToken) -> Result<(), String> { if self.conn.is_some() { return Ok(()); @@ -489,72 +510,121 @@ pub async fn execute_sql_file_paths( return Err(error); } }; + let constraints_disabled = request.skip_relational_constraints && mysql_executor.is_some(); + if constraints_disabled { + let executor = mysql_executor.as_mut().expect("missing MySQL executor while disabling constraints"); + if let Err(error) = executor.set_foreign_key_checks(state, &token, false).await { + emit(sql_file_execution_error_progress(&request.execution_id, started_at, &progress, error.clone())); + return Err(error); + } + executor.constraints_disabled = true; + } let file_count = file_paths.len(); let mut prev_statement_index = 0usize; let mut prev_success_count = 0usize; let mut prev_failure_count = 0usize; let mut prev_affected_rows = 0u64; - for (file_index, file_path) in file_paths.iter().enumerate() { - let file_name = file_path.file_name().map(|n| n.to_string_lossy().to_string()).unwrap_or_default(); + let import_result = async { + for (file_index, file_path) in file_paths.iter().enumerate() { + let file_name = file_path.file_name().map(|n| n.to_string_lossy().to_string()).unwrap_or_default(); - // Emit a file-boundary progress event so the frontend knows which file is - // currently executing and can display a "File N/M" indicator. - if file_count > 1 { - emit(SqlFileProgress { - execution_id: request.execution_id.clone(), - status: SqlFileStatus::Running, - statement_index: progress.statement_index, - success_count: progress.success_count, - failure_count: progress.failure_count, - affected_rows: progress.affected_rows, - elapsed_ms: started_at.elapsed().as_millis(), - statement_summary: String::new(), - error: None, - file_index: Some(file_index), - file_name: Some(file_name.clone()), + // Emit a file-boundary progress event so the frontend knows which file is + // currently executing and can display a "File N/M" indicator. + if file_count > 1 { + emit(SqlFileProgress { + execution_id: request.execution_id.clone(), + status: SqlFileStatus::Running, + statement_index: progress.statement_index, + success_count: progress.success_count, + failure_count: progress.failure_count, + affected_rows: progress.affected_rows, + elapsed_ms: started_at.elapsed().as_millis(), + statement_summary: String::new(), + error: None, + file_index: Some(file_index), + file_name: Some(file_name.clone()), + }); + } + + let mut splitter = StreamingSqlFileSplitter::new(database_type, options); + let mut pending_statements = Vec::with_capacity(SQL_FILE_STATEMENT_BATCH_SIZE); + let normalize_mysql_binary_literals = import_target.as_ref().is_some_and(|target| { + crate::sql::is_mysql_compatible_import_target(&target.db_type, target.driver_profile.as_deref()) }); - } + let mut decoder = + match SqlFileStreamDecoder::open_for_target(file_path, normalize_mysql_binary_literals).await { + Ok(decoder) => decoder, + Err(error) => { + emit(sql_file_execution_error_progress( + &request.execution_id, + started_at, + &progress, + error.clone(), + )); + return Err(error); + } + }; - let mut splitter = StreamingSqlFileSplitter::new(database_type, options); - let mut pending_statements = Vec::with_capacity(SQL_FILE_STATEMENT_BATCH_SIZE); - let normalize_mysql_binary_literals = import_target.as_ref().is_some_and(|target| { - crate::sql::is_mysql_compatible_import_target(&target.db_type, target.driver_profile.as_deref()) - }); - let mut decoder = match SqlFileStreamDecoder::open_for_target(file_path, normalize_mysql_binary_literals).await - { - Ok(decoder) => decoder, - Err(error) => { - emit(sql_file_execution_error_progress(&request.execution_id, started_at, &progress, error.clone())); - return Err(error); - } - }; - - loop { - let chunk = match decoder.next_chunk().await { - Ok(chunk) => chunk, - Err(error) => { - emit(sql_file_progress( - &request.execution_id, - SqlFileStatus::Error, - progress.statement_index, - progress.success_count, - progress.failure_count, - progress.affected_rows, - started_at, - "", - Some(error.clone()), - )); - return Err(error); + loop { + let chunk = match decoder.next_chunk().await { + Ok(chunk) => chunk, + Err(error) => { + emit(sql_file_progress( + &request.execution_id, + SqlFileStatus::Error, + progress.statement_index, + progress.success_count, + progress.failure_count, + progress.affected_rows, + started_at, + "", + Some(error.clone()), + )); + return Err(error); + } + }; + let Some(chunk) = chunk else { + break; + }; + if token.is_cancelled() { + emit_sql_file_terminal_progress(request, &token, started_at, &progress, &mut emit); + return Ok(()); } - }; - let Some(chunk) = chunk else { - break; - }; - if token.is_cancelled() { - emit_sql_file_terminal_progress(request, &token, started_at, &progress, &mut emit); - return Ok(()); + let mut next_statements = splitter.push_chunk(&chunk); + if let Some(filter) = restore_filter.as_mut() { + if let Err(error) = filter_restore_statements( + &mut next_statements, + filter, + request.selected_tables.as_deref().unwrap_or_default(), + ) { + emit(sql_file_execution_error_progress( + &request.execution_id, + started_at, + &progress, + error.clone(), + )); + return Err(error); + } + } + pending_statements.extend(next_statements); + if pending_statements.len() < SQL_FILE_STATEMENT_BATCH_SIZE { + continue; + } + execute_sql_file_statement_batch( + state, + request, + &token, + started_at, + &mut pending_statements, + import_target.as_ref(), + mysql_executor.as_mut(), + &mut progress, + &mut emit, + ) + .await?; } - let mut next_statements = splitter.push_chunk(&chunk); + + let mut next_statements = splitter.finish(); if let Some(filter) = restore_filter.as_mut() { if let Err(error) = filter_restore_statements( &mut next_statements, @@ -571,9 +641,6 @@ pub async fn execute_sql_file_paths( } } pending_statements.extend(next_statements); - if pending_statements.len() < SQL_FILE_STATEMENT_BATCH_SIZE { - continue; - } execute_sql_file_statement_batch( state, request, @@ -586,57 +653,48 @@ pub async fn execute_sql_file_paths( &mut emit, ) .await?; - } - let mut next_statements = splitter.finish(); - if let Some(filter) = restore_filter.as_mut() { - if let Err(error) = filter_restore_statements( - &mut next_statements, - filter, - request.selected_tables.as_deref().unwrap_or_default(), - ) { - emit(sql_file_execution_error_progress(&request.execution_id, started_at, &progress, error.clone())); - return Err(error); + // After each file, emit a per-file summary with diff-based counters so + // the frontend can build a per-file breakdown table. + if file_count > 1 { + emit(SqlFileProgress { + execution_id: request.execution_id.clone(), + status: SqlFileStatus::StatementDone, + statement_index: progress.statement_index - prev_statement_index, + success_count: progress.success_count - prev_success_count, + failure_count: progress.failure_count - prev_failure_count, + affected_rows: progress.affected_rows - prev_affected_rows, + elapsed_ms: started_at.elapsed().as_millis(), + statement_summary: String::new(), + error: None, + file_index: Some(file_index), + file_name: Some(file_name), + }); + prev_statement_index = progress.statement_index; + prev_success_count = progress.success_count; + prev_failure_count = progress.failure_count; + prev_affected_rows = progress.affected_rows; + } + if file_index + 1 < file_count && request.part_cooldown_ms > 0 { + tokio::select! { + _ = token.cancelled() => { + emit_sql_file_terminal_progress(request, &token, started_at, &progress, &mut emit); + return Ok(()); + } + _ = tokio::time::sleep(Duration::from_millis(request.part_cooldown_ms)) => {} + } } } - pending_statements.extend(next_statements); - execute_sql_file_statement_batch( - state, - request, - &token, - started_at, - &mut pending_statements, - import_target.as_ref(), - mysql_executor.as_mut(), - &mut progress, - &mut emit, - ) - .await?; - - // After each file, emit a per-file summary with diff-based counters so - // the frontend can build a per-file breakdown table. - if file_count > 1 { - emit(SqlFileProgress { - execution_id: request.execution_id.clone(), - status: SqlFileStatus::StatementDone, - statement_index: progress.statement_index - prev_statement_index, - success_count: progress.success_count - prev_success_count, - failure_count: progress.failure_count - prev_failure_count, - affected_rows: progress.affected_rows - prev_affected_rows, - elapsed_ms: started_at.elapsed().as_millis(), - statement_summary: String::new(), - error: None, - file_index: Some(file_index), - file_name: Some(file_name), - }); - prev_statement_index = progress.statement_index; - prev_success_count = progress.success_count; - prev_failure_count = progress.failure_count; - prev_affected_rows = progress.affected_rows; + emit_sql_file_terminal_progress(request, &token, started_at, &progress, &mut emit); + Ok(()) + } + .await; + if constraints_disabled { + if let Some(executor) = mysql_executor.as_mut() { + let _ = executor.set_foreign_key_checks(state, &token, true).await; } } - emit_sql_file_terminal_progress(request, &token, started_at, &progress, &mut emit); - Ok(()) + import_result } pub async fn inspect_sql_file_tables(file_path: &Path) -> Result, String> { @@ -2589,6 +2647,8 @@ mod tests { file_path: path.to_string_lossy().to_string(), continue_on_error: true, selected_tables: None, + part_cooldown_ms: 0, + skip_relational_constraints: false, }; let mut progress = Vec::new(); diff --git a/crates/dbx-core/src/sql_file_zip_package.rs b/crates/dbx-core/src/sql_file_zip_package.rs new file mode 100644 index 0000000000..b75614744d --- /dev/null +++ b/crates/dbx-core/src/sql_file_zip_package.rs @@ -0,0 +1,178 @@ +use std::io::{Read, Write}; +use std::path::{Component, Path, PathBuf}; + +use serde::Deserialize; + +pub const SQL_FILE_ZIP_MAX_PARTS: usize = 10_000; +pub const SQL_FILE_ZIP_MAX_MANIFEST_BYTES: u64 = 1024 * 1024; +pub const SQL_FILE_ZIP_MAX_TOTAL_BYTES: u64 = 4 * 1024 * 1024 * 1024; +pub const SQL_FILE_ZIP_MAX_PART_BYTES: u64 = 1024 * 1024 * 1024; + +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase")] +struct SqlZipManifest { + format: String, + source_file_name: String, + parts: Vec, +} + +#[derive(Debug, Clone)] +pub struct SqlFileZipPackage { + pub source_file_name: String, + pub part_names: Vec, + pub total_bytes: u64, +} + +pub fn inspect_sql_file_zip_package(path: &Path) -> Result { + let file = std::fs::File::open(path).map_err(|error| format!("Failed to open SQL ZIP package: {error}"))?; + let mut archive = zip::ZipArchive::new(file).map_err(|error| format!("Invalid SQL ZIP package: {error}"))?; + let manifest_bytes = read_zip_entry(&mut archive, "manifest.json", SQL_FILE_ZIP_MAX_MANIFEST_BYTES)?; + let manifest: SqlZipManifest = + serde_json::from_slice(&manifest_bytes).map_err(|error| format!("Invalid SQL ZIP manifest: {error}"))?; + if manifest.format != "dbx-sql-export-parts-v1" { + return Err("Unsupported SQL ZIP package format".to_string()); + } + if manifest.parts.is_empty() || manifest.parts.len() > SQL_FILE_ZIP_MAX_PARTS { + return Err("SQL ZIP manifest has an invalid number of parts".to_string()); + } + let mut seen = std::collections::HashSet::new(); + let mut total_bytes = 0_u64; + for part in &manifest.parts { + validate_part_name(part)?; + if !seen.insert(part) { + return Err(format!("SQL ZIP manifest contains duplicate part: {part}")); + } + let entry = archive.by_name(part).map_err(|_| format!("SQL ZIP package is missing manifest part: {part}"))?; + if entry.is_dir() || entry.size() > SQL_FILE_ZIP_MAX_PART_BYTES { + return Err(format!("SQL ZIP part exceeds limits: {part}")); + } + total_bytes = total_bytes.checked_add(entry.size()).ok_or("SQL ZIP package size overflow")?; + if total_bytes > SQL_FILE_ZIP_MAX_TOTAL_BYTES { + return Err("SQL ZIP package exceeds total extracted size limit".to_string()); + } + } + Ok(SqlFileZipPackage { source_file_name: manifest.source_file_name, part_names: manifest.parts, total_bytes }) +} + +pub fn extract_sql_file_zip_package(path: &Path, destination: &Path) -> Result { + let package = inspect_sql_file_zip_package(path)?; + std::fs::create_dir_all(destination) + .map_err(|error| format!("Failed to create SQL ZIP extraction directory: {error}"))?; + let file = std::fs::File::open(path).map_err(|error| format!("Failed to open SQL ZIP package: {error}"))?; + let mut archive = zip::ZipArchive::new(file).map_err(|error| format!("Invalid SQL ZIP package: {error}"))?; + for (index, part) in package.part_names.iter().enumerate() { + let entry = archive.by_name(part).map_err(|_| format!("SQL ZIP package is missing manifest part: {part}"))?; + let output = destination.join(format!("{:05}-{}", index + 1, part)); + let mut target = + std::fs::File::create(&output).map_err(|error| format!("Failed to create extracted SQL part: {error}"))?; + // `entry.size()` is declared metadata; the zip reader does not clamp decompressed output + // to it, so the copy itself must be bounded to keep a lying header from writing past the + // inspected limit before the mismatch check runs. + let declared = entry.size(); + let mut bounded = entry.take(declared + 1); + let copied = + std::io::copy(&mut bounded, &mut target).map_err(|error| format!("Failed to extract SQL part: {error}"))?; + if copied != declared { + return Err(format!("Failed to fully extract SQL part: {part}")); + } + target.flush().map_err(|error| format!("Failed to flush extracted SQL part: {error}"))?; + } + Ok(package) +} + +fn read_zip_entry(archive: &mut zip::ZipArchive, name: &str, max_bytes: u64) -> Result, String> { + let entry = archive.by_name(name).map_err(|_| format!("SQL ZIP package is missing {name}"))?; + if entry.is_dir() || entry.size() > max_bytes { + return Err(format!("SQL ZIP package {name} exceeds limits")); + } + let mut bytes = Vec::with_capacity(entry.size() as usize); + let declared = entry.size(); + let mut bounded = entry.take(max_bytes + 1); + bounded.read_to_end(&mut bytes).map_err(|error| format!("Failed to read SQL ZIP {name}: {error}"))?; + if bytes.len() as u64 != declared { + return Err(format!("SQL ZIP {name} does not match its declared size")); + } + Ok(bytes) +} + +fn validate_part_name(name: &str) -> Result<(), String> { + let path = Path::new(name); + if name.trim().is_empty() + || !name.ends_with(".sql") + || path.is_absolute() + || path.components().any(|component| !matches!(component, Component::Normal(_))) + { + return Err(format!("Unsafe SQL ZIP manifest part: {name}")); + } + Ok(()) +} + +pub fn extracted_sql_zip_paths(destination: &Path, package: &SqlFileZipPackage) -> Vec { + package + .part_names + .iter() + .enumerate() + .map(|(index, part)| destination.join(format!("{:05}-{}", index + 1, part))) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use zip::write::SimpleFileOptions; + + fn write_package(path: &Path, parts: &[(&str, &str)], manifest_parts: &[&str]) { + let file = std::fs::File::create(path).unwrap(); + let mut zip = zip::ZipWriter::new(file); + let options = SimpleFileOptions::default(); + for (name, contents) in parts { + zip.start_file(*name, options).unwrap(); + zip.write_all(contents.as_bytes()).unwrap(); + } + let manifest = serde_json::json!({ "format": "dbx-sql-export-parts-v1", "sourceFileName": "source.sql", "parts": manifest_parts }); + zip.start_file("manifest.json", options).unwrap(); + zip.write_all(serde_json::to_string(&manifest).unwrap().as_bytes()).unwrap(); + zip.finish().unwrap(); + } + + #[test] + fn preserves_manifest_order_when_extracting_parts() { + let dir = tempfile::tempdir().unwrap(); + let package_path = dir.path().join("package.zip"); + write_package(&package_path, &[("z.sql", "SELECT 2;\n"), ("a.sql", "SELECT 1;\n")], &["a.sql", "z.sql"]); + let destination = dir.path().join("parts"); + let package = extract_sql_file_zip_package(&package_path, &destination).unwrap(); + let paths = extracted_sql_zip_paths(&destination, &package); + assert_eq!(std::fs::read_to_string(&paths[0]).unwrap(), "SELECT 1;\n"); + assert_eq!(std::fs::read_to_string(&paths[1]).unwrap(), "SELECT 2;\n"); + } + + #[test] + fn rejects_unsafe_manifest_paths() { + let dir = tempfile::tempdir().unwrap(); + let package_path = dir.path().join("package.zip"); + write_package(&package_path, &[("safe.sql", "SELECT 1;\n")], &["../safe.sql"]); + assert!(inspect_sql_file_zip_package(&package_path).unwrap_err().contains("Unsafe")); + } + + #[test] + fn rejects_parts_whose_declared_size_understates_content() { + let dir = tempfile::tempdir().unwrap(); + let package_path = dir.path().join("package.zip"); + write_package(&package_path, &[("a.sql", &"SELECT 1;\n".repeat(64))], &["a.sql"]); + // Shrink the central-directory uncompressed size of the first entry so the declared + // metadata understates the real decompressed content: the bounded copy must stop at the + // declared size instead of inflating whatever the stream contains onto disk. + let mut bytes = std::fs::read(&package_path).unwrap(); + let eocd = bytes.windows(4).rposition(|window| window == b"PK\x05\x06").unwrap(); + let directory_offset = u32::from_le_bytes(bytes[eocd + 16..eocd + 20].try_into().unwrap()) as usize; + assert_eq!(&bytes[directory_offset..directory_offset + 4], b"PK\x01\x02"); + let declared_size_at = directory_offset + 24; + bytes[declared_size_at..declared_size_at + 4].copy_from_slice(&1u32.to_le_bytes()); + std::fs::write(&package_path, bytes).unwrap(); + + let destination = dir.path().join("parts"); + let error = extract_sql_file_zip_package(&package_path, &destination).unwrap_err(); + assert!(error.contains("Failed to fully extract SQL part"), "{error}"); + } +} diff --git a/crates/dbx-core/tests/live_mysql57.rs b/crates/dbx-core/tests/live_mysql57.rs index 195823e5f6..bdff3caa61 100644 --- a/crates/dbx-core/tests/live_mysql57.rs +++ b/crates/dbx-core/tests/live_mysql57.rs @@ -1288,6 +1288,8 @@ INSERT INTO install_check (id) VALUES (1), (2); .into_owned(), continue_on_error: false, selected_tables: None, + part_cooldown_ms: 0, + skip_relational_constraints: false, }; let _ = execute_sql_statement( @@ -1374,6 +1376,8 @@ INSERT INTO children (parent_id) VALUES (LAST_INSERT_ID()); .into_owned(), continue_on_error: false, selected_tables: None, + part_cooldown_ms: 0, + skip_relational_constraints: false, }; let _ = execute_sql_statement( @@ -1444,6 +1448,8 @@ async fn live_sql_file_import_preserves_raw_mysql_binary_literal_bytes() { file_path: std::env::temp_dir().join(format!("mysql-binary-dump-{suffix}.sql")).to_string_lossy().into_owned(), continue_on_error: false, selected_tables: None, + part_cooldown_ms: 0, + skip_relational_constraints: false, }; tokio::fs::write(&request.file_path, script).await.unwrap(); diff --git a/crates/dbx-core/tests/live_mysql_database_export.rs b/crates/dbx-core/tests/live_mysql_database_export.rs index a89f0bf0e3..92fe4a5ec4 100644 --- a/crates/dbx-core/tests/live_mysql_database_export.rs +++ b/crates/dbx-core/tests/live_mysql_database_export.rs @@ -76,6 +76,8 @@ async fn live_mysql_selected_table_restore_preserves_unselected_tables() { let request = SqlFileRequest { execution_id: format!("restore-{suffix}-{compressed}"), connection_id: connection_id.clone(), database: database.clone(), file_path: path.display().to_string(), continue_on_error: false, selected_tables: Some(vec![SqlFileTable { database: Some(database.clone()), name: "chosen".into() }]), + part_cooldown_ms: 0, + skip_relational_constraints: false, }; execute_sql_file_path(&state, &request, &path, CancellationToken::new(), std::time::Instant::now(), |_| {}).await?; let rows = execute_sql_statement(&state, &connection_id, &database, "SELECT id, LENGTH(body) FROM chosen", None, None).await?; @@ -90,6 +92,8 @@ async fn live_mysql_selected_table_restore_preserves_unselected_tables() { let request = SqlFileRequest { execution_id: format!("invalid-{suffix}"), connection_id: connection_id.clone(), database: database.clone(), file_path: path.display().to_string(), continue_on_error: true, selected_tables: Some(vec![SqlFileTable { database: None, name: "chosen".into() }]), + part_cooldown_ms: 0, + skip_relational_constraints: false, }; let mut events = Vec::new(); assert!(execute_sql_file_path(&state, &request, &path, CancellationToken::new(), std::time::Instant::now(), |event| events.push(event)).await.is_err()); @@ -171,6 +175,8 @@ async fn live_mysql_database_export_restores_dependent_views() { file_path: file_path.to_string_lossy().to_string(), continue_on_error: false, selected_tables: None, + part_cooldown_ms: 0, + skip_relational_constraints: false, }; execute_sql_file_path( &state, diff --git a/crates/dbx-core/tests/live_postgres_all_schema_export.rs b/crates/dbx-core/tests/live_postgres_all_schema_export.rs index 06f390cf53..ab087b3cb8 100644 --- a/crates/dbx-core/tests/live_postgres_all_schema_export.rs +++ b/crates/dbx-core/tests/live_postgres_all_schema_export.rs @@ -150,6 +150,8 @@ async fn live_postgres_all_schema_export_restores_one_sql_file() { file_path: export_path.display().to_string(), continue_on_error: false, selected_tables: None, + part_cooldown_ms: 0, + skip_relational_constraints: false, }, &export_path, CancellationToken::new(), diff --git a/crates/dbx-core/tests/live_sqlserver_completion.rs b/crates/dbx-core/tests/live_sqlserver_completion.rs index 6715ee35c9..c46d89f0b7 100644 --- a/crates/dbx-core/tests/live_sqlserver_completion.rs +++ b/crates/dbx-core/tests/live_sqlserver_completion.rs @@ -1295,6 +1295,8 @@ async fn live_sqlserver_sql_file_import_executes_go_batches() { file_path: "fixture.sql".to_string(), continue_on_error: false, selected_tables: None, + part_cooldown_ms: 0, + skip_relational_constraints: false, }; let done_seen = AtomicBool::new(false); diff --git a/crates/dbx-web/src/routes/sql_file.rs b/crates/dbx-web/src/routes/sql_file.rs index 9bf1952720..9c216d6bee 100644 --- a/crates/dbx-web/src/routes/sql_file.rs +++ b/crates/dbx-web/src/routes/sql_file.rs @@ -5,7 +5,7 @@ use std::time::Duration; use axum::extract::{Multipart, Path as AxumPath, State}; use axum::response::sse::{Event, Sse}; use axum::Json; -use dbx_core::sql::{SqlFileProgress, SqlFileRequest, SqlFileStatus}; +use dbx_core::sql::{self, SqlFileProgress, SqlFileRequest, SqlFileStatus}; use dbx_core::sql_file_import::{ execute_sql_file_paths, sql_file_error_progress, sql_file_progress as build_sql_file_progress, SqlFileProgressEmitter, @@ -95,6 +95,31 @@ pub async fn preview_sql_file( let file_path = safe_uploaded_sql_path(&tmp_dir, &file_name)?; std::fs::write(&file_path, &data).map_err(|e| AppError::from(e.to_string()))?; + if file_name.to_ascii_lowercase().ends_with(".zip") { + let extraction_dir = tmp_dir.join(format!("package-{}", Uuid::new_v4())); + let package = dbx_core::sql_file_zip_package::extract_sql_file_zip_package(&file_path, &extraction_dir) + .map_err(AppError::from)?; + let paths = dbx_core::sql_file_zip_package::extracted_sql_zip_paths(&extraction_dir, &package) + .into_iter() + .map(|path| path.to_string_lossy().to_string()) + .collect::>(); + let content = + sql::decode_sql_file_bytes(&std::fs::read(&paths[0]).map_err(|e| AppError::from(e.to_string()))?) + .map_err(AppError::from)?; + let preview: String = content.chars().take(20_000).collect(); + let bootstrap_analysis = dbx_core::sql_file_import::mysql_like_sql_file_bootstrap_analysis(&content); + return Ok(Json(serde_json::json!({ + "fileName": file_name, + "filePath": file_path.to_string_lossy(), + "sizeBytes": data.len(), + "preview": preview, + "canExecuteWithoutSelectedDatabase": bootstrap_analysis.can_execute_without_selected_database, + "establishesDatabaseContext": bootstrap_analysis.establishes_database_context, + "packageFilePaths": paths, + "packagePartCount": package.part_names.len(), + }))); + } + let size_bytes = data.len() as u64; let content = dbx_core::sql_file_import::read_sql_file_preview(&file_path, 1_000_000).await.map_err(AppError::from)?; @@ -199,6 +224,7 @@ pub async fn execute_sql_file( }) .await; + cleanup_sql_file_package_paths(&file_paths); cleanup_sql_file_execution(&state_clone, &req.execution_id).await; }); @@ -211,6 +237,21 @@ fn send_sql_file_progress(tx: &broadcast::Sender, progress: SqlFileProgr } } +fn cleanup_sql_file_package_paths(file_paths: &[PathBuf]) { + let mut directories = std::collections::HashSet::new(); + for path in file_paths { + let Some(parent) = path.parent() else { + continue; + }; + if parent.file_name().and_then(|name| name.to_str()).is_some_and(|name| name.starts_with("package-")) { + directories.insert(parent.to_path_buf()); + } + } + for directory in directories { + let _ = std::fs::remove_dir_all(directory); + } +} + async fn cleanup_sql_file_execution(state: &WebState, execution_id: &str) { state.remove_sse_channel(execution_id).await; state.sql_file_executions.write().await.remove(execution_id); @@ -245,14 +286,19 @@ fn cleanup_sql_file_uploads_older_than(tmp_dir: &Path, max_age: Duration) { }; for entry in entries.flatten() { let path = entry.path(); - let expired = entry - .metadata() - .ok() - .filter(|metadata| metadata.is_file()) - .and_then(|metadata| metadata.modified().ok()) - .and_then(|modified| modified.elapsed().ok()) - .is_some_and(|age| age >= max_age); - if expired { + let Ok(metadata) = entry.metadata() else { + continue; + }; + // Extracted `package-*` directories count too: a preview that never reaches execution + // leaves them behind, and they must expire with the uploaded files. + let expired = + metadata.modified().ok().and_then(|modified| modified.elapsed().ok()).is_some_and(|age| age >= max_age); + if !expired { + continue; + } + if metadata.is_dir() { + let _ = std::fs::remove_dir_all(path); + } else { let _ = std::fs::remove_file(path); } } diff --git a/src-tauri/src/commands/sql_file.rs b/src-tauri/src/commands/sql_file.rs index 76e94bdec7..9cf5a9dec6 100644 --- a/src-tauri/src/commands/sql_file.rs +++ b/src-tauri/src/commands/sql_file.rs @@ -41,6 +41,40 @@ pub async fn inspect_sql_file_tables( pub async fn preview_sql_file(file_path: String) -> Result { let path = PathBuf::from(&file_path); let metadata = tokio::fs::metadata(&path).await.map_err(|e| e.to_string())?; + if path + .extension() + .and_then(|extension| extension.to_str()) + .is_some_and(|extension| extension.eq_ignore_ascii_case("zip")) + { + sweep_stale_sql_zip_packages(); + let extraction_dir = std::env::temp_dir().join(format!("dbx-sql-package-{}", uuid::Uuid::new_v4())); + let (package, extracted_paths) = tokio::task::spawn_blocking({ + let path = path.clone(); + let extraction_dir = extraction_dir.clone(); + move || { + let package = dbx_core::sql_file_zip_package::extract_sql_file_zip_package(&path, &extraction_dir)?; + let paths = dbx_core::sql_file_zip_package::extracted_sql_zip_paths(&extraction_dir, &package) + .into_iter() + .map(|part| part.to_string_lossy().to_string()) + .collect::>(); + Ok::<_, String>((package, paths)) + } + }) + .await + .map_err(|error| format!("Failed to extract SQL ZIP package: {error}"))??; + let prefix = read_sql_file_preview(PathBuf::from(&extracted_paths[0]).as_path(), 1_000_000).await?; + let bootstrap_analysis = mysql_like_sql_file_bootstrap_analysis(&prefix); + return Ok(SqlFilePreview { + file_name: path.file_name().and_then(|name| name.to_str()).unwrap_or("package.zip").to_string(), + file_path, + size_bytes: metadata.len(), + preview: prefix.chars().take(20_000).collect(), + can_execute_without_selected_database: bootstrap_analysis.can_execute_without_selected_database, + establishes_database_context: bootstrap_analysis.establishes_database_context, + package_file_paths: Some(extracted_paths), + package_part_count: Some(package.part_names.len()), + }); + } let prefix = read_sql_file_preview(&path, 1_000_000).await?; let bootstrap_analysis = mysql_like_sql_file_bootstrap_analysis(&prefix); let preview = prefix.chars().take(20_000).collect(); @@ -52,6 +86,8 @@ pub async fn preview_sql_file(file_path: String) -> Result= MAX_AGE); + if expired { + let _ = std::fs::remove_dir_all(entry.path()); + } + } +} + #[tauri::command] pub async fn cancel_sql_file_execution(execution_id: String) -> Result { let executions = sql_file_executions().read().await;