github/dbx
4
0 Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-11-15 22:53:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(mysql): prevent stale export pool cleanup

Browse Source
This commit is contained in:
Guoyu Su authored and GitHub committed 2026-08-25 22:50:40 +08:00
1 parent aa91a0aaa9
commit 0658d18922
8 files changed
+400 -66
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/dbx-core/src/database_export.rs
+88 -48
View File
@@ -3,7 +3,8 @@ use serde_json::Value;
use std::collections::{HashMap, HashSet};
use std::future::Future;
use std::io::{BufWriter, Write};
use std::sync::RwLock;
use std::pin::Pin;
use std::sync::{Arc, RwLock};
use std::time::Duration;
use tokio_util::sync::CancellationToken;
@@ -1586,10 +1587,10 @@ pub async fn clear_export_cancelled(export_id: &str) {
/// not accept a cancellation token; polling here keeps the export task
/// responsive and dropping the pending future follows the same bounded
/// prefetch cancellation behavior used below.
async fn await_export_operation<T, F>(export_id: &str, operation: F) -> Result<T, String>
where
F: Future<Output = Result<T, String>>,
{
async fn await_export_operation<T>(
export_id: &str,
operation: Pin<Box<dyn Future<Output = Result<T, String>> + Send + '_>>,
) -> Result<T, String> {
tokio::pin!(operation);
loop {
if is_export_cancelled_now(export_id) {
@@ -1609,6 +1610,43 @@ where
}
}
struct AbortExportTaskOnDrop(tokio::task::AbortHandle);
impl Drop for AbortExportTaskOnDrop {
fn drop(&mut self) {
self.0.abort();
}
}
async fn get_export_table_ddl_isolated(
state: Arc<crate::connection::AppState>,
connection_id: String,
database: String,
schema: String,
table: String,
) -> Result<String, String> {
let task = tokio::spawn(async move {
crate::schema::get_table_relation_export_ddl_core(&state, &connection_id, &database, &schema, &table, None)
.await
});
let _abort_on_drop = AbortExportTaskOnDrop(task.abort_handle());
task.await.map_err(|error| format!("Database export metadata task failed: {error}"))?
}
async fn get_export_table_columns_isolated(
state: Arc<crate::connection::AppState>,
connection_id: String,
database: String,
schema: String,
table: String,
) -> Result<Vec<crate::db::ColumnInfo>, String> {
let task = tokio::spawn(async move {
crate::schema::get_columns_core(&state, &connection_id, &database, &schema, &table).await
});
let _abort_on_drop = AbortExportTaskOnDrop(task.abort_handle());
task.await.map_err(|error| format!("Database export metadata task failed: {error}"))?
}
fn snapshot_batch_cancelled(db_type: &DatabaseType, export_id: &str) -> bool {
*db_type != DatabaseType::Postgres && is_export_cancelled_now(export_id)
}
@@ -2106,7 +2144,7 @@ mod windows_export_destination {
}
pub async fn export_database_sql_core(
state: &crate::connection::AppState,
state: &Arc<crate::connection::AppState>,
request: &DatabaseExportRequest,
on_progress: impl Fn(ExportProgress) + Sync,
) -> Result<(), String> {
@@ -2126,7 +2164,7 @@ pub async fn export_database_sql_core(
}
async fn export_database_sql_core_inner(
state: &crate::connection::AppState,
state: &Arc<crate::connection::AppState>,
request: &DatabaseExportRequest,
on_progress: impl Fn(ExportProgress) + Sync,
) -> Result<(), String> {
@@ -2275,7 +2313,7 @@ async fn export_database_sql_core_inner(
let postgres_sequences = if request.include_structure && matches!(db_type, DatabaseType::Postgres) {
match await_export_operation(
&request.export_id,
list_postgres_export_sequences(
Box::pin(list_postgres_export_sequences(
state,
&pool_key,
&request.schema,
@@ -2283,7 +2321,7 @@ async fn export_database_sql_core_inner(
&request.excluded_tables,
request.include_objects,
request.fail_on_error,
),
)),
)
.await
{
@@ -2416,8 +2454,8 @@ async fn export_database_sql_core_inner(
use futures::StreamExt;
let prefetch_targets: Vec<(usize, String)> =
tables.iter().enumerate().map(|(index, table_info)| (index, table_info.name.clone())).collect();
let mut prefetch_stream =
futures::stream::iter(prefetch_targets.into_iter().map(|(index, table_name)| async move {
let mut prefetch_stream = futures::stream::iter(prefetch_targets.into_iter().map(|(index, table_name)| {
Box::pin(async move {
if is_export_cancelled_now(&request.export_id) {
return (index, PrefetchedTableMetadata { ddl: None, columns: None });
}
@@ -2425,14 +2463,13 @@ async fn export_database_sql_core_inner(
Some(
await_export_operation(
&request.export_id,
crate::schema::get_table_relation_export_ddl_core(
state,
&request.connection_id,
&request.database,
&request.schema,
&table_name,
None,
),
Box::pin(get_export_table_ddl_isolated(
state.clone(),
request.connection_id.clone(),
request.database.clone(),
request.schema.clone(),
table_name.clone(),
)),
)
.await,
)
@@ -2446,13 +2483,13 @@ async fn export_database_sql_core_inner(
Some(
await_export_operation(
&request.export_id,
crate::schema::get_columns_core(
state,
&request.connection_id,
&request.database,
&request.schema,
&table_name,
),
Box::pin(get_export_table_columns_isolated(
state.clone(),
request.connection_id.clone(),
request.database.clone(),
request.schema.clone(),
table_name.clone(),
)),
)
.await,
)
@@ -2460,8 +2497,9 @@ async fn export_database_sql_core_inner(
None
};
(index, PrefetchedTableMetadata { ddl, columns })
}))
.buffer_unordered(database_export_metadata_prefetch_concurrency(db_type));
})
}))
.buffer_unordered(database_export_metadata_prefetch_concurrency(db_type));
while let Some((index, metadata)) = prefetch_stream.next().await {
if metadata
.ddl
@@ -2597,14 +2635,13 @@ async fn export_database_sql_core_inner(
None => {
await_export_operation(
&request.export_id,
crate::schema::get_table_relation_export_ddl_core(
state,
&request.connection_id,
&request.database,
&request.schema,
table_name,
None,
),
Box::pin(get_export_table_ddl_isolated(
state.clone(),
request.connection_id.clone(),
request.database.clone(),
request.schema.clone(),
table_name.clone(),
)),
)
.await
}
@@ -2645,13 +2682,13 @@ async fn export_database_sql_core_inner(
None => {
await_export_operation(
&request.export_id,
crate::schema::get_columns_core(
state,
&request.connection_id,
&request.database,
&request.schema,
table_name,
),
Box::pin(get_export_table_columns_isolated(
state.clone(),
request.connection_id.clone(),
request.database.clone(),
request.schema.clone(),
table_name.clone(),
)),
)
.await
}
@@ -3124,10 +3161,13 @@ mod tests {
clear_export_cancelled(&export_id).await;
let task_export_id = export_id.clone();
let task = tokio::spawn(async move {
await_export_operation(&task_export_id, async {
tokio::time::sleep(std::time::Duration::from_secs(30)).await;
Ok::<_, String>(())
})
await_export_operation(
&task_export_id,
Box::pin(async {
tokio::time::sleep(std::time::Duration::from_secs(30)).await;
Ok::<_, String>(())
}),
)
.await
});
@@ -3145,7 +3185,7 @@ mod tests {
async fn await_export_operation_preserves_completed_metadata() {
let export_id = format!("complete-metadata-{}", uuid::Uuid::new_v4());
clear_export_cancelled(&export_id).await;
let result = await_export_operation(&export_id, async { Ok::<_, String>(42_u32) }).await;
let result = await_export_operation(&export_id, Box::pin(async { Ok::<_, String>(42_u32) })).await;
assert_eq!(result, Ok(42));
}