diff --git a/crates/dbx-core/src/lib.rs b/crates/dbx-core/src/lib.rs index ffe6263d3..87b429caa 100644 --- a/crates/dbx-core/src/lib.rs +++ b/crates/dbx-core/src/lib.rs @@ -34,6 +34,7 @@ pub mod sql; pub mod sql_analysis; pub mod sql_dialect; pub mod sql_editability; +pub mod sql_file_import; pub mod storage; pub mod table_import; pub mod table_structure_sql; diff --git a/crates/dbx-core/src/sql_file_import.rs b/crates/dbx-core/src/sql_file_import.rs new file mode 100644 index 000000000..e3cdf2d38 --- /dev/null +++ b/crates/dbx-core/src/sql_file_import.rs @@ -0,0 +1,497 @@ +use std::time::Instant; + +use tokio_util::sync::CancellationToken; + +use crate::connection::AppState; +use crate::models::connection::DatabaseType; +use crate::query::execute_sql_statement; +use crate::sql::{ + optimize_sql_file_import_statements, statement_summary, SqlFileImportStatement, SqlFileImportStatementKind, + SqlFileProgress, SqlFileRequest, SqlFileStatus, SqlParsingOptions, SqlStatementSplitter, +}; + +#[derive(Debug, Clone)] +struct SqlFileImportTarget { + db_type: DatabaseType, + driver_profile: Option, +} + +#[derive(Debug)] +struct StatementErrorDecision { + progress: Vec, + failure_count: usize, + result: Result, +} + +pub async fn execute_sql_file_content( + state: &AppState, + request: &SqlFileRequest, + file_content: &str, + token: CancellationToken, + started_at: Instant, + mut emit: 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; + + 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 = 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(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; + if execute_statement_with_progress( + state, + request, + &token, + started_at, + next_statement_index, + &planned_statement, + &mut success_count, + &mut failure_count, + &mut affected_rows, + &mut emit, + ) + .await? + { + return Ok(()); + } + 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(()) +} + +pub 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, +) -> 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, + } +} + +pub fn sql_file_error_progress(execution_id: &str, started_at: Instant, error: String) -> SqlFileProgress { + sql_file_progress(execution_id, SqlFileStatus::Error, 0, 0, 0, 0, started_at, "", Some(error)) +} + +async fn sql_file_import_target(state: &AppState, connection_id: &str) -> Option { + 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( + state: &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, + emit: &mut impl FnMut(SqlFileProgress), +) -> Result { + if token.is_cancelled() { + let summary = statement_summary(&statement.sql); + emit(sql_file_progress( + &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(sql_file_progress( + &request.execution_id, + SqlFileStatus::Running, + statement_index, + *success_count, + *failure_count, + *affected_rows, + started_at, + &summary, + None, + )); + *success_count += statement.source_statement_count; + emit(sql_file_progress( + &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(sql_file_progress( + &request.execution_id, + SqlFileStatus::Running, + statement_index, + *success_count, + *failure_count, + *affected_rows, + started_at, + &summary, + None, + )); + + match execute_sql_statement( + state, + &request.connection_id, + &request.database, + &statement.sql, + None, + Some(token.clone()), + ) + .await + { + Ok(result) => { + *success_count += statement.source_statement_count; + *affected_rows += result.affected_rows; + emit(sql_file_progress( + &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( + state, + request, + token, + started_at, + statement_index + 1 - statement.source_statement_count, + statement, + success_count, + failure_count, + affected_rows, + emit, + ) + .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 { + emit(progress); + } + decision.result + } + } +} + +async fn execute_merged_statement_fallback_with_progress( + state: &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, + emit: &mut impl FnMut(SqlFileProgress), +) -> Result { + for (offset, source_sql) in statement.source_sqls.iter().enumerate() { + let statement_index = first_statement_index + offset; + if token.is_cancelled() { + emit(sql_file_progress( + &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(sql_file_progress( + &request.execution_id, + SqlFileStatus::Running, + statement_index, + *success_count, + *failure_count, + *affected_rows, + started_at, + &summary, + None, + )); + + match execute_sql_statement( + state, + &request.connection_id, + &request.database, + source_sql, + None, + Some(token.clone()), + ) + .await + { + Ok(result) => { + *success_count += 1; + *affected_rows += result.affected_rows; + emit(sql_file_progress( + &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 { + emit(progress); + } + if decision.result? { + return Ok(true); + } + } + } + } + + Ok(false) +} + +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) } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[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()); + } +} diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index 6d8ca074c..7635b3965 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -65,6 +65,7 @@ async fn main() { password_hash: RwLock::new(password_hash), sessions: RwLock::new(HashSet::new()), sse_channels: RwLock::new(HashMap::new()), + sql_file_executions: RwLock::new(HashMap::new()), login_rate_limit: tokio::sync::Mutex::new(state::LoginRateLimit { fail_count: 0, locked_until: None }), }); diff --git a/crates/dbx-web/src/routes/connection.rs b/crates/dbx-web/src/routes/connection.rs index 912f8cb68..3f4d2ec2c 100644 --- a/crates/dbx-web/src/routes/connection.rs +++ b/crates/dbx-web/src/routes/connection.rs @@ -206,6 +206,7 @@ mod tests { password_hash: RwLock::new(None), sessions: RwLock::new(HashSet::new()), sse_channels: RwLock::new(HashMap::new()), + sql_file_executions: RwLock::new(HashMap::new()), login_rate_limit: Mutex::new(LoginRateLimit { fail_count: 0, locked_until: None }), }); (state, dir) diff --git a/crates/dbx-web/src/routes/sql_file.rs b/crates/dbx-web/src/routes/sql_file.rs index cf6a9540e..d1b3c4a53 100644 --- a/crates/dbx-web/src/routes/sql_file.rs +++ b/crates/dbx-web/src/routes/sql_file.rs @@ -4,29 +4,23 @@ use std::sync::Arc; use axum::extract::{Multipart, Path as AxumPath, State}; use axum::response::sse::{Event, Sse}; use axum::Json; -use dbx_core::query; use dbx_core::sql; -use dbx_core::sql::{SqlFileImportStatement, SqlFileImportStatementKind}; +use dbx_core::sql::{SqlFileProgress, SqlFileRequest, SqlFileStatus}; +use dbx_core::sql_file_import::{ + execute_sql_file_content, sql_file_error_progress, sql_file_progress as build_sql_file_progress, +}; use futures::stream::Stream; use serde::Deserialize; +use tokio::sync::broadcast; +use tokio_util::sync::CancellationToken; use crate::error::AppError; use crate::state::WebState; -#[derive(Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct SqlFileExecuteRequest { - pub execution_id: String, - pub connection_id: String, - pub database: String, - pub file_path: String, - pub continue_on_error: bool, -} - #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct SqlFileExecuteWrapper { - pub request: SqlFileExecuteRequest, + pub request: SqlFileRequest, } #[derive(Deserialize)] @@ -70,49 +64,40 @@ pub async fn execute_sql_file( ) -> Result, AppError> { let req = body.request; let execution_id = req.execution_id.clone(); + let file_path = validated_uploaded_sql_path(&state.data_dir, &req.file_path)?; + let token = CancellationToken::new(); + { + let mut executions = state.sql_file_executions.write().await; + if executions.contains_key(&execution_id) { + return Err(AppError(format!("SQL file execution '{execution_id}' already exists"))); + } + executions.insert(execution_id.clone(), token.clone()); + } let (tx, _) = tokio::sync::broadcast::channel::(256); state.sse_channels.write().await.insert(execution_id.clone(), tx.clone()); let app = state.app.clone(); let state_clone = state.clone(); - let file_path = validated_uploaded_sql_path(&state.data_dir, &req.file_path)?; - tokio::spawn(async move { + let started_at = std::time::Instant::now(); match std::fs::metadata(&file_path) { Ok(meta) if meta.len() > 200 * 1024 * 1024 => { - let progress = dbx_core::sql::SqlFileProgress { - execution_id: req.execution_id.clone(), - status: dbx_core::sql::SqlFileStatus::Error, - statement_index: 0, - success_count: 0, - failure_count: 0, - affected_rows: 0, - elapsed_ms: 0, - statement_summary: String::new(), - error: Some(format!("File too large: {} bytes (max {} bytes)", meta.len(), 200 * 1024 * 1024)), - }; - if let Ok(json) = serde_json::to_string(&progress) { - let _ = tx.send(json); - } + send_sql_file_progress( + &tx, + sql_file_error_progress( + &req.execution_id, + started_at, + format!("File too large: {} bytes (max {} bytes)", meta.len(), 200 * 1024 * 1024), + ), + ); + cleanup_sql_file_execution(&state_clone, &req.execution_id).await; return; } Err(e) => { - let progress = dbx_core::sql::SqlFileProgress { - execution_id: req.execution_id.clone(), - status: dbx_core::sql::SqlFileStatus::Error, - statement_index: 0, - success_count: 0, - failure_count: 0, - affected_rows: 0, - elapsed_ms: 0, - statement_summary: String::new(), - error: Some(e.to_string()), - }; - if let Ok(json) = serde_json::to_string(&progress) { - let _ = tx.send(json); - } + send_sql_file_progress(&tx, sql_file_error_progress(&req.execution_id, started_at, e.to_string())); + cleanup_sql_file_execution(&state_clone, &req.execution_id).await; return; } _ => {} @@ -124,264 +109,37 @@ pub async fn execute_sql_file( }) { Ok(content) => content, Err(e) => { - let progress = dbx_core::sql::SqlFileProgress { - execution_id: req.execution_id.clone(), - status: dbx_core::sql::SqlFileStatus::Error, - statement_index: 0, - success_count: 0, - failure_count: 0, - affected_rows: 0, - elapsed_ms: 0, - statement_summary: String::new(), - error: Some(e.to_string()), - }; - if let Ok(json) = serde_json::to_string(&progress) { - let _ = tx.send(json); - } + send_sql_file_progress(&tx, sql_file_error_progress(&req.execution_id, started_at, e.to_string())); + cleanup_sql_file_execution(&state_clone, &req.execution_id).await; return; } }; - // Send started - let started = dbx_core::sql::SqlFileProgress { - execution_id: req.execution_id.clone(), - status: dbx_core::sql::SqlFileStatus::Started, - statement_index: 0, - success_count: 0, - failure_count: 0, - affected_rows: 0, - elapsed_ms: 0, - statement_summary: String::new(), - error: None, - }; - if let Ok(json) = serde_json::to_string(&started) { - let _ = tx.send(json); - } - - 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 planned_statements = sql::optimize_sql_file_import_statements( - &statements, - import_target.as_ref().map(|(db_type, _)| *db_type), - import_target.as_ref().and_then(|(_, driver_profile)| driver_profile.as_deref()), + send_sql_file_progress( + &tx, + build_sql_file_progress(&req.execution_id, SqlFileStatus::Started, 0, 0, 0, 0, started_at, "", None), ); - let start = std::time::Instant::now(); - let mut success_count = 0usize; - let mut failure_count = 0usize; - let mut total_affected: u64 = 0; - let mut statement_index = 0usize; - for planned_statement in planned_statements { - let next_statement_index = statement_index + planned_statement.source_statement_count; - let summary = sql::statement_summary(&planned_statement.sql); + let _ = execute_sql_file_content(&app, &req, &file_content, token, started_at, |progress| { + send_sql_file_progress(&tx, progress); + }) + .await; - // Send running - let running = dbx_core::sql::SqlFileProgress { - execution_id: req.execution_id.clone(), - status: dbx_core::sql::SqlFileStatus::Running, - statement_index: next_statement_index, - success_count, - failure_count, - affected_rows: total_affected, - elapsed_ms: start.elapsed().as_millis(), - statement_summary: summary.clone(), - error: None, - }; - if let Ok(json) = serde_json::to_string(&running) { - let _ = tx.send(json); - } - - if planned_statement.kind == SqlFileImportStatementKind::Skip { - success_count += planned_statement.source_statement_count; - let done = dbx_core::sql::SqlFileProgress { - execution_id: req.execution_id.clone(), - status: dbx_core::sql::SqlFileStatus::StatementDone, - statement_index: next_statement_index, - success_count, - failure_count, - affected_rows: total_affected, - elapsed_ms: start.elapsed().as_millis(), - statement_summary: summary, - error: None, - }; - if let Ok(json) = serde_json::to_string(&done) { - let _ = tx.send(json); - } - statement_index = next_statement_index; - continue; - } - - match query::execute_sql_statement( - &app, - &req.connection_id, - &req.database, - &planned_statement.sql, - None, - None, - ) - .await - { - Ok(result) => { - success_count += planned_statement.source_statement_count; - total_affected += result.affected_rows; - let done = dbx_core::sql::SqlFileProgress { - execution_id: req.execution_id.clone(), - status: dbx_core::sql::SqlFileStatus::StatementDone, - statement_index: next_statement_index, - success_count, - failure_count, - affected_rows: total_affected, - elapsed_ms: start.elapsed().as_millis(), - statement_summary: summary, - error: None, - }; - if let Ok(json) = serde_json::to_string(&done) { - let _ = tx.send(json); - } - } - Err(e) => { - let stopped = execute_web_merged_statement_fallback( - &app, - &req, - &tx, - start, - statement_index + 1, - &planned_statement, - &mut success_count, - &mut failure_count, - &mut total_affected, - e, - ) - .await; - if stopped || !req.continue_on_error { - break; - } - } - } - statement_index = next_statement_index; - } - - // Send final done - let final_done = dbx_core::sql::SqlFileProgress { - execution_id: req.execution_id.clone(), - status: dbx_core::sql::SqlFileStatus::Done, - statement_index, - success_count, - failure_count, - affected_rows: total_affected, - elapsed_ms: start.elapsed().as_millis(), - statement_summary: String::new(), - error: None, - }; - if let Ok(json) = serde_json::to_string(&final_done) { - let _ = tx.send(json); - } - - state_clone.remove_sse_channel(&req.execution_id).await; + cleanup_sql_file_execution(&state_clone, &req.execution_id).await; }); Ok(Json(serde_json::json!({ "executionId": execution_id }))) } -async fn execute_web_merged_statement_fallback( - app: &Arc, - req: &SqlFileExecuteRequest, - tx: &tokio::sync::broadcast::Sender, - start: std::time::Instant, - first_statement_index: usize, - statement: &SqlFileImportStatement, - success_count: &mut usize, - failure_count: &mut usize, - total_affected: &mut u64, - batch_error: String, -) -> bool { - if statement.source_statement_count <= 1 { - let summary = sql::statement_summary(&statement.sql); - *failure_count += 1; - let failed = dbx_core::sql::SqlFileProgress { - execution_id: req.execution_id.clone(), - status: dbx_core::sql::SqlFileStatus::StatementFailed, - statement_index: first_statement_index, - success_count: *success_count, - failure_count: *failure_count, - affected_rows: *total_affected, - elapsed_ms: start.elapsed().as_millis(), - statement_summary: summary, - error: Some(batch_error), - }; - if let Ok(json) = serde_json::to_string(&failed) { - let _ = tx.send(json); - } - return false; +fn send_sql_file_progress(tx: &broadcast::Sender, progress: SqlFileProgress) { + if let Ok(json) = serde_json::to_string(&progress) { + let _ = tx.send(json); } +} - for (offset, source_sql) in statement.source_sqls.iter().enumerate() { - let statement_index = first_statement_index + offset; - let summary = sql::statement_summary(source_sql); - let running = dbx_core::sql::SqlFileProgress { - execution_id: req.execution_id.clone(), - status: dbx_core::sql::SqlFileStatus::Running, - statement_index, - success_count: *success_count, - failure_count: *failure_count, - affected_rows: *total_affected, - elapsed_ms: start.elapsed().as_millis(), - statement_summary: summary.clone(), - error: None, - }; - if let Ok(json) = serde_json::to_string(&running) { - let _ = tx.send(json); - } - - match query::execute_sql_statement(app, &req.connection_id, &req.database, source_sql, None, None).await { - Ok(result) => { - *success_count += 1; - *total_affected += result.affected_rows; - let done = dbx_core::sql::SqlFileProgress { - execution_id: req.execution_id.clone(), - status: dbx_core::sql::SqlFileStatus::StatementDone, - statement_index, - success_count: *success_count, - failure_count: *failure_count, - affected_rows: *total_affected, - elapsed_ms: start.elapsed().as_millis(), - statement_summary: summary, - error: None, - }; - if let Ok(json) = serde_json::to_string(&done) { - let _ = tx.send(json); - } - } - Err(error) => { - *failure_count += 1; - let failed = dbx_core::sql::SqlFileProgress { - execution_id: req.execution_id.clone(), - status: dbx_core::sql::SqlFileStatus::StatementFailed, - statement_index, - success_count: *success_count, - failure_count: *failure_count, - affected_rows: *total_affected, - elapsed_ms: start.elapsed().as_millis(), - statement_summary: summary, - error: Some(error), - }; - if let Ok(json) = serde_json::to_string(&failed) { - let _ = tx.send(json); - } - if !req.continue_on_error { - return true; - } - } - } - } - - false +async fn cleanup_sql_file_execution(state: &WebState, execution_id: &str) { + state.remove_sse_channel(execution_id).await; + state.sql_file_executions.write().await.remove(execution_id); } fn safe_uploaded_sql_path(tmp_dir: &Path, file_name: &str) -> Result { @@ -422,9 +180,13 @@ pub async fn cancel_sql_file( State(state): State>, Json(req): Json, ) -> Json { - // Remove the channel to stop the execution loop - state.sse_channels.write().await.remove(&req.execution_id); - Json(serde_json::json!({ "cancelled": true })) + let executions = state.sql_file_executions.read().await; + if let Some(token) = executions.get(&req.execution_id) { + token.cancel(); + Json(serde_json::json!({ "cancelled": true })) + } else { + Json(serde_json::json!({ "cancelled": false })) + } } #[cfg(test)] diff --git a/crates/dbx-web/src/state.rs b/crates/dbx-web/src/state.rs index ee0aeabd7..20f3a855c 100644 --- a/crates/dbx-web/src/state.rs +++ b/crates/dbx-web/src/state.rs @@ -3,6 +3,7 @@ use std::collections::{HashMap, HashSet}; use std::path::PathBuf; use std::sync::Arc; use tokio::sync::{broadcast, Mutex, RwLock}; +use tokio_util::sync::CancellationToken; pub struct LoginRateLimit { pub fail_count: u32, @@ -15,6 +16,7 @@ pub struct WebState { pub password_hash: RwLock>, pub sessions: RwLock>, pub sse_channels: RwLock>>, + pub sql_file_executions: RwLock>, pub login_rate_limit: Mutex, } diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index ecd6c4729..f3d85028b 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -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>, diff --git a/src-tauri/src/commands/sql_file.rs b/src-tauri/src/commands/sql_file.rs index d7ec967fe..442e0f7b4 100644 --- a/src-tauri/src/commands/sql_file.rs +++ b/src-tauri/src/commands/sql_file.rs @@ -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>> = std::sync::LazyLock::new(|| RwLock::new(HashMap::new())); -#[derive(Debug)] -struct StatementErrorDecision { - progress: Vec, - failure_count: usize, - result: Result, -} - -#[derive(Debug, Clone)] -struct SqlFileImportTarget { - db_type: DatabaseType, - driver_profile: Option, -} - #[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 { - 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>, - 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 { - 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>, - 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 { - 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 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, -) -> 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()); - } }