diff --git a/Cargo.lock b/Cargo.lock index a92322f35..311926253 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1842,6 +1842,7 @@ dependencies = [ "csv", "deadpool-postgres", "duckdb", + "encoding_rs", "font-kit", "futures", "iana-time-zone", diff --git a/crates/dbx-core/Cargo.toml b/crates/dbx-core/Cargo.toml index 1ec64841a..e96fd66ad 100644 --- a/crates/dbx-core/Cargo.toml +++ b/crates/dbx-core/Cargo.toml @@ -12,6 +12,7 @@ serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" regex = "1" percent-encoding = "2" +encoding_rs = "0.8" log = "0.4" tokio = { version = "1", features = ["full"] } tokio-util = { version = "0.7", features = ["compat"] } diff --git a/crates/dbx-core/src/sql.rs b/crates/dbx-core/src/sql.rs index 2a988978b..58cc14bd4 100644 --- a/crates/dbx-core/src/sql.rs +++ b/crates/dbx-core/src/sql.rs @@ -73,6 +73,39 @@ pub struct SqlFileProgress { pub error: Option, } +pub fn decode_sql_file_bytes(bytes: &[u8]) -> Result { + if bytes.starts_with(&[0xEF, 0xBB, 0xBF]) { + return std::str::from_utf8(&bytes[3..]).map(|text| text.to_string()).map_err(|_| sql_file_encoding_error()); + } + + if bytes.starts_with(&[0xFF, 0xFE]) { + return decode_sql_file_with_encoding(&bytes[2..], encoding_rs::UTF_16LE); + } + + if bytes.starts_with(&[0xFE, 0xFF]) { + return decode_sql_file_with_encoding(&bytes[2..], encoding_rs::UTF_16BE); + } + + if let Ok(text) = std::str::from_utf8(bytes) { + return Ok(text.strip_prefix('\u{feff}').unwrap_or(text).to_string()); + } + + decode_sql_file_with_encoding(bytes, encoding_rs::GBK) +} + +fn decode_sql_file_with_encoding(bytes: &[u8], encoding: &'static encoding_rs::Encoding) -> Result { + let (text, had_errors) = encoding.decode_without_bom_handling(bytes); + if had_errors { + return Err(sql_file_encoding_error()); + } + Ok(text.into_owned()) +} + +fn sql_file_encoding_error() -> String { + "Unsupported SQL file encoding. Save the file as UTF-8, UTF-8 with BOM, UTF-16 with BOM, or GBK, then try again." + .to_string() +} + #[derive(Default)] pub struct SqlStatementSplitter { buffer: String, @@ -1003,7 +1036,7 @@ mod tests { use crate::models::connection::DatabaseType; use super::{ - find_statement_at_cursor_for_database, prepare_sql_file_statement, split_sql_script, + decode_sql_file_bytes, find_statement_at_cursor_for_database, prepare_sql_file_statement, split_sql_script, split_sql_statements_for_database, starts_with_executable_sql_keyword, starts_with_executable_sql_keyword_for_database, SqlFileStatementAction, SqlStatementSplitter, }; @@ -1016,6 +1049,32 @@ mod tests { ); } + #[test] + fn decodes_utf8_bom_sql_file_bytes_without_bom_statement_prefix() { + let sql = decode_sql_file_bytes(b"\xEF\xBB\xBFCREATE TABLE t(id int);").unwrap(); + + assert_eq!(sql, "CREATE TABLE t(id int);"); + } + + #[test] + fn decodes_gbk_sql_file_bytes_before_execution() { + let bytes = b"INSERT INTO t VALUES ('\xD6\xD0\xCE\xC4');"; + let sql = decode_sql_file_bytes(bytes).unwrap(); + + assert_eq!(sql, "INSERT INTO t VALUES ('中文');"); + } + + #[test] + fn decodes_utf16le_bom_sql_file_bytes() { + let bytes = [ + 0xFF, 0xFE, b'S', 0x00, b'E', 0x00, b'L', 0x00, b'E', 0x00, b'C', 0x00, b'T', 0x00, b' ', 0x00, b'1', 0x00, + b';', 0x00, + ]; + let sql = decode_sql_file_bytes(&bytes).unwrap(); + + assert_eq!(sql, "SELECT 1;"); + } + #[test] fn ignores_semicolons_inside_quotes_and_comments() { let sql = "\ diff --git a/crates/dbx-web/src/routes/sql_file.rs b/crates/dbx-web/src/routes/sql_file.rs index 710558419..3874de37b 100644 --- a/crates/dbx-web/src/routes/sql_file.rs +++ b/crates/dbx-web/src/routes/sql_file.rs @@ -50,7 +50,7 @@ pub async fn preview_sql_file( std::fs::write(&file_path, &data).map_err(|e| AppError(e.to_string()))?; let size_bytes = data.len() as u64; - let content = String::from_utf8_lossy(&data); + let content = sql::decode_sql_file_bytes(&data).map_err(AppError)?; let preview: String = content.chars().take(5000).collect(); return Ok(Json(serde_json::json!({ @@ -118,8 +118,11 @@ pub async fn execute_sql_file( _ => {} } - let file_content = match std::fs::read_to_string(&file_path) { - Ok(c) => c, + let file_content = match std::fs::read(&file_path).and_then(|bytes| { + sql::decode_sql_file_bytes(&bytes) + .map_err(|message| std::io::Error::new(std::io::ErrorKind::InvalidData, message)) + }) { + Ok(content) => content, Err(e) => { let progress = dbx_core::sql::SqlFileProgress { execution_id: req.execution_id.clone(), @@ -155,11 +158,14 @@ pub async fn execute_sql_file( let _ = tx.send(json); } - 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 statements = import_target + .as_ref() + .map(|(db_type, _)| sql::split_sql_statements_for_database(&file_content, *db_type)) + .unwrap_or_else(|| sql::split_sql_statements(&file_content)); let start = std::time::Instant::now(); let mut success_count = 0usize; let mut failure_count = 0usize; diff --git a/src-tauri/src/commands/sql_file.rs b/src-tauri/src/commands/sql_file.rs index 91e3f296e..62b3850aa 100644 --- a/src-tauri/src/commands/sql_file.rs +++ b/src-tauri/src/commands/sql_file.rs @@ -4,7 +4,6 @@ use std::sync::Arc; use std::time::Instant; use tauri::{AppHandle, Emitter, State}; -use tokio::io::{AsyncBufReadExt, BufReader}; use tokio::sync::RwLock; use tokio_util::sync::CancellationToken; @@ -13,8 +12,8 @@ use crate::commands::query::execute_sql_statement; use dbx_core::models::connection::DatabaseType; pub use dbx_core::sql::{ - prepare_sql_file_statement, statement_summary, SqlFilePreview, SqlFileProgress, SqlFileRequest, - SqlFileStatementAction, SqlFileStatus, SqlStatementSplitter, + decode_sql_file_bytes, prepare_sql_file_statement, statement_summary, SqlFilePreview, SqlFileProgress, + SqlFileRequest, SqlFileStatementAction, SqlFileStatus, SqlParsingOptions, SqlStatementSplitter, }; static SQL_FILE_EXECUTIONS: std::sync::LazyLock>> = @@ -46,11 +45,8 @@ struct SqlFileSummary { 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())?; - let mut file = tokio::fs::File::open(&path).await.map_err(|e| e.to_string())?; - let mut buffer = vec![0; 4096]; - let bytes_read = tokio::io::AsyncReadExt::read(&mut file, &mut buffer).await.map_err(|e| e.to_string())?; - buffer.truncate(bytes_read); - let preview = String::from_utf8_lossy(&buffer).to_string(); + let bytes = tokio::fs::read(&path).await.map_err(|e| e.to_string())?; + let preview = decode_sql_file_bytes(&bytes)?.chars().take(5000).collect(); Ok(SqlFilePreview { file_name: path.file_name().and_then(|name| name.to_str()).unwrap_or("script.sql").to_string(), @@ -106,8 +102,8 @@ async fn execute_sql_file_inner( let mut failure_count = 0; let mut affected_rows = 0; - let file = match tokio::fs::File::open(&request.file_path).await { - Ok(file) => file, + let file_bytes = match tokio::fs::read(&request.file_path).await { + Ok(bytes) => bytes, Err(error) => { let error = error.to_string(); emit_file_io_error_progress( @@ -123,12 +119,30 @@ async fn execute_sql_file_inner( return Err(error); } }; - let mut reader = BufReader::new(file); - let mut splitter = SqlStatementSplitter::default(); - let mut line = String::new(); + let file_content = match decode_sql_file_bytes(&file_bytes) { + Ok(content) => content, + Err(error) => { + emit_file_io_error_progress( + app, + &request.execution_id, + statement_index, + success_count, + failure_count, + affected_rows, + started_at, + error.clone(), + ); + return Err(error); + } + }; let import_target = sql_file_import_target(state.inner().as_ref(), &request.connection_id).await; + let options = + import_target.as_ref().map(|target| SqlParsingOptions::for_database_type(target.db_type)).unwrap_or_default(); + let mut splitter = SqlStatementSplitter::with_options(options); + let mut statements = splitter.push_chunk(&file_content); + statements.extend(splitter.finish()); - loop { + for statement in statements { if token.is_cancelled() { emit_progress( app, @@ -145,51 +159,6 @@ async fn execute_sql_file_inner( return Ok(()); } - line.clear(); - let bytes_read = match reader.read_line(&mut line).await { - Ok(bytes_read) => bytes_read, - Err(error) => { - let error = error.to_string(); - emit_file_io_error_progress( - app, - &request.execution_id, - statement_index, - success_count, - failure_count, - affected_rows, - started_at, - error.clone(), - ); - return Err(error); - } - }; - if bytes_read == 0 { - break; - } - - for statement in splitter.push_chunk(&line) { - statement_index += 1; - if execute_statement_with_progress( - app, - state, - request, - &token, - started_at, - statement_index, - &statement, - import_target.as_ref(), - &mut success_count, - &mut failure_count, - &mut affected_rows, - ) - .await? - { - return Ok(()); - } - } - } - - for statement in splitter.finish() { statement_index += 1; if execute_statement_with_progress( app,