feat(mongodb): bson dumps and database backup/restore

This commit is contained in:
Hand Sonic authored and GitHub committed 2026-09-16 22:46:48 +08:00
1 parent ecc1135de3
commit f5673ca0a1
54 files changed
+6355 -106

No files matched your search

+19
View File
@@ -988,6 +988,25 @@ async fn main() {
.route("/mongo/export/progress/{exportId}", get(routes::mongodb_import_export::export_progress))
.route("/mongo/export/download/{exportId}", get(routes::mongodb_import_export::export_download))
.route("/mongo/export/cancel", post(routes::mongodb_import_export::cancel_export))
.route("/mongo/dump/catalog", post(routes::mongodb_dump::catalog))
.route(
"/mongo/dump/source",
post(routes::mongodb_dump::prepare_source).layer(DefaultBodyLimit::max(
routes::table_import::import_request_body_limit_for_upload(web_body_limit_bytes()),
)),
)
.route("/mongo/dump/source/release", post(routes::mongodb_dump::release_source))
.route("/mongo/dump/upload-limit", get(routes::mongodb_dump::upload_limit))
.route(
"/mongo/dump/source/upload",
post(routes::mongodb_dump::upload_restore_source).layer(DefaultBodyLimit::max(
routes::table_import::import_request_body_limit_for_upload(web_body_limit_bytes()),
)),
)
.route("/mongo/dump/export", post(routes::mongodb_dump::start_dump))
.route("/mongo/dump/restore", post(routes::mongodb_dump::start_restore))
.route("/mongo/dump/progress/{taskId}", get(routes::mongodb_dump::progress))
.route("/mongo/dump/cancel", post(routes::mongodb_dump::cancel))
// History
.route("/history", get(routes::history::load_history).delete(routes::history::clear_history))
.route("/history/save", post(routes::history::save_history))
+1
View File
@@ -17,6 +17,7 @@ pub mod jdbc;
pub mod layout;
pub mod mcp_policy;
pub mod mongo;
pub mod mongodb_dump;
pub mod mongodb_import_export;
#[cfg(feature = "mq-admin")]
pub mod mq;
+402
View File
@@ -0,0 +1,402 @@
use std::path::{Component, Path as FilePath, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use axum::{
extract::{Multipart, Path, State},
response::Sse,
Json,
};
use dbx_core::{
mongodb_dump::{
self, MongoDatabaseDumpProgress, MongoDatabaseDumpRequest, MongoDatabaseRestoreRequest, MongoDumpFormat,
MongoRestoreSourceRequest,
},
transfer,
};
use serde::Deserialize;
use tokio::io::AsyncWriteExt;
use super::export_download::export_download_filename;
use crate::{
error::AppError,
sse::{TransferProgressChannel, TransferReplayEventKind},
state::{WebExportFile, WebState},
};
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CatalogRequest {
connection_id: String,
database: String,
}
pub async fn catalog(
State(state): State<Arc<WebState>>,
Json(request): Json<CatalogRequest>,
) -> Result<Json<serde_json::Value>, AppError> {
let catalog = mongodb_dump::inspect_mongodb_database_dump(&state.app, &request.connection_id, &request.database)
.await
.map_err(AppError::from)?;
Ok(Json(serde_json::to_value(catalog).map_err(|e| AppError::from(e.to_string()))?))
}
fn upload_path(root: &FilePath, filename: &str) -> Result<PathBuf, AppError> {
let path = FilePath::new(filename);
if filename.is_empty()
|| filename.contains(['\\', ':', '\0'])
|| path.components().any(|c| !matches!(c, Component::Normal(_)))
{
return Err(AppError::from("Invalid dump upload path".to_string()));
}
if path.components().count() > 4 {
return Err(AppError::from("Dump directory nesting is too deep".to_string()));
}
Ok(root.join(path))
}
pub async fn prepare_source(
State(state): State<Arc<WebState>>,
multipart: Multipart,
) -> Result<Json<serde_json::Value>, AppError> {
receive_source(state, multipart, false).await
}
pub async fn upload_restore_source(
State(state): State<Arc<WebState>>,
multipart: Multipart,
) -> Result<Json<serde_json::Value>, AppError> {
receive_source(state, multipart, true).await
}
pub async fn upload_limit() -> Json<usize> {
Json(crate::web_body_limit_bytes())
}
async fn field_text(mut field: axum::extract::multipart::Field<'_>, limit: usize) -> Result<String, AppError> {
let mut bytes = Vec::new();
while let Some(chunk) = field.chunk().await.map_err(|e| AppError::bad_request(e.to_string()))? {
if bytes.len().saturating_add(chunk.len()) > limit {
return Err(AppError::bad_request("Dump metadata field is too large"));
}
bytes.extend_from_slice(&chunk);
}
String::from_utf8(bytes).map_err(|e| AppError::bad_request(e.to_string()))
}
async fn receive_source(
state: Arc<WebState>,
mut multipart: Multipart,
acquiring: bool,
) -> Result<Json<serde_json::Value>, AppError> {
let tmp = state.data_dir.join("tmp");
tokio::fs::create_dir_all(&tmp).await.map_err(|e| AppError::from(e.to_string()))?;
let upload =
tempfile::Builder::new().prefix("mongo-upload-").tempdir_in(&tmp).map_err(|e| AppError::from(e.to_string()))?;
let mut format = None;
let mut gzip = false;
let mut paths = Vec::new();
let mut bytes = 0usize;
let mut manifest: Option<Vec<mongodb_dump::MongoDumpFile>> = None;
let mut request: Option<MongoDatabaseRestoreRequest> = None;
let mut expected_upload = std::collections::HashMap::<String, u64>::new();
while let Some(mut field) = multipart.next_field().await.map_err(|e| AppError::from(e.to_string()))? {
match field.name().unwrap_or("") {
"file" => {
if paths.len() >= 100_000 {
return Err(AppError::from("Too many dump files".to_string()));
}
let filename =
field.file_name().ok_or_else(|| AppError::from("Missing dump filename".to_string()))?.to_string();
if format.is_none() {
return Err(AppError::bad_request("Dump format must precede files"));
}
if !acquiring
&& format == Some(MongoDumpFormat::Directory)
&& !filename.ends_with(".metadata.json")
&& !filename.ends_with(".metadata.json.gz")
{
return Err(AppError::bad_request(
"Directory preview accepts metadata only; upload BSON after confirmation",
));
}
if acquiring && request.is_none() {
return Err(AppError::bad_request("Restore plan must precede files"));
}
let expected_size = if acquiring {
Some(
expected_upload
.remove(&filename)
.ok_or_else(|| AppError::bad_request("File is not in the selected restore plan"))?,
)
} else {
None
};
let mut file_bytes = 0u64;
let path = upload_path(upload.path(), &filename)?;
tokio::fs::create_dir_all(path.parent().unwrap()).await.map_err(|e| AppError::from(e.to_string()))?;
let mut output = tokio::fs::OpenOptions::new()
.write(true)
.create_new(true)
.open(&path)
.await
.map_err(|e| AppError::from(e.to_string()))?;
while let Some(chunk) = field.chunk().await.map_err(|e| AppError::from(e.to_string()))? {
if let Some(request) = &request {
if transfer::is_cancelled(&request.task_id).await {
return Err(AppError::bad_request("MongoDB dump/restore cancelled"));
}
}
bytes = bytes.saturating_add(chunk.len());
file_bytes = file_bytes.saturating_add(chunk.len() as u64);
if expected_size.is_some_and(|size| file_bytes > size) {
return Err(AppError::bad_request("File size changed after preview"));
}
let limit = if !acquiring && format == Some(MongoDumpFormat::Directory) {
32 * 1024 * 1024
} else {
crate::web_body_limit_bytes()
};
if bytes > limit {
let mut error =
AppError::from(format!("Dump upload exceeds the configured {} byte limit", limit));
error.status = axum::http::StatusCode::PAYLOAD_TOO_LARGE;
return Err(error);
}
output.write_all(&chunk).await.map_err(|e| AppError::from(e.to_string()))?;
}
output.flush().await.map_err(|e| AppError::from(e.to_string()))?;
if expected_size.is_some_and(|size| file_bytes != size) {
return Err(AppError::bad_request("File size changed after preview"));
}
paths.push((filename, path));
}
"format" => {
let value = field_text(field, 32).await?;
format = Some(
serde_json::from_value::<MongoDumpFormat>(serde_json::Value::String(value))
.map_err(|e| AppError::from(e.to_string()))?,
);
}
"gzip" => {
let value = field_text(field, 16).await?;
gzip = value.parse::<bool>().map_err(|_| AppError::from("Invalid gzip option".to_string()))?;
}
"manifest" if !acquiring => {
let value = field_text(field, 16 * 1024 * 1024).await?;
manifest = Some(serde_json::from_str(&value).map_err(|e| AppError::bad_request(e.to_string()))?);
}
"request" if acquiring => {
let value = field_text(field, 1024 * 1024).await?;
if request.is_some() {
return Err(AppError::bad_request("Duplicate restore plan"));
}
request = Some(serde_json::from_str(&value).map_err(|e| AppError::bad_request(e.to_string()))?);
expected_upload = mongodb_dump::mongodb_restore_upload_files(request.as_ref().unwrap())
.map_err(AppError::from)?
.into_iter()
.map(|f| (f.path, f.size_bytes))
.collect();
}
_ => return Err(AppError::from("Unexpected dump upload field".to_string())),
}
}
let format = format.ok_or_else(|| AppError::from("Dump format is required".to_string()))?;
if paths.is_empty() && manifest.is_none() {
return Err(AppError::from("No dump files uploaded".to_string()));
}
if acquiring {
if !expected_upload.is_empty() {
return Err(AppError::bad_request("Selected restore files are missing"));
}
let request = request.ok_or_else(|| AppError::bad_request("Missing restore plan"))?;
if transfer::is_cancelled(&request.task_id).await {
return Err(AppError::bad_request("MongoDB dump/restore cancelled"));
}
let preview = mongodb_dump::attach_mongodb_restore_upload(&request, upload).await.map_err(AppError::from)?;
return Ok(Json(serde_json::to_value(preview).map_err(|e| AppError::from(e.to_string()))?));
}
if format == MongoDumpFormat::Directory {
let manifest = manifest.ok_or_else(|| AppError::bad_request("Directory preview requires a file manifest"))?;
let expected = manifest
.iter()
.filter(|f| f.path.ends_with(".metadata.json") || f.path.ends_with(".metadata.json.gz"))
.collect::<Vec<_>>();
if expected.len() != paths.len()
|| paths.iter().any(|(name, path)| {
!expected
.iter()
.any(|f| f.path == *name && std::fs::metadata(path).is_ok_and(|m| m.len() == f.size_bytes))
})
{
return Err(AppError::bad_request("Uploaded metadata does not match the file manifest"));
}
let preview =
mongodb_dump::prepare_mongodb_directory_catalog(upload, manifest, gzip).await.map_err(AppError::from)?;
return Ok(Json(serde_json::to_value(preview).map_err(|e| AppError::from(e.to_string()))?));
}
let path = match format {
MongoDumpFormat::Archive if paths.len() == 1 => paths[0].1.clone(),
MongoDumpFormat::Archive => return Err(AppError::from("Archive restore accepts one file".to_string())),
MongoDumpFormat::Directory => {
if paths.iter().any(|(name, _)| FilePath::new(name).components().count() < 2) {
return Err(AppError::from("Directory uploads require database-relative paths".to_string()));
}
upload.path().to_owned()
}
};
let preview = mongodb_dump::prepare_owned_mongodb_restore_source(
MongoRestoreSourceRequest { path: path.to_string_lossy().into_owned(), format, gzip },
upload,
)
.await
.map_err(AppError::from)?;
Ok(Json(serde_json::to_value(preview).map_err(|e| AppError::from(e.to_string()))?))
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SourceRef {
source_ref: String,
}
pub async fn release_source(Json(request): Json<SourceRef>) -> Json<serde_json::Value> {
Json(serde_json::json!({ "released": mongodb_dump::release_mongodb_restore_source(&request.source_ref) }))
}
#[derive(Deserialize)]
pub struct DumpWrapper {
request: MongoDatabaseDumpRequest,
}
#[derive(Deserialize)]
pub struct RestoreWrapper {
request: MongoDatabaseRestoreRequest,
}
enum Task {
Dump(MongoDatabaseDumpRequest, WebExportFile),
Restore(MongoDatabaseRestoreRequest),
}
async fn start_task(state: Arc<WebState>, task: Task) -> Result<Json<serde_json::Value>, AppError> {
let id = match &task {
Task::Dump(request, _) => &request.task_id,
Task::Restore(request) => &request.task_id,
}
.clone();
let key = format!("mongo-database:{id}");
let channel = Arc::new(TransferProgressChannel::new());
{
let mut channels = state.transfer_progress_channels.write().await;
if channels.contains_key(&key) {
return Err(AppError::from("Database dump/restore task already exists".to_string()));
}
channels.insert(key.clone(), channel.clone());
}
let task_id = id.clone();
dbx_core::export_runtime::spawn_export_task(async move {
let mut latest = MongoDatabaseDumpProgress::initial(&task_id);
let mut callback = |progress: MongoDatabaseDumpProgress| {
latest = progress;
if latest.status == "running" {
if let Ok(json) = serde_json::to_string(&latest) {
channel.send(json, TransferReplayEventKind::Progress);
}
}
};
let result = match task {
Task::Dump(request, file) => {
let result = mongodb_dump::dump_mongodb_database(
&state.app,
&request,
|id| {
let id = id.to_string();
Box::pin(async move { transfer::is_cancelled(&id).await })
},
&mut callback,
)
.await;
if result.is_ok() {
state.export_files.write().await.insert(task_id.clone(), file);
}
result
}
Task::Restore(request) => {
mongodb_dump::restore_mongodb_database(
&state.app,
&request,
|id| {
let id = id.to_string();
Box::pin(async move { transfer::is_cancelled(&id).await })
},
&mut callback,
)
.await
}
};
if let Err(error) = result {
latest.error_message = Some(error);
if latest.status == "running" {
latest.status = "error".into();
latest.phase = "done".into();
}
}
if let Ok(json) = serde_json::to_string(&latest) {
channel.send(json, TransferReplayEventKind::Terminal);
}
transfer::clear_cancelled(&task_id).await;
tokio::spawn(async move {
tokio::time::sleep(Duration::from_secs(60)).await;
state.transfer_progress_channels.write().await.remove(&key);
});
});
Ok(Json(serde_json::json!({ "taskId": id })))
}
pub async fn start_dump(
State(state): State<Arc<WebState>>,
Json(body): Json<DumpWrapper>,
) -> Result<Json<serde_json::Value>, AppError> {
let mut request = body.request;
if request.format != MongoDumpFormat::Archive {
return Err(AppError::from("Web database exports use MongoDB archive format".to_string()));
}
let tmp = state.data_dir.join("tmp");
tokio::fs::create_dir_all(&tmp).await.map_err(|e| AppError::from(e.to_string()))?;
let ext = if request.gzip { "archive.gz" } else { "archive" };
let filename = export_download_filename(&request.file_path, &request.database, ext);
request.file_path =
tmp.join(format!("mongo_database_{}.{ext}", uuid::Uuid::new_v4())).to_string_lossy().into_owned();
let file = WebExportFile { file_path: request.file_path.clone(), download_filename: filename, format: ext.into() };
start_task(state, Task::Dump(request, file)).await
}
pub async fn start_restore(
State(state): State<Arc<WebState>>,
Json(body): Json<RestoreWrapper>,
) -> Result<Json<serde_json::Value>, AppError> {
start_task(state, Task::Restore(body.request)).await
}
pub async fn progress(
State(state): State<Arc<WebState>>,
Path(id): Path<String>,
) -> Result<Sse<impl futures::Stream<Item = Result<axum::response::sse::Event, std::convert::Infallible>>>, AppError> {
let channel = state
.transfer_progress_channels
.read()
.await
.get(&format!("mongo-database:{id}"))
.cloned()
.ok_or_else(|| AppError::from("Database task not found".to_string()))?;
Ok(crate::sse::sse_from_transfer_channel(channel))
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CancelRequest {
task_id: String,
}
pub async fn cancel(Json(request): Json<CancelRequest>) -> Json<serde_json::Value> {
transfer::set_cancelled(&request.task_id).await;
Json(serde_json::json!({ "cancelled": true }))
}
@@ -17,13 +17,40 @@ use futures::stream::Stream;
use futures::StreamExt;
use serde::Deserialize;
use tokio::io::AsyncWriteExt;
use tokio_util::io::ReaderStream;
use crate::error::AppError;
use crate::routes::export_download::{attachment_content_disposition, export_download_filename};
use crate::sse::{TransferProgressChannel, TransferReplayEventKind};
use crate::state::{WebExportFile, WebState};
const MONGO_IMPORT_PROGRESS_TTL: Duration = Duration::from_secs(30);
struct RemoveFileOnDrop(PathBuf);
impl Drop for RemoveFileOnDrop {
fn drop(&mut self) {
let _ = std::fs::remove_file(&self.0);
}
}
struct ExportDownloadStream {
// Fields are dropped in declaration order, so Windows closes the file before deleting it.
chunks: ReaderStream<tokio::fs::File>,
_cleanup: RemoveFileOnDrop,
}
impl Stream for ExportDownloadStream {
type Item = std::io::Result<Bytes>;
fn poll_next(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
std::pin::Pin::new(&mut self.get_mut().chunks).poll_next(cx)
}
}
fn initial_import_progress(import_id: &str, started_at: Instant) -> MongoImportProgress {
MongoImportProgress {
import_id: import_id.to_string(),
@@ -153,11 +180,11 @@ pub async fn preview_import(
if let Some((source_ref, file_path)) = uploaded_file {
let file_path_str = file_path.to_string_lossy().to_string();
let format = match format.or_else(|| mongodb_import_export::format_from_path(&file_path_str).ok()) {
let format = match format {
Some(format) => format,
None => {
cleanup_uploaded_import_path(&file_path).await;
return Err(AppError::from("Unsupported MongoDB import file type".to_string()));
return Err(AppError::from("MongoDB import format is required".to_string()));
}
};
let preview = mongodb_import_export::preview_mongodb_import_file(&MongoImportPreviewRequest {
@@ -269,7 +296,7 @@ pub async fn execute_import(
let state_clone = state.clone();
tokio::spawn(async move {
let tx_clone = tx.clone();
let _result = mongodb_import_export::import_mongodb_file_core(
let result = mongodb_import_export::import_mongodb_file_core(
&app,
&req,
|id| {
@@ -279,6 +306,17 @@ pub async fn execute_import(
|progress| send_import_progress(&tx_clone, &progress),
)
.await;
if let Err(error) = result {
let current = tx.borrow().clone();
if let Ok(mut progress) = serde_json::from_str::<MongoImportProgress>(&current) {
if progress.status == MongoImportStatus::Running {
progress.phase = MongoImportPhase::Done;
progress.status = MongoImportStatus::Error;
progress.error_message = Some(error.display_message());
send_import_progress(&tx, &progress);
}
}
}
cleanup_uploaded_import_source(&req.file_path).await;
schedule_import_progress_cleanup(state_clone, req.import_id.clone());
});
@@ -402,26 +440,35 @@ pub async fn start_export(
let ext = match req.format {
MongoExportFormat::Csv => "csv",
MongoExportFormat::Ndjson => "ndjson",
MongoExportFormat::Bson if req.gzip => "bson.gz",
MongoExportFormat::Bson => "bson",
};
let tmp_file = tmp_dir.join(format!("mongo_export_{export_id}.{ext}"));
let tmp_file = tmp_dir.join(format!("mongo_export_{}.{ext}", uuid::Uuid::new_v4()));
let file_path = tmp_file.to_string_lossy().to_string();
let download_filename = export_download_filename(&req.file_path, &req.collection, ext);
req.file_path = file_path.clone();
state
.export_files
.write()
.await
.insert(export_id.clone(), WebExportFile { file_path, download_filename, format: ext.to_string() });
let export_file = WebExportFile { file_path, download_filename, format: ext.to_string() };
let channel_key = format!("mongo-export:{export_id}");
let tx = {
let mut channels = state.sse_channels.write().await;
channels.entry(export_id.clone()).or_insert_with(|| tokio::sync::broadcast::channel::<String>(256).0).clone()
let mut channels = state.transfer_progress_channels.write().await;
let channel = Arc::new(TransferProgressChannel::new());
channels.insert(channel_key.clone(), channel.clone());
channel
};
let app = state.app.clone();
let state_clone = state.clone();
dbx_core::export_runtime::spawn_export_task(async move {
let mut latest = MongoExportProgress {
export_id: req.export_id.clone(),
status: MongoExportStatus::Running,
documents_read: 0,
bytes_written: 0,
total_documents: None,
error_message: None,
elapsed_ms: 0,
};
let result = mongodb_import_export::export_mongodb_query_core(
&app,
&req,
@@ -430,32 +477,38 @@ pub async fn start_export(
Box::pin(async move { transfer::is_cancelled(&id).await })
},
|progress| {
if let Ok(json) = serde_json::to_string(&progress) {
let _ = tx.send(json);
latest = progress;
if latest.status == MongoExportStatus::Running {
if let Ok(json) = serde_json::to_string(&latest) {
tx.send(json, TransferReplayEventKind::Progress);
}
}
},
)
.await;
if let Err(error) = result {
let _ = tokio::fs::remove_file(&req.file_path).await;
state_clone.export_files.write().await.remove(&req.export_id);
let progress = MongoExportProgress {
export_id: req.export_id.clone(),
status: MongoExportStatus::Error,
documents_read: 0,
bytes_written: 0,
total_documents: None,
error_message: Some(error),
elapsed_ms: 0,
};
if let Ok(json) = serde_json::to_string(&progress) {
let _ = tx.send(json);
match result {
Ok(_) => {
// The download must exist before a terminal event can reach the browser.
state_clone.export_files.write().await.insert(req.export_id.clone(), export_file);
}
Err(error) => {
let _ = tokio::fs::remove_file(&req.file_path).await;
if latest.status != MongoExportStatus::Cancelled {
latest.status = MongoExportStatus::Error;
}
latest.error_message = Some(error);
}
}
if let Ok(json) = serde_json::to_string(&latest) {
tx.send(json, TransferReplayEventKind::Terminal);
}
transfer::clear_cancelled(&req.export_id).await;
tokio::time::sleep(Duration::from_secs(5)).await;
state_clone.remove_sse_channel(&req.export_id).await;
tokio::spawn(async move {
tokio::time::sleep(Duration::from_secs(30)).await;
state_clone.transfer_progress_channels.write().await.remove(&channel_key);
});
});
Ok(Json(serde_json::json!({ "exportId": export_id })))
@@ -465,12 +518,14 @@ pub async fn export_progress(
State(state): State<Arc<WebState>>,
Path(export_id): Path<String>,
) -> Result<Sse<impl Stream<Item = Result<axum::response::sse::Event, std::convert::Infallible>>>, AppError> {
let tx = {
let mut channels = state.sse_channels.write().await;
channels.entry(export_id).or_insert_with(|| tokio::sync::broadcast::channel::<String>(256).0).clone()
};
let rx = tx.subscribe();
Ok(crate::sse::sse_from_channel(rx))
let channel = state
.transfer_progress_channels
.read()
.await
.get(&format!("mongo-export:{export_id}"))
.cloned()
.ok_or_else(|| AppError::from("Export not found".to_string()))?;
Ok(crate::sse::sse_from_transfer_channel(channel))
}
pub async fn cancel_export(
@@ -491,17 +546,27 @@ pub async fn export_download(
.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()))?;
let _ = tokio::fs::remove_file(&export_file.file_path).await;
export_file_response(export_file).await
}
pub(crate) async fn export_file_response(export_file: WebExportFile) -> Result<Response, AppError> {
let file = tokio::fs::File::open(&export_file.file_path).await.map_err(|e| AppError::from(e.to_string()))?;
let content_type = match export_file.format.as_str() {
"csv" => "text/csv; charset=utf-8",
_ => "application/x-ndjson; charset=utf-8",
"ndjson" => "application/x-ndjson; charset=utf-8",
"bson.gz" | "archive.gz" => "application/gzip",
"bson" => "application/bson",
_ => "application/octet-stream",
};
let download = ExportDownloadStream {
chunks: ReaderStream::new(file),
_cleanup: RemoveFileOnDrop(PathBuf::from(&export_file.file_path)),
};
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))
.body(Body::from_stream(download))
.map_err(|error| AppError::from(error.to_string()))
}
@@ -509,6 +574,152 @@ pub async fn export_download(
mod tests {
use super::*;
#[tokio::test]
async fn download_keeps_file_until_body_is_consumed_or_abandoned() {
for extension in ["bson", "bson.gz"] {
for consume in [true, false] {
let directory = tempfile::tempdir().unwrap();
let path = directory.path().join(format!("records.{extension}"));
let bytes = vec![42u8; 100_000];
std::fs::write(&path, &bytes).unwrap();
let response = export_file_response(WebExportFile {
file_path: path.to_str().unwrap().into(),
download_filename: format!("records.{extension}"),
format: extension.into(),
})
.await
.unwrap_or_else(|error| panic!("{}", error.message));
assert!(path.exists(), "download file was removed before streaming");
assert_eq!(
response.headers()[header::CONTENT_TYPE],
if extension == "bson" { "application/bson" } else { "application/gzip" }
);
assert!(response.headers()[header::CONTENT_DISPOSITION]
.to_str()
.unwrap()
.contains(&format!("records.{extension}")));
if consume {
assert_eq!(axum::body::to_bytes(response.into_body(), bytes.len()).await.unwrap(), bytes);
} else {
let mut stream = response.into_body().into_data_stream();
assert!(!stream.next().await.unwrap().unwrap().is_empty());
drop(stream);
}
assert!(!path.exists(), "download file leaked after the response ended");
}
}
}
#[tokio::test]
async fn export_failure_and_cancellation_are_replayed_to_late_subscribers() {
use axum::response::IntoResponse;
use dbx_core::{connection::AppState, storage::Storage};
for cancelled in [false, true] {
let directory = tempfile::tempdir().unwrap();
let storage = Storage::open(&directory.path().join("storage.db")).await.unwrap();
let state = Arc::new(WebState::for_tests(Arc::new(AppState::new(storage)), directory.path().into()));
let id = uuid::Uuid::new_v4().to_string();
let expected = if cancelled { MongoExportStatus::Cancelled } else { MongoExportStatus::Error };
if cancelled {
transfer::set_cancelled(&id).await;
}
let Json(started) = start_export(
State(state.clone()),
Json(StartExportRequest {
request: MongoExportRequest {
export_id: id.clone(),
connection_id: "missing".into(),
database: "test".into(),
collection: "empty".into(),
filter: None,
projection: None,
sort: None,
collation: None,
format: MongoExportFormat::Bson,
gzip: false,
include_header: false,
file_path: "empty.bson".into(),
execution_id: None,
},
}),
)
.await
.unwrap_or_else(|error| panic!("{}", error.message));
assert_eq!(started["exportId"], id);
let key = format!("mongo-export:{id}");
tokio::time::timeout(Duration::from_secs(5), async {
loop {
let latest = state.transfer_progress_channels.read().await.get(&key).unwrap().latest();
if latest.is_some_and(|json| {
serde_json::from_str::<MongoExportProgress>(&json).unwrap().status == expected
}) {
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.unwrap();
assert!(!state.export_files.read().await.contains_key(&id));
let response = export_progress(State(state.clone()), Path(id))
.await
.unwrap_or_else(|error| panic!("{}", error.message));
let mut stream = response.into_response().into_body().into_data_stream();
let chunk = tokio::time::timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap();
assert!(String::from_utf8_lossy(&chunk)
.contains(&format!("\"status\":{}", serde_json::to_string(&expected).unwrap())));
}
}
#[tokio::test]
async fn early_import_failure_finishes_progress_and_cleans_source() {
use dbx_core::{connection::AppState, storage::Storage};
let directory = tempfile::tempdir().unwrap();
let storage = Storage::open(&directory.path().join("storage.db")).await.unwrap();
let state = Arc::new(WebState::for_tests(Arc::new(AppState::new(storage)), directory.path().into()));
let upload_dir = import_upload_dir(directory.path());
std::fs::create_dir_all(&upload_dir).unwrap();
let path = upload_dir.join("empty.bson");
std::fs::write(&path, []).unwrap();
let id = uuid::Uuid::new_v4().to_string();
let Json(started) = execute_import(
State(state.clone()),
Json(ExecuteImportWrapper {
request: MongoImportRequest {
import_id: id.clone(),
connection_id: "missing".into(),
database: "test".into(),
collection: "empty".into(),
file_path: path.to_str().unwrap().into(),
source_ref: None,
format: mongodb_import_export::MongoImportFormat::Bson,
parse_options: Default::default(),
batch_size: 1,
execution_id: None,
},
}),
)
.await
.unwrap_or_else(|error| panic!("{}", error.message));
assert_eq!(started["importId"], id);
tokio::time::timeout(Duration::from_secs(5), async {
loop {
let progress = state.table_import_channels.read().await.get(&id).unwrap().borrow().clone();
let progress: MongoImportProgress = serde_json::from_str(&progress).unwrap();
if progress.status == MongoImportStatus::Error && !path.exists() {
assert_eq!(progress.phase, MongoImportPhase::Done);
assert!(progress.error_message.unwrap().contains("Batch size"));
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.unwrap();
}
#[test]
fn source_ref_rejects_path_traversal() {
let data_dir = std::env::temp_dir().join(format!("dbx-mongo-import-data-{}", uuid::Uuid::new_v4()));