From 754c8b4a53b98e3945902be4ec19f7c8d8d90de5 Mon Sep 17 00:00:00 2001 From: zipg Date: Tue, 21 Jul 2026 21:08:13 +0800 Subject: [PATCH] fix(export): use inactivity timeout for result export --- crates/dbx-core/src/db/postgres.rs | 28 +++- crates/dbx-core/src/query.rs | 83 +++++++++- crates/dbx-core/src/query_result_export.rs | 145 +++++------------- crates/dbx-core/tests/live_mysql57.rs | 82 +++++++++- .../live_postgres_query_result_export.rs | 117 +++++++++++++- 5 files changed, 336 insertions(+), 119 deletions(-) diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 135620e23..b119e0ff9 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -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; diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index a73ae7822..9cea00269 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -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( + stream_future: F, + timeout: Option, + progress_clock: Arc, + cancel_token: Option<&CancellationToken>, + timeout_message: String, +) -> Result +where + F: Future>, +{ + 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() diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs index d8ec5812b..0cbb8c931 100644 --- a/crates/dbx-core/src/query_result_export.rs +++ b/crates/dbx-core/src/query_result_export.rs @@ -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) -> Option { } } -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( - stream_future: F, - timeout: Option, - progress_clock: Arc, - cancel_token: Option<&CancellationToken>, - timeout_message: String, -) -> Result -where - F: Future>, -{ - 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::>(), @@ -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(); diff --git a/crates/dbx-core/tests/live_mysql57.rs b/crates/dbx-core/tests/live_mysql57.rs index 68f060851..7a81162f1 100644 --- a/crates/dbx-core/tests/live_mysql57.rs +++ b/crates/dbx-core/tests/live_mysql57.rs @@ -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::().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::>().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() { diff --git a/crates/dbx-core/tests/live_postgres_query_result_export.rs b/crates/dbx-core/tests/live_postgres_query_result_export.rs index dda513b8d..f9cfc1175 100644 --- a/crates/dbx-core/tests/live_postgres_query_result_export.rs +++ b/crates/dbx-core/tests/live_postgres_query_result_export.rs @@ -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); +}