From d1286047439c61aa81b44700e4d214f382d60d97 Mon Sep 17 00:00:00 2001 From: miracle Date: Fri, 26 Jun 2026 21:41:56 +0800 Subject: [PATCH] =?UTF-8?q?fix(connection):=20=E4=BF=AE=E5=A4=8D=E6=95=B0?= =?UTF-8?q?=E6=8D=AE=E5=BA=93=E8=BF=9E=E6=8E=A5=E8=B6=85=E6=97=B6=E5=90=8E?= =?UTF-8?q?=E6=97=A0=E6=B3=95=E5=8F=AF=E9=9D=A0=E6=81=A2=E5=A4=8D=E7=9A=84?= =?UTF-8?q?=E9=97=AE=E9=A2=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(core): implement PIP-0001 database connection timeout recovery - Add DbOperationBudget struct for unified execution budget model - Add deadpool create/recycle timeout for PostgreSQL pools - Add checkout_postgres_client helper replacing 14 bare pool.get() calls with timeout + cancel token support - Fix PostgreSQL TLS cancel: rebuild TLS connector from stored context, fallback to NoTls on failure - Add timeout to MySQL kill_query and kill_query_with_opts (5s) - Extend is_connection_error to match deadpool pool timeout messages - Unify keepalive default to 30s (was 60s Rust / 0s frontend mismatch) - Inject PostgreSQL TCP keepalive params (idle=30s, interval=10s, retries=3) - Inject MySQL tcp_keepalive=30s in OptsBuilder - Make MySQL inactive_connection_ttl configurable from idle_timeout_secs - Add stage-based structured logging for pool.checkout/health.check/cancel - Store PostgresCancelContext in AppState for TLS cancel reconstruction - Clean up cancel contexts on pool removal/drain - Add unit tests for DbOperationBudget and pool timeout error detection * fix(core): align timeout recovery with PIP-0001 * fix(core): handle timeout recovery review issues * fix(core): complete connection timeout recovery * fix(core): fall back postgres cancel connector build * refactor: revert comment translations to English in connection.rs and query.rs - Revert DbOperationBudget doc comments to English (query.rs) - Revert PostgresCancelContext and TLS cancel comments to English (connection.rs) - Keep Chinese error-matching strings in is_connection_error unchanged --------- Co-authored-by: onceMisery --- .../connection/ConnectionDialog.vue | 6 +- .../__tests__/connectionStore.timeout.spec.ts | 120 ++++++ .../__tests__/queryStore.cancel.spec.ts | 51 +++ apps/desktop/src/stores/connectionStore.ts | 26 +- apps/desktop/src/stores/queryStore.ts | 20 +- crates/dbx-core/src/connection.rs | 65 ++- crates/dbx-core/src/db/mysql.rs | 232 +++++++++-- crates/dbx-core/src/db/postgres.rs | 388 +++++++++++++---- crates/dbx-core/src/models/connection.rs | 2 +- crates/dbx-core/src/query.rs | 391 ++++++++++++++++-- crates/dbx-core/src/schema.rs | 2 +- 11 files changed, 1134 insertions(+), 169 deletions(-) create mode 100644 apps/desktop/src/stores/__tests__/connectionStore.timeout.spec.ts create mode 100644 apps/desktop/src/stores/__tests__/queryStore.cancel.spec.ts diff --git a/apps/desktop/src/components/connection/ConnectionDialog.vue b/apps/desktop/src/components/connection/ConnectionDialog.vue index 8e8d7d54b2..46e5c53cb2 100644 --- a/apps/desktop/src/components/connection/ConnectionDialog.vue +++ b/apps/desktop/src/components/connection/ConnectionDialog.vue @@ -135,7 +135,7 @@ const defaultForm = (): ConnectionForm => ({ connect_timeout_secs: 10, query_timeout_secs: 30, idle_timeout_secs: 60, - keepalive_interval_secs: 0, + keepalive_interval_secs: 30, ssl: false, ca_cert_path: "", client_cert_path: "", @@ -908,7 +908,7 @@ watch( connect_timeout_secs: config.connect_timeout_secs || 10, query_timeout_secs: config.query_timeout_secs ?? 30, idle_timeout_secs: config.idle_timeout_secs ?? 60, - keepalive_interval_secs: config.keepalive_interval_secs ?? 0, + keepalive_interval_secs: config.keepalive_interval_secs ?? 30, ssl: config.ssl || false, ca_cert_path: config.ca_cert_path || "", client_cert_path: config.client_cert_path || "", @@ -1523,7 +1523,7 @@ function connectionConfigForSubmit(id: string): ConnectionConfig { const idleTimeout = Number(config.idle_timeout_secs); config.idle_timeout_secs = Number.isFinite(idleTimeout) && idleTimeout >= 0 ? idleTimeout : 60; const keepaliveInterval = Number(config.keepalive_interval_secs); - config.keepalive_interval_secs = Number.isFinite(keepaliveInterval) && keepaliveInterval >= 0 ? keepaliveInterval : 0; + config.keepalive_interval_secs = Number.isFinite(keepaliveInterval) && keepaliveInterval >= 0 ? keepaliveInterval : 30; if (config.db_type === "manticoresearch") { config.url_params = ""; } diff --git a/apps/desktop/src/stores/__tests__/connectionStore.timeout.spec.ts b/apps/desktop/src/stores/__tests__/connectionStore.timeout.spec.ts new file mode 100644 index 0000000000..20a8f62493 --- /dev/null +++ b/apps/desktop/src/stores/__tests__/connectionStore.timeout.spec.ts @@ -0,0 +1,120 @@ +import { createPinia, setActivePinia } from "pinia"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import type { ConnectionConfig, TreeNode } from "@/types/database"; + +function installLocalStorage() { + const data = new Map(); + vi.stubGlobal("localStorage", { + getItem: vi.fn((key: string) => data.get(key) ?? null), + setItem: vi.fn((key: string, value: string) => data.set(key, value)), + removeItem: vi.fn((key: string) => data.delete(key)), + }); +} + +function postgresConnection(overrides: Partial = {}): ConnectionConfig { + return { + id: "pg-1", + name: "Postgres", + db_type: "postgres", + host: "127.0.0.1", + port: 5432, + username: "postgres", + password: "", + database: "app", + read_only: false, + ...overrides, + } as ConnectionConfig; +} + +describe("connectionStore timeout recovery", () => { + beforeEach(() => { + vi.useFakeTimers(); + vi.resetModules(); + vi.unstubAllGlobals(); + installLocalStorage(); + setActivePinia(createPinia()); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it("times out connected health checks and falls back to reconnect", async () => { + const checkConnectionHealth = vi.fn(() => new Promise(() => undefined)); + const connectDb = vi.fn().mockResolvedValue(undefined); + + vi.doMock("@/lib/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/api", () => ({ + checkConnectionHealth, + connectDb, + deleteSchemaCachePrefix: vi.fn().mockResolvedValue(undefined), + saveConnections: vi.fn().mockResolvedValue(undefined), + saveSidebarLayout: vi.fn().mockResolvedValue(undefined), + })); + + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + const connection = postgresConnection({ connect_timeout_secs: 10 }); + store.connections = [connection]; + store.connectedIds.add(connection.id); + + const ensure = store.ensureConnected(connection.id); + await vi.advanceTimersByTimeAsync(5001); + await ensure; + + expect(checkConnectionHealth).toHaveBeenCalledWith(connection.id); + expect(connectDb).toHaveBeenCalledWith(connection); + expect(store.connectedIds.has(connection.id)).toBe(true); + }, 10_000); + + it("normalizes missing keepalive interval to 30 seconds", async () => { + vi.doMock("@/lib/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/api", () => ({ + deleteSchemaCachePrefix: vi.fn().mockResolvedValue(undefined), + saveConnections: vi.fn().mockResolvedValue(undefined), + saveSidebarLayout: vi.fn().mockResolvedValue(undefined), + })); + + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + await store.addConnection(postgresConnection({ keepalive_interval_secs: undefined })); + + expect(store.connections[0]?.keepalive_interval_secs).toBe(30); + }); + + it("clears connection node loading when health check timeout forces reconnect failure", async () => { + const checkConnectionHealth = vi.fn(() => new Promise(() => undefined)); + const connectDb = vi.fn().mockRejectedValue(new Error("reconnect failed")); + + vi.doMock("@/lib/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/api", () => ({ + checkConnectionHealth, + connectDb, + deleteSchemaCachePrefix: vi.fn().mockResolvedValue(undefined), + saveConnections: vi.fn().mockResolvedValue(undefined), + saveSidebarLayout: vi.fn().mockResolvedValue(undefined), + })); + + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + const connection = postgresConnection({ connect_timeout_secs: 10 }); + const node: TreeNode = { + id: connection.id, + label: connection.name, + type: "connection", + connectionId: connection.id, + isLoading: true, + children: [], + }; + store.connections = [connection]; + store.connectedIds.add(connection.id); + store.treeNodes = [node]; + + const ensure = store.ensureConnected(connection.id).catch((error) => error); + await vi.advanceTimersByTimeAsync(5001); + const error = await ensure; + + expect(error).toBeInstanceOf(Error); + expect(node.isLoading).toBe(false); + }, 10_000); +}); diff --git a/apps/desktop/src/stores/__tests__/queryStore.cancel.spec.ts b/apps/desktop/src/stores/__tests__/queryStore.cancel.spec.ts new file mode 100644 index 0000000000..08b041c0b2 --- /dev/null +++ b/apps/desktop/src/stores/__tests__/queryStore.cancel.spec.ts @@ -0,0 +1,51 @@ +import { createPinia, setActivePinia } from "pinia"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +function installLocalStorage() { + const data = new Map(); + vi.stubGlobal("localStorage", { + getItem: vi.fn((key: string) => data.get(key) ?? null), + setItem: vi.fn((key: string, value: string) => data.set(key, value)), + removeItem: vi.fn((key: string) => data.delete(key)), + }); +} + +describe("queryStore cancel timeout recovery", () => { + beforeEach(() => { + vi.useFakeTimers(); + vi.resetModules(); + vi.unstubAllGlobals(); + installLocalStorage(); + setActivePinia(createPinia()); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + it("clears cancelling state when cancelQuery does not return", async () => { + const cancelQuery = vi.fn(() => new Promise(() => undefined)); + + vi.doMock("@/lib/api", () => ({ cancelQuery })); + + const { useQueryStore } = await import("@/stores/queryStore"); + const store = useQueryStore(); + const tabId = store.createTab("pg-1", "app", "query_1"); + store.setExecutingWithId(tabId, "exec-1"); + + const cancel = store.cancelTabExecution(tabId); + expect(store.tabs[0]?.isCancelling).toBe(true); + + await vi.advanceTimersByTimeAsync(10_001); + const result = await cancel; + + expect(result).toBe(false); + expect(cancelQuery).toHaveBeenCalledWith("exec-1"); + expect(store.tabs[0]).toMatchObject({ + isExecuting: false, + isCancelling: false, + executionId: undefined, + }); + expect(store.tabs[0]?.result?.rows[0]?.[0]).toContain("Cancel request timed out"); + }, 15_000); +}); diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index c31b6969f5..2414cc1fe4 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -59,9 +59,11 @@ import { REDIS_SCAN_PAGE_SIZE_DEFAULT } from "@/lib/redisKeyPattern"; const PINNED_TREE_NODES_STORAGE_KEY = "dbx-pinned-tree-nodes"; const ACTIVE_CONNECTION_STORAGE_KEY = "dbx-active-connection"; const CONNECTION_HEALTH_CHECK_TTL_MS = 2000; +const CONNECTION_HEALTH_CHECK_TIMEOUT_MS = 5000; const METADATA_LOAD_MIN_TIMEOUT_MS = 15_000; const METADATA_LOAD_DISABLED_QUERY_TIMEOUT_MS = 60_000; const DISCONNECT_REQUEST_TIMEOUT_MS = 5_000; +const DEFAULT_KEEPALIVE_INTERVAL_SECS = 30; const MONGO_LEGACY_DRIVER_PROFILE = "mongodb-legacy"; const MONGO_LEGACY_DRIVER_LABEL = "MongoDB (Legacy)"; const SUPERSEDED_CONNECTION_ATTEMPT_MESSAGE = "Connection attempt was superseded by a newer attempt"; @@ -304,6 +306,25 @@ export const useConnectionStore = defineStore("connection", () => { return Math.max(METADATA_LOAD_MIN_TIMEOUT_MS, boundedTimeoutSecs * 1000); } + async function withConnectionHealthTimeout(connectionId: string, promise: Promise): Promise { + let timer: ReturnType | undefined; + try { + return await Promise.race([ + promise, + new Promise((_, reject) => { + timer = setTimeout(() => { + reject(new Error(`Connection health check timed out after ${Math.ceil(CONNECTION_HEALTH_CHECK_TIMEOUT_MS / 1000)}s.`)); + }, CONNECTION_HEALTH_CHECK_TIMEOUT_MS); + }), + ]); + } catch (error) { + clearConnectionNodeLoading(connectionId); + throw error; + } finally { + if (timer) clearTimeout(timer); + } + } + async function withMetadataLoadTimeout(connectionId: string, promise: Promise, label: string): Promise { const timeoutMs = metadataLoadTimeoutMs(getConfig(connectionId)); let timer: ReturnType | undefined; @@ -486,7 +507,7 @@ export const useConnectionStore = defineStore("connection", () => { connect_timeout_secs: config.connect_timeout_secs || 10, query_timeout_secs: config.query_timeout_secs ?? 30, idle_timeout_secs: config.idle_timeout_secs ?? 60, - keepalive_interval_secs: config.keepalive_interval_secs ?? 0, + keepalive_interval_secs: config.keepalive_interval_secs ?? DEFAULT_KEEPALIVE_INTERVAL_SECS, }; } @@ -1230,7 +1251,7 @@ export const useConnectionStore = defineStore("connection", () => { if (hasRecentConnectionHealthCheck(connectionId)) return; // Optimistic: verify backend pool is actually healthy try { - await api.checkConnectionHealth(connectionId); + await withConnectionHealthTimeout(connectionId, api.checkConnectionHealth(connectionId)); markConnectionHealthChecked(connectionId); return; } catch { @@ -1273,6 +1294,7 @@ export const useConnectionStore = defineStore("connection", () => { return; } recordConnectionError(connectionId, e); + clearConnectionNodeLoading(connectionId); throw e; } finally { if (connectInFlight.get(connectionId) === connectPromise) { diff --git a/apps/desktop/src/stores/queryStore.ts b/apps/desktop/src/stores/queryStore.ts index 832a7ba0dd..342d458c7d 100644 --- a/apps/desktop/src/stores/queryStore.ts +++ b/apps/desktop/src/stores/queryStore.ts @@ -49,6 +49,7 @@ const STORAGE_KEY = "dbx-open-tabs"; const ACTIVE_TAB_KEY = "dbx-active-tab"; const ORACLE_LIKE_METADATA_TYPES = new Set(["oracle", "dameng", "oceanbase-oracle"]); const BACKGROUND_CLIENT_SESSION_SUFFIXES = ["count", "explain", "export"] as const; +const CANCEL_QUERY_TIMEOUT_MS = 10_000; interface BuildQueryResultExportRequestOptions { exportId: string; @@ -98,6 +99,20 @@ async function withFrontendQueryTimeout(promise: Promise, timeoutSecs: num } } +async function withCancelQueryTimeout(promise: Promise): Promise { + let timer: ReturnType | undefined; + try { + return await Promise.race([ + promise, + new Promise((_, reject) => { + timer = setTimeout(() => reject(new Error("Cancel request timed out after 10s.")), CANCEL_QUERY_TIMEOUT_MS); + }), + ]); + } finally { + if (timer) clearTimeout(timer); + } +} + function normalizeOracleLikeMetadataIdentifier(dbType: string, identifier: string | undefined, quoted?: boolean) { if (!identifier || quoted || !ORACLE_LIKE_METADATA_TYPES.has(dbType)) return identifier; return identifier.toUpperCase(); @@ -1946,7 +1961,7 @@ export const useQueryStore = defineStore("query", () => { if (!executionId) return false; tab.isCancelling = true; try { - const canceled = await api.cancelQuery(executionId); + const canceled = await withCancelQueryTimeout(api.cancelQuery(executionId)); if (!canceled) { const current = tabs.value.find((t) => t.id === id); if (current && current.executionId === executionId) { @@ -1962,7 +1977,10 @@ export const useQueryStore = defineStore("query", () => { if (tab) useConnectionStore().recordConnectionLostError(tab.connectionId, e); const current = tabs.value.find((t) => t.id === id); if (current && current.executionId === executionId) { + current.isExecuting = false; current.isCancelling = false; + current.executionId = undefined; + current.queryExecutionStartedAt = undefined; current.result = toErrorResult(e); } return false; diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index f59b4da5c1..42c0ee0656 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -33,7 +33,7 @@ pub const JDBC_PLUGIN_NOT_INSTALLED: &str = pub const PRESTOSQL_JDBC_DRIVER_CLASS: &str = "io.prestosql.jdbc.PrestoDriver"; const DEFAULT_AGENT_CONNECT_TIMEOUT_SECS: u64 = 30; const ACCESS_AGENT_CONNECT_TIMEOUT_SECS: u64 = 30; -const POOL_CLOSE_TIMEOUT_SECS: u64 = 5; +const POOL_CLOSE_TIMEOUT_SECS: u64 = 3; #[cfg(feature = "duckdb-bundled")] mod duckdb_types { @@ -136,6 +136,9 @@ pub struct AppState { pub plugins: PluginRegistry, pub agent_manager: crate::agent_manager::AgentManager, pub nacos_registry: crate::nacos::NacosAdminRegistry, + /// PostgreSQL TLS cancel context, keyed by pool_key. + /// Used to reconstruct a TLS connector compatible with the original connection when cancelling. + postgres_cancel_contexts: Arc>>, #[cfg(feature = "mq-admin")] pub mq_registry: crate::mq::MqAdminRegistry, } @@ -214,6 +217,7 @@ pub async fn connect_mysql_metadata_pool( max_connections: usize, ) -> Result<(db::mysql::MySqlPool, MysqlMode), String> { let url = connection_url_for_endpoint(db_config, host, port); + let idle_timeout_secs = Some(db_config.idle_timeout_secs); if db_config.needs_bare_mysql() { return match db::mysql::connect_bare_with_pool_limit(&url, connect_timeout, max_connections).await { Ok(pool) => Ok((pool, MysqlMode::Bare)), @@ -233,11 +237,12 @@ pub async fn connect_mysql_metadata_pool( }; } - match db::mysql::connect_with_ca_cert_and_pool_limit( + match db::mysql::connect_with_ca_cert_pool_limit_and_idle( &url, Some(&db_config.ca_cert_path), connect_timeout, max_connections, + idle_timeout_secs, ) .await { @@ -251,11 +256,12 @@ pub async fn connect_mysql_metadata_pool( log::info!( "MySQL metadata connection without a default database failed ({err}); retrying with configured default database." ); - let pool = db::mysql::connect_with_ca_cert_and_pool_limit( + let pool = db::mysql::connect_with_ca_cert_pool_limit_and_idle( &fallback_url, Some(&config.ca_cert_path), connect_timeout, max_connections, + idle_timeout_secs, ) .await?; let mode = detect_ob_oracle_mode(config, &pool).await; @@ -365,6 +371,7 @@ impl AppState { app_version, ), nacos_registry: crate::nacos::NacosAdminRegistry::new(), + postgres_cancel_contexts: Arc::new(RwLock::new(HashMap::new())), #[cfg(feature = "mq-admin")] mq_registry: crate::mq::MqAdminRegistry::new(), } @@ -420,9 +427,10 @@ impl AppState { self.stop_keepalive_task(&pool_key).await; self.pool_activity.write().await.insert(pool_key.clone(), PoolActivity::now()); self.start_keepalive_task(&pool_key, &pool, config).await; + let previous_key = pool_key.clone(); let previous = self.connections.write().await.insert(pool_key, pool); if let Some(pool) = previous { - close_pool_kind(pool).await; + close_pool_kind_with_timeout(previous_key, pool).await; } } @@ -461,7 +469,7 @@ impl AppState { config: &ConnectionConfig, ) -> Result<(), String> { if let Err(err) = self.ensure_current_connection_attempt(connection_id, Some(attempt)).await { - close_pool_kind(pool).await; + close_pool_kind_with_timeout(pool_key, pool).await; return Err(err); } self.insert_connection_pool(pool_key, pool, config).await; @@ -495,6 +503,7 @@ impl AppState { let connections = self.connections.clone(); let keepalive_tasks = self.keepalive_tasks.clone(); let pool_activity = self.pool_activity.clone(); + let cancel_contexts = self.postgres_cancel_contexts.clone(); let running_queries = self.running_queries.clone(); let idle_timeout = Duration::from_secs(idle_timeout_secs.max(1)); let handle = tokio::spawn(async move { @@ -517,6 +526,7 @@ impl AppState { ); keepalive_tasks.write().await.remove(&key); pool_activity.write().await.remove(&key); + cancel_contexts.write().await.remove(&key); let removed = connections.write().await.remove(&key); if let Some(pool) = removed { close_pool_kind_with_timeout(key, pool).await; @@ -533,6 +543,7 @@ impl AppState { log::warn!("Connection keepalive failed for '{key}': {err}; invalidating pool"); keepalive_tasks.write().await.remove(&key); pool_activity.write().await.remove(&key); + cancel_contexts.write().await.remove(&key); let removed = connections.write().await.remove(&key); if let Some(pool) = removed { close_pool_kind_with_timeout(key, pool).await; @@ -546,6 +557,7 @@ impl AppState { ); keepalive_tasks.write().await.remove(&key); pool_activity.write().await.remove(&key); + cancel_contexts.write().await.remove(&key); let removed = connections.write().await.remove(&key); if let Some(pool) = removed { close_pool_kind_with_timeout(key, pool).await; @@ -582,6 +594,11 @@ impl AppState { self.pool_activity.write().await.insert(pool_key.to_string(), PoolActivity::now()); } + /// Get the PostgreSQL TLS cancel context (used to reconstruct the TLS connector when cancelling a query). + pub async fn get_postgres_cancel_context(&self, pool_key: &str) -> Option { + self.postgres_cancel_contexts.read().await.get(pool_key).cloned() + } + pub fn pool_activity_touch(&self, pool_key: &str) -> PoolActivityTouch { PoolActivityTouch { pool_key: pool_key.to_string(), @@ -678,7 +695,14 @@ impl AppState { | DatabaseType::Gaussdb | DatabaseType::Kwdb | DatabaseType::Questdb - | DatabaseType::OpenGauss => PoolKind::Postgres(db::postgres::connect(&url, connect_timeout).await?), + | DatabaseType::OpenGauss => { + let pg_pool = db::postgres::connect(&url, connect_timeout).await?; + // Build TLS cancel context for reconstructing TLS connection during cancel + if let Some(ctx) = db::postgres::build_postgres_cancel_context(&url) { + self.postgres_cancel_contexts.write().await.insert(pool_key.clone(), ctx); + } + PoolKind::Postgres(pg_pool) + } DatabaseType::Sqlite => { let extensions = db::sqlite::sqlite_extension_specs_from_url_params(db_config.url_params.as_deref()) .into_iter() @@ -1005,7 +1029,7 @@ impl AppState { }; if let Err(err) = self.ensure_current_connection_attempt(connection_id, connection_attempt).await { - close_pool_kind(pool).await; + close_pool_kind_with_timeout(pool_key.clone(), pool).await; return Err(err); } self.insert_connection_pool(pool_key.clone(), pool, &db_config).await; @@ -1316,9 +1340,10 @@ impl AppState { self.stop_keepalive_task(pool_key).await; self.pool_activity.write().await.remove(pool_key); + self.postgres_cancel_contexts.write().await.remove(pool_key); let removed = self.connections.write().await.remove(pool_key); if let Some(pool) = removed { - close_pool_kind(pool).await; + close_pool_kind_with_timeout(pool_key.to_string(), pool).await; true } else { false @@ -1347,9 +1372,10 @@ impl AppState { } else { self.stop_keepalive_task(&pool_key).await; self.pool_activity.write().await.remove(&pool_key); + self.postgres_cancel_contexts.write().await.remove(&pool_key); let removed = self.connections.write().await.remove(&pool_key); if let Some(pool) = removed { - close_pool_kind(pool).await; + close_pool_kind_with_timeout(pool_key.clone(), pool).await; } } self.get_or_create_pool_for_session(connection_id, database, client_session_id).await @@ -1376,9 +1402,10 @@ impl AppState { } self.stop_keepalive_task(&pool_key).await; self.pool_activity.write().await.remove(&pool_key); + self.postgres_cancel_contexts.write().await.remove(&pool_key); let removed = self.connections.write().await.remove(&pool_key); if let Some(pool) = removed { - close_pool_kind(pool).await; + close_pool_kind_with_timeout(pool_key, pool).await; Ok(true) } else { Ok(false) @@ -1388,9 +1415,10 @@ impl AppState { pub async fn remove_pool_by_key(&self, pool_key: &str) -> bool { self.stop_keepalive_task(pool_key).await; self.pool_activity.write().await.remove(pool_key); + self.postgres_cancel_contexts.write().await.remove(pool_key); let removed = self.connections.write().await.remove(pool_key); if let Some(pool) = removed { - close_pool_kind(pool).await; + close_pool_kind_with_timeout(pool_key.to_string(), pool).await; true } else { false @@ -1418,21 +1446,23 @@ impl AppState { self.stop_keepalive_tasks(&keys_to_remove).await; { let mut activity = self.pool_activity.write().await; + let mut cancel_contexts = self.postgres_cancel_contexts.write().await; for key in &keys_to_remove { activity.remove(key); + cancel_contexts.remove(key); } } let mut conns = self.connections.write().await; let mut removed = Vec::with_capacity(keys_to_remove.len()); for key in keys_to_remove { if let Some(pool) = conns.remove(&key) { - removed.push(pool); + removed.push((key, pool)); } } drop(conns); let closed = !removed.is_empty(); - for pool in removed { - close_pool_kind(pool).await; + for (key, pool) in removed { + close_pool_kind_with_timeout(key, pool).await; } Ok(closed) } @@ -1720,11 +1750,14 @@ impl AppState { } } let mut conns = self.connections.write().await; + let mut removed = Vec::with_capacity(dead_keys.len()); for key in &dead_keys { if let Some(pool) = conns.remove(key) { - close_pool_kind(pool).await; + removed.push((key.clone(), pool)); } } + drop(conns); + close_removed_pools(removed).await; } // Re-establish SSH tunnels that have died @@ -1766,8 +1799,10 @@ impl AppState { self.stop_keepalive_tasks(&keys_to_remove).await; { let mut activity = self.pool_activity.write().await; + let mut cancel_contexts = self.postgres_cancel_contexts.write().await; for key in &keys_to_remove { activity.remove(key); + cancel_contexts.remove(key); } } let mut conns = self.connections.write().await; diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index dcfb8acc67..bf46b00984 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -10,6 +10,7 @@ use std::path::PathBuf; use std::sync::atomic::AtomicBool; use std::time::Duration; use std::time::Instant; +use tokio_util::sync::CancellationToken; use crate::models::connection::DatabaseType; use crate::sql::starts_with_executable_sql_keyword; @@ -386,16 +387,26 @@ pub async fn connect_with_ca_cert_and_pool_limit( ca_cert_path: Option<&str>, fallback_timeout: Duration, max_connections: usize, +) -> Result { + connect_with_ca_cert_pool_limit_and_idle(url, ca_cert_path, fallback_timeout, max_connections, None).await +} + +pub async fn connect_with_ca_cert_pool_limit_and_idle( + url: &str, + ca_cert_path: Option<&str>, + fallback_timeout: Duration, + max_connections: usize, + idle_timeout_secs: Option, ) -> Result { let timeout = super::parse_connect_timeout_with_fallback(url, fallback_timeout); - let pool = create_pool(url, ca_cert_path, max_connections)?; + let pool = create_pool(url, ca_cert_path, max_connections, idle_timeout_secs)?; let result = verify_pool_connection(&pool, timeout).await; if let Err(ref e) = result { if mysql_error_should_retry_without_ssl(e) { if let Some(fallback_url) = ssl_fallback_url(url) { log::info!("SSL handshake failed, retrying with ssl-mode=disabled"); - let fallback_pool = create_pool(&fallback_url, None, max_connections)?; + let fallback_pool = create_pool(&fallback_url, None, max_connections, idle_timeout_secs)?; return match verify_pool_connection(&fallback_pool, timeout).await { Ok(()) => Ok(fallback_pool), Err(e) => Err(e), @@ -413,7 +424,12 @@ struct MySqlTlsFiles { sslkey: Option, } -fn create_pool(url: &str, ca_cert_path: Option<&str>, max_connections: usize) -> Result { +fn create_pool( + url: &str, + ca_cert_path: Option<&str>, + max_connections: usize, + idle_timeout_secs: Option, +) -> Result { let tls_url = mysql_tls_url(url)?; let opts = mysql_async::Opts::from_url(&mysql_async_url(&tls_url.url)).map_err(|e| format!("Invalid MySQL URL: {e}"))?; @@ -422,14 +438,17 @@ fn create_pool(url: &str, ca_cert_path: Option<&str>, max_connections: usize) -> // Single-connection pools (max_connections == 1) are client session pools that // must preserve session state (e.g. TEMPORARY TABLEs) across queries. // Disable COM_RESET_CONNECTION for these pools to avoid clearing that state. + let inactive_ttl = + idle_timeout_secs.filter(|&s| s >= 30).map(Duration::from_secs).unwrap_or(Duration::from_secs(300)); let pool_opts = mysql_async::PoolOpts::new() .with_constraints(mysql_async::PoolConstraints::new(1, max_connections).unwrap()) - .with_inactive_connection_ttl(Duration::from_secs(300)) + .with_inactive_connection_ttl(inactive_ttl) .with_reset_connection(max_connections > 1); let mut builder = mysql_async::OptsBuilder::from_opts(opts) .stmt_cache_size(0) .prefer_socket(false) .pool_opts(Some(pool_opts)) + .tcp_keepalive(Some(30u32)) .setup(mysql_setup_queries(url)); if let Some(ssl_opts) = mysql_ssl_opts(base_ssl_opts, url, ca_cert_path, &tls_url.files)? { builder = builder.ssl_opts(ssl_opts); @@ -992,12 +1011,12 @@ pub async fn connect_bare_with_pool_limit( max_connections: usize, ) -> Result { let timeout = super::parse_connect_timeout_with_fallback(url, fallback_timeout); - let pool = create_pool(url, None, max_connections)?; + let pool = create_pool(url, None, max_connections, None)?; verify_pool_connection(&pool, timeout).await.map(|_| pool) } pub async fn list_databases(pool: &MySqlPool) -> Result, String> { - let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; let result = match conn.query_iter("SELECT SCHEMA_NAME FROM information_schema.SCHEMATA ORDER BY SCHEMA_NAME").await { Ok(result) => result, @@ -1018,7 +1037,7 @@ pub async fn list_databases(pool: &MySqlPool) -> Result, Strin } pub async fn list_databases_show(pool: &MySqlPool) -> Result, String> { - let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; let result = conn.query_iter("SHOW DATABASES").await.map_err(|e| e.to_string())?; let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; Ok(database_infos_from_names(rows.iter().map(|row| get_str(row, 0)), true)) @@ -1057,7 +1076,7 @@ pub async fn list_tables_filtered( object_types: Option<&[String]>, ) -> Result, String> { let sql = list_tables_sql(database, filter, limit, offset, object_types); - let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; let result = match conn.query_iter(&sql).await { Ok(result) => result, Err(err) => { @@ -1198,7 +1217,7 @@ pub async fn completion_assistant_search( request.object_kinds.clone() }; let pattern = mysql_completion_like_pattern(&request.mask, request.match_mode.as_ref()); - let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; let mut candidates = Vec::new(); if kinds @@ -1421,7 +1440,7 @@ fn table_comment_sql(database: &str, table: &str) -> String { pub async fn get_table_comment(pool: &MySqlPool, database: &str, table: &str) -> Result, String> { let sql = table_comment_sql(database, table); - let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; let result = conn.query_iter(&sql).await.map_err(|e| e.to_string())?; let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; Ok(rows @@ -1444,7 +1463,7 @@ async fn list_table_status_show(pool: &MySqlPool, database: &str) -> Result = result.collect_and_drop().await.map_err(|e| e.to_string())?; Ok(rows @@ -1467,7 +1486,7 @@ async fn list_table_status_show(pool: &MySqlPool, database: &str) -> Result Result, String> { let sql = show_tables_sql(database, true); - let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; let rows: Vec = match conn.query_iter(&sql).await { Ok(result) => result.collect_and_drop().await.map_err(|e| e.to_string())?, Err(_) => { @@ -1589,7 +1608,7 @@ fn row_to_object(row: &mysql_async::Row, database: &str) -> ObjectInfo { } pub async fn list_objects(pool: &MySqlPool, database: &str) -> Result, String> { - let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; let tables_sql = list_tables_objects_sql(database); let result = conn.query_iter(&tables_sql).await.map_err(|e| e.to_string())?; @@ -1625,7 +1644,7 @@ pub async fn list_object_statistics(pool: &MySqlPool, database: &str) -> Result< ORDER BY TABLE_NAME", quote_value(database), ); - let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; let result = conn.query_iter(&sql).await.map_err(|e| e.to_string())?; let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; Ok(rows @@ -1672,7 +1691,7 @@ pub async fn list_table_objects_show(pool: &MySqlPool, database: &str) -> Result } async fn list_routine_objects(pool: &MySqlPool, database: &str) -> Result, String> { - let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; let routines_sql = list_routines_sql(database); let result = conn.query_iter(&routines_sql).await.map_err(|e| e.to_string())?; let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; @@ -1680,7 +1699,7 @@ async fn list_routine_objects(pool: &MySqlPool, database: &str) -> Result Result, String> { - let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; let mut objects = Vec::new(); let routines_sql = list_routines_sql(database); @@ -1991,31 +2010,115 @@ fn skip_mysql_quoted(sql: &str, start: usize, quote: u8) -> usize { /// Get a connection from the pool with a health check. If the connection is dead /// (e.g. after app was backgrounded), it tries again with a fresh connection. pub async fn get_conn_with_health_check(pool: &MySqlPool) -> Result { - let timeout = crate::db::connection_timeout(); - let mut conn = get_conn_with_timeout(pool, timeout).await?; - match ping_conn_with_timeout(&mut conn, timeout).await { - Ok(()) => Ok(conn), - _ => { - let _ = tokio::time::timeout(timeout, conn.disconnect()).await; - let mut conn = get_conn_with_timeout(pool, timeout).await?; - ping_conn_with_timeout(&mut conn, timeout).await?; + get_conn_with_health_check_with_timeout(pool, super::connection_timeout()).await +} + +pub async fn get_conn_with_health_check_with_timeout( + pool: &MySqlPool, + timeout: Duration, +) -> Result { + get_conn_with_health_check_with_cancel(pool, timeout, timeout, None).await +} + +pub async fn get_conn_with_health_check_with_cancel( + pool: &MySqlPool, + timeout: Duration, + cleanup_timeout: Duration, + cancel_token: Option<&CancellationToken>, +) -> Result { + let start = Instant::now(); + let mut conn = get_conn_with_timeout_and_cancel(pool, timeout, cancel_token).await?; + match ping_conn_with_timeout_and_cancel(&mut conn, timeout, cancel_token).await { + Ok(()) => { + log::debug!( + "[db:health.check:done] elapsed_ms={} timeout_ms={}", + start.elapsed().as_millis(), + timeout.as_millis() + ); + Ok(conn) + } + Err(err) if err == crate::query::QUERY_CANCELED => { + let _ = tokio::time::timeout(cleanup_timeout, conn.disconnect()).await; + Err(err) + } + Err(err) => { + log::warn!( + "[db:health.check:error] elapsed_ms={} timeout_ms={} error={}; retrying", + start.elapsed().as_millis(), + timeout.as_millis(), + err + ); + let _ = tokio::time::timeout(cleanup_timeout, conn.disconnect()).await; + let mut conn = get_conn_with_timeout_and_cancel(pool, timeout, cancel_token).await?; + if let Err(err) = ping_conn_with_timeout_and_cancel(&mut conn, timeout, cancel_token).await { + if err == crate::query::QUERY_CANCELED { + let _ = tokio::time::timeout(cleanup_timeout, conn.disconnect()).await; + } + return Err(err); + } + log::info!( + "[db:health.check:recovered] elapsed_ms={} timeout_ms={}", + start.elapsed().as_millis(), + timeout.as_millis() + ); Ok(conn) } } } -async fn get_conn_with_timeout(pool: &MySqlPool, timeout: Duration) -> Result { +async fn get_conn_with_timeout_and_cancel( + pool: &MySqlPool, + timeout: Duration, + cancel_token: Option<&CancellationToken>, +) -> Result { + let get_future = async { + tokio::time::timeout(timeout, pool.get_conn()) + .await + .map_err(|_| "MySQL get connection timed out".to_string())? + .map_err(|e| e.to_string()) + }; + + match cancel_token { + Some(token) => { + tokio::select! { + biased; + _ = token.cancelled() => Err(crate::query::canceled_error()), + result = get_future => result, + } + } + None => get_future.await, + } +} + +pub async fn get_conn_with_timeout(pool: &MySqlPool, timeout: Duration) -> Result { tokio::time::timeout(timeout, pool.get_conn()) .await .map_err(|_| "MySQL get connection timed out".to_string())? .map_err(|e| e.to_string()) } -async fn ping_conn_with_timeout(conn: &mut mysql_async::Conn, timeout: Duration) -> Result<(), String> { - tokio::time::timeout(timeout, conn.ping()) - .await - .map_err(|_| "MySQL ping timed out".to_string())? - .map_err(|e| e.to_string()) +async fn ping_conn_with_timeout_and_cancel( + conn: &mut mysql_async::Conn, + timeout: Duration, + cancel_token: Option<&CancellationToken>, +) -> Result<(), String> { + let ping_future = async { + tokio::time::timeout(timeout, conn.ping()) + .await + .map_err(|_| "MySQL ping timed out".to_string())? + .map_err(|e| e.to_string()) + }; + + match cancel_token { + Some(token) => { + tokio::select! { + biased; + _ = token.cancelled() => Err(crate::query::canceled_error()), + result = ping_future => result, + } + } + None => ping_future.await, + } } async fn execute_result_set_with_text_protocol_on_conn( @@ -2234,13 +2337,61 @@ async fn stream_query_rows_prepared( } pub async fn kill_query(pool: &MySqlPool, connection_id: u32) -> Result<(), String> { - let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; - conn.query_drop(format!("KILL QUERY {connection_id}")).await.map_err(|e| e.to_string()) + let start = Instant::now(); + let timeout = super::connection_timeout(); + let mut conn = tokio::time::timeout(timeout, pool.get_conn()) + .await + .map_err(|_| { + log::warn!( + "[db:cancel:error] elapsed_ms={} timeout_ms={} error=MySQL kill connection checkout timed out", + start.elapsed().as_millis(), + timeout.as_millis() + ); + "MySQL kill connection checkout timed out".to_string() + })? + .map_err(|e| e.to_string())?; + tokio::time::timeout(timeout, conn.query_drop(format!("KILL QUERY {connection_id}"))) + .await + .map_err(|_| { + log::warn!( + "[db:cancel:error] elapsed_ms={} timeout_ms={} error=MySQL KILL QUERY timed out", + start.elapsed().as_millis(), + timeout.as_millis() + ); + "MySQL KILL QUERY timed out".to_string() + })? + .map_err(|e| e.to_string())?; + log::info!("[db:cancel:done] elapsed_ms={} timeout_ms={}", start.elapsed().as_millis(), timeout.as_millis()); + Ok(()) } pub async fn kill_query_with_opts(opts: mysql_async::Opts, connection_id: u32) -> Result<(), String> { - let mut conn = mysql_async::Conn::new(opts).await.map_err(|e| e.to_string())?; - conn.query_drop(format!("KILL QUERY {connection_id}")).await.map_err(|e| e.to_string()) + let start = Instant::now(); + let timeout = super::connection_timeout(); + let mut conn = tokio::time::timeout(timeout, mysql_async::Conn::new(opts)) + .await + .map_err(|_| { + log::warn!( + "[db:cancel:error] elapsed_ms={} timeout_ms={} error=MySQL kill connection timed out", + start.elapsed().as_millis(), + timeout.as_millis() + ); + "MySQL kill connection timed out".to_string() + })? + .map_err(|e| e.to_string())?; + tokio::time::timeout(timeout, conn.query_drop(format!("KILL QUERY {connection_id}"))) + .await + .map_err(|_| { + log::warn!( + "[db:cancel:error] elapsed_ms={} timeout_ms={} error=MySQL KILL QUERY execution timed out", + start.elapsed().as_millis(), + timeout.as_millis() + ); + "MySQL KILL QUERY execution timed out".to_string() + })? + .map_err(|e| e.to_string())?; + log::info!("[db:cancel:done] elapsed_ms={} timeout_ms={}", start.elapsed().as_millis(), timeout.as_millis()); + Ok(()) } pub async fn execute_query_on_conn_with_max_rows( @@ -2400,7 +2551,7 @@ pub async fn list_indexes(pool: &MySqlPool, database: &str, table: &str) -> Resu quote_value(database), quote_value(table), ); - let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; let result = conn.query_iter(&sql).await.map_err(|e| e.to_string())?; let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; @@ -2438,7 +2589,7 @@ pub async fn list_foreign_keys(pool: &MySqlPool, database: &str, table: &str) -> quote_value(database), quote_value(table), ); - let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; let result = conn.query_iter(&sql).await.map_err(|e| e.to_string())?; let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; @@ -2465,7 +2616,7 @@ pub async fn list_triggers(pool: &MySqlPool, database: &str, table: &str) -> Res quote_value(database), quote_value(table), ); - let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?; let result = conn.query_iter(&sql).await.map_err(|e| e.to_string())?; let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; @@ -2483,6 +2634,7 @@ pub async fn list_triggers(pool: &MySqlPool, database: &str, table: &str) -> Res #[cfg(test)] mod tests { use super::*; + use crate::db::connection_timeout; use mysql_async::consts::ColumnFlags; #[test] @@ -3044,7 +3196,7 @@ mod tests { #[test] fn parse_connect_timeout_ignores_out_of_range() { - let default = crate::db::connection_timeout(); + let default = connection_timeout(); let url = "mysql://host:3306/db?connect_timeout=999"; assert_eq!(crate::db::parse_connect_timeout(url), default); let url2 = "mysql://host:3306/db?connect_timeout=0"; @@ -3053,14 +3205,14 @@ mod tests { #[test] fn parse_connect_timeout_returns_default_when_missing() { - let default = crate::db::connection_timeout(); + let default = connection_timeout(); let url = "mysql://host:3306/db?ssl-mode=preferred&charset=utf8mb4"; assert_eq!(crate::db::parse_connect_timeout(url), default); } #[test] fn parse_connect_timeout_returns_default_when_no_query() { - let default = crate::db::connection_timeout(); + let default = connection_timeout(); let url = "mysql://host:3306/db"; assert_eq!(crate::db::parse_connect_timeout(url), default); } diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 47df3e1a7d..8bf4dd5ddc 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -21,6 +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::sql::starts_with_executable_sql_keyword; use crate::types::{ ColumnInfo, CompletionAssistantCandidate, CompletionAssistantCandidateKind, CompletionAssistantMatchMode, @@ -718,8 +719,7 @@ pub async fn stream_query_rows( cancelled: &AtomicBool, mut on_row: impl FnMut(&[serde_json::Value]) -> Result<(), String>, ) -> Result { - let client = - pool.get().await.map_err(|e| format!("PostgreSQL connection failed: {}", pg_pool_error_to_string(e)))?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; match stream_query_rows_on_client(&client, sql, max_rows, cancelled, &mut on_row).await { Ok(rows) => Ok(rows), Err(error) if should_retry_postgres_text_query_message(&error.to_ascii_lowercase()) => { @@ -797,7 +797,8 @@ async fn stream_query_rows_text_on_client( } pub async fn connect(url: &str, fallback_timeout: Duration) -> Result { - let postgres_url = postgres_connection_url(url)?; + let url_with_keepalive = inject_postgres_keepalive_params(url); + let postgres_url = postgres_connection_url(&url_with_keepalive)?; let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); let timeout = super::parse_connect_timeout_with_fallback(url, fallback_timeout); @@ -823,6 +824,8 @@ pub async fn connect(url: &str, fallback_timeout: Duration) -> Result Result, - sslkey: Option, - sslrootcert: Option, +pub struct PostgresSslFiles { + pub sslcert: Option, + pub sslkey: Option, + pub sslrootcert: Option, +} + +/// TLS context info, used to reconstruct the TLS connector when cancelling a query. +#[derive(Debug, Clone)] +pub struct PostgresCancelContext { + pub ssl_files: PostgresSslFiles, + pub accepts_invalid_certs: bool, + pub verifies_hostname: bool, + pub ssl_mode: SslMode, +} + +/// Build a TLS cancel context from the connection URL. +/// Returns None if URL parsing fails or sslmode=disable (no TLS cancel needed). +pub fn build_postgres_cancel_context(url: &str) -> Option { + let postgres_url = postgres_connection_url(url).ok()?; + let pg_config = tokio_postgres::Config::from_str(&postgres_url.url).ok()?; + if pg_config.get_ssl_mode() == SslMode::Disable { + return None; + } + Some(PostgresCancelContext { + ssl_files: postgres_url.ssl_files, + accepts_invalid_certs: postgres_url.accepts_invalid_certs, + verifies_hostname: postgres_url.verifies_hostname, + ssl_mode: pg_config.get_ssl_mode(), + }) +} + +/// Reconstruct a TLS connector from the cancel context, used for TLS connection cancellation. +fn make_rustls_connect_from_context( + ctx: &PostgresCancelContext, +) -> Result { + // Build a minimal pg_config solely for ssl_mode determination + let mut pg_config = tokio_postgres::Config::new(); + pg_config.ssl_mode(ctx.ssl_mode); + let tls_config = postgres_tls_config(&pg_config, &ctx.ssl_files, ctx.accepts_invalid_certs, ctx.verifies_hostname)?; + Ok(tokio_postgres_rustls::MakeRustlsConnect::new(tls_config)) } #[derive(Debug, Clone, PartialEq, Eq)] @@ -857,6 +896,27 @@ struct PostgresConnectionUrl { verifies_hostname: bool, } +/// Inject TCP keepalive parameters into the PostgreSQL URL (only when the user has not explicitly specified them). +/// Default parameters shorten half-open connection detection time, suitable for desktop/VPN/NAT environments. +fn inject_postgres_keepalive_params(url: &str) -> String { + let (base, fragment) = url.split_once('#').map_or((url, ""), |(base, fragment)| (base, fragment)); + let query = base.split('?').nth(1); + let has_keepalives = query + .map(|q| q.split('&').any(|p| p.split('=').next().is_some_and(|k| k.eq_ignore_ascii_case("keepalives")))) + .unwrap_or(false); + if has_keepalives { + return url.to_string(); // User has explicitly configured keepalive + } + let separator = if base.contains('?') { "&" } else { "?" }; + let injected = + format!("{base}{separator}keepalives=1&keepalives_idle=30&keepalives_interval=10&keepalives_retries=3"); + if fragment.is_empty() { + injected + } else { + format!("{injected}#{fragment}") + } +} + fn postgres_connection_url(url: &str) -> Result { let Some(query_start) = url.find('?') else { let pg_config = @@ -1144,7 +1204,7 @@ fn validate_postgres_ssl_paths(url: &str) -> Result<(), String> { } pub async fn list_databases(pool: &Pool) -> Result, String> { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let stmt = client .prepare_cached( "SELECT datname FROM pg_database \ @@ -1176,7 +1236,7 @@ pub async fn list_tables_filtered( if crate::sql::fuzzy_filter_enabled(filter) { like_fuzzy_pattern(filter) } else { String::new() }; let limit_param = limit.and_then(|value| i64::try_from(value).ok()); let offset_param = offset.and_then(|value| i64::try_from(value).ok()).unwrap_or(0); - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let stmt = client.prepare_cached(postgres_tables_sql()).await.map_err(|e| e.to_string())?; let rows = client .query(&stmt, &[&schema, &filter_pattern, &fuzzy_filter_pattern, &limit_param, &offset_param]) @@ -1207,7 +1267,7 @@ pub async fn completion_assistant_search( request.object_kinds.clone() }; let pattern = postgres_completion_like_pattern(&request.mask, request.match_mode.as_ref()); - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let mut candidates = Vec::new(); if kinds.iter().any(|kind| matches!(kind, CompletionAssistantObjectKind::Schema)) { @@ -1390,7 +1450,7 @@ fn postgres_completion_like_pattern(value: &str, mode: Option<&CompletionAssista pub async fn get_table_comment(pool: &Pool, schema: &str, table: &str) -> Result, String> { let schema = if schema.is_empty() { "public" } else { schema }; - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let stmt = client.prepare_cached(postgres_table_comment_sql()).await.map_err(|e| e.to_string())?; let rows = client.query(&stmt, &[&schema, &table]).await.map_err(|e| e.to_string())?; Ok(rows.first().and_then(|row| row.try_get::<_, Option>(0).ok().flatten()).filter(|s| !s.is_empty())) @@ -1532,7 +1592,7 @@ fn list_objects_sql(include_timestamps: bool) -> &'static str { } pub async fn list_objects(pool: &Pool, schema: &str) -> Result, String> { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let stmt = client.prepare_cached(list_objects_sql(true)).await.map_err(|e| e.to_string())?; let rows = match client.query(&stmt, &[&schema]).await { Ok(rows) => rows, @@ -1559,7 +1619,7 @@ pub async fn list_objects(pool: &Pool, schema: &str) -> Result, pub async fn list_object_statistics(pool: &Pool, schema: &str) -> Result, String> { let schema = if schema.is_empty() { "public" } else { schema }; - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let stmt = client .prepare_cached( "SELECT c.relname, \ @@ -1589,7 +1649,7 @@ pub async fn list_schemas(pool: &Pool) -> Result, String> { } pub async fn list_schema_infos(pool: &Pool) -> Result, String> { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let stmt = client .prepare_cached( "SELECT n.nspname AS schema_name, d.description AS schema_comment \ @@ -1736,7 +1796,7 @@ async fn get_columns_with_sql( pub async fn get_columns(pool: &Pool, schema: &str, table: &str) -> Result, String> { let schema = if schema.is_empty() { "public" } else { schema }; - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; match get_columns_with_sql(&client, POSTGRES_COLUMNS_SQL, schema, table).await { Ok(columns) => Ok(columns), Err(primary_error) => match get_columns_with_sql(&client, POSTGRES_COLUMNS_COMPAT_SQL, schema, table).await { @@ -1783,10 +1843,10 @@ pub async fn execute_query_with_max_rows( let row_limit = query_result_row_limit(max_rows); if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; execute_select_query(&client, sql, start, row_limit).await } else { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let affected = client.execute(sql, &[]).await.map_err(pg_error_to_string)?; clear_postgres_caches_after_ddl(pool, Some(&client), sql); @@ -1809,14 +1869,17 @@ pub async fn execute_query_with_max_rows_and_cancel( sql: &str, max_rows: Option, cancel_token: Option, - timeout_duration: Option, + budget: DbOperationBudget, + cancel_context: Option, ) -> Result { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, cancel_token.as_ref(), budget.checkout_timeout).await?; let pg_cancel_token = client.cancel_token(); wait_postgres_query( pg_cancel_token, + cancel_context, cancel_token, - timeout_duration, + budget.query_timeout, + budget.cancel_timeout, execute_query_with_max_rows_inner(&client, sql, max_rows), ) .await @@ -1834,7 +1897,7 @@ pub async fn execute_query_with_schema_and_max_rows( ) -> Result { let start = Instant::now(); let checkout_start = Instant::now(); - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; log::info!( "[postgres][execute_with_schema:pool:done] elapsed_ms={} total_ms={} schema={}", checkout_start.elapsed().as_millis(), @@ -1850,10 +1913,13 @@ pub async fn execute_query_with_schema_and_max_rows( } let set_schema_start = Instant::now(); - client - .execute(&format!("SET search_path TO {}, public", pg_quote_ident(schema)), &[]) - .await - .map_err(pg_error_to_string)?; + execute_postgres_infra_statement( + &client, + &format!("SET search_path TO {}, public", pg_quote_ident(schema)), + super::connection_timeout(), + "schema.set", + ) + .await?; log::info!( "[postgres][execute_with_schema:set-search-path:done] elapsed_ms={} total_ms={}", set_schema_start.elapsed().as_millis(), @@ -1872,16 +1938,8 @@ pub async fn execute_query_with_schema_and_max_rows( result.is_ok() ); - // Always reset search_path so the connection is clean when returned to the pool - let reset_start = Instant::now(); - let _ = client.execute("RESET search_path", &[]).await; - log::info!( - "[postgres][execute_with_schema:reset-search-path:done] elapsed_ms={} total_ms={}", - reset_start.elapsed().as_millis(), - start.elapsed().as_millis() - ); - - result + let reset_result = reset_postgres_search_path(&client, super::connection_timeout(), start).await; + merge_postgres_query_and_reset_result(result, reset_result) } pub async fn execute_query_with_schema_and_max_rows_and_cancel( @@ -1890,11 +1948,12 @@ pub async fn execute_query_with_schema_and_max_rows_and_cancel( sql: &str, max_rows: Option, cancel_token: Option, - timeout_duration: Option, + budget: DbOperationBudget, + cancel_context: Option, ) -> Result { let start = Instant::now(); let checkout_start = Instant::now(); - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, cancel_token.as_ref(), budget.checkout_timeout).await?; log::info!( "[postgres][execute_with_schema:pool:done] elapsed_ms={} total_ms={} schema={}", checkout_start.elapsed().as_millis(), @@ -1909,18 +1968,23 @@ pub async fn execute_query_with_schema_and_max_rows_and_cancel( let pg_cancel_token = client.cancel_token(); return wait_postgres_query( pg_cancel_token, + cancel_context, cancel_token, - timeout_duration, + budget.query_timeout, + budget.cancel_timeout, execute_query_with_max_rows_inner(&client, sql, max_rows), ) .await; } let set_schema_start = Instant::now(); - client - .execute(&format!("SET search_path TO {}, public", pg_quote_ident(schema)), &[]) - .await - .map_err(pg_error_to_string)?; + execute_postgres_infra_statement( + &client, + &format!("SET search_path TO {}, public", pg_quote_ident(schema)), + budget.recycle_timeout, + "schema.set", + ) + .await?; log::info!( "[postgres][execute_with_schema:set-search-path:done] elapsed_ms={} total_ms={}", set_schema_start.elapsed().as_millis(), @@ -1931,8 +1995,10 @@ pub async fn execute_query_with_schema_and_max_rows_and_cancel( let pg_cancel_token = client.cancel_token(); let result = wait_postgres_query( pg_cancel_token, + cancel_context, cancel_token, - timeout_duration, + budget.query_timeout, + budget.cancel_timeout, execute_query_with_max_rows_inner(&client, sql, max_rows), ) .await; @@ -1946,21 +2012,84 @@ pub async fn execute_query_with_schema_and_max_rows_and_cancel( result.is_ok() ); - let reset_start = Instant::now(); - let _ = client.execute("RESET search_path", &[]).await; - log::info!( - "[postgres][execute_with_schema:reset-search-path:done] elapsed_ms={} total_ms={}", - reset_start.elapsed().as_millis(), - start.elapsed().as_millis() - ); + let reset_result = reset_postgres_search_path(&client, budget.cleanup_timeout, start).await; + merge_postgres_query_and_reset_result(result, reset_result) +} - result +async fn reset_postgres_search_path( + client: &deadpool_postgres::Client, + timeout_duration: Duration, + start: Instant, +) -> Result<(), String> { + let reset_start = Instant::now(); + match execute_postgres_infra_statement(client, "RESET search_path", timeout_duration, "schema.reset").await { + Ok(_) => { + log::info!( + "[postgres][execute_with_schema:reset-search-path:done] elapsed_ms={} total_ms={}", + reset_start.elapsed().as_millis(), + start.elapsed().as_millis() + ); + Ok(()) + } + Err(err) => { + log::warn!( + "[postgres][execute_with_schema:reset-search-path:error] elapsed_ms={} total_ms={} error={}", + reset_start.elapsed().as_millis(), + start.elapsed().as_millis(), + err + ); + Err(postgres_schema_reset_cleanup_error(err)) + } + } +} + +fn merge_postgres_query_and_reset_result( + query_result: Result, + reset_result: Result<(), String>, +) -> Result { + match (query_result, reset_result) { + (Ok(result), Ok(())) => Ok(result), + (Err(query_err), Ok(())) => Err(query_err), + (Ok(_), Err(reset_err)) => Err(reset_err), + (Err(query_err), Err(reset_err)) => Err(format!("{query_err}; {reset_err}")), + } +} + +fn postgres_schema_reset_cleanup_error(err: String) -> String { + format!("PostgreSQL schema.reset cleanup failed: {err}") +} + +pub(crate) async fn execute_postgres_infra_statement( + client: &deadpool_postgres::Client, + sql: &str, + timeout_duration: Duration, + stage: &str, +) -> Result { + tokio::time::timeout(timeout_duration, client.execute(sql, &[])) + .await + .map_err(|_| format!("PostgreSQL {stage} timed out after {} seconds", timeout_duration.as_secs()))? + .map_err(pg_error_to_string) +} + +pub(crate) async fn wait_postgres_operation( + pg_cancel_token: tokio_postgres::CancelToken, + cancel_context: Option, + timeout_duration: Option, + cancel_timeout: Duration, + future: F, +) -> Result +where + F: Future>, +{ + wait_postgres_query(pg_cancel_token, cancel_context, None, timeout_duration, cancel_timeout, future).await } async fn wait_postgres_query( pg_cancel_token: tokio_postgres::CancelToken, + cancel_context: Option, cancel_token: Option, timeout_duration: Option, + cancel_timeout: Duration, future: F, ) -> Result where @@ -1971,13 +2100,13 @@ where tokio::select! { biased; _ = token.cancelled() => { - cancel_postgres_query(pg_cancel_token).await; + cancel_postgres_query(pg_cancel_token, cancel_context.as_ref(), cancel_timeout).await; Err(crate::query::canceled_error()) } result = tokio::time::timeout(duration, future) => match result { Ok(result) => result, Err(_) => { - cancel_postgres_query(pg_cancel_token).await; + cancel_postgres_query(pg_cancel_token, cancel_context.as_ref(), cancel_timeout).await; Err(format!("Query timed out after {} seconds", duration.as_secs())) } }, @@ -1986,7 +2115,7 @@ where (None, Some(duration)) => match tokio::time::timeout(duration, future).await { Ok(result) => result, Err(_) => { - cancel_postgres_query(pg_cancel_token).await; + cancel_postgres_query(pg_cancel_token, cancel_context.as_ref(), cancel_timeout).await; Err(format!("Query timed out after {} seconds", duration.as_secs())) } }, @@ -1994,7 +2123,7 @@ where tokio::select! { biased; _ = token.cancelled() => { - cancel_postgres_query(pg_cancel_token).await; + cancel_postgres_query(pg_cancel_token, cancel_context.as_ref(), cancel_timeout).await; Err(crate::query::canceled_error()) } result = future => result, @@ -2004,14 +2133,103 @@ where } } -async fn cancel_postgres_query(pg_cancel_token: tokio_postgres::CancelToken) { - match tokio::time::timeout(Duration::from_secs(2), pg_cancel_token.cancel_query(NoTls)).await { +/// PostgreSQL pool checkout with timeout and cancel token support. +/// When the checkout phase is stuck, the cancel token can terminate the wait early. +/// The timeout error message includes "checkout timed out" to ensure is_connection_error can classify it correctly. +pub async fn checkout_postgres_client( + pool: &Pool, + cancel_token: Option<&CancellationToken>, + checkout_timeout: Duration, +) -> Result { + let start = Instant::now(); + let get_future = async { + tokio::time::timeout(checkout_timeout, pool.get()) + .await + .map_err(|_| { + let elapsed = start.elapsed().as_millis(); + log::warn!( + "[db:pool.checkout:error] elapsed_ms={} timeout_ms={} error=checkout timed out", + elapsed, + checkout_timeout.as_millis() + ); + format!("PostgreSQL connection pool checkout timed out ({}s)", checkout_timeout.as_secs()) + })? + .map_err(|e| { + let elapsed = start.elapsed().as_millis(); + let err = pg_pool_error_to_string(e); + log::warn!( + "[db:pool.checkout:error] elapsed_ms={} timeout_ms={} error={}", + elapsed, + checkout_timeout.as_millis(), + err + ); + format!("PostgreSQL connection pool checkout failed: {err}") + }) + }; + + let result = match cancel_token { + Some(token) => tokio::select! { + biased; + _ = token.cancelled() => { + log::info!( + "[db:pool.checkout:cancelled] elapsed_ms={} timeout_ms={}", + start.elapsed().as_millis(), + checkout_timeout.as_millis() + ); + return Err(crate::query::canceled_error()); + } + result = get_future => result, + }, + None => get_future.await, + }; + if result.is_ok() { + log::debug!( + "[db:pool.checkout:done] elapsed_ms={} timeout_ms={}", + start.elapsed().as_millis(), + checkout_timeout.as_millis() + ); + } + result +} + +async fn cancel_postgres_query( + pg_cancel_token: tokio_postgres::CancelToken, + cancel_context: Option<&PostgresCancelContext>, + cancel_timeout: Duration, +) { + let cancel_timeout = postgres_cancel_attempt_timeout(cancel_timeout, cancel_context); + if let Some(ctx) = cancel_context { + match make_rustls_connect_from_context(ctx) { + Ok(tls) => match tokio::time::timeout(cancel_timeout, pg_cancel_token.cancel_query(tls)).await { + Ok(Ok(())) => return, + Ok(Err(err)) => { + log::warn!("Failed to send PostgreSQL TLS cancel request: {err}"); + return; + } + Err(_) => { + log::warn!("Timed out sending PostgreSQL TLS cancel request ({}s)", cancel_timeout.as_secs()); + return; + } + }, + Err(err) => { + log::warn!("Failed to build TLS connector for cancel: {err}; falling back to NoTls cancel"); + } + } + } + match tokio::time::timeout(cancel_timeout, pg_cancel_token.cancel_query(NoTls)).await { Ok(Ok(())) => {} Ok(Err(err)) => log::warn!("Failed to send PostgreSQL cancel request: {err}"), - Err(_) => log::warn!("Timed out sending PostgreSQL cancel request"), + Err(_) => log::warn!("Timed out sending PostgreSQL cancel request ({}s)", cancel_timeout.as_secs()), } } +fn postgres_cancel_attempt_timeout( + cancel_timeout: Duration, + _cancel_context: Option<&PostgresCancelContext>, +) -> Duration { + cancel_timeout +} + fn is_transaction_recovery_statement(sql: &str) -> bool { starts_with_executable_sql_keyword(sql, &["ROLLBACK", "ABORT", "COMMIT", "END"]) } @@ -2140,7 +2358,7 @@ async fn list_indexes_with_sql( } pub async fn list_indexes(pool: &Pool, schema: &str, table: &str) -> Result, String> { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; match list_indexes_with_sql(&client, POSTGRES_INDEXES_SQL, schema, table).await { Ok(indexes) => Ok(indexes), Err(primary_error) => match list_indexes_with_sql(&client, POSTGRES_INDEXES_COMPAT_SQL, schema, table).await { @@ -2160,7 +2378,7 @@ pub async fn list_indexes(pool: &Pool, schema: &str, table: &str) -> Result Result, String> { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let stmt = client .prepare_cached( "SELECT fk.constraint_name, fk.column_name, \ @@ -2201,7 +2419,7 @@ pub async fn list_foreign_keys(pool: &Pool, schema: &str, table: &str) -> Result } pub async fn list_triggers(pool: &Pool, schema: &str, table: &str) -> Result, String> { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let stmt = client .prepare_cached( "SELECT trigger_name, event_manipulation, action_timing \ @@ -2225,7 +2443,7 @@ pub async fn list_triggers(pool: &Pool, schema: &str, table: &str) -> Result Result, String> { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; // Use pg_proc + pg_get_functiondef() instead of information_schema.routines // for reliable function definition retrieval (information_schema.routines.routine_definition // is NULL for non-SQL functions like plpgsql) @@ -2267,7 +2485,7 @@ pub async fn list_functions(pool: &Pool, schema: &str) -> Result Result, String> { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; // Use pg_class + pg_sequence + pg_namespace instead of pg_sequences view // for better compatibility and permission handling let stmt = client @@ -2327,7 +2545,7 @@ pub async fn list_sequences(pool: &Pool, schema: &str, with_last_values: bool) - } pub async fn list_rules(pool: &Pool, schema: &str) -> Result, String> { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let stmt = client .prepare_cached( "SELECT schemaname, tablename, rulename, definition \ @@ -2350,7 +2568,7 @@ pub async fn list_rules(pool: &Pool, schema: &str) -> Result, Stri } pub async fn list_owners(pool: &Pool, schema: &str) -> Result, String> { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let stmt = client.prepare_cached(POSTGRES_OWNERS_SQL).await.map_err(|e| e.to_string())?; let rows = client.query(&stmt, &[&schema]).await.map_err(|e| e.to_string())?; @@ -2374,14 +2592,14 @@ pub async fn execute_batch(pool: &Pool, statements: &[String]) -> Result<(), Str if combined.is_empty() { return Ok(()); } - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; client.batch_execute(&combined).await.map_err(pg_error_to_string)?; clear_postgres_caches_after_ddl(pool, Some(&client), &combined); Ok(()) } pub async fn terminate_current_user_database_backends(pool: &Pool, database: &str) -> Result { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; client .execute( "SELECT pg_terminate_backend(pid) \ @@ -2417,7 +2635,7 @@ fn invalidates_postgres_statement_cache(sql: &str) -> bool { /// `COPY table (col1, col2) TO STDOUT (FORMAT CSV, HEADER)`. /// Returns the raw COPY output bytes. pub async fn copy_out(pool: &Pool, sql: &str) -> Result, String> { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let stream = client.copy_out(sql).await.map_err(pg_error_to_string)?; tokio::pin!(stream); let mut result = Vec::new(); @@ -2431,7 +2649,7 @@ pub async fn copy_out(pool: &Pool, sql: &str) -> Result, String> { /// `COPY table (col1, col2) FROM STDIN (FORMAT CSV)`. /// `data` is the raw input in the format specified by the COPY command. pub async fn copy_in(pool: &Pool, sql: &str, data: &[u8]) -> Result<(), String> { - let client = pool.get().await.map_err(|e| e.to_string())?; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; let sink = client.copy_in::(sql).await.map_err(pg_error_to_string)?; let mut sink = Box::pin(sink); sink.as_mut().send(bytes::Bytes::copy_from_slice(data)).await.map_err(pg_error_to_string)?; @@ -2794,6 +3012,38 @@ mod tests { assert!(error.contains("sslkey")); } + #[test] + fn inject_postgres_keepalive_params_preserves_url_fragment() { + let url = "postgres://localhost/app?sslmode=require#read-only"; + + assert_eq!( + inject_postgres_keepalive_params(url), + "postgres://localhost/app?sslmode=require&keepalives=1&keepalives_idle=30&keepalives_interval=10&keepalives_retries=3#read-only" + ); + } + + #[test] + fn postgres_cancel_attempt_timeout_is_single_budget() { + assert_eq!(postgres_cancel_attempt_timeout(Duration::from_secs(5), None), Duration::from_secs(5)); + assert_eq!( + postgres_cancel_attempt_timeout( + Duration::from_secs(5), + Some(&PostgresCancelContext { + ssl_files: PostgresSslFiles::default(), + accepts_invalid_certs: true, + verifies_hostname: false, + ssl_mode: SslMode::Require, + }) + ), + Duration::from_secs(5) + ); + } + + #[test] + fn postgres_cancel_context_omits_disabled_ssl_mode() { + assert!(build_postgres_cancel_context("postgres://localhost/app?sslmode=disable").is_none()); + } + #[test] fn postgres_tls_accepts_invalid_certs_for_require_sslmode() { let pg_config = tokio_postgres::Config::from_str("postgres://localhost/db?sslmode=require").unwrap(); diff --git a/crates/dbx-core/src/models/connection.rs b/crates/dbx-core/src/models/connection.rs index 0ef642f4f5..b66ec2cdda 100644 --- a/crates/dbx-core/src/models/connection.rs +++ b/crates/dbx-core/src/models/connection.rs @@ -211,7 +211,7 @@ pub fn default_idle_timeout_secs() -> u64 { } pub fn default_keepalive_interval_secs() -> u64 { - 60 + 30 } fn default_proxy_port() -> u16 { diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index f1239da5de..554d0549d9 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -20,7 +20,7 @@ use tokio_util::sync::CancellationToken; use crate::connection::{AppState, PoolKind}; use crate::database_capabilities; use crate::db; -use crate::models::connection::DatabaseType; +use crate::models::connection::{ConnectionConfig, DatabaseType}; #[cfg(feature = "duckdb-bundled")] use crate::sql::starts_with_duckdb_result_sql_keyword; use crate::sql::{split_sql_batches, split_sql_statements}; @@ -38,6 +38,54 @@ pub enum PoolErrorAction { ReconnectAndRetry, } +/// Unified database operation execution budget. +/// query_timeout = None only means SQL execution has no upper limit; +/// checkout/connect/recycle/cancel/cleanup always have hard upper limits and cannot be disabled. +#[derive(Debug, Clone)] +pub struct DbOperationBudget { + pub checkout_timeout: Duration, + pub connect_timeout: Duration, + pub recycle_timeout: Duration, + pub query_timeout: Option, + pub cancel_timeout: Duration, + pub cleanup_timeout: Duration, +} + +impl DbOperationBudget { + /// Build an execution budget from connection config. + /// checkout/connect/recycle use connect_timeout_secs (clamped to 1s minimum, 300s maximum). + /// query_timeout follows resolve_query_timeout semantics (Some(0) -> None). + /// cancel/cleanup are fixed values and cannot be disabled. + pub fn from_config(connect_timeout_secs: u64, query_timeout_secs: Option) -> Self { + let infra_timeout = Duration::from_secs(connect_timeout_secs.clamp(1, 300)); + Self { + checkout_timeout: infra_timeout, + connect_timeout: infra_timeout, + recycle_timeout: infra_timeout, + query_timeout: resolve_query_timeout(query_timeout_secs), + cancel_timeout: Duration::from_secs(5), + cleanup_timeout: Duration::from_secs(3), + } + } + + pub fn from_connection_config(config: &ConnectionConfig) -> Self { + Self::from_config(config.effective_connect_timeout_secs(), Some(config.query_timeout_secs)) + } + + /// Use global default values (when no connection config is available). + pub fn with_defaults() -> Self { + let default_infra = db::connection_timeout(); + Self { + checkout_timeout: default_infra, + connect_timeout: default_infra, + recycle_timeout: default_infra, + query_timeout: Some(QUERY_TIMEOUT), + cancel_timeout: Duration::from_secs(5), + cleanup_timeout: Duration::from_secs(3), + } + } +} + /// Check read-only protection for a connection, blocking write SQL statements. /// Only clones the connection name when read-only mode is active, avoiding /// unnecessary allocations otherwise. @@ -656,6 +704,7 @@ pub fn is_connection_error(err: &str) -> bool { || lower.contains("broken pipe") || lower.contains("reset by peer") || lower.contains("timed out") + || (lower.contains("pool") && lower.contains("timeout")) || lower.contains("closed") || lower.contains("关闭的连接") || lower.contains("连接已关闭") @@ -680,6 +729,10 @@ fn is_agent_rpc_timeout_error(lower: &str) -> bool { lower.starts_with("agent rpc call timed out ") } +fn is_schema_reset_cleanup_error(lower: &str) -> bool { + lower.contains("schema.reset cleanup failed") +} + fn should_discard_agent_pool_after_error(err: &str) -> bool { let lower = err.to_lowercase(); is_dbx_query_timeout_error(&lower) @@ -695,6 +748,7 @@ pub fn pool_error_action(db_type: Option, err: &str) -> PoolErrorA let lower = err.to_lowercase(); if db::sqlserver::is_driver_panic_error(err) || (is_dbx_query_timeout_error(&lower) && should_discard_pool_after_query_timeout(db_type)) + || is_schema_reset_cleanup_error(&lower) || (db_type.is_some_and(|db_type| database_capabilities::is_agent_type(&db_type)) && should_discard_agent_pool_after_error(err) && !is_connection_error(err)) @@ -824,6 +878,23 @@ fn resolve_query_timeout(timeout_secs: Option) -> Option { } } +async fn operation_budget_for_pool_key( + state: &AppState, + pool_key: &str, + query_timeout: Option, +) -> DbOperationBudget { + let mut budget = configured_operation_budget_for_pool_key(state, pool_key).await; + budget.query_timeout = query_timeout; + budget +} + +async fn configured_operation_budget_for_pool_key(state: &AppState, pool_key: &str) -> DbOperationBudget { + let configs = state.configs.read().await; + crate::connection::config_for_pool_key(pool_key, &configs) + .map(DbOperationBudget::from_connection_config) + .unwrap_or_else(DbOperationBudget::with_defaults) +} + #[allow(clippy::too_many_arguments)] pub async fn do_execute( state: &AppState, @@ -851,6 +922,7 @@ pub async fn do_execute( let conn_name = config.filter(|c| c.read_only).map(|c| c.name.clone()); (attached, conn_name) }; + let operation_budget = operation_budget_for_pool_key(state, pool_key, query_timeout).await; if let Some(name) = conn_name_if_readonly { crate::query_execution_sql::check_read_only(sql, &name)?; } @@ -889,7 +961,21 @@ pub async fn do_execute( let bare = *mode == crate::connection::MysqlMode::Bare; let max_rows = options.max_rows; drop(connections); - let mut conn = db::mysql::get_conn_with_health_check(&p).await?; + let mut conn = match db::mysql::get_conn_with_health_check_with_cancel( + &p, + operation_budget.checkout_timeout, + operation_budget.cleanup_timeout, + cancel_token.as_ref(), + ) + .await + { + Ok(conn) => conn, + Err(err) if err == QUERY_CANCELED => { + state.remove_pool_by_key(pool_key).await; + return Err(err); + } + Err(err) => return Err(err), + }; let connection_id = conn.id(); if let Some(ref execution_id) = options.execution_id { let kill_opts = conn.opts().clone(); @@ -913,7 +999,7 @@ pub async fn do_execute( let p = p.clone(); let schema = schema.map(|s| s.to_string()); let max_rows = options.max_rows; - let query_timeout = query_timeout; + let cancel_context = state.get_postgres_cancel_context(pool_key).await; drop(connections); if let Some(schema) = schema { db::postgres::execute_query_with_schema_and_max_rows_and_cancel( @@ -922,12 +1008,20 @@ pub async fn do_execute( sql, max_rows, cancel_token, - query_timeout, + operation_budget.clone(), + cancel_context, ) .await } else { - db::postgres::execute_query_with_max_rows_and_cancel(&p, sql, max_rows, cancel_token, query_timeout) - .await + db::postgres::execute_query_with_max_rows_and_cancel( + &p, + sql, + max_rows, + cancel_token, + operation_budget.clone(), + cancel_context, + ) + .await } } PoolKind::Sqlite(p) => { @@ -1553,12 +1647,21 @@ async fn execute_multi_mysql( options: QueryExecutionOptions, ) -> Result, String> { let query_timeout = resolve_query_timeout(options.timeout_secs); + let operation_budget = operation_budget_for_pool_key(state, pool_key, query_timeout).await; let bare = mode == crate::connection::MysqlMode::Bare; let max_rows = options.max_rows; - let mut conn = match db::mysql::get_conn_with_health_check(pool).await { + let mut conn = match db::mysql::get_conn_with_health_check_with_cancel( + pool, + operation_budget.checkout_timeout, + operation_budget.cleanup_timeout, + cancel_token.as_ref(), + ) + .await + { Ok(conn) => conn, Err(err) => { if matches!(pool_error_action(db_type, &err), PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry) + || err == QUERY_CANCELED { state.remove_pool_by_key(pool_key).await; } @@ -1850,6 +1953,7 @@ pub async fn execute_statements_in_transaction( let start = std::time::Instant::now(); let db_type = connection_database_type(state, connection_id).await; + let operation_budget = configured_operation_budget_for_pool_key(state, &pool_key).await; // Clone the pool handle within the lock, then drop it before any async work. let path = { @@ -1886,8 +1990,13 @@ pub async fn execute_statements_in_transaction( }; let result = match path { - Some(TxPath::Pg(pool)) => exec_tx_pg_inner(pool, statements, schema, start).await, - Some(TxPath::Mysql(pool, _bare)) => exec_tx_mysql_inner(pool, statements, start).await, + Some(TxPath::Pg(pool)) => { + let cancel_context = state.get_postgres_cancel_context(&pool_key).await; + exec_tx_pg_inner(pool, statements, schema, start, operation_budget.clone(), cancel_context).await + } + Some(TxPath::Mysql(pool, _bare)) => { + exec_tx_mysql_inner(pool, statements, start, operation_budget.clone()).await + } Some(TxPath::Sqlite(pool)) => exec_tx_sqlite_inner(pool, statements, start).await, Some(TxPath::Explicit) => { let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await; @@ -1926,24 +2035,41 @@ async fn exec_tx_pg_inner( statements: &[String], schema: Option<&str>, start: std::time::Instant, + budget: DbOperationBudget, + cancel_context: Option, ) -> Result { - let mut client = pool.get().await.map_err(|e| format!("Failed to acquire connection: {}", e))?; + let mut client = db::postgres::checkout_postgres_client(&pool, None, budget.checkout_timeout) + .await + .map_err(|e| format!("Failed to acquire connection: {}", e))?; let had_schema = schema.is_some(); if let Some(s) = schema { - client - .execute(&format!("SET search_path TO {}", db::postgres::pg_quote_ident(s)), &[]) - .await - .map_err(|e| format!("SET search_path failed: {}", e))?; + db::postgres::execute_postgres_infra_statement( + &client, + &format!("SET search_path TO {}", db::postgres::pg_quote_ident(s)), + budget.recycle_timeout, + "schema.set", + ) + .await + .map_err(|e| format!("SET search_path failed: {}", e))?; } - let tx_result = exec_tx_pg_statements(&mut client, statements).await; + let tx_result = exec_tx_pg_statements(&mut client, statements, &budget, cancel_context).await; // Always reset search_path so the connection is clean when returned to the pool - if had_schema { - let _ = client.execute("RESET search_path", &[]).await; - } + let reset_result = if had_schema { + db::postgres::execute_postgres_infra_statement( + &client, + "RESET search_path", + budget.cleanup_timeout, + "schema.reset", + ) + .await + .map_err(|err| format!("PostgreSQL schema.reset cleanup failed: {err}")) + } else { + Ok(0) + }; - match tx_result { - Ok(total_affected) => Ok(db::QueryResult { + match (tx_result, reset_result) { + (Ok(total_affected), Ok(_)) => Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: vec![], @@ -1954,23 +2080,42 @@ async fn exec_tx_pg_inner( session_id: None, has_more: false, }), - Err(e) => Err(e), + (Err(e), Ok(_)) => Err(e), + (Ok(_), Err(reset_err)) => Err(reset_err), + (Err(e), Err(reset_err)) => Err(format!("{e}; {reset_err}")), } } -async fn exec_tx_pg_statements(client: &mut deadpool_postgres::Client, statements: &[String]) -> Result { - let tx = client.transaction().await.map_err(|e| format!("Failed to begin transaction: {}", e))?; +async fn exec_tx_pg_statements( + client: &mut deadpool_postgres::Client, + statements: &[String], + budget: &DbOperationBudget, + cancel_context: Option, +) -> Result { + let tx = tokio::time::timeout(budget.recycle_timeout, client.transaction()) + .await + .map_err(|_| { + format!("Failed to begin transaction: timed out after {} seconds", budget.recycle_timeout.as_secs()) + })? + .map_err(|e| format!("Failed to begin transaction: {}", e))?; let mut total_affected: u64 = 0; for (i, sql) in statements.iter().enumerate() { - match tx.execute(sql, &[]).await { - Ok(affected) => total_affected += affected, - Err(e) => { - // Transaction auto-rollbacks on drop - return Err(format!("Statement {} failed: {}", i + 1, e)); - } - } + let pg_cancel_token = tx.client().cancel_token(); + let affected = db::postgres::wait_postgres_operation( + pg_cancel_token, + cancel_context.clone(), + budget.query_timeout, + budget.cancel_timeout, + async { tx.execute(sql, &[]).await.map_err(|e| e.to_string()) }, + ) + .await + .map_err(|e| format!("Statement {} failed: {}", i + 1, e))?; + total_affected += affected; } - tx.commit().await.map_err(|e| format!("COMMIT failed: {}", e))?; + tokio::time::timeout(budget.cleanup_timeout, tx.commit()) + .await + .map_err(|_| format!("COMMIT timed out after {} seconds", budget.cleanup_timeout.as_secs()))? + .map_err(|e| format!("COMMIT failed: {}", e))?; Ok(total_affected) } @@ -1978,20 +2123,28 @@ async fn exec_tx_mysql_inner( pool: mysql_async::Pool, statements: &[String], start: std::time::Instant, + budget: DbOperationBudget, ) -> Result { - let mut conn = db::mysql::get_conn_with_health_check(&pool).await?; - conn.query_drop("START TRANSACTION").await.map_err(|e| format!("Failed to begin transaction: {}", e))?; + let mut conn = db::mysql::get_conn_with_health_check_with_timeout(&pool, budget.checkout_timeout).await?; + mysql_query_drop_with_timeout( + &mut conn, + "START TRANSACTION", + budget.recycle_timeout, + "Failed to begin transaction", + ) + .await?; let mut total_affected: u64 = 0; for (i, sql) in statements.iter().enumerate() { - match conn.query_iter(sql).await { - Ok(result) => total_affected += result.affected_rows(), + match mysql_query_iter_with_timeout(&mut conn, sql, budget.query_timeout).await { + Ok(affected) => total_affected += affected, Err(e) => { - let _ = conn.query_drop("ROLLBACK").await; + let _ = mysql_query_drop_with_timeout(&mut conn, "ROLLBACK", budget.cleanup_timeout, "ROLLBACK failed") + .await; return Err(format!("Statement {} failed: {}", i + 1, e)); } } } - conn.query_drop("COMMIT").await.map_err(|e| format!("COMMIT failed: {}", e))?; + mysql_query_drop_with_timeout(&mut conn, "COMMIT", budget.cleanup_timeout, "COMMIT failed").await?; Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), @@ -2005,6 +2158,33 @@ async fn exec_tx_mysql_inner( }) } +async fn mysql_query_drop_with_timeout( + conn: &mut mysql_async::Conn, + sql: &str, + timeout_duration: Duration, + context: &str, +) -> Result<(), String> { + tokio::time::timeout(timeout_duration, conn.query_drop(sql)) + .await + .map_err(|_| format!("{context}: timed out after {} seconds", timeout_duration.as_secs()))? + .map_err(|e| format!("{context}: {e}")) +} + +async fn mysql_query_iter_with_timeout( + conn: &mut mysql_async::Conn, + sql: &str, + timeout_duration: Option, +) -> Result { + match timeout_duration { + Some(timeout_duration) => tokio::time::timeout(timeout_duration, conn.query_iter(sql)) + .await + .map_err(|_| format!("Query timed out after {} seconds", timeout_duration.as_secs()))? + .map(|result| result.affected_rows()) + .map_err(|e| e.to_string()), + None => conn.query_iter(sql).await.map(|result| result.affected_rows()).map_err(|e| e.to_string()), + } +} + async fn exec_tx_sqlite_inner( pool: db::sqlite::SqliteHandle, statements: &[String], @@ -2175,6 +2355,55 @@ mod tests { use super::*; use crate::models::connection::{default_redis_key_separator, ConnectionConfig, DatabaseType}; + fn test_connection_config(db_type: DatabaseType) -> ConnectionConfig { + ConnectionConfig { + id: "conn-1".to_string(), + name: "Connection".to_string(), + db_type, + driver_profile: None, + driver_label: None, + url_params: None, + host: "localhost".to_string(), + port: 0, + username: String::new(), + password: String::new(), + database: None, + visible_databases: None, + visible_schemas: None, + attached_databases: Vec::new(), + color: None, + transport_layers: Vec::new(), + connect_timeout_secs: 10, + query_timeout_secs: 30, + idle_timeout_secs: 60, + keepalive_interval_secs: 30, + 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: 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, + } + } + #[tokio::test] async fn wait_for_query_returns_cancelled_when_token_is_cancelled() { let token = CancellationToken::new(); @@ -2252,6 +2481,71 @@ mod tests { assert_eq!(follow_up.rows, vec![vec![serde_json::json!(1)]]); } + #[test] + fn db_operation_budget_from_config() { + let budget = DbOperationBudget::from_config(10, Some(30)); + assert_eq!(budget.checkout_timeout, Duration::from_secs(10)); + assert_eq!(budget.connect_timeout, Duration::from_secs(10)); + assert_eq!(budget.recycle_timeout, Duration::from_secs(10)); + assert_eq!(budget.query_timeout, Some(Duration::from_secs(30))); + assert_eq!(budget.cancel_timeout, Duration::from_secs(5)); + assert_eq!(budget.cleanup_timeout, Duration::from_secs(3)); + } + + #[test] + fn db_operation_budget_from_connection_config_uses_connection_settings() { + let mut config = test_connection_config(DatabaseType::Postgres); + config.connect_timeout_secs = 12; + config.query_timeout_secs = 0; + + let budget = DbOperationBudget::from_connection_config(&config); + + assert_eq!(budget.checkout_timeout, Duration::from_secs(12)); + assert_eq!(budget.connect_timeout, Duration::from_secs(12)); + assert_eq!(budget.recycle_timeout, Duration::from_secs(12)); + assert_eq!(budget.query_timeout, None); + assert_eq!(budget.cancel_timeout, Duration::from_secs(5)); + assert_eq!(budget.cleanup_timeout, Duration::from_secs(3)); + } + + #[test] + fn db_operation_budget_query_timeout_zero_means_no_limit() { + let budget = DbOperationBudget::from_config(10, Some(0)); + assert_eq!(budget.query_timeout, None); + // Infrastructure timeouts still have hard limits + assert_eq!(budget.checkout_timeout, Duration::from_secs(10)); + assert_eq!(budget.cancel_timeout, Duration::from_secs(5)); + } + + #[test] + fn db_operation_budget_query_timeout_zero_keeps_transaction_infra_limits() { + let mut config = test_connection_config(DatabaseType::Mysql); + config.connect_timeout_secs = 7; + config.query_timeout_secs = 0; + + let budget = DbOperationBudget::from_connection_config(&config); + + assert_eq!(budget.query_timeout, None); + assert_eq!(budget.checkout_timeout, Duration::from_secs(7)); + assert_eq!(budget.recycle_timeout, Duration::from_secs(7)); + assert_eq!(budget.cleanup_timeout, Duration::from_secs(3)); + } + + #[test] + fn db_operation_budget_clamps_infra_timeout() { + let budget = DbOperationBudget::from_config(0, Some(30)); + assert_eq!(budget.checkout_timeout, Duration::from_secs(1)); // clamped to min 1s + let budget = DbOperationBudget::from_config(600, Some(30)); + assert_eq!(budget.checkout_timeout, Duration::from_secs(300)); // clamped to max 300s + } + + #[test] + fn db_operation_budget_with_defaults() { + let budget = DbOperationBudget::with_defaults(); + assert_eq!(budget.checkout_timeout, db::connection_timeout()); + assert_eq!(budget.query_timeout, Some(QUERY_TIMEOUT)); + } + #[test] fn is_connection_error_detects_english_messages() { assert!(is_connection_error("connection reset")); @@ -2300,6 +2594,20 @@ mod tests { assert!(!is_connection_error("os error 13")); } + #[test] + fn is_connection_error_detects_deadpool_pool_timeouts() { + // deadpool-postgres PoolError::Timeout messages (contain "pool" + "timeout" but not "timed out") + assert!(is_connection_error("pool wait timeout")); + assert!(is_connection_error("pool create timeout")); + assert!(is_connection_error("pool recycle timeout")); + // checkout helper timeout messages + assert!(is_connection_error("PostgreSQL connection pool checkout timed out (5s)")); + assert!(is_connection_error("MySQL get connection timed out")); + assert!(is_connection_error("MySQL ping timed out")); + assert!(is_connection_error("MySQL kill connection checkout timed out")); + assert!(is_connection_error("MySQL KILL QUERY timed out")); + } + #[test] fn pool_error_action_discards_sqlserver_driver_panic_without_retry() { let err = format!("{} the current client will be rebuilt.", db::sqlserver::SQLSERVER_DRIVER_PANIC_ERROR_PREFIX); @@ -2322,6 +2630,15 @@ mod tests { assert_eq!(pool_error_action(Some(DatabaseType::DuckDb), err), PoolErrorAction::Keep); } + #[test] + fn pool_error_action_discards_schema_reset_cleanup_without_retry() { + let err = "PostgreSQL schema.reset cleanup failed: PostgreSQL schema.reset timed out after 3 seconds"; + + assert_eq!(pool_error_action(Some(DatabaseType::Postgres), err), PoolErrorAction::Discard); + assert_eq!(pool_error_action(Some(DatabaseType::OpenGauss), err), PoolErrorAction::Discard); + assert!(should_discard_pool_after_error(Some(DatabaseType::Postgres), err)); + } + #[test] fn pool_error_action_reconnects_connection_errors() { let err = "connection reset by peer"; diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 2e2bc94cae..4148e1361e 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -3186,7 +3186,7 @@ async fn mysql_object_source( ) -> Result { use mysql_async::prelude::*; let sql = mysql_object_source_sql(name, kind); - let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let mut conn = db::mysql::get_conn_with_timeout(pool, db::connection_timeout()).await?; let result = conn.query_iter(&sql).await.map_err(|e| e.to_string())?; let rows: Vec = result.collect_and_drop().await.map_err(|e| e.to_string())?; let row = rows.first().ok_or("Object source not found")?;