From fa76c742c8f3fa0a8274e98b0ca0619038e0789f Mon Sep 17 00:00:00 2001
From: auth0 <zorokakao@gmail.com>
Date: Sun, 12 Jul 2026 22:59:27 +0800
Subject: [PATCH] fix(mysql): disconnect failed export streams

---
 crates/dbx-core/src/query_result_export.rs | 54 ++++++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs
index 8a0f742a4..799f2a61f 100644
--- a/crates/dbx-core/src/query_result_export.rs
+++ b/crates/dbx-core/src/query_result_export.rs
@@ -35,6 +35,18 @@ const STREAMING_PAGINATION_UNSUPPORTED_ERROR: &str = "当前查询暂不支持
 const AGENT_SESSION_MISSING_ERROR: &str = "查询结果流式导出需要驱动返回结果集会话，但当前驱动未返回 session_id。";
 const STREAM_PROGRESS_TIME_INTERVAL: Duration = Duration::from_secs(1);
 
+async fn disconnect_with_timeout<C, F, Fut>(
+    connection: C,
+    cleanup_timeout: Duration,
+    disconnect: F,
+) -> Result<Result<(), String>, tokio::time::error::Elapsed>
+where
+    F: FnOnce(C) -> Fut,
+    Fut: std::future::Future<Output = Result<(), String>>,
+{
+    tokio::time::timeout(cleanup_timeout, disconnect(connection)).await
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct QueryResultExportRequest {
@@ -941,6 +953,25 @@ async fn try_export_mysql_query_result_stream(
     watcher_done.cancel();
 
     if let Err(error) = stream_result {
+        // A timed-out, cancelled, or failed MySQL result stream may leave an
+        // incomplete protocol packet on the connection. Explicitly disconnect
+        // it so mysql_async cannot recycle the poisoned connection into the pool.
+        match disconnect_with_timeout(conn, operation_budget.cleanup_timeout, |conn| async move {
+            conn.disconnect().await.map_err(|error| error.to_string())
+        })
+        .await
+        {
+            Ok(Ok(())) => {}
+            Ok(Err(disconnect_error)) => {
+                log::warn!(
+                    "Failed to disconnect MySQL export connection {mysql_connection_id} after stream error: {disconnect_error}"
+                );
+            }
+            Err(_) => {
+                log::warn!("Timed out disconnecting MySQL export connection {mysql_connection_id} after stream error");
+            }
+        }
+
         if error == QUERY_CANCELED
             || export_cancelled.load(Ordering::SeqCst)
             || cancel_token.as_ref().is_some_and(|token| token.is_cancelled())
@@ -1411,4 +1442,27 @@ mod tests {
         assert!(safe_keyset_candidate("SELECT * FROM users WHERE active = true").is_none());
         assert!(safe_keyset_candidate("SELECT id, name FROM users").is_none());
     }
+
+    #[tokio::test]
+    async fn failed_mysql_stream_disconnects_connection_without_database_or_xlsx() {
+        let disconnected = Arc::new(AtomicBool::new(false));
+        let disconnected_for_call = disconnected.clone();
+        let result = disconnect_with_timeout((), Duration::from_secs(1), move |_| async move {
+            disconnected_for_call.store(true, Ordering::SeqCst);
+            Ok(())
+        })
+        .await;
+        assert!(matches!(result, Ok(Ok(()))));
+        assert!(disconnected.load(Ordering::SeqCst));
+    }
+
+    #[tokio::test]
+    async fn failed_mysql_stream_disconnect_is_bounded_by_cleanup_timeout() {
+        let result = disconnect_with_timeout((), Duration::from_millis(1), |_| async {
+            tokio::time::sleep(Duration::from_secs(30)).await;
+            Ok(())
+        })
+        .await;
+        assert!(result.is_err());
+    }
 }