From 291890f67e3d3968cc63932b27bebe62fb0ea7b9 Mon Sep 17 00:00:00 2001
From: zipg <geeky@foxmail.com>
Date: Thu, 3 Sep 2026 17:52:15 +0800
Subject: [PATCH] feat(export): respect connection query timeout in table
 export

---
 .../export/ExportProgressDialog.vue           | 10 ++-
 .../__tests__/ExportProgressDialog.spec.ts    | 26 +++++++
 apps/desktop/src/components/grid/DataGrid.vue | 19 ++++-
 .../src/components/layout/ContentArea.vue     |  2 +
 .../src/lib/__tests__/sql/queryError.spec.ts  |  1 +
 apps/desktop/src/lib/sql/queryError.ts        |  2 +-
 crates/dbx-core/src/db/agent_driver.rs        | 70 ++++++++++++++++++-
 crates/dbx-core/src/table_export.rs           | 25 ++++++-
 8 files changed, 148 insertions(+), 7 deletions(-)

diff --git a/apps/desktop/src/components/export/ExportProgressDialog.vue b/apps/desktop/src/components/export/ExportProgressDialog.vue
index cada6aa06c..9f7e6b7900 100644
--- a/apps/desktop/src/components/export/ExportProgressDialog.vue
+++ b/apps/desktop/src/components/export/ExportProgressDialog.vue
@@ -3,12 +3,13 @@ import { computed, onBeforeUnmount, onMounted, ref } from "vue";
 import { useI18n } from "vue-i18n";
 import { Dialog, DialogContent, DialogHeader, DialogTitle, DialogFooter } from "@/components/ui/dialog";
 import { Button } from "@/components/ui/button";
-import { Loader2, CheckCircle2, XCircle, AlertCircle, FolderOpen, Minimize2, X } from "@lucide/vue";
+import { Loader2, CheckCircle2, XCircle, AlertCircle, FolderOpen, Minimize2, Wrench, X } from "@lucide/vue";
 import { useToast } from "@/composables/useToast";
 import { isTauriRuntime } from "@/lib/backend/tauriRuntime";
 import * as api from "@/lib/backend/api";
 import { translateBackendError } from "@/i18n/backend-errors";
 import { formatDataTransferDuration } from "@/composables/useExportTracker";
+import { isQueryTimeoutErrorMessage } from "@/lib/sql/queryError";
 
 const { t } = useI18n();
 const { toast } = useToast();
@@ -38,11 +39,13 @@ const props = defineProps<{
   finishedAt?: number;
   disableCancel?: boolean;
   canMinimize?: boolean;
+  connectionId?: string;
 }>();
 
 const emit = defineEmits<{
   cancel: [];
   minimize: [];
+  changeQueryTimeout: [];
   "update:open": [value: boolean];
 }>();
 
@@ -59,6 +62,7 @@ const displayName = computed(() => {
 });
 const isActive = computed(() => props.status === "Running" || props.status === "Writing");
 const isFinished = computed(() => props.status === "Done" || props.status === "Error" || props.status === "Cancelled");
+const canChangeQueryTimeout = computed(() => props.status === "Error" && !!props.connectionId && !!props.errorMessage && isQueryTimeoutErrorMessage(props.errorMessage));
 const canRevealFile = computed(() => props.status === "Done" && !!props.filePath && isTauriRuntime());
 const progressPercent = computed(() => {
   if (!props.totalRows || props.totalRows <= 0) return 0;
@@ -163,6 +167,10 @@ async function revealExportFile() {
           </Button>
         </template>
         <template v-else-if="isFinished">
+          <Button v-if="canChangeQueryTimeout" variant="outline" size="sm" @click="emit('changeQueryTimeout')">
+            <Wrench class="mr-1 h-3.5 w-3.5" />
+            {{ t("editor.changeQueryTimeout") }}
+          </Button>
           <Button v-if="canRevealFile" size="sm" :disabled="isRevealing" @click="revealExportFile">
             <Loader2 v-if="isRevealing" class="mr-1 h-3.5 w-3.5 animate-spin" />
             <FolderOpen v-else class="mr-1 h-3.5 w-3.5" />
diff --git a/apps/desktop/src/components/export/__tests__/ExportProgressDialog.spec.ts b/apps/desktop/src/components/export/__tests__/ExportProgressDialog.spec.ts
index ff52c6f1be..66c5e2b8ed 100644
--- a/apps/desktop/src/components/export/__tests__/ExportProgressDialog.spec.ts
+++ b/apps/desktop/src/components/export/__tests__/ExportProgressDialog.spec.ts
@@ -80,3 +80,29 @@ describe("ExportProgressDialog display name", () => {
     expect(document.body.textContent).toContain("audit_log (.csv)");
   });
 });
+
+describe("ExportProgressDialog query timeout action", () => {
+  it("offers connection settings for a query timeout", async () => {
+    i18n.global.locale.value = "en";
+    const onChangeQueryTimeout = vi.fn();
+    await mountDialog({
+      ...baseProps,
+      status: "Error",
+      errorMessage: "Agent RPC error (-1): dm.jdbc.driver.DMException: 请求执行超时",
+      connectionId: "dm-1",
+      onChangeQueryTimeout,
+    });
+
+    const button = [...document.body.querySelectorAll("button")].find((element) => element.textContent?.includes("Change query timeout"));
+    expect(button).toBeTruthy();
+    button?.dispatchEvent(new MouseEvent("click", { bubbles: true }));
+    expect(onChangeQueryTimeout).toHaveBeenCalledOnce();
+  });
+
+  it("does not offer timeout settings for unrelated export errors", async () => {
+    i18n.global.locale.value = "en";
+    await mountDialog({ ...baseProps, status: "Error", errorMessage: "Permission denied", connectionId: "dm-1" });
+
+    expect(document.body.textContent).not.toContain("Change query timeout");
+  });
+});
diff --git a/apps/desktop/src/components/grid/DataGrid.vue b/apps/desktop/src/components/grid/DataGrid.vue
index a73bbf36c4..dfbeda648c 100644
--- a/apps/desktop/src/components/grid/DataGrid.vue
+++ b/apps/desktop/src/components/grid/DataGrid.vue
@@ -568,6 +568,7 @@ const emit = defineEmits<{
   "update:whereInput": [value: string];
   "update:orderByInput": [value: string];
   "local-column-filters-change": [value: Record<string, string[]>];
+  changeQueryTimeout: [connectionId: string];
 }>();
 
 const autoRefresh = useDataGridAutoRefresh({
@@ -8155,6 +8156,12 @@ async function cancelActiveExport() {
   await exportCancelHandler.value?.();
 }
 
+function changeExportQueryTimeout() {
+  if (!props.connectionId) return;
+  exportProgressDialog.value = false;
+  emit("changeQueryTimeout", props.connectionId);
+}
+
 const userFacingSql = ref("");
 let userFacingSqlGeneration = 0;
 
@@ -14857,7 +14864,17 @@ function openGridSnapshot() {
     />
     <ImagePreviewDialog v-if="imagePreviewMounted" v-model:open="imagePreviewOpen" :src="imagePreviewSrc" :title="imagePreviewTitle" />
     <component v-if="previewDialogOpen && previewDialogConfig" :is="previewDialogConfig.component" v-model:open="previewDialogOpen" v-bind="previewDialogConfig.props" />
-    <ExportProgressDialog v-if="exportProgressDialogMounted" v-model:open="exportProgressDialog" v-bind="exportProgressState" :disable-cancel="!exportCancelHandler" :can-minimize="exportCanMinimize" @cancel="cancelActiveExport" @minimize="exportProgressDialog = false" />
+    <ExportProgressDialog
+      v-if="exportProgressDialogMounted"
+      v-model:open="exportProgressDialog"
+      v-bind="exportProgressState"
+      :connection-id="connectionId"
+      :disable-cancel="!exportCancelHandler"
+      :can-minimize="exportCanMinimize"
+      @cancel="cancelActiveExport"
+      @minimize="exportProgressDialog = false"
+      @change-query-timeout="changeExportQueryTimeout"
+    />
   </div>
 </template>
 
diff --git a/apps/desktop/src/components/layout/ContentArea.vue b/apps/desktop/src/components/layout/ContentArea.vue
index f4ef53520b..8e8df06995 100644
--- a/apps/desktop/src/components/layout/ContentArea.vue
+++ b/apps/desktop/src/components/layout/ContentArea.vue
@@ -1831,6 +1831,7 @@ defineExpose({
                 @reload="(sql?: string, searchText?: string, whereInput?: string, orderBy?: string, limit?: number, offset?: number, intent?: DataGridReloadIntent) => emit('reload', sql, searchText, whereInput, orderBy, limit, offset, intent)"
                 @paginate="(offset: number, limit: number, whereInput?: string, orderBy?: string) => emit('paginate', offset, limit, whereInput, orderBy)"
                 @sort="(column: string, columnIndex: number, direction: 'asc' | 'desc' | null, whereInput?: string, mode?: DataGridSortMode) => emit('sort', column, columnIndex, direction, whereInput, mode)"
+                @change-query-timeout="(connectionId: string) => emit('openConnectionSettings', connectionId, 'advanced')"
               >
                 <template #result-toolbar-leading="{ compact }">
                   <QueryResultViewSwitcher
@@ -2204,6 +2205,7 @@ defineExpose({
           @reload="(sql?: string, searchText?: string, whereInput?: string, orderBy?: string, limit?: number, offset?: number, intent?: DataGridReloadIntent) => emit('reload', sql, searchText, whereInput, orderBy, limit, offset, intent)"
           @paginate="(offset: number, limit: number, whereInput?: string, orderBy?: string) => emit('paginate', offset, limit, whereInput, orderBy)"
           @sort="(column: string, columnIndex: number, direction: 'asc' | 'desc' | null, whereInput?: string, mode?: DataGridSortMode) => emit('sort', column, columnIndex, direction, whereInput, mode)"
+          @change-query-timeout="(connectionId: string) => emit('openConnectionSettings', connectionId, 'advanced')"
         >
           <template v-if="activeTab.result && isQueryExecutionErrorResult(activeTab.result)" #error-actions="{ errorMessage }">
             <QueryErrorActions
diff --git a/apps/desktop/src/lib/__tests__/sql/queryError.spec.ts b/apps/desktop/src/lib/__tests__/sql/queryError.spec.ts
index 46bc3a856f..956f22eb9a 100644
--- a/apps/desktop/src/lib/__tests__/sql/queryError.spec.ts
+++ b/apps/desktop/src/lib/__tests__/sql/queryError.spec.ts
@@ -44,6 +44,7 @@ describe("isQueryTimeoutErrorMessage", () => {
     expect(isQueryTimeoutErrorMessage("Query timed out after 30 seconds")).toBe(true);
     expect(isQueryTimeoutErrorMessage("查询超时 (60s)，请检查数据库连接是否正常")).toBe(true);
     expect(isQueryTimeoutErrorMessage("查詢逾時 (60s)，請檢查資料庫連線是否正常")).toBe(true);
+    expect(isQueryTimeoutErrorMessage("Agent RPC error (-1): dm.jdbc.driver.DMException: 请求执行超时")).toBe(true);
   });
 
   it("detects statement timeout messages", () => {
diff --git a/apps/desktop/src/lib/sql/queryError.ts b/apps/desktop/src/lib/sql/queryError.ts
index 0e49ed633c..f50162bec4 100644
--- a/apps/desktop/src/lib/sql/queryError.ts
+++ b/apps/desktop/src/lib/sql/queryError.ts
@@ -22,7 +22,7 @@ export function isQueryTimeoutErrorMessage(message: string, backendError?: Backe
     return QUERY_TIMEOUT_STAGES.has(stage);
   }
   const lower = message.toLowerCase();
-  if (lower.includes("query timed out") || lower.includes("查询超时") || lower.includes("查詢逾時")) return true;
+  if (lower.includes("query timed out") || lower.includes("查询超时") || lower.includes("查詢逾時") || lower.includes("请求执行超时") || lower.includes("請求執行逾時")) return true;
   // Agent RPC client-side timeout (tokio::time::timeout in agent_driver.rs). This is the
   // fallback when JDBC setQueryTimeout never fires (unsupported/unresponsive driver), so the
   // backend already treats it as a query timeout (is_agent_rpc_timeout_error in query.rs) —
diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs
index c37b3164e9..6fed9b44e8 100644
--- a/crates/dbx-core/src/db/agent_driver.rs
+++ b/crates/dbx-core/src/db/agent_driver.rs
@@ -2495,6 +2495,22 @@ impl AgentDriverClient {
         self.call_method(AgentMethod::StartTableRead, serde_json::to_value(params).map_err(|e| e.to_string())?).await
     }
 
+    pub async fn start_table_read_with_timeout_and_cancel<T: DeserializeOwned + Send + 'static>(
+        &mut self,
+        params: AgentTableReadStartParams,
+        timeout_duration: Option<Duration>,
+        cancel_token: Option<CancellationToken>,
+    ) -> Result<T, String> {
+        self.invalidate_cached_query();
+        self.call_method_with_timeout_and_cancel(
+            AgentMethod::StartTableRead,
+            serde_json::to_value(params).map_err(|e| e.to_string())?,
+            timeout_duration,
+            cancel_token,
+        )
+        .await
+    }
+
     pub async fn fetch_table_read_page<T: DeserializeOwned + Send + 'static>(
         &mut self,
         session_id: &str,
@@ -2508,6 +2524,23 @@ impl AgentDriverClient {
         .await
     }
 
+    pub async fn fetch_table_read_page_with_timeout_and_cancel<T: DeserializeOwned + Send + 'static>(
+        &mut self,
+        session_id: &str,
+        page_size: usize,
+        timeout_duration: Option<Duration>,
+        cancel_token: Option<CancellationToken>,
+    ) -> Result<T, String> {
+        self.call_method_with_timeout_and_cancel(
+            AgentMethod::FetchTableReadPage,
+            serde_json::to_value(AgentTableReadPageParams { session_id: session_id.to_string(), page_size })
+                .map_err(|e| e.to_string())?,
+            timeout_duration,
+            cancel_token,
+        )
+        .await
+    }
+
     pub async fn close_table_read_session<T: DeserializeOwned + Send + 'static>(
         &mut self,
         session_id: &str,
@@ -4179,7 +4212,7 @@ def respond(req):
 
     session_id = params.get('agentSessionId', '__legacy__')
     with session_lock(session_id):
-        if method == 'execute_query':
+        if method in ('execute_query', 'start_table_read'):
             sql = params.get('sql', '')
             with state_lock:
                 query_count += 1
@@ -4194,6 +4227,9 @@ def respond(req):
             if sql == 'error':
                 write_response(req, error='synthetic query error')
                 return
+            if method == 'start_table_read':
+                write_response(req, {'rows': [], 'session_id': None, 'has_more': False})
+                return
             write_response(req, {'sql': sql, 'count': current_count})
             return
         write_response(req, {'ok': True})
@@ -4297,6 +4333,34 @@ for line in sys.stdin:
         let _ = std::fs::remove_file(script_path);
     }
 
+    #[tokio::test]
+    async fn table_read_uses_the_supplied_rpc_timeout() {
+        let (runtime, script_path) = spawn_stateful_test_runtime("table-read-timeout-test").await;
+        let mut client = AgentDriverClient::shared_session(runtime.clone(), "table-read-session".to_string());
+        let params = AgentTableReadStartParams {
+            sql: "slow table export".to_string(),
+            database: Some("TEST".to_string()),
+            schema: Some("APP".to_string()),
+            page_size: 100,
+            max_rows: 100,
+            fetch_size: Some(100),
+            timeout_secs: Some(1),
+        };
+
+        let error = client
+            .start_table_read_with_timeout_and_cancel::<serde_json::Value>(
+                params,
+                Some(Duration::from_millis(75)),
+                Some(CancellationToken::new()),
+            )
+            .await
+            .unwrap_err();
+        assert!(error.contains("Agent RPC call timed out"));
+
+        runtime.kill();
+        let _ = std::fs::remove_file(script_path);
+    }
+
     #[tokio::test]
     async fn cached_agent_query_handles_hits_errors_ttl_and_invalidation() {
         let (runtime, script_path) = spawn_stateful_test_runtime("query-cache-test").await;
@@ -5016,7 +5080,11 @@ for line in sys.stdin:
     #[test]
     fn exposes_table_read_protocol_wrappers() {
         let _start_table_read = AgentDriverClient::start_table_read::<serde_json::Value>;
+        let _start_table_read_with_timeout_and_cancel =
+            AgentDriverClient::start_table_read_with_timeout_and_cancel::<serde_json::Value>;
         let _fetch_table_read_page = AgentDriverClient::fetch_table_read_page::<serde_json::Value>;
+        let _fetch_table_read_page_with_timeout_and_cancel =
+            AgentDriverClient::fetch_table_read_page_with_timeout_and_cancel::<serde_json::Value>;
         let _close_table_read_session = AgentDriverClient::close_table_read_session::<serde_json::Value>;
     }
 
diff --git a/crates/dbx-core/src/table_export.rs b/crates/dbx-core/src/table_export.rs
index 842275add2..d5c8d4114a 100644
--- a/crates/dbx-core/src/table_export.rs
+++ b/crates/dbx-core/src/table_export.rs
@@ -23,7 +23,9 @@ use crate::database_export::{
 };
 use crate::db::agent_driver::AgentTableReadStartParams;
 use crate::models::connection::DatabaseType;
-use crate::query::{close_query_session, execute_sql_statement_with_options, QueryExecutionOptions};
+use crate::query::{
+    close_query_session, execute_sql_statement_with_options, query_timeout_duration, QueryExecutionOptions,
+};
 use crate::transfer::{
     count_sql_with_where_and_identifier_quote, execute_read_on_pool, execute_read_on_pool_with_max_rows,
     keyset_pagination_sql_with_identifier_quote, pagination_sql_with_filter_order_and_identifier_quote,
@@ -682,6 +684,7 @@ async fn fetch_table_export_batch(
                 let sql = table_cursor_sql(request, sql_context, query_col_names, column_types, primary_keys);
                 let max_rows = request.row_limit.unwrap_or(i32::MAX as usize);
                 let query_timeout = table_export_query_timeout_secs(state, pool_key).await;
+                let rpc_timeout = query_timeout_duration(Some(query_timeout));
                 let params = AgentTableReadStartParams {
                     sql,
                     database: Some(request.database.clone()),
@@ -697,7 +700,14 @@ async fn fetch_table_export_batch(
                 };
                 let client = client.clone();
                 let mut client = client.lock().await;
-                match client.start_table_read::<QueryResult>(params).await {
+                match client
+                    .start_table_read_with_timeout_and_cancel::<QueryResult>(
+                        params,
+                        rpc_timeout,
+                        Some(cancel_token.clone()),
+                    )
+                    .await
+                {
                     Ok(result) => {
                         *cursor_session = result.session_id.clone().map(TableExportCursorSession::Agent);
                         if result.session_id.is_none() && !result.has_more {
@@ -750,7 +760,16 @@ async fn fetch_table_export_batch(
                 };
                 let client = client.clone();
                 let mut client = client.lock().await;
-                match client.fetch_table_read_page::<QueryResult>(&session_id, active_batch_size).await {
+                let query_timeout = table_export_query_timeout_secs(state, pool_key).await;
+                match client
+                    .fetch_table_read_page_with_timeout_and_cancel::<QueryResult>(
+                        &session_id,
+                        active_batch_size,
+                        query_timeout_duration(Some(query_timeout)),
+                        Some(cancel_token.clone()),
+                    )
+                    .await
+                {
                     Ok(result) => {
                         *cursor_session =
                             result.session_id.clone().or(Some(session_id)).map(TableExportCursorSession::Agent);