refactor(sql-file): share import executor
This commit is contained in:
t8y2 committed
1 parent
f5cc81f761
commit
da6074caaa
8 files changed
+568
-845
No files matched your search
@@ -6,9 +6,6 @@ use dbx_core::db;
|
||||
use dbx_core::models::connection::DatabaseType;
|
||||
use dbx_core::sql::split_sql_statements;
|
||||
|
||||
// Re-export core functions for use by other modules (e.g., sql_file.rs)
|
||||
pub use dbx_core::query::execute_sql_statement;
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn execute_query(
|
||||
state: State<'_, Arc<AppState>>,
|
||||
|
||||
@@ -8,31 +8,13 @@ use tokio::sync::RwLock;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::commands::connection::AppState;
|
||||
use crate::commands::query::execute_sql_statement;
|
||||
use dbx_core::models::connection::DatabaseType;
|
||||
use dbx_core::sql_file_import::{execute_sql_file_content, sql_file_error_progress, sql_file_progress};
|
||||
|
||||
pub use dbx_core::sql::{
|
||||
decode_sql_file_bytes, optimize_sql_file_import_statements, statement_summary, SqlFileImportStatement,
|
||||
SqlFileImportStatementKind, SqlFilePreview, SqlFileProgress, SqlFileRequest, SqlFileStatus, SqlParsingOptions,
|
||||
SqlStatementSplitter,
|
||||
};
|
||||
pub use dbx_core::sql::{decode_sql_file_bytes, SqlFilePreview, SqlFileRequest, SqlFileStatus};
|
||||
|
||||
static SQL_FILE_EXECUTIONS: std::sync::LazyLock<RwLock<HashMap<String, CancellationToken>>> =
|
||||
std::sync::LazyLock::new(|| RwLock::new(HashMap::new()));
|
||||
|
||||
#[derive(Debug)]
|
||||
struct StatementErrorDecision {
|
||||
progress: Vec<SqlFileProgress>,
|
||||
failure_count: usize,
|
||||
result: Result<bool, String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct SqlFileImportTarget {
|
||||
db_type: DatabaseType,
|
||||
driver_profile: Option<String>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct SqlFileSummary {
|
||||
@@ -98,352 +80,26 @@ async fn execute_sql_file_inner(
|
||||
token: CancellationToken,
|
||||
started_at: Instant,
|
||||
) -> Result<(), String> {
|
||||
let mut statement_index = 0;
|
||||
let mut success_count = 0;
|
||||
let mut failure_count = 0;
|
||||
let mut affected_rows = 0;
|
||||
|
||||
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,
|
||||
statement_index,
|
||||
success_count,
|
||||
failure_count,
|
||||
affected_rows,
|
||||
started_at,
|
||||
error.clone(),
|
||||
);
|
||||
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,
|
||||
statement_index,
|
||||
success_count,
|
||||
failure_count,
|
||||
affected_rows,
|
||||
started_at,
|
||||
error.clone(),
|
||||
);
|
||||
emit_file_io_error_progress(app, &request.execution_id, 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());
|
||||
|
||||
let planned_statements = optimize_sql_file_import_statements(
|
||||
&statements,
|
||||
import_target.as_ref().map(|target| target.db_type),
|
||||
import_target.as_ref().and_then(|target| target.driver_profile.as_deref()),
|
||||
);
|
||||
|
||||
for planned_statement in planned_statements {
|
||||
if token.is_cancelled() {
|
||||
emit_progress(
|
||||
app,
|
||||
&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;
|
||||
if execute_statement_with_progress(
|
||||
app,
|
||||
state,
|
||||
request,
|
||||
&token,
|
||||
started_at,
|
||||
next_statement_index,
|
||||
&planned_statement,
|
||||
&mut success_count,
|
||||
&mut failure_count,
|
||||
&mut affected_rows,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
return Ok(());
|
||||
}
|
||||
statement_index = next_statement_index;
|
||||
}
|
||||
|
||||
emit_progress(
|
||||
app,
|
||||
&request.execution_id,
|
||||
SqlFileStatus::Done,
|
||||
statement_index,
|
||||
success_count,
|
||||
failure_count,
|
||||
affected_rows,
|
||||
started_at,
|
||||
"",
|
||||
None,
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn sql_file_import_target(state: &AppState, connection_id: &str) -> Option<SqlFileImportTarget> {
|
||||
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<AppState>>,
|
||||
request: &SqlFileRequest,
|
||||
token: &CancellationToken,
|
||||
started_at: Instant,
|
||||
statement_index: usize,
|
||||
statement: &SqlFileImportStatement,
|
||||
success_count: &mut usize,
|
||||
failure_count: &mut usize,
|
||||
affected_rows: &mut u64,
|
||||
) -> Result<bool, String> {
|
||||
if token.is_cancelled() {
|
||||
let summary = statement_summary(&statement.sql);
|
||||
emit_progress(
|
||||
app,
|
||||
&request.execution_id,
|
||||
SqlFileStatus::Cancelled,
|
||||
statement_index,
|
||||
*success_count,
|
||||
*failure_count,
|
||||
*affected_rows,
|
||||
started_at,
|
||||
&summary,
|
||||
None,
|
||||
);
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
if statement.kind == SqlFileImportStatementKind::Skip {
|
||||
let summary = statement_summary(&statement.sql);
|
||||
emit_progress(
|
||||
app,
|
||||
&request.execution_id,
|
||||
SqlFileStatus::Running,
|
||||
statement_index,
|
||||
*success_count,
|
||||
*failure_count,
|
||||
*affected_rows,
|
||||
started_at,
|
||||
&summary,
|
||||
None,
|
||||
);
|
||||
*success_count += statement.source_statement_count;
|
||||
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(&statement.sql);
|
||||
|
||||
emit_progress(
|
||||
app,
|
||||
&request.execution_id,
|
||||
SqlFileStatus::Running,
|
||||
statement_index,
|
||||
*success_count,
|
||||
*failure_count,
|
||||
*affected_rows,
|
||||
started_at,
|
||||
&summary,
|
||||
None,
|
||||
);
|
||||
|
||||
match execute_sql_statement(
|
||||
state.inner().as_ref(),
|
||||
&request.connection_id,
|
||||
&request.database,
|
||||
&statement.sql,
|
||||
None,
|
||||
Some(token.clone()),
|
||||
)
|
||||
execute_sql_file_content(state.inner().as_ref(), request, &file_content, token, started_at, |progress| {
|
||||
let _ = app.emit("sql-file-progress", progress);
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(result) => {
|
||||
*success_count += statement.source_statement_count;
|
||||
*affected_rows += result.affected_rows;
|
||||
emit_progress(
|
||||
app,
|
||||
&request.execution_id,
|
||||
SqlFileStatus::StatementDone,
|
||||
statement_index,
|
||||
*success_count,
|
||||
*failure_count,
|
||||
*affected_rows,
|
||||
started_at,
|
||||
&summary,
|
||||
None,
|
||||
);
|
||||
Ok(false)
|
||||
}
|
||||
Err(error) => {
|
||||
if statement.source_statement_count > 1 && !token.is_cancelled() {
|
||||
return execute_merged_statement_fallback_with_progress(
|
||||
app,
|
||||
state,
|
||||
request,
|
||||
token,
|
||||
started_at,
|
||||
statement_index + 1 - statement.source_statement_count,
|
||||
statement,
|
||||
success_count,
|
||||
failure_count,
|
||||
affected_rows,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let decision = statement_error_decision(
|
||||
&request.execution_id,
|
||||
token,
|
||||
request.continue_on_error,
|
||||
started_at,
|
||||
statement_index,
|
||||
*success_count,
|
||||
*failure_count,
|
||||
*affected_rows,
|
||||
&summary,
|
||||
error,
|
||||
);
|
||||
|
||||
*failure_count = decision.failure_count;
|
||||
for progress in decision.progress {
|
||||
let _ = app.emit("sql-file-progress", progress);
|
||||
}
|
||||
decision.result
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute_merged_statement_fallback_with_progress(
|
||||
app: &AppHandle,
|
||||
state: &State<'_, Arc<AppState>>,
|
||||
request: &SqlFileRequest,
|
||||
token: &CancellationToken,
|
||||
started_at: Instant,
|
||||
first_statement_index: usize,
|
||||
statement: &SqlFileImportStatement,
|
||||
success_count: &mut usize,
|
||||
failure_count: &mut usize,
|
||||
affected_rows: &mut u64,
|
||||
) -> Result<bool, String> {
|
||||
for (offset, source_sql) in statement.source_sqls.iter().enumerate() {
|
||||
let statement_index = first_statement_index + offset;
|
||||
if token.is_cancelled() {
|
||||
emit_progress(
|
||||
app,
|
||||
&request.execution_id,
|
||||
SqlFileStatus::Cancelled,
|
||||
statement_index,
|
||||
*success_count,
|
||||
*failure_count,
|
||||
*affected_rows,
|
||||
started_at,
|
||||
&statement_summary(source_sql),
|
||||
None,
|
||||
);
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
let summary = statement_summary(source_sql);
|
||||
emit_progress(
|
||||
app,
|
||||
&request.execution_id,
|
||||
SqlFileStatus::Running,
|
||||
statement_index,
|
||||
*success_count,
|
||||
*failure_count,
|
||||
*affected_rows,
|
||||
started_at,
|
||||
&summary,
|
||||
None,
|
||||
);
|
||||
|
||||
match execute_sql_statement(
|
||||
state.inner().as_ref(),
|
||||
&request.connection_id,
|
||||
&request.database,
|
||||
source_sql,
|
||||
None,
|
||||
Some(token.clone()),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(result) => {
|
||||
*success_count += 1;
|
||||
*affected_rows += result.affected_rows;
|
||||
emit_progress(
|
||||
app,
|
||||
&request.execution_id,
|
||||
SqlFileStatus::StatementDone,
|
||||
statement_index,
|
||||
*success_count,
|
||||
*failure_count,
|
||||
*affected_rows,
|
||||
started_at,
|
||||
&summary,
|
||||
None,
|
||||
);
|
||||
}
|
||||
Err(error) => {
|
||||
let decision = statement_error_decision(
|
||||
&request.execution_id,
|
||||
token,
|
||||
request.continue_on_error,
|
||||
started_at,
|
||||
statement_index,
|
||||
*success_count,
|
||||
*failure_count,
|
||||
*affected_rows,
|
||||
&summary,
|
||||
error,
|
||||
);
|
||||
|
||||
*failure_count = decision.failure_count;
|
||||
for progress in decision.progress {
|
||||
let _ = app.emit("sql-file-progress", progress);
|
||||
}
|
||||
if decision.result? {
|
||||
return Ok(true);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
fn register_sql_file_execution(
|
||||
@@ -463,68 +119,6 @@ fn remove_sql_file_execution(executions: &mut HashMap<String, CancellationToken>
|
||||
executions.remove(execution_id);
|
||||
}
|
||||
|
||||
fn statement_error_decision(
|
||||
execution_id: &str,
|
||||
token: &CancellationToken,
|
||||
continue_on_error: bool,
|
||||
started_at: Instant,
|
||||
statement_index: usize,
|
||||
success_count: usize,
|
||||
failure_count: usize,
|
||||
affected_rows: u64,
|
||||
summary: &str,
|
||||
error: String,
|
||||
) -> StatementErrorDecision {
|
||||
if token.is_cancelled() {
|
||||
return StatementErrorDecision {
|
||||
progress: vec![sql_file_progress(
|
||||
execution_id,
|
||||
SqlFileStatus::Cancelled,
|
||||
statement_index,
|
||||
success_count,
|
||||
failure_count,
|
||||
affected_rows,
|
||||
started_at,
|
||||
summary,
|
||||
None,
|
||||
)],
|
||||
failure_count,
|
||||
result: Ok(true),
|
||||
};
|
||||
}
|
||||
|
||||
let failure_count = failure_count + 1;
|
||||
let statement_failed = sql_file_progress(
|
||||
execution_id,
|
||||
SqlFileStatus::StatementFailed,
|
||||
statement_index,
|
||||
success_count,
|
||||
failure_count,
|
||||
affected_rows,
|
||||
started_at,
|
||||
summary,
|
||||
Some(error.clone()),
|
||||
);
|
||||
|
||||
if continue_on_error {
|
||||
return StatementErrorDecision { progress: vec![statement_failed], failure_count, result: Ok(false) };
|
||||
}
|
||||
|
||||
let terminal_error = sql_file_progress(
|
||||
execution_id,
|
||||
SqlFileStatus::Error,
|
||||
statement_index,
|
||||
success_count,
|
||||
failure_count,
|
||||
affected_rows,
|
||||
started_at,
|
||||
summary,
|
||||
Some(error.clone()),
|
||||
);
|
||||
|
||||
StatementErrorDecision { progress: vec![statement_failed, terminal_error], failure_count, result: Err(error) }
|
||||
}
|
||||
|
||||
fn emit_progress(
|
||||
app: &AppHandle,
|
||||
execution_id: &str,
|
||||
@@ -553,74 +147,8 @@ fn emit_progress(
|
||||
);
|
||||
}
|
||||
|
||||
fn emit_file_io_error_progress(
|
||||
app: &AppHandle,
|
||||
execution_id: &str,
|
||||
statement_index: usize,
|
||||
success_count: usize,
|
||||
failure_count: usize,
|
||||
affected_rows: u64,
|
||||
started_at: Instant,
|
||||
error: String,
|
||||
) {
|
||||
let _ = app.emit(
|
||||
"sql-file-progress",
|
||||
file_io_error_progress(
|
||||
execution_id,
|
||||
statement_index,
|
||||
success_count,
|
||||
failure_count,
|
||||
affected_rows,
|
||||
started_at,
|
||||
error,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
fn file_io_error_progress(
|
||||
execution_id: &str,
|
||||
statement_index: usize,
|
||||
success_count: usize,
|
||||
failure_count: usize,
|
||||
affected_rows: u64,
|
||||
started_at: Instant,
|
||||
error: String,
|
||||
) -> SqlFileProgress {
|
||||
sql_file_progress(
|
||||
execution_id,
|
||||
SqlFileStatus::Error,
|
||||
statement_index,
|
||||
success_count,
|
||||
failure_count,
|
||||
affected_rows,
|
||||
started_at,
|
||||
"",
|
||||
Some(error),
|
||||
)
|
||||
}
|
||||
|
||||
fn sql_file_progress(
|
||||
execution_id: &str,
|
||||
status: SqlFileStatus,
|
||||
statement_index: usize,
|
||||
success_count: usize,
|
||||
failure_count: usize,
|
||||
affected_rows: u64,
|
||||
started_at: Instant,
|
||||
statement_summary: &str,
|
||||
error: Option<String>,
|
||||
) -> SqlFileProgress {
|
||||
SqlFileProgress {
|
||||
execution_id: execution_id.to_string(),
|
||||
status,
|
||||
statement_index,
|
||||
success_count,
|
||||
failure_count,
|
||||
affected_rows,
|
||||
elapsed_ms: started_at.elapsed().as_millis(),
|
||||
statement_summary: statement_summary.to_string(),
|
||||
error,
|
||||
}
|
||||
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)]
|
||||
@@ -714,14 +242,14 @@ mod execution_tests {
|
||||
|
||||
#[test]
|
||||
fn file_io_errors_build_terminal_error_progress() {
|
||||
let progress = file_io_error_progress("exec-1", 4, 2, 1, 17, Instant::now(), "read failed".to_string());
|
||||
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, 4);
|
||||
assert_eq!(progress.success_count, 2);
|
||||
assert_eq!(progress.failure_count, 1);
|
||||
assert_eq!(progress.affected_rows, 17);
|
||||
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()));
|
||||
}
|
||||
@@ -742,70 +270,4 @@ mod execution_tests {
|
||||
assert!(original.is_cancelled());
|
||||
assert!(!replacement.is_cancelled());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stop_on_error_returns_err_with_terminal_error_progress() {
|
||||
let decision = statement_error_decision(
|
||||
"exec-1",
|
||||
&CancellationToken::new(),
|
||||
false,
|
||||
Instant::now(),
|
||||
3,
|
||||
1,
|
||||
0,
|
||||
5,
|
||||
"bad statement",
|
||||
"syntax error".to_string(),
|
||||
);
|
||||
|
||||
assert_eq!(decision.failure_count, 1);
|
||||
assert_eq!(decision.result, Err("syntax error".to_string()));
|
||||
assert_eq!(decision.progress.len(), 2);
|
||||
assert_eq!(decision.progress[0].status, SqlFileStatus::StatementFailed);
|
||||
assert_eq!(decision.progress[1].status, SqlFileStatus::Error);
|
||||
assert_eq!(decision.progress[1].error, Some("syntax error".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancelled_in_flight_error_does_not_increment_failure_count() {
|
||||
let token = CancellationToken::new();
|
||||
token.cancel();
|
||||
|
||||
let decision = statement_error_decision(
|
||||
"exec-1",
|
||||
&token,
|
||||
false,
|
||||
Instant::now(),
|
||||
2,
|
||||
1,
|
||||
4,
|
||||
9,
|
||||
"slow statement",
|
||||
"Query canceled".to_string(),
|
||||
);
|
||||
|
||||
assert_eq!(decision.failure_count, 4);
|
||||
assert_eq!(decision.result, Ok(true));
|
||||
assert_eq!(decision.progress.len(), 1);
|
||||
assert_eq!(decision.progress[0].status, SqlFileStatus::Cancelled);
|
||||
assert_eq!(decision.progress[0].failure_count, 4);
|
||||
assert_eq!(decision.progress[0].error, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn progress_payload_serializes_camel_case_status() {
|
||||
let progress =
|
||||
sql_file_progress("exec-1", SqlFileStatus::StatementDone, 1, 1, 0, 3, Instant::now(), "select 1", None);
|
||||
|
||||
let value = serde_json::to_value(progress).unwrap();
|
||||
|
||||
assert_eq!(value["executionId"], "exec-1");
|
||||
assert_eq!(value["statementIndex"], 1);
|
||||
assert_eq!(value["successCount"], 1);
|
||||
assert_eq!(value["failureCount"], 0);
|
||||
assert_eq!(value["affectedRows"], 3);
|
||||
assert_eq!(value["statementSummary"], "select 1");
|
||||
assert_eq!(value["status"], "statementDone");
|
||||
assert!(value.get("execution_id").is_none());
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user