refactor(sql-file): share import executor

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 }),
     });
 

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)

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<Json<serde_json::Value>, 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::<String>(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<dbx_core::connection::AppState>,
-    req: &SqlFileExecuteRequest,
-    tx: &tokio::sync::broadcast::Sender<String>,
-    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<String>, 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<PathBuf, AppError> {
@@ -422,9 +180,13 @@ pub async fn cancel_sql_file(
     State(state): State<Arc<WebState>>,
     Json(req): Json<CancelSqlFileRequest>,
 ) -> Json<serde_json::Value> {
-    // 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)]

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<Option<String>>,
     pub sessions: RwLock<HashSet<String>>,
     pub sse_channels: RwLock<HashMap<String, broadcast::Sender<String>>>,
+    pub sql_file_executions: RwLock<HashMap<String, CancellationToken>>,
     pub login_rate_limit: Mutex<LoginRateLimit>,
 }