fix(postgres): correct query and connection timeout handling under SSH tunnel

apps/desktop/src/lib/__tests__/connection/connectionAttemptTimeout.spec.ts

@@ -34,7 +34,7 @@ describe("connectionAttemptTimeout", () => {
     expect(connectionAttemptTimeoutMs({ db_type: "access", connect_timeout_secs: 45, transport_layers: [] })).toBe(47_000);
   });
 
-  it("includes HTTP tunnel connection timeout values", () => {
+  it("adds HTTP tunnel setup and tunneled endpoint probe budgets", () => {
     expect(
       connectionAttemptTimeoutMs({
         db_type: "mysql",
@@ -48,6 +48,25 @@ describe("connectionAttemptTimeout", () => {
           },
         ],
       }),
+    ).toBe(37_000);
+  });
+
+  it("adds SSH setup and tunneled endpoint probe budgets", () => {
+    expect(
+      connectionAttemptTimeoutMs({
+        db_type: "postgres",
+        connect_timeout_secs: 10,
+        transport_layers: [
+          {
+            type: "ssh",
+            id: "ssh",
+            host: "bastion.example.com",
+            port: 22,
+            user: "dbx",
+            connect_timeout_secs: 5,
+          },
+        ],
+      }),
     ).toBe(27_000);
   });
 
@@ -81,7 +100,7 @@ describe("connectionAttemptTimeout", () => {
         },
         (profileId) => (profileId === profile.id ? profile : undefined),
       ),
-    ).toBe(42_000);
+    ).toBe(52_000);
   });
 
   it("keeps disabled shared layers outside the attempt deadline", () => {

apps/desktop/src/lib/__tests__/sql/queryTimeout.spec.ts

@@ -0,0 +1,21 @@
+import { describe, expect, it } from "vitest";
+import { frontendQueryTimeoutSecsForSql, queryTimeoutSecsForConnection } from "@/lib/sql/queryTimeout";
+
+describe("queryTimeout", () => {
+  it("lets PostgreSQL row queries use the backend inactivity timeout", () => {
+    expect(frontendQueryTimeoutSecsForSql("SELECT * FROM sample_records LIMIT 2000", "postgres", 30)).toBe(0);
+    expect(frontendQueryTimeoutSecsForSql("/* page */\nWITH rows AS (SELECT 1) SELECT * FROM rows", "postgres", 30)).toBe(0);
+    expect(frontendQueryTimeoutSecsForSql("UPDATE sample_records SET state = 'ready' RETURNING id", "postgres", 30)).toBe(0);
+  });
+
+  it("keeps the frontend guard for non-row PostgreSQL statements", () => {
+    expect(frontendQueryTimeoutSecsForSql("UPDATE sample_records SET state = 'ready'", "postgres", 30)).toBe(60);
+    expect(frontendQueryTimeoutSecsForSql("INSERT INTO sample_records(note) VALUES ('RETURNING is text')", "postgres", 30)).toBe(60);
+    expect(frontendQueryTimeoutSecsForSql("UPDATE sample_records SET note = 'ready' /* RETURNING */", "postgres", 30)).toBe(60);
+  });
+
+  it("keeps the existing frontend guard for other database types", () => {
+    expect(frontendQueryTimeoutSecsForSql("SELECT * FROM sample_records LIMIT 2000", "mysql", 30)).toBe(60);
+    expect(queryTimeoutSecsForConnection({ query_timeout_secs: undefined })).toBe(60);
+  });
+});

apps/desktop/src/lib/connection/connectionAttemptTimeout.ts

@@ -65,16 +65,22 @@ function resolvedTimeoutLayer(layer: TransportLayerConfig, resolveTunnelProfile?
 export function connectionAttemptTimeoutMs(config: Pick<ConnectionConfig, "connect_timeout_secs" | "transport_layers"> & Partial<Pick<ConnectionConfig, "db_type">>, resolveTunnelProfile?: TunnelProfileResolver): number {
   const baseTimeoutSecs = positiveSeconds(config.connect_timeout_secs, DEFAULT_CONNECT_TIMEOUT_SECS);
   const agentMinTimeoutSecs = config.db_type === "access" ? ACCESS_AGENT_MIN_CONNECT_TIMEOUT_SECS : AGENT_DRIVER_MIN_CONNECT_TIMEOUT_SECS;
-  const timeouts = [DRIVER_STARTUP_FLOOR_TYPES.has(config.db_type as DatabaseType) ? Math.max(baseTimeoutSecs, agentMinTimeoutSecs) : baseTimeoutSecs];
+  let timeoutSecs = DRIVER_STARTUP_FLOOR_TYPES.has(config.db_type as DatabaseType) ? Math.max(baseTimeoutSecs, agentMinTimeoutSecs) : baseTimeoutSecs;
+  let hasEnabledTransportLayer = false;
   for (const unresolvedLayer of config.transport_layers ?? []) {
     const layer = resolvedTimeoutLayer(unresolvedLayer, resolveTunnelProfile);
     if (layer.enabled === false) continue;
+    hasEnabledTransportLayer = true;
     if (layer.type === "ssh" || layer.type === "http_tunnel") {
-      timeouts.push(positiveSeconds(layer.connect_timeout_secs, DEFAULT_CONNECT_TIMEOUT_SECS));
+      timeoutSecs += positiveSeconds(layer.connect_timeout_secs, DEFAULT_CONNECT_TIMEOUT_SECS);
     }
   }
+  // Transport setup, the tunneled endpoint probe, and the database connection
+  // happen sequentially in the backend. Keep the UI guard outside their total
+  // budget so it cannot cancel a connection attempt that is still progressing.
+  if (hasEnabledTransportLayer) timeoutSecs += baseTimeoutSecs;
   const fallbackBuffer = config.db_type === "mongodb" ? MONGO_LEGACY_FALLBACK_TIMEOUT_BUFFER_MS : 0;
-  return Math.ceil(Math.max(...timeouts) * 1000 + CONNECTION_ATTEMPT_TIMEOUT_BUFFER_MS + fallbackBuffer);
+  return Math.ceil(timeoutSecs * 1000 + CONNECTION_ATTEMPT_TIMEOUT_BUFFER_MS + fallbackBuffer);
 }
 
 export function connectionAttemptTimeoutMessage(timeoutMs: number): string {

apps/desktop/src/lib/sql/queryTimeout.ts

@@ -1,8 +1,12 @@
 import { splitSqlStatementRanges } from "@/lib/sql/sqlStatementRanges";
+import { tokenizeSqlSemantic } from "@/lib/sql/semantic/tokens";
 import type { ConnectionConfig, DatabaseType } from "@/types/database";
 
 export const DEFAULT_QUERY_TIMEOUT_SECS = 60;
 
+const POSTGRES_ROW_STATEMENT_KEYWORDS = new Set(["select", "show", "explain", "table", "with"]);
+const POSTGRES_RETURNING_STATEMENT_KEYWORDS = new Set(["insert", "update", "delete", "merge"]);
+
 export function queryTimeoutSecsForConnection(connection?: Pick<ConnectionConfig, "query_timeout_secs"> | null): number {
   const value = Number(connection?.query_timeout_secs);
   return Number.isFinite(value) && value >= 0 ? value : DEFAULT_QUERY_TIMEOUT_SECS;
@@ -11,7 +15,24 @@ export function queryTimeoutSecsForConnection(connection?: Pick<ConnectionConfig
 export function frontendQueryTimeoutSecsForSql(sql: string, databaseType: DatabaseType | undefined, queryTimeoutSecs: number): number {
   if (queryTimeoutSecs === 0) return 0;
 
+  // PostgreSQL applies the configured timeout while the backend receives rows
+  // over a slow SSH tunnel. An absolute browser-side deadline can reject a
+  // successful, steadily progressing page fetch first, so let the backend own
+  // the timeout for row-returning PostgreSQL statements.
+  if (databaseType === "postgres" && postgresQueryMayReturnRows(sql, databaseType)) return 0;
+
   const baseTimeoutSecs = Math.max(queryTimeoutSecs * 2, 60);
   const statementCount = Math.max(splitSqlStatementRanges(sql, databaseType).length, 1);
   return baseTimeoutSecs * statementCount;
 }
+
+function postgresQueryMayReturnRows(sql: string, databaseType: DatabaseType): boolean {
+  return splitSqlStatementRanges(sql, databaseType).some(({ sql: statement }) => {
+    const tokens = tokenizeSqlSemantic(statement, "postgres");
+    const firstKeyword = tokens.find((token) => token.kind === "word" && token.depth === 0)?.normalized;
+    if (!firstKeyword) return false;
+    if (POSTGRES_ROW_STATEMENT_KEYWORDS.has(firstKeyword)) return true;
+    if (!POSTGRES_RETURNING_STATEMENT_KEYWORDS.has(firstKeyword)) return false;
+    return tokens.some((token) => token.kind === "word" && token.depth === 0 && token.normalized === "returning");
+  });
+}