use std::sync::Arc; use axum::extract::{Multipart, Path, State}; use axum::response::sse::{Event, Sse}; use axum::Json; use dbx_core::query; use dbx_core::sql; use futures::stream::Stream; use serde::Deserialize; 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, } #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct CancelSqlFileRequest { pub execution_id: String, } pub async fn preview_sql_file( State(state): State>, mut multipart: Multipart, ) -> Result, AppError> { let tmp_dir = state.data_dir.join("tmp"); std::fs::create_dir_all(&tmp_dir).map_err(|e| AppError(e.to_string()))?; while let Some(field) = multipart.next_field().await.map_err(|e| AppError(e.to_string()))? { let file_name = field.file_name().unwrap_or("upload.sql").to_string(); let data = field.bytes().await.map_err(|e| AppError(e.to_string()))?; let file_path = tmp_dir.join(&file_name); 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 preview: String = content.chars().take(5000).collect(); return Ok(Json(serde_json::json!({ "fileName": file_name, "filePath": file_path.to_string_lossy(), "sizeBytes": size_bytes, "preview": preview, }))); } Err(AppError("No file uploaded".to_string())) } pub async fn execute_sql_file( State(state): State>, Json(body): Json, ) -> Result, AppError> { let req = body.request; let execution_id = req.execution_id.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 = std::path::Path::new(&req.file_path); if !file_path.is_absolute() { return Err(AppError("File path must be absolute".to_string())); } if req.file_path.contains("..") { return Err(AppError("File path must not contain '..'".to_string())); } tokio::spawn(async move { match std::fs::metadata(&req.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); } 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); } return; } _ => {} } let file_content = match std::fs::read_to_string(&req.file_path) { Ok(c) => c, 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); } 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 statements = sql::split_sql_statements(&file_content); let start = std::time::Instant::now(); let mut success_count = 0usize; let mut failure_count = 0usize; let mut total_affected: u64 = 0; for (i, stmt) in statements.iter().enumerate() { let summary = sql::statement_summary(stmt); // Send running let running = dbx_core::sql::SqlFileProgress { execution_id: req.execution_id.clone(), status: dbx_core::sql::SqlFileStatus::Running, statement_index: i, 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); } match query::execute_sql_statement(&app, &req.connection_id, &req.database, stmt, 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: i, 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) => { failure_count += 1; let failed = dbx_core::sql::SqlFileProgress { execution_id: req.execution_id.clone(), status: dbx_core::sql::SqlFileStatus::StatementFailed, statement_index: i, success_count, failure_count, affected_rows: total_affected, elapsed_ms: start.elapsed().as_millis(), statement_summary: summary, error: Some(e), }; if let Ok(json) = serde_json::to_string(&failed) { let _ = tx.send(json); } if !req.continue_on_error { break; } } } } // Send final done let final_done = dbx_core::sql::SqlFileProgress { execution_id: req.execution_id.clone(), status: dbx_core::sql::SqlFileStatus::Done, statement_index: statements.len(), 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; }); Ok(Json(serde_json::json!({ "executionId": execution_id }))) } pub async fn sql_file_progress( State(state): State>, Path(execution_id): Path, ) -> Result>>, AppError> { let channels = state.sse_channels.read().await; let tx = channels.get(&execution_id).ok_or_else(|| AppError("Execution not found".to_string()))?; let rx = tx.subscribe(); drop(channels); Ok(crate::sse::sse_from_channel(rx)) } 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 })) }