fix(export): use inactivity timeout for SQL Server streams

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<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,
@@ -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::<Result<(), String>>(),
+            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::<Result<(), String>>().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()));
+    }
 }

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);
+}