From f4c7c1d764064929b7fea5131accf3582bef9ae8 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Sat, 16 May 2026 12:13:40 +0800 Subject: [PATCH] fix: handle MySQL backup SQL import --- crates/dbx-core/src/sql.rs | 148 +++++++++++++++++++++++++++- docs/content/docs/sql-file.cn.mdx | 2 +- docs/content/docs/sql-file.mdx | 2 +- src-tauri/src/commands/sql_file.rs | 63 +++++++++++- src-web/src/routes/sql_file.rs | 42 +++++++- src/components/sidebar/TreeItem.vue | 39 +++++++- src/i18n/locales/en.ts | 4 + src/i18n/locales/zh-CN.ts | 4 + src/lib/createDatabaseSql.ts | 43 ++++++++ tests/createDatabaseSql.test.ts | 34 +++++++ 10 files changed, 371 insertions(+), 10 deletions(-) create mode 100644 src/lib/createDatabaseSql.ts create mode 100644 tests/createDatabaseSql.test.ts diff --git a/crates/dbx-core/src/sql.rs b/crates/dbx-core/src/sql.rs index ebf8f19ad..0b4dad8df 100644 --- a/crates/dbx-core/src/sql.rs +++ b/crates/dbx-core/src/sql.rs @@ -1,5 +1,7 @@ use serde::{Deserialize, Serialize}; +use crate::models::connection::DatabaseType; + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct SqlFileRequest { @@ -31,6 +33,12 @@ pub enum SqlFileStatus { Cancelled, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SqlFileStatementAction { + Execute(String), + Skip, +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct SqlFileProgress { @@ -247,6 +255,28 @@ pub fn statement_summary(statement: &str) -> String { collapsed.chars().take(MAX_LEN).collect() } +pub fn prepare_sql_file_statement( + statement: &str, + db_type: &DatabaseType, + driver_profile: Option<&str>, +) -> SqlFileStatementAction { + let statement = statement.trim(); + let Some(body) = mysql_executable_comment_body(statement) else { + return SqlFileStatementAction::Execute(statement.to_string()); + }; + + if !is_mysql_compatible_import_target(db_type, driver_profile) { + return SqlFileStatementAction::Skip; + } + + let body = body.trim(); + if body.is_empty() || is_mysql_key_toggle_statement(body) { + return SqlFileStatementAction::Skip; + } + + SqlFileStatementAction::Execute(body.to_string()) +} + pub fn starts_with_executable_sql_keyword(sql: &str, keywords: &[&str]) -> bool { let Some(token) = first_executable_sql_token(sql) else { return false; @@ -254,6 +284,81 @@ pub fn starts_with_executable_sql_keyword(sql: &str, keywords: &[&str]) -> bool keywords.iter().any(|keyword| token.eq_ignore_ascii_case(keyword)) } +fn is_mysql_compatible_import_target(db_type: &DatabaseType, driver_profile: Option<&str>) -> bool { + matches!(db_type, DatabaseType::Mysql | DatabaseType::Doris | DatabaseType::StarRocks | DatabaseType::Goldendb) + || driver_profile.map(|profile| profile.to_ascii_lowercase()).is_some_and(|profile| { + matches!( + profile.as_str(), + "mariadb" | "tidb" | "oceanbase" | "custom_mysql" | "doris" | "starrocks" | "selectdb" | "goldendb" + ) + }) +} + +fn mysql_executable_comment_body(statement: &str) -> Option<&str> { + let bytes = statement.as_bytes(); + let start = leading_mysql_executable_comment_start(statement)?; + let body_start = if bytes.get(start + 2) == Some(&b'!') { start + 3 } else { start + 4 }; + let mut body_start = body_start; + while body_start < bytes.len() && (bytes[body_start].is_ascii_digit() || bytes[body_start].is_ascii_whitespace()) { + body_start += 1; + } + + let close = find_block_comment_close(bytes, body_start)?; + if has_executable_sql(&statement[close + 2..]) { + return None; + } + + Some(&statement[body_start..close]) +} + +fn leading_mysql_executable_comment_start(statement: &str) -> Option { + let bytes = statement.as_bytes(); + let mut i = 0; + + while i < bytes.len() { + while i < bytes.len() && bytes[i].is_ascii_whitespace() { + i += 1; + } + + if i + 1 < bytes.len() && bytes[i] == b'-' && bytes[i + 1] == b'-' { + i += 2; + while i < bytes.len() && bytes[i] != b'\n' { + i += 1; + } + continue; + } + + if i + 1 < bytes.len() && bytes[i] == b'/' && bytes[i + 1] == b'*' { + if i + 2 < bytes.len() && (bytes[i + 2] == b'!' || (i + 3 < bytes.len() && &bytes[i + 2..i + 4] == b"M!")) { + return Some(i); + } + + let close = find_block_comment_close(bytes, i + 2)?; + i = close + 2; + continue; + } + + return None; + } + + None +} + +fn find_block_comment_close(bytes: &[u8], mut start: usize) -> Option { + while start + 1 < bytes.len() { + if bytes[start] == b'*' && bytes[start + 1] == b'/' { + return Some(start); + } + start += 1; + } + None +} + +fn is_mysql_key_toggle_statement(statement: &str) -> bool { + let upper = statement.split_whitespace().collect::>().join(" ").to_ascii_uppercase(); + upper.starts_with("ALTER TABLE ") && (upper.ends_with(" ENABLE KEYS") || upper.ends_with(" DISABLE KEYS")) +} + fn first_executable_sql_token(sql: &str) -> Option<&str> { let bytes = sql.as_bytes(); let mut i = 0; @@ -399,7 +504,12 @@ fn split_sql_script(sql: &str) -> Result, String> { #[cfg(test)] mod tests { - use super::{split_sql_script, starts_with_executable_sql_keyword, SqlStatementSplitter}; + use crate::models::connection::DatabaseType; + + use super::{ + prepare_sql_file_statement, split_sql_script, starts_with_executable_sql_keyword, SqlFileStatementAction, + SqlStatementSplitter, + }; #[test] fn splits_semicolon_delimited_statements() { @@ -504,6 +614,42 @@ mod tests { assert!(starts_with_executable_sql_keyword("/*M! SELECT 1 */", &["SELECT"])); } + #[test] + fn prepares_mysql_executable_comments_for_mysql_compatible_imports() { + assert_eq!( + prepare_sql_file_statement( + "/*!40101 SET character_set_client = @saved_cs_client */", + &DatabaseType::Mysql, + None + ), + SqlFileStatementAction::Execute("SET character_set_client = @saved_cs_client".to_string()) + ); + } + + #[test] + fn skips_mysql_key_toggle_comments_for_mysql_compatible_imports() { + assert_eq!( + prepare_sql_file_statement(" /*!40000 ALTER TABLE `dd_admin` ENABLE KEYS */", &DatabaseType::Mysql, None), + SqlFileStatementAction::Skip + ); + assert_eq!( + prepare_sql_file_statement("/*!40000 ALTER TABLE `dd_admin` DISABLE KEYS */", &DatabaseType::Mysql, None), + SqlFileStatementAction::Skip + ); + } + + #[test] + fn skips_mysql_executable_comments_for_non_mysql_imports() { + assert_eq!( + prepare_sql_file_statement( + "/*!40101 SET character_set_client = @saved_cs_client */", + &DatabaseType::Postgres, + None + ), + SqlFileStatementAction::Skip + ); + } + #[test] fn split_batches_by_go() { assert_eq!(super::split_sql_batches("SELECT 1\nGO\nSELECT 2"), vec!["SELECT 1", "SELECT 2"]); diff --git a/docs/content/docs/sql-file.cn.mdx b/docs/content/docs/sql-file.cn.mdx index 22782194c..e46a50b66 100644 --- a/docs/content/docs/sql-file.cn.mdx +++ b/docs/content/docs/sql-file.cn.mdx @@ -46,7 +46,7 @@ DBX 会按流式方式读取 SQL 文件，并尽量正确识别语句边界， - 字符串中的分号不会被当作语句结束 - 行注释和块注释中的分号不会被当作语句结束 - PostgreSQL 的 dollar-quoted 函数体会保持为同一条语句 -- MySQL executable comments 会作为可执行语句处理 +- MySQL executable comments 会在导入前按目标连接处理：MySQL 兼容连接会展开可执行内容，非 MySQL 连接会跳过；`ENABLE/DISABLE KEYS` 这类恢复优化指令会跳过 ## 错误处理 diff --git a/docs/content/docs/sql-file.mdx b/docs/content/docs/sql-file.mdx index bfbd08c92..2b1b28ebc 100644 --- a/docs/content/docs/sql-file.mdx +++ b/docs/content/docs/sql-file.mdx @@ -46,7 +46,7 @@ DBX reads SQL files as a stream and tries to detect statement boundaries correct - Semicolons inside strings are not treated as statement endings - Semicolons inside line and block comments are ignored - PostgreSQL dollar-quoted function bodies are kept together -- MySQL executable comments are treated as executable statements +- MySQL executable comments are prepared for the target connection before import: MySQL-compatible connections unwrap executable content, non-MySQL connections skip it, and restore-only `ENABLE/DISABLE KEYS` directives are skipped ## Error Handling diff --git a/src-tauri/src/commands/sql_file.rs b/src-tauri/src/commands/sql_file.rs index f73d7d696..91e3f296e 100644 --- a/src-tauri/src/commands/sql_file.rs +++ b/src-tauri/src/commands/sql_file.rs @@ -10,9 +10,11 @@ use tokio_util::sync::CancellationToken; use crate::commands::connection::AppState; use crate::commands::query::execute_sql_statement; +use dbx_core::models::connection::DatabaseType; pub use dbx_core::sql::{ - statement_summary, SqlFilePreview, SqlFileProgress, SqlFileRequest, SqlFileStatus, SqlStatementSplitter, + prepare_sql_file_statement, statement_summary, SqlFilePreview, SqlFileProgress, SqlFileRequest, + SqlFileStatementAction, SqlFileStatus, SqlStatementSplitter, }; static SQL_FILE_EXECUTIONS: std::sync::LazyLock>> = @@ -25,6 +27,12 @@ struct StatementErrorDecision { result: Result, } +#[derive(Debug, Clone)] +struct SqlFileImportTarget { + db_type: DatabaseType, + driver_profile: Option, +} + #[cfg(test)] #[derive(Debug, Clone, PartialEq, Eq)] struct SqlFileSummary { @@ -118,6 +126,7 @@ async fn execute_sql_file_inner( let mut reader = BufReader::new(file); let mut splitter = SqlStatementSplitter::default(); let mut line = String::new(); + let import_target = sql_file_import_target(state.inner().as_ref(), &request.connection_id).await; loop { if token.is_cancelled() { @@ -168,6 +177,7 @@ async fn execute_sql_file_inner( started_at, statement_index, &statement, + import_target.as_ref(), &mut success_count, &mut failure_count, &mut affected_rows, @@ -189,6 +199,7 @@ async fn execute_sql_file_inner( started_at, statement_index, &statement, + import_target.as_ref(), &mut success_count, &mut failure_count, &mut affected_rows, @@ -214,6 +225,13 @@ async fn execute_sql_file_inner( Ok(()) } +async fn sql_file_import_target(state: &AppState, connection_id: &str) -> Option { + let configs = state.configs.read().await; + configs + .get(connection_id) + .map(|config| SqlFileImportTarget { db_type: config.db_type, driver_profile: config.driver_profile.clone() }) +} + async fn execute_statement_with_progress( app: &AppHandle, state: &State<'_, Arc>, @@ -222,13 +240,13 @@ async fn execute_statement_with_progress( started_at: Instant, statement_index: usize, statement: &str, + import_target: Option<&SqlFileImportTarget>, success_count: &mut usize, failure_count: &mut usize, affected_rows: &mut u64, ) -> Result { - let summary = statement_summary(statement); - if token.is_cancelled() { + let summary = statement_summary(statement); emit_progress( app, &request.execution_id, @@ -244,6 +262,43 @@ async fn execute_statement_with_progress( return Ok(true); } + let statement_action = import_target + .map(|target| prepare_sql_file_statement(statement, &target.db_type, target.driver_profile.as_deref())) + .unwrap_or_else(|| SqlFileStatementAction::Execute(statement.to_string())); + let executable_statement = match statement_action { + SqlFileStatementAction::Execute(statement) => statement, + SqlFileStatementAction::Skip => { + let summary = statement_summary(statement); + emit_progress( + app, + &request.execution_id, + SqlFileStatus::Running, + statement_index, + *success_count, + *failure_count, + *affected_rows, + started_at, + &summary, + None, + ); + *success_count += 1; + emit_progress( + app, + &request.execution_id, + SqlFileStatus::StatementDone, + statement_index, + *success_count, + *failure_count, + *affected_rows, + started_at, + &summary, + None, + ); + return Ok(false); + } + }; + let summary = statement_summary(&executable_statement); + emit_progress( app, &request.execution_id, @@ -261,7 +316,7 @@ async fn execute_statement_with_progress( state.inner().as_ref(), &request.connection_id, &request.database, - statement, + &executable_statement, None, Some(token.clone()), ) diff --git a/src-web/src/routes/sql_file.rs b/src-web/src/routes/sql_file.rs index 9b3720fd3..125933b53 100644 --- a/src-web/src/routes/sql_file.rs +++ b/src-web/src/routes/sql_file.rs @@ -5,6 +5,7 @@ use axum::response::sse::{Event, Sse}; use axum::Json; use dbx_core::query; use dbx_core::sql; +use dbx_core::sql::SqlFileStatementAction; use futures::stream::Stream; use serde::Deserialize; @@ -160,13 +161,29 @@ pub async fn execute_sql_file( } let statements = sql::split_sql_statements(&file_content); + let import_target = { + let configs = app.configs.read().await; + configs.get(&req.connection_id).map(|config| (config.db_type, config.driver_profile.clone())) + }; let start = std::time::Instant::now(); let mut success_count = 0usize; let mut failure_count = 0usize; let mut total_affected: u64 = 0; for (i, stmt) in statements.iter().enumerate() { - let summary = sql::statement_summary(stmt); + let statement_action = import_target + .as_ref() + .map(|(db_type, driver_profile)| { + sql::prepare_sql_file_statement(stmt, db_type, driver_profile.as_deref()) + }) + .unwrap_or_else(|| SqlFileStatementAction::Execute(stmt.to_string())); + let (stmt_to_execute, summary, should_execute) = match statement_action { + SqlFileStatementAction::Execute(statement) => { + let summary = sql::statement_summary(&statement); + (statement, summary, true) + } + SqlFileStatementAction::Skip => (String::new(), sql::statement_summary(stmt), false), + }; // Send running let running = dbx_core::sql::SqlFileProgress { @@ -184,7 +201,28 @@ pub async fn execute_sql_file( let _ = tx.send(json); } - match query::execute_sql_statement(&app, &req.connection_id, &req.database, stmt, None, None).await { + if !should_execute { + success_count += 1; + let done = dbx_core::sql::SqlFileProgress { + execution_id: req.execution_id.clone(), + status: dbx_core::sql::SqlFileStatus::StatementDone, + statement_index: i, + success_count, + failure_count, + affected_rows: total_affected, + elapsed_ms: start.elapsed().as_millis(), + statement_summary: summary, + error: None, + }; + if let Ok(json) = serde_json::to_string(&done) { + let _ = tx.send(json); + } + continue; + } + + match query::execute_sql_statement(&app, &req.connection_id, &req.database, &stmt_to_execute, None, None) + .await + { Ok(result) => { success_count += 1; total_affected += result.affected_rows; diff --git a/src/components/sidebar/TreeItem.vue b/src/components/sidebar/TreeItem.vue index 922ab8713..26d8ef3d5 100644 --- a/src/components/sidebar/TreeItem.vue +++ b/src/components/sidebar/TreeItem.vue @@ -88,6 +88,7 @@ import { } from "@/lib/databaseCapabilities"; import { sidebarSelectionCopyAction, treeNodeRowAction, treeNodeRowDoubleClickAction } from "@/lib/treeNodeClick"; import { formatCsv, formatJson, formatSqlInsert } from "@/lib/exportFormats"; +import { buildCreateDatabaseSql, supportsCreateDatabaseCharset } from "@/lib/createDatabaseSql"; import { hexToRgba } from "@/lib/color"; import DangerConfirmDialog from "@/components/editor/DangerConfirmDialog.vue"; import { isTauriRuntime } from "@/lib/tauriRuntime"; @@ -604,6 +605,8 @@ const duplicateTableName = ref(""); const showCreateDatabaseDialog = ref(false); const createDatabaseName = ref(""); +const createDatabaseCharset = ref("utf8mb4"); +const createDatabaseCollation = ref("utf8mb4_unicode_ci"); const showDropDatabaseConfirm = ref(false); const showCreateSchemaDialog = ref(false); const createSchemaName = ref(""); @@ -682,6 +685,11 @@ const canCreateDatabase = computed(() => { return props.node.type === "connection" && supportsDatabaseCreation(config?.db_type); }); +const canSetCreateDatabaseCharset = computed(() => { + const config = props.node.connectionId ? connectionStore.getConfig(props.node.connectionId) : undefined; + return supportsCreateDatabaseCharset(config?.db_type, config?.driver_profile); +}); + const canDropDatabase = computed(() => { const config = props.node.connectionId ? connectionStore.getConfig(props.node.connectionId) : undefined; return props.node.type === "database" && supportsDatabaseCreation(config?.db_type); @@ -786,6 +794,8 @@ function buildDropSchemaSql(): string { function openCreateDatabaseDialog() { createDatabaseName.value = ""; + createDatabaseCharset.value = "utf8mb4"; + createDatabaseCollation.value = "utf8mb4_unicode_ci"; showCreateDatabaseDialog.value = true; } @@ -796,7 +806,14 @@ async function confirmCreateDatabase() { showCreateDatabaseDialog.value = false; try { await connectionStore.ensureConnected(node.connectionId); - const sql = `CREATE DATABASE ${quoteIdent(name)};`; + const config = connectionStore.getConfig(node.connectionId); + const sql = buildCreateDatabaseSql({ + databaseType: config?.db_type, + driverProfile: config?.driver_profile, + name, + charset: createDatabaseCharset.value, + collation: createDatabaseCollation.value, + }); await api.executeQuery(node.connectionId, "", sql); toast(t("contextMenu.createDatabaseSuccess", { name }), 3000); await connectionStore.loadDatabases(node.connectionId, { force: true }); @@ -2037,6 +2054,26 @@ const isDragging = computed(() => dragState.active && dragState.draggedId === pr :placeholder="t('contextMenu.createDatabaseNamePlaceholder')" @keydown.enter.prevent="confirmCreateDatabase" /> +
+
+ + +
+
+ + +
+