diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs index cc75be08e..49db4b22b 100644 --- a/crates/dbx-core/src/query_result_export.rs +++ b/crates/dbx-core/src/query_result_export.rs @@ -1,8 +1,9 @@ use serde::{Deserialize, Serialize}; use std::fs::File; +use std::future::Future; use std::io::{BufWriter, Seek, Write}; use std::sync::{ - atomic::{AtomicBool, Ordering}, + atomic::{AtomicBool, AtomicU64, Ordering}, Arc, }; use std::time::{Duration, Instant}; @@ -192,6 +193,85 @@ 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, @@ -1309,6 +1389,9 @@ 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() { Some(token) => { @@ -1378,15 +1461,20 @@ async fn try_export_sqlserver_query_result_stream( } } } + // Mark only after the row is fully written so local XLSX work never consumes + // the next database inactivity window. + progress_clock_for_stream.mark(); Ok(()) }, ); - match query_export_timeout(request.timeout_secs) { - Some(timeout) => tokio::time::timeout(timeout, stream_future) - .await - .map_err(|_| format!("Query timed out after {} seconds", timeout.as_secs()))??, - None => stream_future.await?, - }; + 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?; drop(client); if rows_exported != last_progress_rows { @@ -1593,4 +1681,101 @@ mod tests { .await; assert!(result.is_err()); } + + #[tokio::test] + async fn sqlserver_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::>(), + Some(Duration::from_millis(20)), + progress_clock, + None, + "query timeout".to_string(), + ) + .await; + + assert_eq!(result, Err("query timeout".to_string())); + } + + #[tokio::test] + async fn sqlserver_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( + async move { + for row in 1..=5 { + tokio::time::sleep(Duration::from_millis(50)).await; + progress_clock_for_stream.mark(); + assert!(row <= 5); + } + Ok::<_, String>(5_u8) + }, + Some(Duration::from_millis(150)), + progress_clock, + None, + "query timeout".to_string(), + ) + .await; + + assert_eq!(result, Ok(5)); + } + + #[tokio::test] + async fn sqlserver_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( + async move { + std::thread::sleep(Duration::from_millis(50)); + progress_clock_for_stream.mark(); + Ok::<_, String>(()) + }, + Some(Duration::from_millis(20)), + progress_clock, + None, + "query timeout".to_string(), + ) + .await; + + assert_eq!(result, Ok(())); + } + + #[tokio::test] + async fn sqlserver_stream_timeout_zero_disables_idle_timeout() { + let progress_clock = Arc::new(StreamProgressClock::new()); + let result = await_stream_with_progress_timeout( + async { + tokio::time::sleep(Duration::from_millis(20)).await; + Ok::<_, String>(()) + }, + None, + progress_clock, + None, + "query timeout".to_string(), + ) + .await; + + assert_eq!(result, Ok(())); + } + + #[tokio::test] + async fn sqlserver_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(); + let task = tokio::spawn(async move { + await_stream_with_progress_timeout( + async { std::future::pending::>().await }, + Some(Duration::from_secs(1)), + progress_clock, + Some(&cancel_token_for_task), + "query timeout".to_string(), + ) + .await + }); + tokio::task::yield_now().await; + cancel_token.cancel(); + + assert_eq!(task.await.unwrap(), Err(QUERY_CANCELED.to_string())); + } } diff --git a/crates/dbx-core/tests/live_sqlserver_query_result_export.rs b/crates/dbx-core/tests/live_sqlserver_query_result_export.rs new file mode 100644 index 000000000..14bc09dcf --- /dev/null +++ b/crates/dbx-core/tests/live_sqlserver_query_result_export.rs @@ -0,0 +1,131 @@ +use dbx_core::connection::{AppState, PoolKind}; +use dbx_core::models::connection::DatabaseType; +use dbx_core::query_result_export::{export_query_result_core, ExportStatus, QueryResultExportRequest}; +use dbx_core::storage::Storage; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +fn live_sqlserver_config(id: &str, database: &str) -> dbx_core::models::connection::ConnectionConfig { + dbx_core::models::connection::ConnectionConfig { + id: id.to_string(), + name: id.to_string(), + db_type: DatabaseType::SqlServer, + driver_profile: None, + driver_label: None, + url_params: None, + agent_java_options: Vec::new(), + host: std::env::var("DBX_LIVE_SQLSERVER_HOST").unwrap_or_else(|_| "127.0.0.1".to_string()), + port: std::env::var("DBX_LIVE_SQLSERVER_PORT").ok().and_then(|value| value.parse().ok()).unwrap_or(1433), + username: std::env::var("DBX_LIVE_SQLSERVER_USER").unwrap_or_else(|_| "sa".to_string()), + password: std::env::var("DBX_LIVE_SQLSERVER_PASSWORD").expect("DBX_LIVE_SQLSERVER_PASSWORD"), + database: Some(database.to_string()), + visible_databases: None, + visible_schemas: None, + attached_databases: Vec::new(), + init_script: None, + color: None, + transport_layers: Vec::new(), + connect_timeout_secs: 10, + query_timeout_secs: 30, + idle_timeout_secs: 60, + keepalive_interval_secs: 0, + ssl: false, + ca_cert_path: String::new(), + client_cert_path: String::new(), + client_key_path: String::new(), + sysdba: false, + oracle_connection_type: None, + connection_string: None, + redis_connection_mode: None, + redis_sentinel_master: String::new(), + redis_sentinel_nodes: String::new(), + redis_sentinel_username: String::new(), + redis_sentinel_password: String::new(), + redis_sentinel_tls: false, + redis_cluster_nodes: String::new(), + redis_key_separator: dbx_core::models::connection::default_redis_key_separator(), + redis_scan_page_size: None, + etcd_endpoints: String::new(), + gbase_server: String::new(), + informix_server: String::new(), + external_config: None, + jdbc_driver_class: None, + jdbc_driver_paths: Vec::new(), + one_time: false, + read_only: false, + is_production: false, + production_databases: vec![], + database_info: None, + } +} + +#[tokio::test] +#[ignore = "requires DBX_LIVE_SQLSERVER_HOST/PORT/USER/PASSWORD pointing at SQL Server"] +async fn live_sqlserver_xlsx_export_can_outlive_query_timeout_while_rows_keep_arriving() { + let database = std::env::var("DBX_LIVE_SQLSERVER_DATABASE").unwrap_or_else(|_| "tempdb".to_string()); + let host = std::env::var("DBX_LIVE_SQLSERVER_HOST").unwrap_or_else(|_| "127.0.0.1".to_string()); + let port = std::env::var("DBX_LIVE_SQLSERVER_PORT").ok().and_then(|value| value.parse().ok()).unwrap_or(1433); + let user = std::env::var("DBX_LIVE_SQLSERVER_USER").unwrap_or_else(|_| "sa".to_string()); + let password = std::env::var("DBX_LIVE_SQLSERVER_PASSWORD").expect("DBX_LIVE_SQLSERVER_PASSWORD"); + let client = + dbx_core::db::sqlserver::connect(&host, port, &user, &password, Some(&database), None, Duration::from_secs(10)) + .await + .expect("connect SQL Server"); + + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let dir = std::env::temp_dir().join(format!("dbx-live-sqlserver-xlsx-{suffix}")); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + let connection_id = "live-sqlserver-xlsx-export"; + let pool_key = format!("{connection_id}:{database}"); + state.configs.write().await.insert(connection_id.to_string(), live_sqlserver_config(connection_id, &database)); + state.connections.write().await.insert(pool_key, PoolKind::SqlServer(Arc::new(tokio::sync::Mutex::new(client)))); + + let file_path = dir.join("result.xlsx"); + let sql = "WITH numbers AS (\ + SELECT TOP (130000) ROW_NUMBER() OVER (ORDER BY (SELECT NULL)) AS id \ + FROM sys.all_objects AS first_source CROSS JOIN sys.all_objects AS second_source\ + ) SELECT id, REPLICATE(N'x', 64) AS payload FROM numbers ORDER BY id"; + let request = QueryResultExportRequest { + export_id: format!("live-sqlserver-xlsx-{suffix}"), + connection_id: connection_id.to_string(), + database: database.clone(), + schema: Some("dbo".to_string()), + sql: sql.to_string(), + query_base_sql: sql.to_string(), + database_type: DatabaseType::SqlServer, + use_agent_cursor: false, + file_path: file_path.to_string_lossy().to_string(), + format: "xlsx".to_string(), + include_sql_sheet: false, + page_size: 5000, + row_limit: Some(200_000), + total_rows: Some(130_000), + timeout_secs: Some(1), + keyset_optimization_enabled: true, + client_session_id: None, + execution_id: Some(format!("live-sqlserver-xlsx-{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(); + + result.expect("stream 130,000 rows to XLSX"); + assert!(elapsed > Duration::from_secs(1), "export should outlive configured timeout: {elapsed:?}"); + assert_eq!(rows_exported.load(Ordering::Relaxed), 130_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); +}