diff --git a/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java b/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java index a7eff9772c..8d87747727 100644 --- a/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java +++ b/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java @@ -175,6 +175,13 @@ public final class JdbcExecutor { try (ResultSet rs = stmt.getResultSet()) { result = readResultSet(rs, elapsed, effectiveMaxRows, valueReader); } + // `elapsed` stops at `stmt.execute()`. Retrieving and + // converting the rows above is frequently the dominant + // cost of a statement, so report the whole duration the + // same way the native drivers do. Without this a + // fetch-heavy result shows up in the UI as a few + // milliseconds while the caller waited for seconds. + result.setExecution_time_ms(System.currentTimeMillis() - start); } else { result = new QueryResult( Collections.emptyList(), @@ -464,7 +471,11 @@ public final class JdbcExecutor { ); createdSession = session; targetSessions.put(sessionId, session); - return readSessionPage(targetSessions, session, options.getPageSize(), elapsed); + QueryPageResult page = readSessionPage(targetSessions, session, options.getPageSize(), elapsed); + // Same reason as `execute`: pulling the first page's rows is + // part of the duration the caller waited for. + page.setExecution_time_ms(System.currentTimeMillis() - start); + return page; } catch (Exception e) { if (createdSession != null) { closeSession(targetSessions, createdSession.id); @@ -711,8 +722,14 @@ public final class JdbcExecutor { throw new IllegalArgumentException(missingMessage); } synchronized (session) { + long fetchStart = System.currentTimeMillis(); try { - return readSessionPage(targetSessions, session, pageSize, 0L); + QueryPageResult page = readSessionPage(targetSessions, session, pageSize, 0L); + // A later page reads its rows after `stmt.execute()` has long + // returned, so it has to measure the fetch itself instead of + // reporting 0ms while the caller waits. + page.setExecution_time_ms(System.currentTimeMillis() - fetchStart); + return page; } catch (RuntimeException | Error error) { closeSession(targetSessions, sessionId); throw error; diff --git a/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java b/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java index 11637f24bf..f15999d71a 100644 --- a/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java +++ b/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java @@ -119,6 +119,26 @@ class JdbcExecutorTest { assertEquals(2, fixture.getColumnTypeCalls()); } + @Test + void executeCountsRowRetrievalInExecutionTime() { + // `stmt.execute()` returns before the rows are consumed, so a driver + // that only timed that call reported a handful of milliseconds for a + // statement whose fetch took seconds. Row retrieval has to be part of + // the duration the caller waited for. + QueryResult result = JdbcExecutor.INSTANCE.execute( + executionConnection(true, -1, null, new AtomicInteger(), slowResultSet(30L), null), + "SELECT * FROM slow_table", + "", + schema -> "" + ); + + assertEquals(Arrays.asList(Arrays.asList(1, "Ada"), Arrays.asList(2, "Grace")), result.getRows()); + assertTrue( + result.getExecution_time_ms() >= 25L, + "row retrieval must count toward execution_time_ms but was " + result.getExecution_time_ms() + "ms" + ); + } + @Test void executeReturnsMultipleStatementWarningsForNoResultStatements() { SQLWarning first = new SQLWarning("identity value is 443", "S0003", 7998); @@ -432,6 +452,69 @@ class JdbcExecutorTest { return rows[rowIndex][columnIndex - 1]; } + /** Result set whose cursor advances slowly, emulating a remote fetch. */ + private static ResultSet slowResultSet(long perRowDelayMs) { + Object[][] rows = { + {1, "Ada"}, + {2, "Grace"} + }; + String[] labels = {"id", "name"}; + int[] sqlTypes = {Types.INTEGER, Types.VARCHAR}; + String[] typeNames = {"INTEGER", "VARCHAR"}; + AtomicInteger cursor = new AtomicInteger(-1); + + InvocationHandler metaHandler = (Object unused, Method method, Object[] args) -> { + switch (method.getName()) { + case "getColumnCount": + return labels.length; + case "getColumnLabel": + return labels[(Integer) args[0] - 1]; + case "getColumnType": + return sqlTypes[(Integer) args[0] - 1]; + case "getColumnTypeName": + return typeNames[(Integer) args[0] - 1]; + default: + return defaultValue(method.getReturnType()); + } + }; + ResultSetMetaData metadata = (ResultSetMetaData) Proxy.newProxyInstance( + ResultSetMetaData.class.getClassLoader(), + new Class[]{ResultSetMetaData.class}, + metaHandler + ); + + InvocationHandler resultSetHandler = (Object unused, Method method, Object[] args) -> { + switch (method.getName()) { + case "getMetaData": + return metadata; + case "next": + int nextIndex = cursor.incrementAndGet(); + if (nextIndex >= rows.length) { + return false; + } + Thread.sleep(perRowDelayMs); + return true; + case "getObject": + case "getInt": + case "getString": + Object cell = currentCell(rows, cursor.get(), (Integer) args[0]); + if (method.getName().equals("getInt")) { + return ((Number) cell).intValue(); + } + return method.getName().equals("getObject") ? cell : String.valueOf(cell); + case "wasNull": + return false; + default: + return defaultValue(method.getReturnType()); + } + }; + return (ResultSet) Proxy.newProxyInstance( + ResultSet.class.getClassLoader(), + new Class[]{ResultSet.class}, + resultSetHandler + ); + } + private static ResultSet resultSet(Object objectValue, StringSupplier stringSupplier, boolean wasNull) { byte[] bytesValue = objectValue instanceof byte[] ? (byte[]) objectValue : null; return resultSet(objectValue, bytesValue, stringSupplier, wasNull, null); diff --git a/apps/desktop/src/components/editor/QueryEditor.vue b/apps/desktop/src/components/editor/QueryEditor.vue index a62050d26a..4d12996f81 100644 --- a/apps/desktop/src/components/editor/QueryEditor.vue +++ b/apps/desktop/src/components/editor/QueryEditor.vue @@ -7783,6 +7783,19 @@ watch( { deep: true }, ); +// The editor component is reused across tabs, so a tab switch can change the +// connection/database without remounting it. +watch([() => props.connectionId, () => props.database, () => props.catalog, () => props.clientSessionId], () => warmActiveTabConnection()); + +// Restored tabs mount before their connection is established, so the warm-up +// above is skipped and never retried when connecting finishes later. +watch( + () => (props.connectionId ? connectionStore.connectedIds.has(props.connectionId) : false), + (connected) => { + if (connected) warmActiveTabConnection(); + }, +); + // A content-search jump. Fires on id change (a new or reused tab, or a repeat // click on the same result), and on mount when the request is already pending. watch( @@ -8086,12 +8099,24 @@ function resumeQueryEditorBackgroundWork() { editorIsActive = true; registerTableReferenceDropListener(); scheduleSemanticDiagnostics(); + // Warm the database driver/pool while the user is still reading or typing, so + // the first Run does not pay pool creation or external-driver startup. + warmActiveTabConnection(); if (view.value) schedulePreviewContextRefresh(view.value); restoreEditorSelection(undefined, !props.initialViewport); restoreEditorFocus(); restoreEditorViewport(); } +function warmActiveTabConnection() { + if (!props.connectionId) return; + connectionStore.warmConnection(props.connectionId, { + database: props.database, + catalog: props.catalog, + clientSessionId: props.clientSessionId, + }); +} + onActivated(resumeQueryEditorBackgroundWork); onDeactivated(pauseQueryEditorBackgroundWork); @@ -8099,6 +8124,7 @@ onDeactivated(pauseQueryEditorBackgroundWork); onMounted(() => { if (typeof window === "undefined") return; window.addEventListener(BEFORE_TAB_SWITCH_EVENT, captureEditorStateBeforeTabSwitch); + warmActiveTabConnection(); }); onBeforeUnmount(() => { diff --git a/apps/desktop/src/composables/__tests__/useNavigationTargets.store.spec.ts b/apps/desktop/src/composables/__tests__/useNavigationTargets.store.spec.ts index bfb1d96ffd..9bd6856490 100644 --- a/apps/desktop/src/composables/__tests__/useNavigationTargets.store.spec.ts +++ b/apps/desktop/src/composables/__tests__/useNavigationTargets.store.spec.ts @@ -9,6 +9,7 @@ const mocks = vi.hoisted(() => ({ activeConnectionId: "", getConfig: vi.fn((connectionId: string) => ({ id: connectionId, db_type: "postgres" })), ensureConnected: vi.fn(), + warmConnection: vi.fn(), connectionIdentifierQuote: vi.fn(() => undefined), // 与真实 store 一致：setTableMeta 记录连接元数据代次 metadataGenerationFor: vi.fn(() => 0), diff --git a/apps/desktop/src/composables/__tests__/useSidebarDataOpenRuntime.spec.ts b/apps/desktop/src/composables/__tests__/useSidebarDataOpenRuntime.spec.ts index 20fbb44581..64e9ad020b 100644 --- a/apps/desktop/src/composables/__tests__/useSidebarDataOpenRuntime.spec.ts +++ b/apps/desktop/src/composables/__tests__/useSidebarDataOpenRuntime.spec.ts @@ -15,6 +15,7 @@ const mocks = vi.hoisted(() => ({ openDataTabsNextToActive: false, metadataGeneration: 0, ensureConnected: vi.fn(), + warmConnection: vi.fn(), executeTabSql: vi.fn(), loadTableMetadata: vi.fn(), buildTableSelectSql: vi.fn(), @@ -26,6 +27,7 @@ vi.mock("@/stores/connectionStore", () => ({ useConnectionStore: () => ({ getConfig: () => ({ id: "connection-1", db_type: mocks.databaseType }), ensureConnected: mocks.ensureConnected, + warmConnection: mocks.warmConnection, connectionIdentifierQuote: () => undefined, metadataGenerationFor: () => mocks.metadataGeneration, }), @@ -237,6 +239,19 @@ describe("useSidebarDataOpenRuntime", () => { expect(mocks.tabs).toHaveLength(2); }); + it("prewarms the tab session pool before the first page is fetched", async () => { + await useSidebarDataOpenRuntime().openData(tableNode); + + const createdTab = mocks.tabs[0]!; + // The tab executes on its own session pool; warming it up front keeps the + // connection setup off the critical path of the visible loading spinner. + expect(mocks.warmConnection).toHaveBeenCalledWith("connection-1", { + database: "app", + catalog: undefined, + clientSessionId: createdTab.id, + }); + }); + it("keeps adjacent placement for every tab created in always-new mode", async () => { mocks.dataTabReuseMode = "always-new"; mocks.openDataTabsNextToActive = true; diff --git a/apps/desktop/src/composables/useSidebarDataOpenRuntime.ts b/apps/desktop/src/composables/useSidebarDataOpenRuntime.ts index 54c979831a..3405be2494 100644 --- a/apps/desktop/src/composables/useSidebarDataOpenRuntime.ts +++ b/apps/desktop/src/composables/useSidebarDataOpenRuntime.ts @@ -17,6 +17,7 @@ import { tableOpenPageLimit } from "@/lib/table/tableOpenPageLimit"; import { tableDataLargeValuePreviewOptions } from "@/lib/dataGrid/dataGridLargeValues"; import { canActivateExistingDataTableTab } from "@/lib/tabs/dataTabActivation"; import { beginDataTabNavigation, endDataTabNavigation, isCurrentDataTabNavigation } from "@/lib/tabs/dataTabNavigationGeneration"; +import { dataTabExecutionDatabase } from "@/lib/table/dataTabExecutionDatabase"; const DATA_TAB_METADATA_TTL_MS = TABLE_METADATA_CACHE_TTL_MS; @@ -60,6 +61,21 @@ export function useSidebarDataOpenRuntime() { }); lastPhaseAt = now; }; + // Each data tab executes on its own session-scoped pool + // (`:[:catalog:..]:session:`), so its first Run + // pays a fresh database connection. On a remote link that setup costs + // hundreds of milliseconds of round trips and lands on the critical path of + // the page fetch the user is waiting for. Start it in the background as soon + // as the tab exists (and again once the connection is up) so it overlaps the + // table-metadata load instead of the visible loading spinner. + const warmTabSessionPool = (targetTabId: string) => { + if (!config) return; + connectionStore.warmConnection(node.connectionId, { + database: dataTabExecutionDatabase(config, node.database, node.catalog), + catalog: node.catalog, + clientSessionId: targetTabId, + }); + }; openDataLog("info", "start", { traceId, type: node.type, @@ -215,6 +231,7 @@ export function useSidebarDataOpenRuntime() { })(); openDataLog("info", "tab-created", { traceId, tabId, elapsed: elapsed() }); logPhase("tab-created", { tabId }); + warmTabSessionPool(tabId); // 接管该 tab：登记本次导航代次，作废在途的 openTableTarget/openData 旧代次 // ——旧导航即使目标身份相同（同表不同 whereInput）也不得再落地 const navigationToken = beginDataTabNavigation(tabId); @@ -312,6 +329,7 @@ export function useSidebarDataOpenRuntime() { } openDataLog("info", "ensure-connected:done", { traceId, elapsed: elapsed() }); logPhase("ensure-connected", { tabId }); + warmTabSessionPool(tabId); if (!config) throw new Error("Connection config not found"); const limit = tableOpenPageLimit(settingsStore.editorSettings.tableOpenPageSize); diff --git a/apps/desktop/src/lib/backend/api.ts b/apps/desktop/src/lib/backend/api.ts index 7848fcb9b4..8559bc5eb0 100644 --- a/apps/desktop/src/lib/backend/api.ts +++ b/apps/desktop/src/lib/backend/api.ts @@ -70,6 +70,7 @@ export const sessionCredentialStatus = forward("sessionCredentialStatus"); export const forgetSessionCredential = forward("forgetSessionCredential"); export const replaceNacosSessionCredential = forward("replaceNacosSessionCredential"); export const checkConnectionHealth = forward("checkConnectionHealth"); +export const prewarmConnection = forward("prewarmConnection"); export const connectionIdentifierQuote = forward("connectionIdentifierQuote"); export const closeDatabaseConnection = forward("closeDatabaseConnection"); export const refreshConnections = forward("refreshConnections"); diff --git a/apps/desktop/src/lib/backend/http.ts b/apps/desktop/src/lib/backend/http.ts index 3126827a36..ed4d2b1094 100644 --- a/apps/desktop/src/lib/backend/http.ts +++ b/apps/desktop/src/lib/backend/http.ts @@ -488,6 +488,10 @@ export async function checkConnectionHealth(connectionId: string): Promise return post("/api/connection/check-health", { connectionId }); } +export async function prewarmConnection(connectionId: string, database?: string, catalog?: string, clientSessionId?: string): Promise { + return post("/api/connection/prewarm", { connectionId, database, catalog, clientSessionId }); +} + export async function connectionIdentifierQuote(connectionId: string, database?: string): Promise { const quote = await post("/api/connection/identifier-quote", { connectionId, diff --git a/apps/desktop/src/lib/backend/tauri.ts b/apps/desktop/src/lib/backend/tauri.ts index 531f3a2bf9..e192f432e2 100644 --- a/apps/desktop/src/lib/backend/tauri.ts +++ b/apps/desktop/src/lib/backend/tauri.ts @@ -1335,6 +1335,10 @@ export async function checkConnectionHealth(connectionId: string): Promise return invokeBackend("check_connection_health", { connectionId }); } +export async function prewarmConnection(connectionId: string, database?: string, catalog?: string, clientSessionId?: string): Promise { + return invokeBackend("prewarm_connection", { connectionId, database, catalog, clientSessionId }); +} + export async function connectionIdentifierQuote(connectionId: string, database?: string): Promise { const quote = await invoke("connection_identifier_quote", { connectionId, diff --git a/apps/desktop/src/stores/__tests__/connectionStore.prewarm.spec.ts b/apps/desktop/src/stores/__tests__/connectionStore.prewarm.spec.ts new file mode 100644 index 0000000000..7680045f09 --- /dev/null +++ b/apps/desktop/src/stores/__tests__/connectionStore.prewarm.spec.ts @@ -0,0 +1,114 @@ +import { createPinia, setActivePinia } from "pinia"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { ConnectionConfig } 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(): ConnectionConfig { + return { + id: "pg-1", + name: "Postgres", + db_type: "postgres", + host: "127.0.0.1", + port: 5432, + username: "postgres", + password: "", + database: "app", + } as ConnectionConfig; +} + +function installApiMocks(prewarmConnection: ReturnType) { + vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/backend/api", () => ({ + checkConnectionHealth: vi.fn().mockResolvedValue(undefined), + deleteSchemaCachePrefix: vi.fn().mockResolvedValue(undefined), + getColumns: vi.fn().mockResolvedValue([]), + listInstalledAgents: vi.fn().mockResolvedValue([]), + listObjects: vi.fn().mockResolvedValue([]), + listTables: vi.fn().mockResolvedValue([]), + loadSchemaCache: vi.fn().mockResolvedValue(null), + prewarmConnection, + saveConnections: vi.fn().mockResolvedValue(undefined), + saveSchemaCache: vi.fn().mockResolvedValue(undefined), + saveSidebarLayout: vi.fn().mockResolvedValue(undefined), + })); +} + +async function createStore() { + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + const connection = postgresConnection(); + store.connections = [connection]; + store.connectedIds.add(connection.id); + return store; +} + +describe("connectionStore driver prewarm", () => { + beforeEach(() => { + vi.resetModules(); + vi.unstubAllGlobals(); + installLocalStorage(); + setActivePinia(createPinia()); + }); + + it("warms the driver once per pool target and skips unconnected connections", async () => { + const prewarmConnection = vi.fn().mockResolvedValue(undefined); + installApiMocks(prewarmConnection); + const store = await createStore(); + + store.warmConnection("pg-1", { database: "app" }); + store.warmConnection("pg-1", { database: "app" }); + store.warmConnection("pg-1", { database: "app", clientSessionId: "tab-1" }); + store.warmConnection("unknown", { database: "app" }); + await Promise.resolve(); + await Promise.resolve(); + + expect(prewarmConnection).toHaveBeenCalledTimes(2); + expect(prewarmConnection).toHaveBeenCalledWith("pg-1", "app", undefined, undefined); + expect(prewarmConnection).toHaveBeenCalledWith("pg-1", "app", undefined, "tab-1"); + }); + + it("does not surface warm-up failures and retries with a fresh attempt", async () => { + const prewarmConnection = vi.fn().mockRejectedValueOnce(new Error("driver unavailable")).mockResolvedValue(undefined); + installApiMocks(prewarmConnection); + const store = await createStore(); + + expect(() => store.warmConnection("pg-1", { database: "app" })).not.toThrow(); + await new Promise((resolve) => setTimeout(resolve, 0)); + store.warmConnection("pg-1", { database: "app" }); + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(prewarmConnection).toHaveBeenCalledTimes(2); + }); + + it("forgets warm-up bookkeeping when a connection is torn down", async () => { + const prewarmConnection = vi.fn().mockResolvedValue(undefined); + installApiMocks(prewarmConnection); + const store = await createStore(); + const waitForCall = () => new Promise((resolve) => setTimeout(resolve, 0)); + + store.warmConnection("pg-1", { database: "app", clientSessionId: "tab-1" }); + await waitForCall(); + expect(prewarmConnection).toHaveBeenCalledTimes(1); + + // Inside the TTL window the second request is deduplicated. + store.warmConnection("pg-1", { database: "app", clientSessionId: "tab-1" }); + await waitForCall(); + expect(prewarmConnection).toHaveBeenCalledTimes(1); + + // The pool is gone after a teardown, so the next warm-up must run again + // instead of trusting the stale timestamp. + store.markConnectionOffline("pg-1"); + store.connectedIds.add("pg-1"); + store.warmConnection("pg-1", { database: "app", clientSessionId: "tab-1" }); + await waitForCall(); + expect(prewarmConnection).toHaveBeenCalledTimes(2); + }); +}); diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index 6e5af9ed70..9d95bf14be 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -209,6 +209,9 @@ const ACTIVE_CONNECTION_STORAGE_KEY = "dbx-active-connection"; const SIDEBAR_TABLE_NAME_FILTERS_STORAGE_KEY = "dbx-sidebar-table-name-filters"; const CONNECTION_HEALTH_CHECK_TTL_MS = 2000; const CONNECTION_HEALTH_CHECK_TIMEOUT_MS = 5000; +/** How long a successful driver/pool warm-up is trusted before the next tab + * activation warms the same pool target again. */ +const CONNECTION_PREWARM_TTL_MS = 30_000; const METADATA_LOAD_MIN_TIMEOUT_MS = 15_000; const METADATA_LOAD_DISABLED_QUERY_TIMEOUT_MS = 60_000; const DISCONNECT_REQUEST_TIMEOUT_MS = 5_000; @@ -495,6 +498,11 @@ export const useConnectionStore = defineStore("connection", () => { const databaseCompatibilityModes = ref>({}); const databaseCompatibilityRefreshes = new Map>(); const lastConnectionHealthCheckAt = ref>({}); + /** Per pool-target timestamp of the last successful driver/pool warm-up (see + * `warmConnection`). Fire-and-forget: a skipped warm-up only means the pool is + * created by the first real request instead. */ + const lastConnectionPrewarmAt = ref>({}); + const connectionPrewarmInFlight = new Map>(); const agentDrivers = ref([]); let agentDriversRefreshPromise: Promise | null = null; let localAgentDriversRefreshPromise: Promise | null = null; @@ -1318,6 +1326,58 @@ export const useConnectionStore = defineStore("connection", () => { return typeof checkedAt === "number" && Date.now() - checkedAt < CONNECTION_HEALTH_CHECK_TTL_MS; } + /** + * Forgets warm-up bookkeeping for a connection. The backend pool is gone once a + * connection is torn down, so a remembered timestamp must not make + * `warmConnection` skip a warm-up that the next execution would then pay for. + */ + function clearConnectionPrewarmState(connectionId: string) { + const prefix = `${connectionId}\u0000`; + for (const key of Object.keys(lastConnectionPrewarmAt.value)) { + if (key.startsWith(prefix)) delete lastConnectionPrewarmAt.value[key]; + } + for (const key of connectionPrewarmInFlight.keys()) { + if (key.startsWith(prefix)) connectionPrewarmInFlight.delete(key); + } + } + + function connectionPrewarmKey(connectionId: string, target: { database?: string; catalog?: string; clientSessionId?: string }) { + return [connectionId, target.database ?? "", target.catalog ?? "", target.clientSessionId ?? ""].join("\u0000"); + } + + /** + * Warm the driver and connection pool for a connection a tab is about to use. + * + * Opening a SQL editor tab used to do nothing until the user pressed Run, so + * the first execution paid pool creation, tunnel setup, and — for externally + * driven databases such as Oracle — JDBC driver/agent startup. Those seconds + * were visible in the loading indicator but are outside the statement timer, + * which is why the summary could read "23ms" after a multi-second wait. + * Warming in the background moves that cost off the critical path. Failures + * are ignored here because the real request reports them with full context. + */ + function warmConnection(connectionId: string, target: { database?: string; catalog?: string; clientSessionId?: string } = {}) { + if (!connectionId || !connectedIds.value.has(connectionId)) return; + const key = connectionPrewarmKey(connectionId, target); + const warmedAt = lastConnectionPrewarmAt.value[key]; + if (typeof warmedAt === "number" && Date.now() - warmedAt < CONNECTION_PREWARM_TTL_MS) return; + if (connectionPrewarmInFlight.has(key)) return; + const promise = api + .prewarmConnection(connectionId, target.database, target.catalog, target.clientSessionId) + .then(() => { + lastConnectionPrewarmAt.value[key] = Date.now(); + }) + .catch(() => { + // A failed warm-up is not a user-facing error: the next real request + // rebuilds the pool and surfaces any genuine failure itself. + }) + .finally(() => { + connectionPrewarmInFlight.delete(key); + }); + connectionPrewarmInFlight.set(key, promise); + void promise; + } + function clearConnectionNodeLoading(connectionId: string) { const node = findConnectionNode(connectionId); if (node) node.isLoading = false; @@ -4382,6 +4442,7 @@ export const useConnectionStore = defineStore("connection", () => { clearPrimaryVisibleObjectNames(connectionId); clearConnectionIdentifierQuote(connectionId); clearConnectionHealthCheck(connectionId); + clearConnectionPrewarmState(connectionId); if (activeConnectionId.value === connectionId) activeConnectionId.value = null; invalidateCompletionCache(connectionId); cancelObjectDdlLoadsForConnection(connectionId); @@ -4403,6 +4464,7 @@ export const useConnectionStore = defineStore("connection", () => { node.isLoading = false; } clearConnectionHealthCheck(connectionId); + clearConnectionPrewarmState(connectionId); } /** @@ -4426,6 +4488,7 @@ export const useConnectionStore = defineStore("connection", () => { clearConnectionIdentifierQuote(connectionId); forgetSuccessfulLocalConnectionAttempt(connectionId); clearConnectionHealthCheck(connectionId); + clearConnectionPrewarmState(connectionId); const node = findConnectionNode(connectionId); if (node) { node.isLoading = false; @@ -9714,6 +9777,7 @@ export const useConnectionStore = defineStore("connection", () => { hasSessionCredential, closeDatabaseConnection, ensureConnected, + warmConnection, loadConnectedConnectionRootForSidebarSearch, isTreeNodeChildrenLoaded, canUseLoadedTreeNodeToggle, diff --git a/crates/dbx-core/src/connection/mod.rs b/crates/dbx-core/src/connection/mod.rs index 6df51a209d..eb447678f9 100644 --- a/crates/dbx-core/src/connection/mod.rs +++ b/crates/dbx-core/src/connection/mod.rs @@ -58,6 +58,12 @@ const DEFAULT_AGENT_CONNECT_TIMEOUT_SECS: u64 = 30; const ACCESS_AGENT_CONNECT_TIMEOUT_SECS: u64 = 30; const POOL_CLOSE_TIMEOUT_SECS: u64 = 3; const HEALTH_CHECK_POOL_ACQUIRE_TIMEOUT: Duration = Duration::from_millis(500); +/// Upper bound for the "is this checked-out connection still alive" query that +/// follows a successful health checkout. Windows keeps retransmitting on a +/// half-open TCP connection for ~21s before the read fails, so a probe without +/// its own budget would make `check_connection_health` (and therefore +/// `ensureConnected`) hang for the whole OS retry window. +const HEALTH_CHECK_PROBE_TIMEOUT: Duration = Duration::from_secs(5); pub(crate) const METADATA_POOL_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(5); pub(crate) const METADATA_POOL_DEFAULT_LIMIT: usize = 6; pub(crate) const METADATA_POOL_SQLSERVER_LIMIT: usize = 1; @@ -4006,7 +4012,6 @@ impl AppState { } PoolKind::Postgres(pool) => { let pool = pool.clone(); - let timeout = crate::db::connection_timeout(); match db::postgres::checkout_postgres_client_classified( &pool, None, @@ -4014,20 +4019,33 @@ impl AppState { ) .await { - Ok(client) => match tokio::time::timeout(timeout, client.simple_query("SELECT 1")).await { - Ok(Ok(_)) => false, - Ok(Err(err)) => { - log::warn!("PostgreSQL connection pool '{pool_key}' is stale: {err}"); - true + Ok(client) => { + match tokio::time::timeout(HEALTH_CHECK_PROBE_TIMEOUT, client.simple_query("SELECT 1")) + .await + { + Ok(Ok(_)) => false, + Ok(Err(err)) => { + log::warn!("PostgreSQL connection pool '{pool_key}' is stale: {err}"); + true + } + Err(_) => { + log::warn!( + "PostgreSQL connection pool '{pool_key}' is stale: health check timed out" + ); + true + } } - Err(_) => { - log::warn!("PostgreSQL connection pool '{pool_key}' is stale: health check timed out"); - true - } - }, - Err(err) if err.is_pool_saturation() => { + } + // The 500 ms probe budget is intentionally shorter than a foreground checkout. A timeout + // while waiting, creating, or recycling is inconclusive: slow remote handshakes, pool + // re-creation after a keepalive eviction, and active metadata exports can legitimately + // exceed it. Removing the pool here would start competing reconnects while useful work is + // still running, which is exactly how a sub-second probe turns into a multi-second wait on + // the user's next statement. Keep the pool and let the executor's ReconnectAndRetry path + // decide, matching the MySQL branch above. + Err(err @ db::PoolCheckoutError::Timeout { .. }) => { log::debug!( - "PostgreSQL connection pool '{pool_key}' is busy; skipping health probe: {err}" + "PostgreSQL connection pool '{pool_key}' did not finish a health checkout; keeping pool: {err}" ); false } @@ -5162,6 +5180,33 @@ impl AppState { Ok(()) } + /// Warm the driver/pool a tab is about to use, off the user's critical path. + /// + /// The first statement of a session pays costs that the user perceives as + /// "the query is still loading" but that never appear in the reported + /// statement duration: creating the pool, spawning a JDBC/agent driver + /// session (JVM startup for external drivers such as Oracle), opening + /// tunnels, and completing TLS/startup handshakes. `get_or_create_pool_*` + /// performs exactly that work and verifies connectivity before returning, so + /// calling it while the editor is being opened moves those seconds from the + /// first Run to a moment where nobody is waiting on the result. + /// + /// This is deliberately *not* a health probe: it never tears an existing + /// pool down. Use `check_connection_health` when the caller needs a verdict. + pub async fn prewarm_connection_pool( + &self, + connection_id: &str, + database: Option<&str>, + catalog: Option<&str>, + client_session_id: Option<&str>, + ) -> Result<(), String> { + let pool_key = self + .get_or_create_pool_for_session_with_catalog(connection_id, database, catalog, client_session_id) + .await?; + self.touch_pool_activity(&pool_key).await; + Ok(()) + } + pub async fn refresh_connections(&self) { // Clone pool handles under a short-lived read lock, then release it // before performing I/O-heavy health checks to avoid blocking writers. @@ -5188,19 +5233,30 @@ impl AppState { }, PoolKind::Postgres(p) => { match db::postgres::checkout_postgres_client_classified(p, None, timeout).await { - Ok(client) => match tokio::time::timeout(timeout, client.simple_query("SELECT 1")).await { - Ok(Ok(_)) => true, - Ok(Err(e)) => { - log::warn!("PostgreSQL connection pool '{key}' is unhealthy: {e}"); - false + Ok(client) => { + match tokio::time::timeout(HEALTH_CHECK_PROBE_TIMEOUT, client.simple_query("SELECT 1")) + .await + { + Ok(Ok(_)) => true, + Ok(Err(e)) => { + log::warn!("PostgreSQL connection pool '{key}' is unhealthy: {e}"); + false + } + Err(_) => { + log::warn!( + "PostgreSQL connection pool '{key}' is unhealthy: health check timed out" + ); + false + } } - Err(_) => { - log::warn!("PostgreSQL connection pool '{key}' is unhealthy: health check timed out"); - false - } - }, - Err(error) if error.is_pool_saturation() => { - log::debug!("PostgreSQL connection pool '{key}' is busy; skipping health probe: {error}"); + } + // A checkout timeout is inconclusive rather than proof of a dead pool: the budget can be + // consumed by a concurrent create/recycle, and tearing the pool down here would make the + // next foreground statement pay a full reconnect. Mirror `remove_stale_connection_pool`. + Err(error @ db::PoolCheckoutError::Timeout { .. }) => { + log::debug!( + "PostgreSQL connection pool '{key}' did not finish a health checkout; keeping pool: {error}" + ); true } Err(error) => { @@ -7996,6 +8052,55 @@ mod tests { let _ = std::fs::remove_dir_all(dir); } + // A health probe that cannot finish its checkout within the 500 ms budget is + // inconclusive, not proof of a dead pool. Tearing the pool down used to turn + // one slow probe into a full reconnect on the user's next statement, which + // showed up as a loading indicator of several seconds next to a summary that + // only counted the statement itself. + #[tokio::test] + async fn postgres_health_check_keeps_pool_when_checkout_budget_is_exhausted() { + // Accept connections but never answer the PostgreSQL startup packet, so + // every attempt to create a connection runs into its own timeout. + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let mut held = Vec::new(); + while let Ok((socket, _)) = listener.accept().await { + held.push(socket); + } + }); + let mut pg_config = tokio_postgres::Config::new(); + pg_config.host(address.ip().to_string()).port(address.port()).user("health-probe").dbname("health-probe"); + let manager = deadpool_postgres::Manager::new(pg_config, tokio_postgres::NoTls); + let pool = deadpool_postgres::Pool::builder(manager) + .runtime(deadpool_postgres::Runtime::Tokio1) + .max_size(2) + .wait_timeout(Some(Duration::from_millis(200))) + .create_timeout(Some(Duration::from_millis(200))) + .recycle_timeout(Some(Duration::from_millis(200))) + .build() + .expect("build PostgreSQL health probe pool"); + let (state, dir) = test_app_state().await; + state.connections.write().await.insert("conn".to_string(), PoolKind::Postgres(pool.clone())); + + let started = Instant::now(); + assert!(!state.remove_stale_connection_pool("conn").await); + + assert!( + started.elapsed() >= Duration::from_millis(150), + "probe must actually run into the checkout budget, took {:?}", + started.elapsed() + ); + assert!( + matches!(state.connections.read().await.get("conn"), Some(PoolKind::Postgres(_))), + "an inconclusive PostgreSQL probe must not remove the pool" + ); + state.connections.write().await.remove("conn"); + server.abort(); + pool.close(); + let _ = std::fs::remove_dir_all(dir); + } + #[tokio::test] async fn mysql_health_check_keeps_pool_when_connection_creation_exceeds_probe_budget() { let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); diff --git a/crates/dbx-core/src/query/mod.rs b/crates/dbx-core/src/query/mod.rs index 35f6795f49..50a25d7df3 100644 --- a/crates/dbx-core/src/query/mod.rs +++ b/crates/dbx-core/src/query/mod.rs @@ -1720,6 +1720,13 @@ async fn do_execute_typed( let pool_db_type = connection_database_type_for_pool_key(state, pool_key).await; let mysql_catalog_dialect = connection_mysql_catalog_dialect_for_pool_key(state, pool_key).await; let pool = state.pool_handle(pool_key).await.ok_or("Connection not found")?; + // Everything the driver dispatch does for this statement — pool checkout, + // schema/search_path setup, the statement, and the post-statement cleanup — + // is what the user actually waited for. Report that span instead of the + // driver-internal timer, which starts only after a client is in hand and + // therefore hides connection-pool stalls from the summary (#6097 fixed the + // same mismatch for SQL Server's shared-connection lock). + let dispatch_start = std::time::Instant::now(); let mut typed_agent_error = None; #[cfg(feature = "duckdb-sidecar")] @@ -2232,6 +2239,10 @@ async fn do_execute_typed( PoolKind::Consul(_) => Err("SQL execution is not supported for Consul connections".to_string()), }; result + .map(|mut result| { + result.execution_time_ms = result.execution_time_ms.max(dispatch_start.elapsed().as_millis()); + result + }) .map_err(|error| { #[cfg(feature = "duckdb-sidecar")] if let Some(duckdb_error) = typed_duckdb_error { @@ -2721,6 +2732,10 @@ pub(crate) async fn execute_sql_statement_with_options_typed_with_outcome( let db_type = connection_database_type(state, connection_id).await; let invalidate = crate::object_cache::sql_may_change_object_metadata(sql, db_type); let mut transaction = None; + // Same contract as the multi-statement entry point: a single statement is + // reported with the duration of the whole request, including creating or + // reconnecting the pool, so the summary matches the loading indicator. + let request_start = std::time::Instant::now(); let result = execute_sql_statement_with_options_typed_inner( state, connection_id, @@ -2735,7 +2750,10 @@ pub(crate) async fn execute_sql_statement_with_options_typed_with_outcome( if invalidate { crate::object_cache::invalidate_connection_object_cache(&state.storage, connection_id).await; } - result.map(|result| (result, transaction)) + result.map(|mut result| { + result.execution_time_ms = result.execution_time_ms.max(request_start.elapsed().as_millis()); + (result, transaction) + }) } async fn recover_postgres_create_table_after_connection_error( @@ -3183,6 +3201,11 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed( ) -> Result, QueryExecutionError> { let db_type = connection_database_type(state, connection_id).await; let invalidate = crate::object_cache::sql_may_change_object_metadata(sql, db_type); + // A run that produced a single statement result is reported back as one + // duration, so make it the duration of the whole request: pool creation or + // reconnection for a cold pool happens outside the statement dispatcher and + // would otherwise be invisible while the loading indicator counts it. + let request_start = std::time::Instant::now(); let result = execute_multi_core_with_options_for_client_and_progress_typed_inner( state, connection_id, @@ -3197,7 +3220,14 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed( if invalidate { crate::object_cache::invalidate_connection_object_cache(&state.storage, connection_id).await; } - result + result.map(|mut results| { + if results.len() == 1 { + if let Some(item) = results.first_mut() { + item.result.execution_time_ms = item.result.execution_time_ms.max(request_start.elapsed().as_millis()); + } + } + results + }) } async fn execute_multi_core_with_options_for_client_and_progress_typed_inner( diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index 548e8226c2..a58d8efe07 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -389,6 +389,7 @@ async fn main() { .route("/connection/final-proxy-port", post(routes::connection::connection_final_proxy_port)) .route("/connection/disconnect", post(routes::connection::disconnect_db)) .route("/connection/check-health", post(routes::connection::check_connection_health)) + .route("/connection/prewarm", post(routes::connection::prewarm_connection)) .route("/connection/session-credential-status", post(routes::connection::session_credential_status)) .route("/connection/forget-session-credential", post(routes::connection::forget_session_credential)) .route( diff --git a/crates/dbx-web/src/routes/connection.rs b/crates/dbx-web/src/routes/connection.rs index c755847797..b34b276a56 100644 --- a/crates/dbx-web/src/routes/connection.rs +++ b/crates/dbx-web/src/routes/connection.rs @@ -106,6 +106,15 @@ pub struct DisconnectRequest { pub client_attempt: Option, } +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct PrewarmConnectionRequest { + pub connection_id: String, + pub database: Option, + pub catalog: Option, + pub client_session_id: Option, +} + #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct CloseDatabaseConnectionRequest { @@ -602,6 +611,21 @@ pub async fn check_connection_health( Ok(Json(())) } +pub async fn prewarm_connection( + State(state): State>, + Json(body): Json, +) -> Result, AppError> { + let database = body.database.as_deref().filter(|value| !value.is_empty()); + let catalog = body.catalog.as_deref().filter(|value| !value.is_empty()); + let client_session_id = body.client_session_id.as_deref().filter(|value| !value.is_empty()); + state + .app + .prewarm_connection_pool(&body.connection_id, database, catalog, client_session_id) + .await + .map_err(AppError::from)?; + Ok(Json(())) +} + /// 查询连接在本次运行期是否已输入并暂存密码（`save_password=false`）。 /// 供前端决定是否需要弹密码框；仅返回布尔状态，不泄露密码本身。 /// 按当前登录会话（owner）查询，不会暴露其他会话的凭据状态。 diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index 0a8b02cb46..b93248c24e 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -2268,6 +2268,27 @@ pub async fn check_connection_health(state: State<'_, Arc>, connection state.check_connection_health(&connection_id).await } +/// Warm the driver and connection pool for a connection a tab is opening, so the +/// first Run does not pay pool creation, tunnel setup, or external-driver (JDBC +/// agent) startup while the user waits. Never removes an existing pool. +#[tauri::command] +pub async fn prewarm_connection( + state: State<'_, Arc>, + connection_id: String, + database: Option, + catalog: Option, + client_session_id: Option, +) -> Result<(), String> { + state + .prewarm_connection_pool( + &connection_id, + database.as_deref().filter(|value| !value.is_empty()), + catalog.as_deref().filter(|value| !value.is_empty()), + client_session_id.as_deref().filter(|value| !value.is_empty()), + ) + .await +} + #[tauri::command] pub async fn connection_identifier_quote( state: State<'_, Arc>, diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 73b259fd17..0b4741e848 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -1870,6 +1870,7 @@ pub fn run() { commands::connection::clear_all_session_credentials, commands::connection::refresh_connections, commands::connection::check_connection_health, + commands::connection::prewarm_connection, commands::connection::connection_identifier_quote, commands::connection::connection_database_info, commands::connection::save_connection_database_info,