fix(mysql): stream large SQL file execution

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<bool, String>,
 }
 
+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<String, String> {
+    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<tokio::fs::File>,
+    decoder: encoding_rs::Decoder,
+    pending_bytes: Vec<u8>,
+    reached_eof: bool,
+}
+
+impl SqlFileStreamDecoder {
+    async fn open(file_path: &Path) -> Result<Self, String> {
+        Self::open_with_detection_limit(file_path, None).await
+    }
+
+    async fn open_with_detection_limit(file_path: &Path, detection_limit: Option<usize>) -> Result<Self, String> {
+        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<Option<String>, 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<usize>,
+) -> 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<DatabaseType>, 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<String> {
+        match self {
+            Self::Statements(splitter) => splitter.push_chunk(chunk),
+            Self::SqlServerBatches(splitter) => splitter.push_chunk(chunk),
+        }
+    }
+
+    fn finish(self) -> Vec<String> {
+        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<String> {
+        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<String> {
+        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<String>) {
+        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<String>) {
+        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<String>,
+    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<DatabaseType>) -> Vec<String> {
     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!(

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");