use std::sync::{ atomic::{AtomicBool, Ordering}, Arc, }; use axum::body::Body; use axum::extract::{Path, State}; use axum::http::{header, StatusCode}; use axum::response::sse::{Event, Sse}; use axum::response::Response; use axum::Json; use dbx_core::database_export::{ self, DatabaseExportOutputCompression, DatabaseExportRequest, ExportProgress, ExportStatus, }; use futures::stream::Stream; use serde::Deserialize; use crate::error::AppError; use crate::routes::export_download::attachment_content_disposition; use crate::state::{WebExportFile, WebState}; #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct StartExportRequest { pub request: DatabaseExportRequest, } #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct CancelExportRequest { pub export_id: String, } pub async fn start_database_export( State(state): State>, Json(body): Json, ) -> Result, AppError> { let mut req = body.request; let export_id = req.export_id.clone(); // Web mode: write to a server-side temp file instead of the client-supplied // path (which is a relative placeholder and would land in the server process // CWD, unreachable by the browser). The file is served back via the download // endpoint after the export completes. let tmp_dir = state.data_dir.join("tmp"); std::fs::create_dir_all(&tmp_dir).map_err(|e| AppError::from(e.to_string()))?; let extension = if req.output_compression == DatabaseExportOutputCompression::Gzip { "sql.gz" } else { "sql" }; let tmp_file = tmp_dir.join(format!("database_export_{export_id}.{extension}")); let tmp_file_path = tmp_file.to_string_lossy().to_string(); req.file_path = tmp_file_path.clone(); let download_filename = format!("{}.{}", sanitize_export_filename(&req.database), extension); let content_format = extension.to_string(); // Store export file mapping for download state.export_files.write().await.insert( export_id.clone(), WebExportFile { file_path: tmp_file_path.clone(), download_filename, format: content_format }, ); let (tx, _) = tokio::sync::broadcast::channel::(256); state.sse_channels.write().await.insert(export_id.clone(), tx.clone()); let app = state.app.clone(); let state_clone = state.clone(); let cancelled = Arc::new(AtomicBool::new(false)); let cancelled_progress = cancelled.clone(); // Exports interleave async fetches with synchronous row formatting and // buffered disk writes; run them off the async workers (see spawn_export_task). dbx_core::export_runtime::spawn_export_task(async move { let result = database_export::export_database_sql_core(&app, &req, |progress| { if matches!(progress.status, ExportStatus::Cancelled) { cancelled_progress.store(true, Ordering::SeqCst); } if let Ok(json) = serde_json::to_string(&progress) { let _ = tx.send(json); } }) .await; // `export_database_sql_core` signals cancellation either by emitting a // Cancelled progress and returning Ok(()), or by returning // Err("Export cancelled"). Treat both as cancellation. let is_cancelled = cancelled.load(Ordering::SeqCst) || result.as_ref().err().map(|e| e.contains("cancelled")).unwrap_or(false); match (&result, is_cancelled) { (_, true) => { // Cancelled: drop temp file + mapping, no Error progress. let _ = tokio::fs::remove_file(&req.file_path).await; state_clone.export_files.write().await.remove(&req.export_id); } (Err(e), false) => { let _ = tokio::fs::remove_file(&req.file_path).await; state_clone.export_files.write().await.remove(&req.export_id); let progress = ExportProgress { export_id: req.export_id.clone(), current_object: String::new(), object_index: 0, total_objects: 0, rows_exported: 0, total_rows: None, status: ExportStatus::Error, error: Some(e.clone()), preparing: false, }; if let Ok(json) = serde_json::to_string(&progress) { let _ = tx.send(json); } } (Ok(()), false) => { // Done: keep temp file + mapping for the download endpoint. } } database_export::clear_export_cancelled(&req.export_id).await; state_clone.remove_sse_channel(&req.export_id).await; }); Ok(Json(serde_json::json!({ "exportId": export_id }))) } /// Sanitize a database name into a safe download filename component, mirroring /// the client-side `safeName` logic in DatabaseExportDialog.vue. fn sanitize_export_filename(name: &str) -> String { let replaced = name.replace(['\\', '/', ':', '*', '?', '"', '<', '>', '|'], "_"); let sanitized = replaced.trim(); if sanitized.is_empty() { "database".to_string() } else { sanitized.to_string() } } pub async fn database_export_progress( State(state): State>, Path(export_id): Path, ) -> Result>>, AppError> { let channels = state.sse_channels.read().await; let tx = channels.get(&export_id).ok_or_else(|| AppError::from("Export not found".to_string()))?; let rx = tx.subscribe(); drop(channels); Ok(crate::sse::sse_from_channel(rx)) } pub async fn cancel_database_export( State(_state): State>, Json(req): Json, ) -> Json { database_export::set_export_cancelled(&req.export_id).await; Json(serde_json::json!({ "cancelled": true })) } pub async fn database_export_download( State(state): State>, Path(export_id): Path, ) -> Result { let export_file = state .export_files .write() .await .remove(&export_id) .ok_or_else(|| AppError::from("Export file not found".to_string()))?; let data = tokio::fs::read(&export_file.file_path).await.map_err(|e| AppError::from(e.to_string()))?; // Clean up temp file let _ = tokio::fs::remove_file(&export_file.file_path).await; let content_type = if export_file.format == "sql.gz" { "application/gzip" } else { "application/sql; charset=utf-8" }; Response::builder() .status(StatusCode::OK) .header(header::CONTENT_TYPE, content_type) .header(header::CONTENT_DISPOSITION, attachment_content_disposition(&export_file.download_filename)) .body(Body::from(data)) .map_err(|e| AppError::from(e.to_string())) }