diff --git a/crates/dbx-core/src/sql_file_import.rs b/crates/dbx-core/src/sql_file_import.rs index 9745421541..a8c48dcec4 100644 --- a/crates/dbx-core/src/sql_file_import.rs +++ b/crates/dbx-core/src/sql_file_import.rs @@ -1,5 +1,7 @@ +use std::path::Path; use std::time::Instant; +use tokio::io::{AsyncReadExt, BufReader}; use tokio_util::sync::CancellationToken; use crate::connection::{AppState, PoolKind}; @@ -29,6 +31,23 @@ struct StatementErrorDecision { result: Result, } +const SQL_FILE_READ_CHUNK_BYTES: usize = 256 * 1024; +const SQL_FILE_STATEMENT_BATCH_SIZE: usize = 256; +const SQL_FILE_PREVIEW_ENCODING_SAMPLE_BYTES: usize = 1024 * 1024; + +struct SqlFileExecutionProgress { + statement_index: usize, + success_count: usize, + failure_count: usize, + affected_rows: u64, +} + +impl SqlFileExecutionProgress { + fn new() -> Self { + Self { statement_index: 0, success_count: 0, failure_count: 0, affected_rows: 0 } + } +} + struct MySqlSqlFileExecutor { connection_id: String, database: String, @@ -243,6 +262,7 @@ pub async fn execute_sql_file_content( // MySQL-family imports need one pinned connection so `USE` and session // state survive across the whole file. let mut mysql_executor = MySqlSqlFileExecutor::build(state, request, import_target.as_ref()).await?; + let mut progress = SqlFileExecutionProgress::new(); execute_planned_statements_with_progress( state, request, @@ -250,11 +270,351 @@ pub async fn execute_sql_file_content( started_at, &planned_statements, mysql_executor.as_mut(), + &mut progress, &mut emit, ) + .await?; + emit_sql_file_terminal_progress(request, &token, started_at, &progress, &mut emit); + Ok(()) +} + +pub async fn execute_sql_file_path( + state: &AppState, + request: &SqlFileRequest, + file_path: &Path, + token: CancellationToken, + started_at: Instant, + mut emit: impl FnMut(SqlFileProgress), +) -> Result<(), String> { + let import_target = sql_file_import_target(state, &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 = StreamingSqlFileSplitter::new(import_target.as_ref().map(|target| target.db_type), options); + let mut mysql_executor = MySqlSqlFileExecutor::build(state, request, import_target.as_ref()).await?; + let mut progress = SqlFileExecutionProgress::new(); + let mut pending_statements = Vec::with_capacity(SQL_FILE_STATEMENT_BATCH_SIZE); + let mut decoder = match SqlFileStreamDecoder::open(file_path).await { + Ok(decoder) => decoder, + Err(error) => { + emit(sql_file_error_progress(&request.execution_id, started_at, 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(()); + } + pending_statements.extend(splitter.push_chunk(&chunk)); + if pending_statements.len() >= SQL_FILE_STATEMENT_BATCH_SIZE { + 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?; + } + } + + pending_statements.extend(splitter.finish()); + 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?; + emit_sql_file_terminal_progress(request, &token, started_at, &progress, &mut emit); + Ok(()) +} + +pub async fn read_sql_file_preview(file_path: &Path, max_chars: usize) -> Result { + let mut decoder = + SqlFileStreamDecoder::open_with_detection_limit(file_path, Some(SQL_FILE_PREVIEW_ENCODING_SAMPLE_BYTES)) + .await?; + let mut preview = String::new(); + while preview.chars().count() < max_chars { + let Some(chunk) = decoder.next_chunk().await? else { + break; + }; + preview.push_str(&chunk); + } + Ok(preview.chars().take(max_chars).collect()) +} + +struct SqlFileStreamDecoder { + reader: BufReader, + decoder: encoding_rs::Decoder, + pending_bytes: Vec, + reached_eof: bool, +} + +impl SqlFileStreamDecoder { + async fn open(file_path: &Path) -> Result { + Self::open_with_detection_limit(file_path, None).await + } + + async fn open_with_detection_limit(file_path: &Path, detection_limit: Option) -> Result { + let (encoding, bom_len) = detect_sql_file_encoding(file_path, detection_limit).await?; + let mut file = tokio::fs::File::open(file_path).await.map_err(|error| error.to_string())?; + let mut prefix = [0u8; 3]; + let prefix_len = file.read(&mut prefix).await.map_err(|error| error.to_string())?; + let prefix = &prefix[..prefix_len]; + let mut pending_bytes = prefix[bom_len..].to_vec(); + pending_bytes.reserve(SQL_FILE_READ_CHUNK_BYTES); + Ok(Self { + reader: BufReader::with_capacity(SQL_FILE_READ_CHUNK_BYTES, file), + decoder: encoding.new_decoder_without_bom_handling(), + pending_bytes, + reached_eof: false, + }) + } + + async fn next_chunk(&mut self) -> Result, String> { + if self.reached_eof && self.pending_bytes.is_empty() { + return Ok(None); + } + while !self.reached_eof && self.pending_bytes.len() < SQL_FILE_READ_CHUNK_BYTES { + let mut buffer = vec![0u8; SQL_FILE_READ_CHUNK_BYTES - self.pending_bytes.len()]; + let read = self.reader.read(&mut buffer).await.map_err(|error| error.to_string())?; + if read == 0 { + self.reached_eof = true; + break; + } + self.pending_bytes.extend_from_slice(&buffer[..read]); + } + + let mut output = String::with_capacity( + self.decoder + .max_utf8_buffer_length_without_replacement(self.pending_bytes.len()) + .unwrap_or(self.pending_bytes.len()), + ); + let (result, read) = + self.decoder.decode_to_string_without_replacement(&self.pending_bytes, &mut output, self.reached_eof); + self.pending_bytes.drain(..read); + match result { + encoding_rs::DecoderResult::InputEmpty => Ok((!output.is_empty()).then_some(output)), + encoding_rs::DecoderResult::OutputFull => Ok(Some(output)), + encoding_rs::DecoderResult::Malformed(_, _) => Err("Unsupported or invalid SQL file encoding".to_string()), + } + } +} + +async fn detect_sql_file_encoding( + file_path: &Path, + detection_limit: Option, +) -> Result<(&'static encoding_rs::Encoding, usize), String> { + let mut file = tokio::fs::File::open(file_path).await.map_err(|error| error.to_string())?; + let mut prefix = [0u8; 3]; + let prefix_len = file.read(&mut prefix).await.map_err(|error| error.to_string())?; + let prefix = &prefix[..prefix_len]; + if prefix.starts_with(&[0xEF, 0xBB, 0xBF]) { + return Ok((encoding_rs::UTF_8, 3)); + } + if prefix.starts_with(&[0xFF, 0xFE]) { + return Ok((encoding_rs::UTF_16LE, 2)); + } + if prefix.starts_with(&[0xFE, 0xFF]) { + return Ok((encoding_rs::UTF_16BE, 2)); + } + + // SQL dumps often begin with ASCII comments even when the remaining file + // is GBK. Validate the entire stream as UTF-8 with bounded buffers before + // falling back to the legacy GBK behavior. + let mut decoder = encoding_rs::UTF_8.new_decoder_without_bom_handling(); + let mut input = prefix.to_vec(); + let mut inspected_bytes = prefix.len(); + let mut reached_eof = false; + loop { + let reached_detection_limit = detection_limit.is_some_and(|limit| inspected_bytes >= limit); + if !reached_eof && !reached_detection_limit { + let remaining = + detection_limit.map(|limit| limit.saturating_sub(inspected_bytes)).unwrap_or(SQL_FILE_READ_CHUNK_BYTES); + let mut buffer = vec![0u8; SQL_FILE_READ_CHUNK_BYTES.min(remaining.max(1))]; + let read = file.read(&mut buffer).await.map_err(|error| error.to_string())?; + if read == 0 { + reached_eof = true; + } else { + inspected_bytes += read; + input.extend_from_slice(&buffer[..read]); + } + } + let mut output = String::with_capacity( + decoder.max_utf8_buffer_length_without_replacement(input.len()).unwrap_or(input.len()), + ); + let (result, read) = decoder.decode_to_string_without_replacement(&input, &mut output, reached_eof); + input.drain(..read); + match result { + encoding_rs::DecoderResult::Malformed(_, _) => return Ok((encoding_rs::GBK, 0)), + encoding_rs::DecoderResult::InputEmpty if reached_eof || reached_detection_limit => { + return Ok((encoding_rs::UTF_8, 0)); + } + encoding_rs::DecoderResult::InputEmpty | encoding_rs::DecoderResult::OutputFull => {} + } + } +} + +enum StreamingSqlFileSplitter { + Statements(SqlStatementSplitter), + SqlServerBatches(SqlServerBatchSplitter), +} + +impl StreamingSqlFileSplitter { + fn new(db_type: Option, options: SqlParsingOptions) -> Self { + if db_type == Some(DatabaseType::SqlServer) { + Self::SqlServerBatches(SqlServerBatchSplitter::default()) + } else { + Self::Statements(SqlStatementSplitter::with_options(options)) + } + } + + fn push_chunk(&mut self, chunk: &str) -> Vec { + match self { + Self::Statements(splitter) => splitter.push_chunk(chunk), + Self::SqlServerBatches(splitter) => splitter.push_chunk(chunk), + } + } + + fn finish(self) -> Vec { + match self { + Self::Statements(splitter) => splitter.finish(), + Self::SqlServerBatches(splitter) => splitter.finish(), + } + } +} + +#[derive(Default)] +struct SqlServerBatchSplitter { + batch: String, + partial_line: String, +} + +impl SqlServerBatchSplitter { + fn push_chunk(&mut self, chunk: &str) -> Vec { + self.partial_line.push_str(chunk); + let mut batches = Vec::new(); + while let Some(newline) = self.partial_line.find('\n') { + let line = self.partial_line[..newline].trim_end_matches('\r').to_string(); + self.partial_line.drain(..=newline); + self.push_line(&line, &mut batches); + } + batches + } + + fn finish(mut self) -> Vec { + let mut batches = Vec::new(); + if !self.partial_line.is_empty() { + let line = std::mem::take(&mut self.partial_line); + self.push_line(line.trim_end_matches('\r'), &mut batches); + } + self.push_batch(&mut batches); + batches + } + + fn push_line(&mut self, line: &str, batches: &mut Vec) { + if line.trim().eq_ignore_ascii_case("go") { + self.push_batch(batches); + } else { + self.batch.push_str(line); + self.batch.push('\n'); + } + } + + fn push_batch(&mut self, batches: &mut Vec) { + let batch = self.batch.trim(); + if !batch.is_empty() { + batches.push(batch.to_string()); + } + self.batch.clear(); + } +} + +#[allow(clippy::too_many_arguments)] +async fn execute_sql_file_statement_batch( + state: &AppState, + request: &SqlFileRequest, + token: &CancellationToken, + started_at: Instant, + statements: &mut Vec, + import_target: Option<&SqlFileImportTarget>, + mysql_executor: Option<&mut MySqlSqlFileExecutor>, + progress: &mut SqlFileExecutionProgress, + emit: &mut impl FnMut(SqlFileProgress), +) -> Result<(), String> { + if statements.is_empty() { + return Ok(()); + } + let statements = std::mem::take(statements); + let planned_statements = optimize_sql_file_import_statements( + &statements, + import_target.map(|target| target.db_type), + import_target.and_then(|target| target.driver_profile.as_deref()), + ); + execute_planned_statements_with_progress( + state, + request, + token, + started_at, + &planned_statements, + mysql_executor, + progress, + emit, + ) .await } +fn emit_sql_file_terminal_progress( + request: &SqlFileRequest, + token: &CancellationToken, + started_at: Instant, + progress: &SqlFileExecutionProgress, + emit: &mut impl FnMut(SqlFileProgress), +) { + emit(sql_file_progress( + &request.execution_id, + if token.is_cancelled() { SqlFileStatus::Cancelled } else { SqlFileStatus::Done }, + progress.statement_index, + progress.success_count, + progress.failure_count, + progress.affected_rows, + started_at, + "", + None, + )); +} + fn split_sql_file_import_statements(file_content: &str, db_type: Option) -> Vec { if db_type == Some(DatabaseType::SqlServer) { // GO is a client-side batch delimiter, not T-SQL. SQL Server module DDL @@ -500,30 +860,15 @@ async fn execute_planned_statements_with_progress( started_at: Instant, planned_statements: &[SqlFileImportStatement], mut mysql_executor: Option<&mut MySqlSqlFileExecutor>, + progress: &mut SqlFileExecutionProgress, emit: &mut impl FnMut(SqlFileProgress), ) -> Result<(), String> { - let mut statement_index = 0; - let mut success_count = 0; - let mut failure_count = 0; - let mut affected_rows = 0; - for planned_statement in planned_statements { if token.is_cancelled() { - emit(sql_file_progress( - &request.execution_id, - SqlFileStatus::Cancelled, - statement_index, - success_count, - failure_count, - affected_rows, - started_at, - "", - None, - )); return Ok(()); } - let next_statement_index = statement_index + planned_statement.source_statement_count; + let next_statement_index = progress.statement_index + planned_statement.source_statement_count; if execute_statement_with_progress( state, request, @@ -531,9 +876,9 @@ async fn execute_planned_statements_with_progress( started_at, next_statement_index, planned_statement, - &mut success_count, - &mut failure_count, - &mut affected_rows, + &mut progress.success_count, + &mut progress.failure_count, + &mut progress.affected_rows, mysql_executor.as_deref_mut(), emit, ) @@ -541,20 +886,8 @@ async fn execute_planned_statements_with_progress( { return Ok(()); } - statement_index = next_statement_index; + progress.statement_index = next_statement_index; } - - emit(sql_file_progress( - &request.execution_id, - SqlFileStatus::Done, - statement_index, - success_count, - failure_count, - affected_rows, - started_at, - "", - None, - )); Ok(()) } @@ -910,6 +1243,19 @@ fn statement_error_decision( mod tests { use super::*; use crate::models::connection::DatabaseType; + use std::sync::atomic::{AtomicU64, Ordering}; + + static TEMP_SQL_FILE_SEQUENCE: AtomicU64 = AtomicU64::new(0); + + async fn temporary_sql_file(bytes: &[u8]) -> std::path::PathBuf { + let path = std::env::temp_dir().join(format!( + "dbx-sql-file-{}-{}.sql", + std::process::id(), + TEMP_SQL_FILE_SEQUENCE.fetch_add(1, Ordering::Relaxed) + )); + tokio::fs::write(&path, bytes).await.unwrap(); + path + } #[test] fn sqlserver_sql_file_splits_go_batches_without_sending_delimiters() { @@ -939,6 +1285,47 @@ mod tests { assert_eq!(statements[1], "SELECT 3;"); } + #[test] + fn streaming_sqlserver_splitter_handles_go_across_chunks() { + let mut splitter = SqlServerBatchSplitter::default(); + let mut batches = splitter.push_chunk("CREATE PROCEDURE dbo.demo AS\nBEGIN\nSELECT 1;\nEND\nG"); + batches.extend(splitter.push_chunk("O\r\nSELECT 2;\nGO\n")); + batches.extend(splitter.finish()); + + assert_eq!(batches, vec!["CREATE PROCEDURE dbo.demo AS\nBEGIN\nSELECT 1;\nEND", "SELECT 2;"]); + } + + #[tokio::test] + async fn streaming_decoder_detects_gbk_after_ascii_prefix() { + let (encoded, _, _) = encoding_rs::GBK.encode("-- Navicat dump\nINSERT INTO t VALUES ('中文');"); + let path = temporary_sql_file(encoded.as_ref()).await; + let mut decoder = SqlFileStreamDecoder::open(&path).await.unwrap(); + let mut decoded = String::new(); + while let Some(chunk) = decoder.next_chunk().await.unwrap() { + decoded.push_str(&chunk); + } + tokio::fs::remove_file(path).await.unwrap(); + + assert_eq!(decoded, "-- Navicat dump\nINSERT INTO t VALUES ('中文');"); + } + + #[tokio::test] + async fn streaming_decoder_preserves_utf16le_bom_files() { + let mut bytes = vec![0xFF, 0xFE]; + for unit in "SELECT '中文';".encode_utf16() { + bytes.extend_from_slice(&unit.to_le_bytes()); + } + let path = temporary_sql_file(&bytes).await; + let mut decoder = SqlFileStreamDecoder::open(&path).await.unwrap(); + let mut decoded = String::new(); + while let Some(chunk) = decoder.next_chunk().await.unwrap() { + decoded.push_str(&chunk); + } + tokio::fs::remove_file(path).await.unwrap(); + + assert_eq!(decoded, "SELECT '中文';"); + } + #[test] fn non_sqlserver_sql_file_keeps_statement_splitting_behavior() { assert_eq!( diff --git a/crates/dbx-core/tests/live_mysql57.rs b/crates/dbx-core/tests/live_mysql57.rs index c643dc7bc9..30d3e37d04 100644 --- a/crates/dbx-core/tests/live_mysql57.rs +++ b/crates/dbx-core/tests/live_mysql57.rs @@ -6,7 +6,7 @@ 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::sql_file_import::execute_sql_file_path; use dbx_core::storage::Storage; use dbx_core::table_import::parse_xlsx_file; use tokio_util::sync::CancellationToken; @@ -402,7 +402,10 @@ INSERT INTO install_check (id) VALUES (1), (2); execution_id: format!("exec-{suffix}"), connection_id: config.id.clone(), database: String::new(), - file_path: "issue-2356-mysql-install.sql".to_string(), + file_path: std::env::temp_dir() + .join(format!("issue-2356-mysql-install-{suffix}.sql")) + .to_string_lossy() + .into_owned(), continue_on_error: false, }; @@ -416,10 +419,11 @@ INSERT INTO install_check (id) VALUES (1), (2); ) .await; - let result = execute_sql_file_content( + tokio::fs::write(&request.file_path, &script).await.unwrap(); + let result = execute_sql_file_path( &state, &request, - &script, + std::path::Path::new(&request.file_path), CancellationToken::new(), std::time::Instant::now(), |_| {}, @@ -444,6 +448,7 @@ INSERT INTO install_check (id) VALUES (1), (2); ) .await; let _ = std::fs::remove_file(db_path); + let _ = std::fs::remove_file(&request.file_path); result.expect("SQL file import should succeed"); let verify = verify.expect("verify imported rows"); diff --git a/src-tauri/src/commands/sql_file.rs b/src-tauri/src/commands/sql_file.rs index 540d882f45..410329e5c1 100644 --- a/src-tauri/src/commands/sql_file.rs +++ b/src-tauri/src/commands/sql_file.rs @@ -8,9 +8,12 @@ use tokio::sync::RwLock; use tokio_util::sync::CancellationToken; use crate::commands::connection::{ensure_connection_writable, AppState}; -use dbx_core::sql_file_import::{execute_sql_file_content, sql_file_error_progress, sql_file_progress}; +use dbx_core::sql_file_import::{ + execute_sql_file_path, mysql_like_sql_file_can_execute_without_selected_database, read_sql_file_preview, + sql_file_progress, +}; -pub use dbx_core::sql::{decode_sql_file_bytes, SqlFilePreview, SqlFileRequest, SqlFileStatus}; +pub use dbx_core::sql::{SqlFilePreview, SqlFileRequest, SqlFileStatus}; static SQL_FILE_EXECUTIONS: OnceLock>> = OnceLock::new(); @@ -31,17 +34,16 @@ 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 bytes = tokio::fs::read(&path).await.map_err(|e| e.to_string())?; - let content = decode_sql_file_bytes(&bytes)?; - let preview = content.chars().take(20_000).collect(); + let prefix = read_sql_file_preview(&path, 1_000_000).await?; + let can_execute_without_selected_database = mysql_like_sql_file_can_execute_without_selected_database(&prefix); + let preview = prefix.chars().take(20_000).collect(); Ok(SqlFilePreview { file_name: path.file_name().and_then(|name| name.to_str()).unwrap_or("script.sql").to_string(), file_path, size_bytes: metadata.len(), preview, - can_execute_without_selected_database: - dbx_core::sql_file_import::mysql_like_sql_file_can_execute_without_selected_database(&content), + can_execute_without_selected_database, }) } @@ -88,25 +90,16 @@ async fn execute_sql_file_inner( token: CancellationToken, started_at: Instant, ) -> Result<(), String> { - 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(app, &request.execution_id, started_at, error.clone()); - return Err(error); - } - }; - let file_content = match decode_sql_file_bytes(&file_bytes) { - Ok(content) => content, - Err(error) => { - emit_file_io_error_progress(app, &request.execution_id, started_at, error.clone()); - return Err(error); - } - }; - - execute_sql_file_content(state.inner().as_ref(), request, &file_content, token, started_at, |progress| { - let _ = app.emit("sql-file-progress", progress); - }) + execute_sql_file_path( + state.inner().as_ref(), + request, + PathBuf::from(&request.file_path).as_path(), + token, + started_at, + |progress| { + let _ = app.emit("sql-file-progress", progress); + }, + ) .await } @@ -156,10 +149,6 @@ fn emit_progress( ); } -fn emit_file_io_error_progress(app: &AppHandle, execution_id: &str, started_at: Instant, error: String) { - let _ = app.emit("sql-file-progress", sql_file_error_progress(execution_id, started_at, error)); -} - #[cfg(test)] async fn run_statements_for_test( statements: Vec, @@ -250,19 +239,6 @@ mod execution_tests { } #[test] - fn file_io_errors_build_terminal_error_progress() { - let progress = sql_file_error_progress("exec-1", Instant::now(), "read failed".to_string()); - - assert_eq!(progress.execution_id, "exec-1"); - assert_eq!(progress.status, SqlFileStatus::Error); - assert_eq!(progress.statement_index, 0); - assert_eq!(progress.success_count, 0); - assert_eq!(progress.failure_count, 0); - assert_eq!(progress.affected_rows, 0); - assert_eq!(progress.statement_summary, ""); - assert_eq!(progress.error, Some("read failed".to_string())); - } - #[test] fn duplicate_execution_id_is_rejected_without_replacing_token() { let mut executions = HashMap::new();