xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: decode sql files before execution

This commit is contained in:
t8y2 2026-05-26 15:43:07 +08:00
parent be0c2b4775
commit 195e6b218f
5 changed files with 100 additions and 64 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
Cargo.lock generated
View File

@ -1842,6 +1842,7 @@ dependencies = [
"csv",
"deadpool-postgres",
"duckdb",
"encoding_rs",
"font-kit",
"futures",
"iana-time-zone",

1
crates/dbx-core/Cargo.toml
View File

@ -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"] }

61
crates/dbx-core/src/sql.rs
View File

@ -73,6 +73,39 @@ pub struct SqlFileProgress {
pub error: Option<String>,
}
pub fn decode_sql_file_bytes(bytes: &[u8]) -> Result<String, String> {
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<String, String> {
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 = "\

14
crates/dbx-web/src/routes/sql_file.rs
View File

@ -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;

87
src-tauri/src/commands/sql_file.rs
View File

@ -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<RwLock<HashMap<String, CancellationToken>>> =
@ -46,11 +45,8 @@ struct SqlFileSummary {
pub async fn preview_sql_file(file_path: String) -> Result<SqlFilePreview, String> {
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,