xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(export): use inactivity timeout for result export

This commit is contained in:
zipg 2026-07-21 21:08:13 +08:00 committed by GitHub
parent 1b94534558
commit 754c8b4a53
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 336 additions and 119 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
crates/dbx-core/src/db/postgres.rs
View File

@ -21,7 +21,7 @@ use tokio_postgres::{NoTls, Row, SimpleQueryMessage};
use tokio_util::sync::CancellationToken;
use super::file_validator::validate_file_path;
use crate::query::DbOperationBudget;
use crate::query::{await_stream_with_progress_timeout, DbOperationBudget, StreamProgressClock};
use crate::sql::starts_with_executable_sql_keyword;
use crate::types::{
ColumnInfo, CompletionAssistantCandidate, CompletionAssistantCandidateKind, CompletionAssistantMatchMode,
@ -2556,15 +2556,27 @@ pub async fn stream_select_query_with_cancel(
}
let pg_cancel_token = client.cancel_token();
let result = wait_postgres_query(
pg_cancel_token,
cancel_context,
cancel_token,
budget.query_timeout,
budget.cancel_timeout,
stream_select_query_inner(&client, sql, row_limit, &mut on_item),
let query_timeout = budget.query_timeout;
let timeout_error =
format!("Query timed out after {} seconds", query_timeout.map_or(0, |timeout| timeout.as_secs()));
let progress_clock = Arc::new(StreamProgressClock::new());
let progress_clock_for_stream = progress_clock.clone();
let mut on_stream_item = |item| {
on_item(item)?;
progress_clock_for_stream.mark();
Ok(())
};
let result = await_stream_with_progress_timeout(
stream_select_query_inner(&client, sql, row_limit, &mut on_stream_item),
query_timeout,
progress_clock,
cancel_token.as_ref(),
timeout_error.clone(),
)
.await;
if result.as_ref().is_err_and(|error| error == &timeout_error || error == crate::query::QUERY_CANCELED) {
cancel_postgres_query(pg_cancel_token, cancel_context.as_ref(), budget.cancel_timeout).await;
}
if schema_was_set {
let reset_result = reset_postgres_search_path(&client, budget.cleanup_timeout, start).await;

83
crates/dbx-core/src/query.rs
View File

@ -13,7 +13,10 @@ use sqlparser::parser::Parser;
use std::collections::HashSet;
use std::future::Future;
use std::ops::ControlFlow;
use std::sync::Arc;
use std::sync::{
atomic::{AtomicU64, Ordering},
Arc,
};
use std::time::Duration;
#[cfg(feature = "duckdb-bundled")]
use tokio::task::JoinHandle;
@ -1024,6 +1027,84 @@ pub fn canceled_error() -> String {
QUERY_CANCELED.to_string()
}
pub(crate) struct StreamProgressClock {
started_at: tokio::time::Instant,
last_progress_ms: AtomicU64,
}
impl StreamProgressClock {
pub(crate) fn new() -> Self {
Self { started_at: tokio::time::Instant::now(), last_progress_ms: AtomicU64::new(0) }
}
pub(crate) fn mark(&self) {
self.last_progress_ms.store(self.started_at.elapsed().as_millis() as u64, Ordering::Relaxed);
}
fn elapsed_since_progress(&self) -> Duration {
let last_progress_ms = self.last_progress_ms.load(Ordering::Relaxed);
let elapsed_ms = self.started_at.elapsed().as_millis() as u64;
Duration::from_millis(elapsed_ms.saturating_sub(last_progress_ms))
}
}
pub(crate) async fn await_stream_with_progress_timeout<F, T>(
stream_future: F,
timeout: Option<Duration>,
progress_clock: Arc<StreamProgressClock>,
cancel_token: Option<&CancellationToken>,
timeout_message: String,
) -> Result<T, String>
where
F: Future<Output = Result<T, String>>,
{
let Some(timeout) = timeout else {
return match cancel_token {
Some(token) => {
tokio::select! {
biased;
_ = token.cancelled() => Err(canceled_error()),
result = stream_future => result,
}
}
None => stream_future.await,
};
};
tokio::pin!(stream_future);
loop {
// Query timeout is an inactivity budget, not a cap on total stream duration.
let remaining = timeout.saturating_sub(progress_clock.elapsed_since_progress());
if remaining.is_zero() {
return Err(timeout_message.clone());
}
let sleep = tokio::time::sleep(remaining);
tokio::pin!(sleep);
match cancel_token {
Some(token) => {
tokio::select! {
biased;
_ = token.cancelled() => return Err(canceled_error()),
result = &mut stream_future => return result,
_ = &mut sleep => {},
}
}
None => {
tokio::select! {
biased;
result = &mut stream_future => return result,
_ = &mut sleep => {},
}
}
}
if progress_clock.elapsed_since_progress() >= timeout {
return Err(timeout_message);
}
}
}
#[cfg(feature = "duckdb-bundled")]
pub fn duckdb_draining_error() -> String {
DUCKDB_DRAINING_MESSAGE.to_string()

145
crates/dbx-core/src/query_result_export.rs
View File

@ -1,9 +1,8 @@
use serde::{Deserialize, Serialize};
use std::fs::File;
use std::future::Future;
use std::io::{BufWriter, Seek, Write};
use std::sync::{
atomic::{AtomicBool, AtomicU64, Ordering},
atomic::{AtomicBool, Ordering},
Arc,
};
use std::time::{Duration, Instant};
@ -14,8 +13,8 @@ use crate::database_export::is_export_cancelled;
pub use crate::database_export::ExportStatus;
use crate::models::connection::DatabaseType;
use crate::query::{
canceled_error, close_query_session, execute_sql_statement_with_options, operation_budget_for_pool_key,
QueryExecutionOptions, QUERY_CANCELED,
await_stream_with_progress_timeout, canceled_error, close_query_session, execute_sql_statement_with_options,
operation_budget_for_pool_key, QueryExecutionOptions, StreamProgressClock, QUERY_CANCELED,
};
use crate::query_result_sql::{
build_query_pagination_execution_plan, QueryPagination, QueryPaginationExecutionPlanOptions,
@ -193,85 +192,6 @@ fn query_export_timeout(timeout_secs: Option<u64>) -> Option<Duration> {
}
}
struct StreamProgressClock {
started_at: tokio::time::Instant,
last_progress_ms: AtomicU64,
}
impl StreamProgressClock {
fn new() -> Self {
Self { started_at: tokio::time::Instant::now(), last_progress_ms: AtomicU64::new(0) }
}
fn mark(&self) {
self.last_progress_ms.store(self.started_at.elapsed().as_millis() as u64, Ordering::Relaxed);
}
fn elapsed_since_progress(&self) -> Duration {
let last_progress_ms = self.last_progress_ms.load(Ordering::Relaxed);
let elapsed_ms = self.started_at.elapsed().as_millis() as u64;
Duration::from_millis(elapsed_ms.saturating_sub(last_progress_ms))
}
}
async fn await_stream_with_progress_timeout<F, T>(
stream_future: F,
timeout: Option<Duration>,
progress_clock: Arc<StreamProgressClock>,
cancel_token: Option<&CancellationToken>,
timeout_message: String,
) -> Result<T, String>
where
F: Future<Output = Result<T, String>>,
{
let Some(timeout) = timeout else {
return match cancel_token {
Some(token) => {
tokio::select! {
biased;
_ = token.cancelled() => Err(canceled_error()),
result = stream_future => result,
}
}
None => stream_future.await,
};
};
tokio::pin!(stream_future);
loop {
// The query timeout is an inactivity budget, not a cap on total export duration.
// This keeps a stalled server bounded while allowing large local file writes to finish.
let remaining = timeout.saturating_sub(progress_clock.elapsed_since_progress());
if remaining.is_zero() {
return Err(timeout_message.clone());
}
let sleep = tokio::time::sleep(remaining);
tokio::pin!(sleep);
match cancel_token {
Some(token) => {
tokio::select! {
biased;
_ = token.cancelled() => return Err(canceled_error()),
result = &mut stream_future => return result,
_ = &mut sleep => {},
}
}
None => {
tokio::select! {
biased;
result = &mut stream_future => return result,
_ = &mut sleep => {},
}
}
}
if progress_clock.elapsed_since_progress() >= timeout {
return Err(timeout_message);
}
}
}
fn should_fetch_next_page(
use_agent_result_session: bool,
has_more: bool,
@ -1025,6 +945,8 @@ async fn try_export_mysql_query_result_stream(
}
});
let progress_clock = Arc::new(StreamProgressClock::new());
let progress_clock_for_stream = progress_clock.clone();
let stream_future = crate::db::mysql::stream_query_result_on_conn(
&mut conn,
&request.sql,
@ -1093,19 +1015,23 @@ async fn try_export_mysql_query_result_stream(
}
}
}
progress_clock_for_stream.mark();
Ok(())
},
);
let stream_result = match query_timeout {
Some(timeout) => match tokio::time::timeout(timeout, stream_future).await {
Ok(result) => result,
Err(_) => {
let _ = crate::db::mysql::kill_query_with_opts(kill_opts, mysql_connection_id).await;
Err(format!("Query timed out after {} seconds", timeout.as_secs()))
}
},
None => stream_future.await,
};
let timeout_error =
format!("Query timed out after {} seconds", query_timeout.map_or(0, |timeout| timeout.as_secs()));
let stream_result = await_stream_with_progress_timeout(
stream_future,
query_timeout,
progress_clock,
cancel_token.as_ref(),
timeout_error.clone(),
)
.await;
if stream_result.as_ref().is_err_and(|error| error == &timeout_error) {
let _ = crate::db::mysql::kill_query_with_opts(kill_opts, mysql_connection_id).await;
}
watcher_done.cancel();
if let Err(error) = stream_result {
@ -1229,8 +1155,11 @@ async fn try_export_clickhouse_query_result_stream(
None
};
let mut xlsx = None;
let query_timeout = query_export_timeout(request.timeout_secs);
let clickhouse_database = if database.is_empty() { "default" } else { database };
let progress_clock = Arc::new(StreamProgressClock::new());
let progress_clock_for_stream = progress_clock.clone();
let stream_future = crate::db::clickhouse_driver::stream_query_with_max_rows(
&client,
clickhouse_database,
@ -1296,16 +1225,18 @@ async fn try_export_clickhouse_query_result_stream(
}
}
}
progress_clock_for_stream.mark();
Ok(())
},
);
let stream_result = match query_export_timeout(request.timeout_secs) {
Some(timeout) => match tokio::time::timeout(timeout, stream_future).await {
Ok(result) => result,
Err(_) => Err(format!("Query timed out after {} seconds", timeout.as_secs())),
},
None => stream_future.await,
};
let stream_result = await_stream_with_progress_timeout(
stream_future,
query_timeout,
progress_clock,
cancel_token.as_ref(),
format!("Query timed out after {} seconds", query_timeout.map_or(0, |timeout| timeout.as_secs())),
)
.await;
if let Err(error) = stream_result {
if error == QUERY_CANCELED
@ -1389,8 +1320,6 @@ async fn try_export_sqlserver_query_result_stream(
None
};
let mut xlsx = None;
let progress_clock = Arc::new(StreamProgressClock::new());
let progress_clock_for_stream = progress_clock.clone();
let query_timeout = query_export_timeout(request.timeout_secs);
let mut client = match cancel_token.as_ref() {
@ -1404,6 +1333,8 @@ async fn try_export_sqlserver_query_result_stream(
None => client.lock().await,
};
let progress_clock = Arc::new(StreamProgressClock::new());
let progress_clock_for_stream = progress_clock.clone();
let stream_future = crate::db::sqlserver::stream_first_result_set(
&mut client,
&request.sql,
@ -1683,7 +1614,7 @@ mod tests {
}
#[tokio::test]
async fn sqlserver_stream_times_out_when_database_makes_no_progress() {
async fn stream_times_out_when_database_makes_no_progress() {
let progress_clock = Arc::new(StreamProgressClock::new());
let result = await_stream_with_progress_timeout(
std::future::pending::<Result<(), String>>(),
@ -1698,7 +1629,7 @@ mod tests {
}
#[tokio::test]
async fn sqlserver_stream_timeout_resets_after_each_completed_row() {
async fn stream_timeout_resets_after_each_completed_row() {
let progress_clock = Arc::new(StreamProgressClock::new());
let progress_clock_for_stream = progress_clock.clone();
let result = await_stream_with_progress_timeout(
@ -1721,7 +1652,7 @@ mod tests {
}
#[tokio::test]
async fn sqlserver_stream_does_not_count_synchronous_local_writes_as_database_idle_time() {
async fn stream_does_not_count_synchronous_local_writes_as_database_idle_time() {
let progress_clock = Arc::new(StreamProgressClock::new());
let progress_clock_for_stream = progress_clock.clone();
let result = await_stream_with_progress_timeout(
@ -1741,7 +1672,7 @@ mod tests {
}
#[tokio::test]
async fn sqlserver_stream_timeout_zero_disables_idle_timeout() {
async fn stream_timeout_zero_disables_idle_timeout() {
let progress_clock = Arc::new(StreamProgressClock::new());
let result = await_stream_with_progress_timeout(
async {
@ -1759,7 +1690,7 @@ mod tests {
}
#[tokio::test]
async fn sqlserver_stream_cancellation_wins_over_idle_timeout() {
async fn stream_cancellation_wins_over_idle_timeout() {
let progress_clock = Arc::new(StreamProgressClock::new());
let cancel_token = CancellationToken::new();
let cancel_token_for_task = cancel_token.clone();

82
crates/dbx-core/tests/live_mysql57.rs
View File

@ -1,5 +1,6 @@
use std::collections::BTreeSet;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::{Duration, Instant};
use dbx_core::connection::AppState;
use dbx_core::models::connection::{ConnectionConfig, DatabaseType};
@ -195,6 +196,85 @@ async fn live_mysql_query_result_export_xlsx_streams_single_query_without_duplic
assert_eq!(exported_ids, expected_ids);
}
#[tokio::test]
#[ignore = "requires a writable MySQL endpoint for a 650,000-row XLSX export"]
async fn live_mysql_xlsx_export_can_outlive_query_timeout_while_rows_keep_arriving() {
let host = std::env::var("DBX_LIVE_MYSQL_EXPORT_HOST").expect("DBX_LIVE_MYSQL_EXPORT_HOST");
let port = std::env::var("DBX_LIVE_MYSQL_EXPORT_PORT").expect("DBX_LIVE_MYSQL_EXPORT_PORT").parse::<u16>().unwrap();
let user = std::env::var("DBX_LIVE_MYSQL_EXPORT_USER").expect("DBX_LIVE_MYSQL_EXPORT_USER");
let password = std::env::var("DBX_LIVE_MYSQL_EXPORT_PASSWORD").expect("DBX_LIVE_MYSQL_EXPORT_PASSWORD");
let database = std::env::var("DBX_LIVE_MYSQL_EXPORT_DATABASE").expect("DBX_LIVE_MYSQL_EXPORT_DATABASE");
let suffix = uuid::Uuid::new_v4().simple().to_string();
let table = format!("dbx_query_export_timeout_{}", &suffix[..8]);
let connection_id = format!("live-mysql-query-export-timeout-{suffix}");
let config = live_mysql_query_export_config(&connection_id, &host, port, &user, &password, &database);
let dir = std::env::temp_dir().join(format!("dbx-live-mysql-query-export-timeout-{suffix}"));
std::fs::create_dir_all(&dir).unwrap();
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
let state = AppState::new(storage);
state.configs.write().await.insert(config.id.clone(), config);
let values = (1..=807).map(|id| format!("({id}, 'row-{id}')")).collect::<Vec<_>>().join(", ");
let cleanup_sql = format!("DROP TABLE IF EXISTS `{table}`");
let create_sql = format!("CREATE TABLE `{table}` (id INT PRIMARY KEY, label VARCHAR(32) NOT NULL)");
let insert_sql = format!("INSERT INTO `{table}` (id, label) VALUES {values}");
let _ = execute_sql_statement(&state, &connection_id, &database, &cleanup_sql, None, None).await;
execute_sql_statement(&state, &connection_id, &database, &create_sql, None, None)
.await
.expect("create live export timeout table");
execute_sql_statement(&state, &connection_id, &database, &insert_sql, None, None)
.await
.expect("insert live export timeout rows");
let file_path = dir.join("result.xlsx");
let sql = format!(
"SELECT a.id * 100000 + b.id AS id, CONCAT(a.label, '-', b.label) AS label \
FROM `{table}` AS a CROSS JOIN `{table}` AS b LIMIT 650000"
);
let request = QueryResultExportRequest {
export_id: format!("live-mysql-query-export-timeout-{suffix}"),
connection_id: connection_id.clone(),
database: database.clone(),
schema: None,
sql: sql.clone(),
query_base_sql: sql,
database_type: DatabaseType::Mysql,
use_agent_cursor: false,
file_path: file_path.to_string_lossy().to_string(),
format: "xlsx".to_string(),
include_sql_sheet: false,
page_size: 10_000,
row_limit: None,
total_rows: Some(650_000),
timeout_secs: Some(1),
keyset_optimization_enabled: false,
client_session_id: None,
execution_id: Some(format!("live-mysql-query-export-timeout-{suffix}")),
date_time_format: None,
};
let rows_exported = AtomicU64::new(0);
let done_seen = AtomicBool::new(false);
let started_at = Instant::now();
let result = export_query_result_core(&state, &request, None, |progress| {
rows_exported.store(progress.rows_exported, Ordering::Relaxed);
if matches!(progress.status, ExportStatus::Done) {
done_seen.store(true, Ordering::Relaxed);
}
})
.await;
let elapsed = started_at.elapsed();
let cleanup_result = execute_sql_statement(&state, &connection_id, &database, &cleanup_sql, None, None).await;
result.expect("stream 650,000 MySQL rows to XLSX");
cleanup_result.expect("cleanup live export timeout table");
assert!(elapsed > Duration::from_secs(1), "export should outlive configured timeout: {elapsed:?}");
assert_eq!(rows_exported.load(Ordering::Relaxed), 650_000);
assert!(done_seen.load(Ordering::Relaxed));
assert!(std::fs::metadata(&file_path).unwrap().len() > 1_000_000);
let _ = std::fs::remove_dir_all(dir);
}
#[tokio::test]
#[ignore = "requires a remote MySQL endpoint"]
async fn live_mysql_call_procedure_returns_select_result_set() {

117
crates/dbx-core/tests/live_postgres_query_result_export.rs
View File

@ -1,9 +1,10 @@
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::{Duration, Instant};
use dbx_core::connection::AppState;
use dbx_core::db::postgres;
use dbx_core::models::connection::{ConnectionConfig, DatabaseType};
use dbx_core::query::execute_sql_statement;
use dbx_core::query_result_export::{export_query_result_core, ExportStatus, QueryResultExportRequest};
use dbx_core::storage::Storage;
@ -153,3 +154,115 @@ async fn live_postgres_query_result_export_uses_single_streamed_query() {
assert!(csv.contains("\"2050\",\"1\""));
assert_eq!(csv.lines().count(), 2051, "unexpected csv row count");
}
#[tokio::test]
#[ignore = "requires DBX_LIVE_POSTGRES_* env vars for a 650,000-row XLSX export"]
async fn live_postgres_xlsx_export_can_outlive_query_timeout_while_rows_keep_arriving() {
let host = std::env::var("DBX_LIVE_POSTGRES_HOST").unwrap_or_else(|_| "127.0.0.1".to_string());
let port = std::env::var("DBX_LIVE_POSTGRES_PORT").ok().and_then(|value| value.parse().ok()).unwrap_or(5432);
let user = std::env::var("DBX_LIVE_POSTGRES_USER").unwrap_or_else(|_| "postgres".to_string());
let password = std::env::var("DBX_LIVE_POSTGRES_PASSWORD").unwrap_or_default();
let database = std::env::var("DBX_LIVE_POSTGRES_DATABASE").unwrap_or_else(|_| "postgres".to_string());
let suffix = uuid::Uuid::new_v4().simple().to_string();
let connection_id = format!("live-postgres-query-export-timeout-{suffix}");
let config = live_postgres_config(&connection_id, &host, port, &user, &password, &database);
let dir = std::env::temp_dir().join(format!("dbx-live-postgres-query-export-timeout-{suffix}"));
std::fs::create_dir_all(&dir).unwrap();
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
let state = AppState::new(storage);
state.configs.write().await.insert(config.id.clone(), config);
let file_path = dir.join("result.xlsx");
let sql = "SELECT i AS id, repeat('x', 64) AS payload FROM generate_series(1, 650000) AS source(i)";
let request = QueryResultExportRequest {
export_id: format!("live-postgres-query-export-timeout-{suffix}"),
connection_id,
database,
schema: Some("public".to_string()),
sql: sql.to_string(),
query_base_sql: sql.to_string(),
database_type: DatabaseType::Postgres,
use_agent_cursor: false,
file_path: file_path.to_string_lossy().to_string(),
format: "xlsx".to_string(),
include_sql_sheet: false,
page_size: 10_000,
row_limit: None,
total_rows: Some(650_000),
timeout_secs: Some(1),
keyset_optimization_enabled: false,
client_session_id: None,
execution_id: Some(format!("live-postgres-query-export-timeout-{suffix}")),
date_time_format: None,
};
let rows_exported = AtomicU64::new(0);
let done_seen = AtomicBool::new(false);
let started_at = Instant::now();
let result = export_query_result_core(&state, &request, None, |progress| {
rows_exported.store(progress.rows_exported, Ordering::Relaxed);
if matches!(progress.status, ExportStatus::Done) {
done_seen.store(true, Ordering::Relaxed);
}
})
.await;
let elapsed = started_at.elapsed();
let file_len = std::fs::metadata(&file_path).map(|metadata| metadata.len()).unwrap_or(0);
let _ = std::fs::remove_dir_all(dir);
result.expect("stream 650,000 PostgreSQL rows to XLSX");
assert!(elapsed > Duration::from_secs(1), "export should outlive configured timeout: {elapsed:?}");
assert_eq!(rows_exported.load(Ordering::Relaxed), 650_000);
assert!(done_seen.load(Ordering::Relaxed));
assert!(file_len > 1_000_000);
}
#[tokio::test]
#[ignore = "requires DBX_LIVE_POSTGRES_* env vars"]
async fn live_postgres_stream_still_times_out_without_progress_and_recovers() {
let host = std::env::var("DBX_LIVE_POSTGRES_HOST").unwrap_or_else(|_| "127.0.0.1".to_string());
let port = std::env::var("DBX_LIVE_POSTGRES_PORT").ok().and_then(|value| value.parse().ok()).unwrap_or(5432);
let user = std::env::var("DBX_LIVE_POSTGRES_USER").unwrap_or_else(|_| "postgres".to_string());
let password = std::env::var("DBX_LIVE_POSTGRES_PASSWORD").unwrap_or_default();
let database = std::env::var("DBX_LIVE_POSTGRES_DATABASE").unwrap_or_else(|_| "postgres".to_string());
let suffix = uuid::Uuid::new_v4().simple().to_string();
let connection_id = format!("live-postgres-query-export-stall-{suffix}");
let config = live_postgres_config(&connection_id, &host, port, &user, &password, &database);
let dir = std::env::temp_dir().join(format!("dbx-live-postgres-query-export-stall-{suffix}"));
std::fs::create_dir_all(&dir).unwrap();
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
let state = AppState::new(storage);
state.configs.write().await.insert(config.id.clone(), config);
let file_path = dir.join("result.csv");
let sql = "SELECT pg_sleep(5), 1 AS id";
let request = QueryResultExportRequest {
export_id: format!("live-postgres-query-export-stall-{suffix}"),
connection_id: connection_id.clone(),
database: database.clone(),
schema: Some("public".to_string()),
sql: sql.to_string(),
query_base_sql: sql.to_string(),
database_type: DatabaseType::Postgres,
use_agent_cursor: false,
file_path: file_path.to_string_lossy().to_string(),
format: "csv".to_string(),
include_sql_sheet: false,
page_size: 100,
row_limit: None,
total_rows: Some(1),
timeout_secs: Some(1),
keyset_optimization_enabled: false,
client_session_id: None,
execution_id: Some(format!("live-postgres-query-export-stall-{suffix}")),
date_time_format: None,
};
let started_at = Instant::now();
let result = export_query_result_core(&state, &request, None, |_| {}).await;
let elapsed = started_at.elapsed();
let recovery = execute_sql_statement(&state, &connection_id, &database, "SELECT 1 AS id", None, None).await;
let _ = std::fs::remove_dir_all(dir);
assert_eq!(result, Err("Query timed out after 1 seconds".to_string()));
assert!(elapsed < Duration::from_secs(5), "stalled query was not cancelled promptly: {elapsed:?}");
assert_eq!(recovery.expect("PostgreSQL connection should recover after export timeout").rows.len(), 1);
}