diff --git a/frontend/src/components/ConnectionModal.i18n.test.tsx b/frontend/src/components/ConnectionModal.i18n.test.tsx index 5d90f90e..f9525803 100644 --- a/frontend/src/components/ConnectionModal.i18n.test.tsx +++ b/frontend/src/components/ConnectionModal.i18n.test.tsx @@ -433,6 +433,7 @@ describe("ConnectionModal i18n", () => { expect(pageText).toContain("首选"); expect(pageText).toContain("必需"); expect(pageText).toContain("跳过验证"); + expect(pageText).toContain("自定义探活 SQL"); }, ); diff --git a/frontend/src/components/ConnectionModal.tsx b/frontend/src/components/ConnectionModal.tsx index 52175322..890c0574 100644 --- a/frontend/src/components/ConnectionModal.tsx +++ b/frontend/src/components/ConnectionModal.tsx @@ -1496,6 +1496,7 @@ const ConnectionModal: React.FC<{ Number(config.keepAliveIntervalMinutes) > 0 ? Number(config.keepAliveIntervalMinutes) : DEFAULT_KEEPALIVE_INTERVAL_MINUTES, + keepAliveSQL: config.keepAliveSQL || "", mysqlTopology: mysqlIsReplica ? "replica" : "single", mysqlReplicaHosts: mysqlReplicaHosts, rocketmqTopology: rocketmqIsCluster ? "cluster" : "single", @@ -2104,6 +2105,7 @@ const ConnectionModal: React.FC<{ timeout: 30, keepAliveEnabled: false, keepAliveIntervalMinutes: DEFAULT_KEEPALIVE_INTERVAL_MINUTES, + keepAliveSQL: "", uri: "", connectionParams: "", includeDatabases: undefined, @@ -2178,6 +2180,7 @@ const ConnectionModal: React.FC<{ httpTunnelPassword: "", keepAliveEnabled: false, keepAliveIntervalMinutes: DEFAULT_KEEPALIVE_INTERVAL_MINUTES, + keepAliveSQL: "", mysqlTopology: "single", rocketmqTopology: "single", mqttTopology: "single", @@ -2228,6 +2231,7 @@ const ConnectionModal: React.FC<{ httpTunnelPassword: "", keepAliveEnabled: false, keepAliveIntervalMinutes: DEFAULT_KEEPALIVE_INTERVAL_MINUTES, + keepAliveSQL: "", mysqlTopology: "single", rocketmqTopology: "single", mqttTopology: "single", diff --git a/frontend/src/components/connectionModal/ConnectionModalNetworkSecuritySection.tsx b/frontend/src/components/connectionModal/ConnectionModalNetworkSecuritySection.tsx index ce292da7..3893f349 100644 --- a/frontend/src/components/connectionModal/ConnectionModalNetworkSecuritySection.tsx +++ b/frontend/src/components/connectionModal/ConnectionModalNetworkSecuritySection.tsx @@ -4,6 +4,11 @@ import { Button, Checkbox, Form, Input, InputNumber, Space, Typography } from "a import { t } from "../../i18n"; import { getStoredSecretPlaceholder } from "../../utils/connectionModalPresentation"; import { noAutoCapInputProps } from "../../utils/inputAutoCap"; +import { + isSingleReadOnlyConnectionQuery, + MAX_CONNECTION_KEEPALIVE_SQL_LENGTH, + supportsConnectionKeepAliveSQL, +} from "../../utils/connectionReadOnly"; const { Text } = Typography; const DEFAULT_KEEPALIVE_INTERVAL_MINUTES = 240; @@ -57,6 +62,13 @@ const ConnectionModalNetworkSecuritySection: React.FC + {keepAliveSQLSupported ? ( + { + const sql = String(value || "").trim(); + if ( + !sql || + !keepAliveEnabled || + isSingleReadOnlyConnectionQuery( + { + type: dbType, + driver: connectionDriver, + oceanBaseProtocol: oceanBaseProtocol, + }, + sql, + ) + ) { + return Promise.resolve(); + } + return Promise.reject( + new Error(t("connection.modal.network.keepAliveSQL.readOnly")), + ); + }, + }, + ]} + style={{ marginTop: 12, marginBottom: 0 }} + > + + + ) : null} ); diff --git a/frontend/src/components/connectionModal/ConnectionModalStep2.tsx b/frontend/src/components/connectionModal/ConnectionModalStep2.tsx index df157236..49ddf914 100644 --- a/frontend/src/components/connectionModal/ConnectionModalStep2.tsx +++ b/frontend/src/components/connectionModal/ConnectionModalStep2.tsx @@ -2590,6 +2590,7 @@ const ConnectionModalStep2: React.FC = (props) => { timeout: 30, keepAliveEnabled: false, keepAliveIntervalMinutes: 240, + keepAliveSQL: "", uri: "", connectionParams: "", restrictDataEdit: false, diff --git a/frontend/src/components/connectionModal/connectionModalConfig.keepalive.test.ts b/frontend/src/components/connectionModal/connectionModalConfig.keepalive.test.ts index edb7dfa7..214422c8 100644 --- a/frontend/src/components/connectionModal/connectionModalConfig.keepalive.test.ts +++ b/frontend/src/components/connectionModal/connectionModalConfig.keepalive.test.ts @@ -18,6 +18,7 @@ const buildBaseValues = () => ({ timeout: 30, keepAliveEnabled: false, keepAliveIntervalMinutes: 240, + keepAliveSQL: "", savePassword: true, uri: "", connectionParams: "", @@ -71,6 +72,7 @@ describe("connectionModalConfig keepalive", () => { ...buildBaseValues(), keepAliveEnabled: true, keepAliveIntervalMinutes: 15, + keepAliveSQL: " SELECT 1 ", }, forPersist: true, translate, @@ -78,6 +80,7 @@ describe("connectionModalConfig keepalive", () => { expect(config.keepAliveEnabled).toBe(true); expect(config.keepAliveIntervalMinutes).toBe(15); + expect(config.keepAliveSQL).toBe("SELECT 1"); }); it("forces file database keepalive off", async () => { @@ -96,6 +99,47 @@ describe("connectionModalConfig keepalive", () => { expect(config.keepAliveEnabled).toBe(false); expect(config.keepAliveIntervalMinutes).toBe(15); + expect(config.keepAliveSQL).toBe(""); + }); + + it("keeps custom SQL while disabled and rejects unsafe SQL when enabled", async () => { + const disabledConfig = await buildConnectionConfig({ + values: { + ...buildBaseValues(), + keepAliveEnabled: false, + keepAliveSQL: " SELECT 1 ", + }, + forPersist: true, + translate, + }); + + expect(disabledConfig.keepAliveSQL).toBe("SELECT 1"); + await expect(buildConnectionConfig({ + values: { + ...buildBaseValues(), + keepAliveEnabled: true, + keepAliveSQL: "DELETE FROM accounts", + }, + forPersist: true, + translate, + })).rejects.toThrow("connection.modal.network.keepAliveSQL.readOnly"); + }); + + it("drops stale custom SQL for unsupported datasource types without blocking save", async () => { + const config = await buildConnectionConfig({ + values: { + ...buildBaseValues(), + type: "redis", + port: 6379, + keepAliveEnabled: true, + keepAliveSQL: "DELETE FROM accounts", + }, + forPersist: true, + translate, + }); + + expect(config.keepAliveEnabled).toBe(true); + expect(config.keepAliveSQL).toBe(""); }); it("persists readOnly only for datasource types that support production guard", async () => { diff --git a/frontend/src/components/connectionModal/connectionModalConfig.ts b/frontend/src/components/connectionModal/connectionModalConfig.ts index 6e3d6286..08206f1e 100644 --- a/frontend/src/components/connectionModal/connectionModalConfig.ts +++ b/frontend/src/components/connectionModal/connectionModalConfig.ts @@ -1,7 +1,10 @@ import type { ConnectionConfig, SavedConnection } from "../../types"; import { deriveLegacyConnectionReadOnlyFlag, + isSingleReadOnlyConnectionQuery, + MAX_CONNECTION_KEEPALIVE_SQL_LENGTH, normalizeConnectionProtectionConfig, + supportsConnectionKeepAliveSQL, supportsConnectionReadOnlyMode, } from "../../utils/connectionReadOnly"; import { resolveConnectionSecretDraft } from "../../utils/connectionSecretDraft"; @@ -794,6 +797,37 @@ export const buildConnectionConfig = async ({ MAX_KEEPALIVE_INTERVAL_MINUTES, ) : DEFAULT_KEEPALIVE_INTERVAL_MINUTES; + const keepAliveSQLInput = String(mergedValues.keepAliveSQL || "").trim(); + const keepAliveSQLSupported = supportsConnectionKeepAliveSQL({ + type, + driver: mergedValues.driver, + oceanBaseProtocol: selectedOceanBaseProtocol, + }); + if ( + keepAliveEnabled && + keepAliveSQLSupported && + keepAliveSQLInput.length > MAX_CONNECTION_KEEPALIVE_SQL_LENGTH + ) { + throw new Error(t("connection.modal.network.keepAliveSQL.maxLength")); + } + if ( + keepAliveEnabled && + keepAliveSQLSupported && + keepAliveSQLInput && + !isSingleReadOnlyConnectionQuery( + { + type, + driver: mergedValues.driver, + oceanBaseProtocol: selectedOceanBaseProtocol, + }, + keepAliveSQLInput, + ) + ) { + throw new Error(t("connection.modal.network.keepAliveSQL.readOnly")); + } + const keepAliveSQL = keepAliveSQLSupported + ? keepAliveSQLInput.slice(0, MAX_CONNECTION_KEEPALIVE_SQL_LENGTH) + : ""; const normalizedConnectionParams = supportsConnectionParamsForType(type) ? type === "oceanbase" ? normalizeOceanBaseConnectionParamsText( @@ -845,6 +879,7 @@ export const buildConnectionConfig = async ({ timeout: Number(mergedValues.timeout || 30), keepAliveEnabled: keepAliveEnabled, keepAliveIntervalMinutes: keepAliveIntervalMinutes, + keepAliveSQL: keepAliveSQL, redisDB: Number.isFinite(Number(mergedValues.redisDB)) ? Math.max(0, Math.trunc(Number(mergedValues.redisDB))) : 0, diff --git a/frontend/src/store.test.ts b/frontend/src/store.test.ts index a5c6c931..e57dd26a 100644 --- a/frontend/src/store.test.ts +++ b/frontend/src/store.test.ts @@ -689,6 +689,7 @@ describe('store appearance persistence', () => { user: 'postgres', keepAliveEnabled: true, keepAliveIntervalMinutes: 0, + keepAliveSQL: ' SELECT 1 ', }, }, ]); @@ -696,6 +697,7 @@ describe('store appearance persistence', () => { const config = useStore.getState().connections[0]?.config; expect(config?.keepAliveEnabled).toBe(true); expect(config?.keepAliveIntervalMinutes).toBe(240); + expect(config?.keepAliveSQL).toBe('SELECT 1'); }); it('keeps StarRocks saved connections as independent datasource type', async () => { diff --git a/frontend/src/store.ts b/frontend/src/store.ts index 4b801758..966b3dbd 100644 --- a/frontend/src/store.ts +++ b/frontend/src/store.ts @@ -107,8 +107,10 @@ import { } from "./utils/sqlFileTabDrafts"; import { deriveLegacyConnectionReadOnlyFlag, + MAX_CONNECTION_KEEPALIVE_SQL_LENGTH, normalizeConnectionProtectionConfig, resolveConnectionProtectionConfig, + supportsConnectionKeepAliveSQL, } from "./utils/connectionReadOnly"; import { DEFAULT_QUERY_EDITOR_EDITOR_HEIGHT_RATIO, @@ -914,6 +916,17 @@ const sanitizeConnectionConfig = (value: unknown): ConnectionConfig => { MIN_KEEPALIVE_INTERVAL_MINUTES, MAX_KEEPALIVE_INTERVAL_MINUTES, ), + keepAliveSQL: supportsConnectionKeepAliveSQL({ + type, + driver: toTrimmedString(raw.driver), + oceanBaseProtocol: + raw.oceanBaseProtocol as ConnectionConfig["oceanBaseProtocol"], + }) + ? toTrimmedString(raw.keepAliveSQL).slice( + 0, + MAX_CONNECTION_KEEPALIVE_SQL_LENGTH, + ) + : "", }; const resolvedProtection = resolveConnectionProtectionConfig(safeConfig); diff --git a/frontend/src/types.ts b/frontend/src/types.ts index b2a6649c..ffb5a766 100644 --- a/frontend/src/types.ts +++ b/frontend/src/types.ts @@ -308,6 +308,7 @@ export interface ConnectionConfig { timeout?: number; keepAliveEnabled?: boolean; keepAliveIntervalMinutes?: number; + keepAliveSQL?: string; redisDB?: number; // Redis database index uri?: string; // Connection URI for copy/paste clickHouseProtocol?: "auto" | "http" | "native"; // ClickHouse connection protocol override diff --git a/frontend/src/utils/connectionReadOnly.test.ts b/frontend/src/utils/connectionReadOnly.test.ts index 3e88d487..2a105cc3 100644 --- a/frontend/src/utils/connectionReadOnly.test.ts +++ b/frontend/src/utils/connectionReadOnly.test.ts @@ -6,10 +6,28 @@ import { isConnectionDataImportRestricted, isConnectionScriptExecutionRestricted, isConnectionStructureEditRestricted, + isSingleReadOnlyConnectionQuery, resolveConnectionProtectionConfig, + supportsConnectionKeepAliveSQL, } from './connectionReadOnly'; describe('connectionReadOnly', () => { + it('accepts only one read-only SQL query for custom keepalive', () => { + const config = { type: 'mysql' } as any; + + expect(supportsConnectionKeepAliveSQL(config)).toBe(true); + expect(isSingleReadOnlyConnectionQuery(config, 'SELECT 1')).toBe(true); + expect(isSingleReadOnlyConnectionQuery(config, 'WITH probe AS (SELECT 1) SELECT * FROM probe')).toBe(true); + expect(isSingleReadOnlyConnectionQuery(config, 'SELECT 1; SELECT 2')).toBe(false); + expect(isSingleReadOnlyConnectionQuery(config, 'DELETE FROM accounts')).toBe(false); + expect(isSingleReadOnlyConnectionQuery(config, '/*!50000 DELETE FROM accounts */ SELECT 1')).toBe(false); + expect(isSingleReadOnlyConnectionQuery(config, 'SELECT /*!50000 SQL_NO_CACHE */ 1')).toBe(false); + expect(isSingleReadOnlyConnectionQuery(config, "SELECT ';' AS probe")).toBe(false); + expect(isSingleReadOnlyConnectionQuery(config, 'SELECT 1 /* ; */')).toBe(false); + expect(isSingleReadOnlyConnectionQuery(config, 'SELECT 1; -- probe')).toBe(true); + expect(supportsConnectionKeepAliveSQL({ type: 'redis' } as any)).toBe(false); + }); + it('maps legacy readOnly connections to the full production protection set', () => { expect(resolveConnectionProtectionConfig({ type: 'postgres', diff --git a/frontend/src/utils/connectionReadOnly.ts b/frontend/src/utils/connectionReadOnly.ts index 4e44e062..50e869b4 100644 --- a/frontend/src/utils/connectionReadOnly.ts +++ b/frontend/src/utils/connectionReadOnly.ts @@ -65,6 +65,8 @@ const CONNECTION_READ_ONLY_TYPES = new Set([ "mongodb", ]); +export const MAX_CONNECTION_KEEPALIVE_SQL_LENGTH = 4096; + const SQL_READ_ONLY_KEYWORDS = new Set([ "select", "with", @@ -280,6 +282,167 @@ const isConnectionReadOnlyStatement = ( return isReadOnlySqlStatement(statement, dialect); }; +const KEEPALIVE_EXECUTABLE_COMMENT_DIALECTS = new Set([ + "mysql", + "mariadb", + "oceanbase", + "diros", + "starrocks", +]); + +const skipKeepAliveLineComment = (text: string, start: number): number => { + let index = start; + while (index < text.length && text[index] !== "\n" && text[index] !== "\r") { + index += 1; + } + return index; +}; + +const skipKeepAliveQuotedText = ( + text: string, + start: number, + quote: string, +): number => { + for (let index = start + 1; index < text.length; index += 1) { + if (text[index] === "\\" && index + 1 < text.length) { + index += 1; + continue; + } + if (text[index] !== quote) continue; + if (index + 1 < text.length && text[index + 1] === quote) { + index += 1; + continue; + } + return index + 1; + } + return text.length; +}; + +const hasExecutableKeepAliveComment = ( + text: string, + dbType: string, +): boolean => { + if (!KEEPALIVE_EXECUTABLE_COMMENT_DIALECTS.has(dbType)) return false; + for (let index = 0; index < text.length;) { + const remaining = text.slice(index); + if ( + remaining.startsWith("/*!") || + remaining.slice(0, 4).toLowerCase() === "/*m!" + ) { + return true; + } + const current = text[index]; + if (current === "'" || current === '"' || current === "`") { + index = skipKeepAliveQuotedText(text, index, current); + continue; + } + if (current === "#") { + index = skipKeepAliveLineComment(text, index + 1); + continue; + } + if (remaining.startsWith("--") && isDashLineCommentStart(remaining, dbType)) { + index = skipKeepAliveLineComment(text, index + 2); + continue; + } + if (remaining.startsWith("/*")) { + const end = remaining.indexOf("*/", 2); + index = end >= 0 ? index + end + 2 : text.length; + continue; + } + index += 1; + } + return false; +}; + +const isKeepAliveWhitespace = (character: string): boolean => + character === " " || + character === "\t" || + character === "\n" || + character === "\r" || + character === "\f" || + character === "\v"; + +const containsOnlyTrailingKeepAliveTrivia = ( + text: string, + dbType: string, +): boolean => { + for (let index = 0; index < text.length;) { + const remaining = text.slice(index); + if (isKeepAliveWhitespace(text[index])) { + index += 1; + continue; + } + if (remaining.startsWith("--") && isDashLineCommentStart(remaining, dbType)) { + index = skipKeepAliveLineComment(text, index + 2); + continue; + } + if (text[index] === "#" && supportsHashLineComment(dbType)) { + index = skipKeepAliveLineComment(text, index + 1); + continue; + } + if (remaining.startsWith("/*")) { + const end = remaining.indexOf("*/", 2); + if (end < 0) return false; + index += end + 2; + continue; + } + return false; + } + return true; +}; + +const hasSafeKeepAliveStatementDelimiter = ( + text: string, + dbType: string, +): boolean => { + const asciiCount = text.split(";").length - 1; + const fullWidthCount = text.split("；").length - 1; + if (asciiCount + fullWidthCount === 0) return true; + if (asciiCount + fullWidthCount !== 1) return false; + const delimiter = asciiCount === 1 ? ";" : "；"; + const index = text.indexOf(delimiter); + return containsOnlyTrailingKeepAliveTrivia( + text.slice(index + delimiter.length), + dbType, + ); +}; + +export const supportsConnectionKeepAliveSQL = ( + config: ConnectionReadOnlyLike, +): boolean => { + const type = resolveConnectionReadOnlyType(config); + return ( + type !== "mongodb" && + type !== "sqlite" && + type !== "duckdb" && + CONNECTION_READ_ONLY_TYPES.has(type) + ); +}; + +export const isSingleReadOnlyConnectionQuery = ( + config: ConnectionReadOnlyLike, + sql: string, +): boolean => { + if (!supportsConnectionKeepAliveSQL(config)) return false; + const dialect = resolveConnectionReadOnlyType(config); + const text = String(sql || ""); + if ( + hasExecutableKeepAliveComment(text, dialect) || + !hasSafeKeepAliveStatementDelimiter(text, dialect) + ) { + return false; + } + const statements = findSqlStatementRanges(text, dialect) + .map((range) => range.text.trim()) + .filter(Boolean); + if (statements.length !== 1) return false; + const keyword = extractLeadingSqlKeyword(statements[0], dialect); + return ( + (keyword === "select" || keyword === "with") && + isConnectionReadOnlyStatement(config, statements[0]) + ); +}; + export const supportsConnectionReadOnlyMode = ( config: ConnectionReadOnlyLike, ): boolean => { diff --git a/frontend/src/utils/connectionRpcConfig.test.ts b/frontend/src/utils/connectionRpcConfig.test.ts index 4972d1ab..666426fe 100644 --- a/frontend/src/utils/connectionRpcConfig.test.ts +++ b/frontend/src/utils/connectionRpcConfig.test.ts @@ -12,6 +12,9 @@ describe('buildRpcConnectionConfig', () => { host: 'db.local', port: '5432' as unknown as number, user: 'postgres', + keepAliveEnabled: true, + keepAliveIntervalMinutes: 15, + keepAliveSQL: 'SELECT 1', useSSH: true, ssh: { host: 'bastion.local', @@ -38,6 +41,9 @@ describe('buildRpcConnectionConfig', () => { expect(result.timeout).toBe(120); expect(result.redisDB).toBe(6); expect(result.database).toBe('app'); + expect(result.keepAliveEnabled).toBe(true); + expect(result.keepAliveIntervalMinutes).toBe(15); + expect(result.keepAliveSQL).toBe('SELECT 1'); }); it('preserves ClickHouse protocol override for RPC calls', () => { diff --git a/frontend/wailsjs/go/models.ts b/frontend/wailsjs/go/models.ts index 717e805d..2842a207 100755 --- a/frontend/wailsjs/go/models.ts +++ b/frontend/wailsjs/go/models.ts @@ -1121,6 +1121,7 @@ export namespace connection { timeout?: number; keepAliveEnabled?: boolean; keepAliveIntervalMinutes?: number; + keepAliveSQL?: string; redisDB?: number; redisSentinelMaster?: string; redisSentinelUser?: string; @@ -1174,6 +1175,7 @@ export namespace connection { this.timeout = source["timeout"]; this.keepAliveEnabled = source["keepAliveEnabled"]; this.keepAliveIntervalMinutes = source["keepAliveIntervalMinutes"]; + this.keepAliveSQL = source["keepAliveSQL"]; this.redisDB = source["redisDB"]; this.redisSentinelMaster = source["redisSentinelMaster"]; this.redisSentinelUser = source["redisSentinelUser"]; diff --git a/internal/app/app.go b/internal/app/app.go index d31450fa..bfc17213 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -50,12 +50,17 @@ var ( ) type cachedDatabase struct { - inst db.Database - lastPing time.Time - config connection.ConnectionConfig - keepAliveEnabled bool - keepAliveInterval time.Duration - keepAliveInFlight bool + inst db.Database + lastPing time.Time + lastKeepAliveAt time.Time + config connection.ConnectionConfig + keepAliveEnabled bool + keepAliveInterval time.Duration + keepAliveSQL string + keepAliveDBType string + keepAliveRevision uint64 + keepAliveInFlight bool + keepAliveInFlightRevision uint64 } type cachedConnectFailure struct { @@ -374,6 +379,7 @@ func normalizeCacheKeyConfig(config connection.ConnectionConfig) connection.Conn // keepalive 仅影响后台保活策略，不应参与物理连接复用键。 normalized.KeepAliveEnabled = false normalized.KeepAliveIntervalMinutes = 0 + normalized.KeepAliveSQL = "" normalized.SavePassword = false if !normalized.UseSSH { @@ -918,13 +924,31 @@ func (a *App) getDatabaseWithPing(config connection.ConnectionConfig, forcePing a.mu.RUnlock() if ok { keepAliveEnabled, keepAliveInterval := resolveConnectionKeepAliveSettings(effectiveConfig) - if entry.keepAliveEnabled != keepAliveEnabled || entry.keepAliveInterval != keepAliveInterval { + keepAliveSQL, keepAliveDBType := resolveConnectionKeepAliveSQL(effectiveConfig) + if entry.keepAliveEnabled != keepAliveEnabled || + entry.keepAliveInterval != keepAliveInterval || + entry.keepAliveSQL != keepAliveSQL || + entry.keepAliveDBType != keepAliveDBType { a.mu.Lock() if cur, exists := a.dbCache[key]; exists && cur.inst == entry.inst { - cur.keepAliveEnabled = keepAliveEnabled - cur.keepAliveInterval = keepAliveInterval - if !keepAliveEnabled { - cur.keepAliveInFlight = false + policyChanged := cur.keepAliveEnabled != keepAliveEnabled || + cur.keepAliveInterval != keepAliveInterval || + cur.keepAliveSQL != keepAliveSQL || + cur.keepAliveDBType != keepAliveDBType + if policyChanged { + wasKeepAliveEnabled := cur.keepAliveEnabled + cur.keepAliveRevision = nextConnectionKeepAliveRevision(cur.keepAliveRevision) + cur.keepAliveEnabled = keepAliveEnabled + cur.keepAliveInterval = keepAliveInterval + cur.keepAliveSQL = keepAliveSQL + cur.keepAliveDBType = keepAliveDBType + if !keepAliveEnabled { + cur.keepAliveInFlight = false + cur.keepAliveInFlightRevision = 0 + cur.lastKeepAliveAt = time.Time{} + } else if !wasKeepAliveEnabled || cur.lastKeepAliveAt.IsZero() { + cur.lastKeepAliveAt = time.Now() + } } a.dbCache[key] = cur entry = cur @@ -1014,9 +1038,30 @@ func (a *App) getDatabaseWithPing(config connection.ConnectionConfig, forcePing now := time.Now() keepAliveEnabled, keepAliveInterval := resolveConnectionKeepAliveSettings(effectiveConfig) + keepAliveSQL, keepAliveDBType := resolveConnectionKeepAliveSQL(effectiveConfig) a.mu.Lock() if existing, exists := a.dbCache[key]; exists && existing.inst != nil { + policyChanged := existing.keepAliveEnabled != keepAliveEnabled || + existing.keepAliveInterval != keepAliveInterval || + existing.keepAliveSQL != keepAliveSQL || + existing.keepAliveDBType != keepAliveDBType + if policyChanged { + wasKeepAliveEnabled := existing.keepAliveEnabled + existing.keepAliveRevision = nextConnectionKeepAliveRevision(existing.keepAliveRevision) + existing.keepAliveEnabled = keepAliveEnabled + existing.keepAliveInterval = keepAliveInterval + existing.keepAliveSQL = keepAliveSQL + existing.keepAliveDBType = keepAliveDBType + if !keepAliveEnabled { + existing.keepAliveInFlight = false + existing.keepAliveInFlightRevision = 0 + existing.lastKeepAliveAt = time.Time{} + } else if !wasKeepAliveEnabled || existing.lastKeepAliveAt.IsZero() { + existing.lastKeepAliveAt = now + } + } + a.dbCache[key] = existing a.mu.Unlock() // Prefer existing cached connection to avoid cache racing duplicates. _ = dbInst.Close() @@ -1028,9 +1073,13 @@ func (a *App) getDatabaseWithPing(config connection.ConnectionConfig, forcePing a.dbCache[key] = cachedDatabase{ inst: dbInst, lastPing: now, + lastKeepAliveAt: now, config: normalizeCacheKeyConfig(effectiveConfig), keepAliveEnabled: keepAliveEnabled, keepAliveInterval: keepAliveInterval, + keepAliveSQL: keepAliveSQL, + keepAliveDBType: keepAliveDBType, + keepAliveRevision: 1, } a.mu.Unlock() diff --git a/internal/app/app_cache_key_test.go b/internal/app/app_cache_key_test.go index d9c01ade..f70c9065 100644 --- a/internal/app/app_cache_key_test.go +++ b/internal/app/app_cache_key_test.go @@ -53,10 +53,12 @@ func TestGetCacheKey_IgnoreKeepAliveSettings(t *testing.T) { Database: "app", KeepAliveEnabled: false, KeepAliveIntervalMinutes: 240, + KeepAliveSQL: "SELECT 1", } modified := base modified.KeepAliveEnabled = true modified.KeepAliveIntervalMinutes = 15 + modified.KeepAliveSQL = "SELECT current_timestamp" left := getCacheKey(base) right := getCacheKey(modified) diff --git a/internal/app/app_keepalive.go b/internal/app/app_keepalive.go index 14476a50..d8dcd6be 100644 --- a/internal/app/app_keepalive.go +++ b/internal/app/app_keepalive.go @@ -2,11 +2,15 @@ package app import ( "context" + "errors" + "strings" "time" + "unicode/utf8" "GoNavi-Wails/internal/connection" "GoNavi-Wails/internal/db" "GoNavi-Wails/internal/logger" + "GoNavi-Wails/internal/sqlaudit" ) const ( @@ -14,12 +18,75 @@ const ( minConnectionKeepAliveIntervalMinutes = 1 maxConnectionKeepAliveIntervalMinutes = 1440 connectionKeepAliveScanInterval = 30 * time.Second + connectionKeepAliveQueryTimeout = 30 * time.Second + maxConnectionKeepAliveSQLLength = 4096 +) + +var ( + errInvalidConnectionKeepAliveSQL = errors.New("custom keep-alive SQL must be one SELECT or WITH statement without write operations") + errConnectionKeepAliveQueryContextUnsupported = errors.New("database driver does not support cancellable custom keep-alive SQL") ) type cachedDatabaseKeepAliveTarget struct { - key string - inst db.Database - config connection.ConnectionConfig + key string + inst db.Database + config connection.ConnectionConfig + sql string + dbType string + revision uint64 +} + +func nextConnectionKeepAliveRevision(current uint64) uint64 { + next := current + 1 + if next == 0 { + return 1 + } + return next +} + +func supportsConnectionKeepAliveSQL(config connection.ConnectionConfig) bool { + dbType := resolveDDLDBType(config) + if dbType == "mongodb" || isFileDatabaseType(dbType) { + return false + } + _, supported := connectionReadOnlySupportedTypes[dbType] + return supported +} + +func resolveConnectionKeepAliveSQL(config connection.ConnectionConfig) (string, string) { + if !config.KeepAliveEnabled || !supportsConnectionKeepAliveSQL(config) { + return "", "" + } + return strings.TrimSpace(config.KeepAliveSQL), resolveDDLDBType(config) +} + +func executeConnectionKeepAlive(ctx context.Context, target cachedDatabaseKeepAliveTarget) error { + if strings.TrimSpace(target.sql) == "" { + return target.inst.Ping() + } + if utf8.RuneCountInString(target.sql) > maxConnectionKeepAliveSQLLength || + !isSafeExplainQuery(target.dbType, target.sql) { + return errInvalidConnectionKeepAliveSQL + } + + if ctx == nil { + ctx = context.Background() + } + queryCtx, cancel := context.WithTimeout(ctx, connectionKeepAliveQueryTimeout) + defer cancel() + queryer, ok := target.inst.(interface { + QueryContext(context.Context, string) ([]map[string]interface{}, []string, error) + }) + if !ok { + return errConnectionKeepAliveQueryContextUnsupported + } + _, _, err := queryer.QueryContext(queryCtx, target.sql) + return err +} + +func isConnectionKeepAlivePolicyError(err error) bool { + return errors.Is(err, errInvalidConnectionKeepAliveSQL) || + errors.Is(err, errConnectionKeepAliveQueryContextUnsupported) } func resolveConnectionKeepAliveSettings(config connection.ConnectionConfig) (bool, time.Duration) { @@ -61,7 +128,7 @@ func (a *App) startConnectionKeepAliveLoop() { case <-ctx.Done(): return case now := <-ticker.C: - a.runConnectionKeepAliveTick(now) + a.runConnectionKeepAliveTickContext(ctx, now) } } }() @@ -86,22 +153,103 @@ func (a *App) stopConnectionKeepAliveLoop() { } func (a *App) runConnectionKeepAliveTick(now time.Time) { - for _, target := range a.collectDueConnectionKeepAliveTargets(now) { - if target.inst == nil { + a.runConnectionKeepAliveTickContext(context.Background(), now) +} + +func (a *App) runConnectionKeepAliveTickContext(ctx context.Context, now time.Time) { + if ctx == nil { + ctx = context.Background() + } + targets := a.collectDueConnectionKeepAliveTargets(now) + for index, target := range targets { + if ctx.Err() != nil { + a.releaseCachedDatabaseKeepAliveTargets(targets[index:]) + return + } + if target.inst == nil || !a.isCachedDatabaseKeepAliveTargetCurrent(target) { continue } - if err := target.inst.Ping(); err != nil { + if err := executeConnectionKeepAlive(ctx, target); err != nil { + if ctx.Err() != nil { + a.releaseCachedDatabaseKeepAliveTargets(targets[index:]) + return + } + if isConnectionKeepAlivePolicyError(err) { + if a.markCachedDatabaseKeepAliveSkipped(target, time.Now()) { + logger.Warnf( + "连接自定义保活配置无效，已跳过本次探活：%s 缓存Key=%s 原因=%s", + formatConnSummary(target.config), + shortCacheKey(target.key), + sqlaudit.RedactError(normalizeErrorMessage(err)), + ) + } + continue + } if closed, summary := a.evictCachedDatabaseAfterKeepAliveFailure(target); closed { logger.Warnf( "连接保活失败，已清理缓存连接：%s 缓存Key=%s 原因=%s", summary, shortCacheKey(target.key), - normalizeErrorMessage(err), + sqlaudit.RedactError(normalizeErrorMessage(err)), ) } continue } - a.markCachedDatabaseKeepAliveSuccess(target.key, target.inst, time.Now()) + a.markCachedDatabaseKeepAliveSuccess(target, time.Now()) + } +} + +func cachedDatabaseKeepAliveTargetOwnsInFlight(entry cachedDatabase, target cachedDatabaseKeepAliveTarget) bool { + return entry.inst == target.inst && + entry.keepAliveInFlight && + entry.keepAliveInFlightRevision == target.revision +} + +func cachedDatabaseKeepAliveTargetMatches(entry cachedDatabase, target cachedDatabaseKeepAliveTarget) bool { + return cachedDatabaseKeepAliveTargetOwnsInFlight(entry, target) && + entry.keepAliveEnabled && + entry.keepAliveRevision == target.revision && + entry.keepAliveSQL == target.sql && + entry.keepAliveDBType == target.dbType +} + +func (a *App) isCachedDatabaseKeepAliveTargetCurrent(target cachedDatabaseKeepAliveTarget) bool { + if a == nil { + return false + } + + a.mu.Lock() + defer a.mu.Unlock() + entry, exists := a.dbCache[target.key] + if !exists { + return false + } + if cachedDatabaseKeepAliveTargetMatches(entry, target) { + return true + } + if cachedDatabaseKeepAliveTargetOwnsInFlight(entry, target) { + entry.keepAliveInFlight = false + entry.keepAliveInFlightRevision = 0 + a.dbCache[target.key] = entry + } + return false +} + +func (a *App) releaseCachedDatabaseKeepAliveTargets(targets []cachedDatabaseKeepAliveTarget) { + if a == nil || len(targets) == 0 { + return + } + + a.mu.Lock() + defer a.mu.Unlock() + for _, target := range targets { + entry, exists := a.dbCache[target.key] + if !exists || !cachedDatabaseKeepAliveTargetOwnsInFlight(entry, target) { + continue + } + entry.keepAliveInFlight = false + entry.keepAliveInFlightRevision = 0 + a.dbCache[target.key] = entry } } @@ -118,23 +266,27 @@ func (a *App) collectDueConnectionKeepAliveTargets(now time.Time) []cachedDataba if entry.inst == nil || !entry.keepAliveEnabled || entry.keepAliveInterval <= 0 || entry.keepAliveInFlight { continue } - if !entry.lastPing.IsZero() && now.Sub(entry.lastPing) < entry.keepAliveInterval { + if !entry.lastKeepAliveAt.IsZero() && now.Sub(entry.lastKeepAliveAt) < entry.keepAliveInterval { continue } entry.keepAliveInFlight = true + entry.keepAliveInFlightRevision = entry.keepAliveRevision a.dbCache[key] = entry targets = append(targets, cachedDatabaseKeepAliveTarget{ - key: key, - inst: entry.inst, - config: entry.config, + key: key, + inst: entry.inst, + config: entry.config, + sql: entry.keepAliveSQL, + dbType: entry.keepAliveDBType, + revision: entry.keepAliveRevision, }) } return targets } -func (a *App) markCachedDatabaseKeepAliveSuccess(key string, inst db.Database, pingedAt time.Time) { +func (a *App) markCachedDatabaseKeepAliveSuccess(target cachedDatabaseKeepAliveTarget, pingedAt time.Time) { if a == nil { return } @@ -142,14 +294,50 @@ func (a *App) markCachedDatabaseKeepAliveSuccess(key string, inst db.Database, p a.mu.Lock() defer a.mu.Unlock() - entry, exists := a.dbCache[key] - if !exists || entry.inst != inst { + entry, exists := a.dbCache[target.key] + if !exists { return } + if cachedDatabaseKeepAliveTargetMatches(entry, target) { + entry.keepAliveInFlight = false + entry.keepAliveInFlightRevision = 0 + entry.lastPing = pingedAt + entry.lastKeepAliveAt = pingedAt + a.dbCache[target.key] = entry + return + } + if cachedDatabaseKeepAliveTargetOwnsInFlight(entry, target) { + entry.keepAliveInFlight = false + entry.keepAliveInFlightRevision = 0 + a.dbCache[target.key] = entry + } +} - entry.keepAliveInFlight = false - entry.lastPing = pingedAt - a.dbCache[key] = entry +func (a *App) markCachedDatabaseKeepAliveSkipped(target cachedDatabaseKeepAliveTarget, attemptedAt time.Time) bool { + if a == nil { + return false + } + + a.mu.Lock() + defer a.mu.Unlock() + + entry, exists := a.dbCache[target.key] + if !exists { + return false + } + if cachedDatabaseKeepAliveTargetMatches(entry, target) { + entry.keepAliveInFlight = false + entry.keepAliveInFlightRevision = 0 + entry.lastKeepAliveAt = attemptedAt + a.dbCache[target.key] = entry + return true + } + if cachedDatabaseKeepAliveTargetOwnsInFlight(entry, target) { + entry.keepAliveInFlight = false + entry.keepAliveInFlightRevision = 0 + a.dbCache[target.key] = entry + } + return false } func (a *App) evictCachedDatabaseAfterKeepAliveFailure(target cachedDatabaseKeepAliveTarget) (bool, string) { @@ -164,10 +352,14 @@ func (a *App) evictCachedDatabaseAfterKeepAliveFailure(target cachedDatabaseKeep a.mu.Lock() entry, exists := a.dbCache[target.key] - if exists && entry.inst == target.inst { + if exists && cachedDatabaseKeepAliveTargetMatches(entry, target) { inst = entry.inst summary = formatConnSummary(entry.config) delete(a.dbCache, target.key) + } else if exists && cachedDatabaseKeepAliveTargetOwnsInFlight(entry, target) { + entry.keepAliveInFlight = false + entry.keepAliveInFlightRevision = 0 + a.dbCache[target.key] = entry } a.mu.Unlock() diff --git a/internal/app/app_keepalive_test.go b/internal/app/app_keepalive_test.go index b030e382..7b094cee 100644 --- a/internal/app/app_keepalive_test.go +++ b/internal/app/app_keepalive_test.go @@ -1,20 +1,33 @@ package app import ( + "context" "errors" "testing" "time" "GoNavi-Wails/internal/connection" + "GoNavi-Wails/internal/db" ) type keepAliveRecordingDB struct { - closed int - pings int - pingErr error + closed int + pings int + pingErr error + queries []string + queryErr error + queryContextCalls int + queryContextDeadline bool + queryContextHook func(context.Context, string) error + connectHook func() } -func (f *keepAliveRecordingDB) Connect(config connection.ConnectionConfig) error { return nil } +func (f *keepAliveRecordingDB) Connect(config connection.ConnectionConfig) error { + if f.connectHook != nil { + f.connectHook() + } + return nil +} func (f *keepAliveRecordingDB) Close() error { f.closed++ return nil @@ -24,7 +37,17 @@ func (f *keepAliveRecordingDB) Ping() error { return f.pingErr } func (f *keepAliveRecordingDB) Query(query string) ([]map[string]interface{}, []string, error) { - return nil, nil, nil + f.queries = append(f.queries, query) + return nil, nil, f.queryErr +} +func (f *keepAliveRecordingDB) QueryContext(ctx context.Context, query string) ([]map[string]interface{}, []string, error) { + f.queryContextCalls++ + _, f.queryContextDeadline = ctx.Deadline() + f.queries = append(f.queries, query) + if f.queryContextHook != nil { + return nil, nil, f.queryContextHook(ctx, query) + } + return nil, nil, f.queryErr } func (f *keepAliveRecordingDB) Exec(query string) (int64, error) { return 0, nil } func (f *keepAliveRecordingDB) GetDatabases() ([]string, error) { return nil, nil } @@ -75,6 +98,305 @@ func TestRunConnectionKeepAliveTick_PingsDueCachedConnection(t *testing.T) { if entry.lastPing.IsZero() { t.Fatal("expected keepalive success to update lastPing") } + if entry.lastKeepAliveAt.IsZero() { + t.Fatal("expected keepalive success to update lastKeepAliveAt") + } +} + +func TestRunConnectionKeepAliveTick_ExecutesCustomReadOnlySQL(t *testing.T) { + app := NewApp() + config := connection.ConnectionConfig{ + Type: "mysql", + Host: "db.local", + Port: 3306, + User: "readonly", + KeepAliveEnabled: true, + KeepAliveSQL: " SELECT 1 ", + } + key := getCacheKey(config) + dbInst := &keepAliveRecordingDB{} + + app.dbCache[key] = cachedDatabase{ + inst: dbInst, + lastPing: time.Now().Add(-5 * time.Hour), + config: normalizeCacheKeyConfig(config), + keepAliveEnabled: true, + keepAliveInterval: 4 * time.Hour, + keepAliveSQL: "SELECT 1", + keepAliveDBType: "mysql", + } + + app.runConnectionKeepAliveTick(time.Now()) + + if dbInst.pings != 0 { + t.Fatalf("expected custom SQL to replace Ping, got %d pings", dbInst.pings) + } + if len(dbInst.queries) != 1 || dbInst.queries[0] != "SELECT 1" { + t.Fatalf("expected trimmed custom SQL once, got %#v", dbInst.queries) + } + if dbInst.queryContextCalls != 1 || !dbInst.queryContextDeadline { + t.Fatalf("expected cancellable QueryContext with deadline, calls=%d deadline=%t", dbInst.queryContextCalls, dbInst.queryContextDeadline) + } +} + +func TestResolveConnectionKeepAliveSQL_NormalizesSupportedSQLConnections(t *testing.T) { + sql, dbType := resolveConnectionKeepAliveSQL(connection.ConnectionConfig{ + Type: "mysql", + KeepAliveEnabled: true, + KeepAliveSQL: " SELECT 1 ", + }) + if sql != "SELECT 1" || dbType != "mysql" { + t.Fatalf("expected normalized MySQL keepalive SQL, sql=%q dbType=%q", sql, dbType) + } + + sql, dbType = resolveConnectionKeepAliveSQL(connection.ConnectionConfig{ + Type: "redis", + KeepAliveEnabled: true, + KeepAliveSQL: "SELECT 1", + }) + if sql != "" || dbType != "" { + t.Fatalf("expected non-SQL datasource to ignore custom SQL, sql=%q dbType=%q", sql, dbType) + } +} + +func TestExecuteConnectionKeepAlive_RejectsUnsafeSQL(t *testing.T) { + tests := []string{ + "DELETE FROM accounts", + "SELECT 1; SELECT 2", + "SELECT 1; DELETE FROM accounts", + "/*!50000 DELETE FROM accounts */ SELECT 1", + "SELECT /*!50000 SQL_NO_CACHE */ 1", + "SELECT ';' AS probe", + "SELECT 1 /* ; */", + } + + for _, query := range tests { + t.Run(query, func(t *testing.T) { + dbInst := &keepAliveRecordingDB{} + err := executeConnectionKeepAlive(context.Background(), cachedDatabaseKeepAliveTarget{ + inst: dbInst, + sql: query, + dbType: "mysql", + }) + if !errors.Is(err, errInvalidConnectionKeepAliveSQL) { + t.Fatalf("expected unsafe SQL rejection, got %v", err) + } + if len(dbInst.queries) != 0 || dbInst.pings != 0 { + t.Fatalf("expected unsafe SQL not to reach database, queries=%#v pings=%d", dbInst.queries, dbInst.pings) + } + }) + } +} + +func TestExecuteConnectionKeepAlive_RequiresCancellableQuery(t *testing.T) { + dbInst := &keepAliveRecordingDB{} + dbWithoutQueryContext := struct{ db.Database }{Database: dbInst} + + err := executeConnectionKeepAlive(context.Background(), cachedDatabaseKeepAliveTarget{ + inst: dbWithoutQueryContext, + sql: "SELECT 1", + dbType: "mysql", + }) + if !errors.Is(err, errConnectionKeepAliveQueryContextUnsupported) { + t.Fatalf("expected drivers without QueryContext to be rejected, got %v", err) + } + if len(dbInst.queries) != 0 || dbInst.pings != 0 { + t.Fatalf("expected no unbounded query fallback, queries=%#v pings=%d", dbInst.queries, dbInst.pings) + } +} + +func TestRunConnectionKeepAliveTick_KeepsHealthyConnectionOnPolicyError(t *testing.T) { + app := NewApp() + config := connection.ConnectionConfig{Type: "mysql", Host: "db.local", Port: 3306, User: "readonly"} + key := getCacheKey(config) + dbInst := &keepAliveRecordingDB{} + + app.dbCache[key] = cachedDatabase{ + inst: dbInst, + config: normalizeCacheKeyConfig(config), + keepAliveEnabled: true, + keepAliveInterval: 4 * time.Hour, + keepAliveSQL: "DELETE FROM accounts", + keepAliveDBType: "mysql", + } + + app.runConnectionKeepAliveTick(time.Now()) + + entry, exists := app.dbCache[key] + if !exists || entry.inst != dbInst { + t.Fatal("expected policy error not to evict the healthy cached connection") + } + if entry.keepAliveInFlight || entry.lastKeepAliveAt.IsZero() { + t.Fatalf("expected skipped policy to finish and advance its schedule, inFlight=%t lastKeepAliveAt=%s", entry.keepAliveInFlight, entry.lastKeepAliveAt) + } + if dbInst.closed != 0 || len(dbInst.queries) != 0 || dbInst.pings != 0 { + t.Fatalf("expected invalid policy not to touch the database, closed=%d queries=%#v pings=%d", dbInst.closed, dbInst.queries, dbInst.pings) + } +} + +func TestRunConnectionKeepAliveTickContext_CancelsCustomQuery(t *testing.T) { + app := NewApp() + config := connection.ConnectionConfig{Type: "postgres", Host: "db.local", Port: 5432, User: "readonly"} + key := getCacheKey(config) + queryStarted := make(chan struct{}) + dbInst := &keepAliveRecordingDB{ + queryContextHook: func(ctx context.Context, _ string) error { + close(queryStarted) + <-ctx.Done() + return ctx.Err() + }, + } + app.dbCache[key] = cachedDatabase{ + inst: dbInst, + config: normalizeCacheKeyConfig(config), + keepAliveEnabled: true, + keepAliveInterval: 4 * time.Hour, + keepAliveSQL: "SELECT 1", + keepAliveDBType: "postgres", + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + defer close(done) + app.runConnectionKeepAliveTickContext(ctx, time.Now()) + }() + <-queryStarted + cancel() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("expected custom keepalive query to stop after loop cancellation") + } + entry, exists := app.dbCache[key] + if !exists || entry.inst != dbInst || entry.keepAliveInFlight { + t.Fatalf("expected cancelled keepalive to preserve and release cached connection, exists=%t inFlight=%t", exists, entry.keepAliveInFlight) + } + if dbInst.closed != 0 { + t.Fatalf("expected shutdown cancellation not to evict the connection, closed=%d", dbInst.closed) + } +} + +func TestCachedDatabaseKeepAliveTarget_SkipsQueuedStalePolicy(t *testing.T) { + app := NewApp() + config := connection.ConnectionConfig{Type: "mysql", Host: "db.local", Port: 3306, User: "readonly"} + key := getCacheKey(config) + dbInst := &keepAliveRecordingDB{} + app.dbCache[key] = cachedDatabase{ + inst: dbInst, + config: normalizeCacheKeyConfig(config), + keepAliveEnabled: true, + keepAliveInterval: 4 * time.Hour, + keepAliveSQL: "SELECT 1", + keepAliveDBType: "mysql", + keepAliveRevision: 1, + } + + targets := app.collectDueConnectionKeepAliveTargets(time.Now()) + if len(targets) != 1 { + t.Fatalf("expected one due target, got %d", len(targets)) + } + entry := app.dbCache[key] + entry.keepAliveRevision = nextConnectionKeepAliveRevision(entry.keepAliveRevision) + entry.keepAliveSQL = "SELECT 2" + app.dbCache[key] = entry + + if app.isCachedDatabaseKeepAliveTargetCurrent(targets[0]) { + t.Fatal("expected queued target to become stale after policy update") + } + entry = app.dbCache[key] + if entry.keepAliveInFlight || entry.keepAliveSQL != "SELECT 2" || entry.keepAliveRevision != 2 { + t.Fatalf("expected stale target release without changing new policy, inFlight=%t sql=%q revision=%d", entry.keepAliveInFlight, entry.keepAliveSQL, entry.keepAliveRevision) + } + if len(dbInst.queries) != 0 { + t.Fatalf("expected stale queued SQL not to execute, queries=%#v", dbInst.queries) + } +} + +func TestRunConnectionKeepAliveTick_DoesNotApplyStalePolicyResult(t *testing.T) { + app := NewApp() + config := connection.ConnectionConfig{Type: "postgres", Host: "db.local", Port: 5432, User: "readonly"} + key := getCacheKey(config) + queryStarted := make(chan struct{}) + finishQuery := make(chan struct{}) + dbInst := &keepAliveRecordingDB{ + queryContextHook: func(_ context.Context, _ string) error { + close(queryStarted) + <-finishQuery + return errors.New("old probe failed") + }, + } + lastKeepAliveAt := time.Now().Add(-5 * time.Hour) + app.dbCache[key] = cachedDatabase{ + inst: dbInst, + lastKeepAliveAt: lastKeepAliveAt, + config: normalizeCacheKeyConfig(config), + keepAliveEnabled: true, + keepAliveInterval: 4 * time.Hour, + keepAliveSQL: "SELECT 1", + keepAliveDBType: "postgres", + keepAliveRevision: 1, + } + + done := make(chan struct{}) + go func() { + defer close(done) + app.runConnectionKeepAliveTick(time.Now()) + }() + <-queryStarted + app.mu.Lock() + entry := app.dbCache[key] + entry.keepAliveRevision = nextConnectionKeepAliveRevision(entry.keepAliveRevision) + entry.keepAliveSQL = "SELECT 2" + app.dbCache[key] = entry + app.mu.Unlock() + close(finishQuery) + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("expected stale custom keepalive query to finish") + } + entry, exists := app.dbCache[key] + if !exists || entry.inst != dbInst { + t.Fatal("expected stale query failure not to evict the current cached connection") + } + if entry.keepAliveInFlight || entry.keepAliveSQL != "SELECT 2" || entry.keepAliveRevision != 2 { + t.Fatalf("expected current policy to remain unchanged, inFlight=%t sql=%q revision=%d", entry.keepAliveInFlight, entry.keepAliveSQL, entry.keepAliveRevision) + } + if !entry.lastKeepAliveAt.Equal(lastKeepAliveAt) { + t.Fatalf("expected stale result not to advance current policy schedule, before=%s after=%s", lastKeepAliveAt, entry.lastKeepAliveAt) + } + if dbInst.closed != 0 { + t.Fatalf("expected stale query failure not to close the current connection, closed=%d", dbInst.closed) + } +} + +func TestRunConnectionKeepAliveTick_RemovesFailedCustomSQLConnection(t *testing.T) { + app := NewApp() + config := connection.ConnectionConfig{Type: "mysql", Host: "db.local", Port: 3306, User: "readonly"} + key := getCacheKey(config) + dbInst := &keepAliveRecordingDB{queryErr: errors.New("token expired")} + + app.dbCache[key] = cachedDatabase{ + inst: dbInst, + lastPing: time.Now().Add(-5 * time.Hour), + config: normalizeCacheKeyConfig(config), + keepAliveEnabled: true, + keepAliveInterval: 4 * time.Hour, + keepAliveSQL: "SELECT 1", + keepAliveDBType: "mysql", + } + + app.runConnectionKeepAliveTick(time.Now()) + + if len(dbInst.queries) != 1 || dbInst.pings != 0 { + t.Fatalf("expected one custom query and no Ping, queries=%#v pings=%d", dbInst.queries, dbInst.pings) + } + if dbInst.closed != 1 || len(app.dbCache) != 0 { + t.Fatalf("expected failed custom keepalive to evict and close connection, closed=%d cache=%d", dbInst.closed, len(app.dbCache)) + } } func TestRunConnectionKeepAliveTick_RemovesFailedCachedConnection(t *testing.T) { @@ -121,6 +443,7 @@ func TestGetDatabaseWithPing_UpdatesCachedKeepAliveSettings(t *testing.T) { User: "postgres", KeepAliveEnabled: true, KeepAliveIntervalMinutes: 15, + KeepAliveSQL: " SELECT current_timestamp ", } key := getCacheKey(config) dbInst := &keepAliveRecordingDB{} @@ -146,4 +469,113 @@ func TestGetDatabaseWithPing_UpdatesCachedKeepAliveSettings(t *testing.T) { if entry.keepAliveInterval != 15*time.Minute { t.Fatalf("expected keepalive interval 15m, got %s", entry.keepAliveInterval) } + if entry.keepAliveSQL != "SELECT current_timestamp" { + t.Fatalf("expected cached custom SQL to be updated, got %q", entry.keepAliveSQL) + } + if entry.keepAliveDBType != "postgres" { + t.Fatalf("expected cached custom SQL dialect postgres, got %q", entry.keepAliveDBType) + } +} + +func TestGetDatabaseWithPing_ConcurrentCacheWinnerReceivesKeepAliveSettings(t *testing.T) { + originalNewDatabaseFunc := newDatabaseFunc + originalDriverRuntimeSupportStatusFunc := driverRuntimeSupportStatusFunc + originalVerifyDriverAgentRevisionFunc := verifyDriverAgentRevisionFunc + defer func() { + newDatabaseFunc = originalNewDatabaseFunc + driverRuntimeSupportStatusFunc = originalDriverRuntimeSupportStatusFunc + verifyDriverAgentRevisionFunc = originalVerifyDriverAgentRevisionFunc + }() + driverRuntimeSupportStatusFunc = func(string) (bool, string) { return true, "" } + verifyDriverAgentRevisionFunc = func(connection.ConnectionConfig) error { return nil } + + app := NewApp() + config := connection.ConnectionConfig{ + Type: "postgres", + Host: "db.local", + Port: 5432, + User: "postgres", + KeepAliveEnabled: true, + KeepAliveIntervalMinutes: 15, + KeepAliveSQL: " SELECT current_timestamp ", + } + key := getCacheKey(config) + winner := &keepAliveRecordingDB{} + created := &keepAliveRecordingDB{} + created.connectHook = func() { + app.mu.Lock() + app.dbCache[key] = cachedDatabase{ + inst: winner, + lastPing: time.Now(), + config: normalizeCacheKeyConfig(config), + } + app.mu.Unlock() + } + newDatabaseFunc = func(string) (db.Database, error) { return created, nil } + + inst, err := app.getDatabaseWithPing(config, false) + if err != nil { + t.Fatalf("expected concurrent cache winner lookup to succeed, got %v", err) + } + if inst != winner { + t.Fatal("expected the concurrent cache winner to be reused") + } + if created.closed != 1 { + t.Fatalf("expected duplicate created connection to be closed once, got %d", created.closed) + } + entry := app.dbCache[key] + if !entry.keepAliveEnabled || entry.keepAliveInterval != 15*time.Minute { + t.Fatalf("expected winner keepalive policy to update, enabled=%t interval=%s", entry.keepAliveEnabled, entry.keepAliveInterval) + } + if entry.keepAliveSQL != "SELECT current_timestamp" || entry.keepAliveDBType != "postgres" { + t.Fatalf("expected winner custom SQL policy to update, sql=%q dbType=%q", entry.keepAliveSQL, entry.keepAliveDBType) + } + if entry.lastKeepAliveAt.IsZero() { + t.Fatal("expected winner keepalive schedule to start from the policy update") + } +} + +func TestGetDatabaseWithPing_ForegroundPingDoesNotDelayCustomKeepAliveSQL(t *testing.T) { + originalDriverRuntimeSupportStatusFunc := driverRuntimeSupportStatusFunc + defer func() { + driverRuntimeSupportStatusFunc = originalDriverRuntimeSupportStatusFunc + }() + driverRuntimeSupportStatusFunc = func(string) (bool, string) { return true, "" } + + app := NewApp() + config := connection.ConnectionConfig{ + Type: "postgres", + Host: "db.local", + Port: 5432, + User: "postgres", + KeepAliveEnabled: true, + KeepAliveIntervalMinutes: 15, + KeepAliveSQL: "SELECT 1", + } + key := getCacheKey(config) + dbInst := &keepAliveRecordingDB{} + lastKeepAliveAt := time.Now().Add(-20 * time.Minute) + app.dbCache[key] = cachedDatabase{ + inst: dbInst, + lastPing: time.Now().Add(-5 * time.Minute), + lastKeepAliveAt: lastKeepAliveAt, + config: normalizeCacheKeyConfig(config), + keepAliveEnabled: true, + keepAliveInterval: 15 * time.Minute, + keepAliveSQL: "SELECT 1", + keepAliveDBType: "postgres", + } + + if _, err := app.getDatabaseWithPing(config, false); err != nil { + t.Fatalf("expected foreground cache health Ping to succeed, got %v", err) + } + entry := app.dbCache[key] + if !entry.lastKeepAliveAt.Equal(lastKeepAliveAt) { + t.Fatalf("expected foreground Ping not to move keepalive schedule, before=%s after=%s", lastKeepAliveAt, entry.lastKeepAliveAt) + } + + app.runConnectionKeepAliveTick(time.Now()) + if dbInst.pings != 1 || len(dbInst.queries) != 1 { + t.Fatalf("expected foreground Ping followed by due custom SQL, pings=%d queries=%#v", dbInst.pings, dbInst.queries) + } } diff --git a/internal/connection/types.go b/internal/connection/types.go index e91bc82d..42014dd3 100644 --- a/internal/connection/types.go +++ b/internal/connection/types.go @@ -114,6 +114,7 @@ type ConnectionConfig struct { Timeout int `json:"timeout,omitempty"` // Connection timeout in seconds (default: 30) KeepAliveEnabled bool `json:"keepAliveEnabled,omitempty"` // Enable background keep-alive ping for long-lived cached connections KeepAliveIntervalMinutes int `json:"keepAliveIntervalMinutes,omitempty"` // Keep-alive ping interval in minutes (default: 240) + KeepAliveSQL string `json:"keepAliveSQL,omitempty"` // Optional single SELECT/WITH probe used instead of the driver ping RedisDB int `json:"redisDB,omitempty"` // Redis database index (0-15) RedisSentinelMaster string `json:"redisSentinelMaster,omitempty"` // Redis Sentinel master name RedisSentinelUser string `json:"redisSentinelUser,omitempty"` // Redis Sentinel auth user diff --git a/shared/i18n/de-DE.json b/shared/i18n/de-DE.json index b9e481e4..f34954eb 100644 --- a/shared/i18n/de-DE.json +++ b/shared/i18n/de-DE.json @@ -3527,6 +3527,10 @@ "connection_modal.network.http_tunnel_disabled_hint": "HTTP Tunnel aktivieren, um Tunnel-Host, Port und Zugangsdaten zu konfigurieren.", "connection_modal.network.http_tunnel_mutex_hint": "HTTP Tunnel kann nicht zusammen mit SSH Tunnel oder Proxy verwendet werden.", "connection_modal.network.http_tunnel_panel_description": "Ein HTTP Tunnel-Gateway für diese Verbindung konfigurieren.", + "connection_modal.network.keepAliveSQL.help": "Leer lassen, um den Treiber-Ping zu verwenden. Es ist nur eine SELECT-/WITH-Anweisung zulässig; verwenden Sie eine einfache Abfrage, die nur wenige Daten zurückgibt, und ein Datenbankkonto mit Leseberechtigung. Dieser Wert wird im Klartext mit der Verbindung gespeichert; keine Zugangsdaten eingeben.", + "connection_modal.network.keepAliveSQL.label": "Benutzerdefiniertes Keepalive-SQL", + "connection_modal.network.keepAliveSQL.maxLength": "Benutzerdefiniertes Keepalive-SQL darf höchstens 4096 Zeichen lang sein.", + "connection_modal.network.keepAliveSQL.readOnly": "Benutzerdefiniertes Keepalive-SQL muss aus genau einer SELECT- oder WITH-Anweisung bestehen.", "connection_modal.network.proxy": "Proxy", "connection_modal.network.proxy_disabled_hint": "Proxy aktivieren, um diese Verbindung über SOCKS5 oder HTTP CONNECT zu leiten.", "connection_modal.network.proxy_panel_description": "Proxy-Einstellungen für diese Verbindung konfigurieren.", diff --git a/shared/i18n/en-US.json b/shared/i18n/en-US.json index 5b7b2e75..2cbf5818 100644 --- a/shared/i18n/en-US.json +++ b/shared/i18n/en-US.json @@ -3527,6 +3527,10 @@ "connection_modal.network.http_tunnel_disabled_hint": "Enable HTTP Tunnel to configure tunnel host, port, and credentials.", "connection_modal.network.http_tunnel_mutex_hint": "HTTP Tunnel cannot be used together with SSH Tunnel or proxy.", "connection_modal.network.http_tunnel_panel_description": "Configure an HTTP Tunnel gateway for this connection.", + "connection_modal.network.keepAliveSQL.help": "Leave blank to use the driver Ping. Only one SELECT/WITH statement is allowed; use a lightweight query that returns little data and a database account with read-only permissions. This value is stored in plain text with the connection; do not include credentials.", + "connection_modal.network.keepAliveSQL.label": "Custom keep-alive SQL", + "connection_modal.network.keepAliveSQL.maxLength": "Custom keep-alive SQL cannot exceed 4096 characters.", + "connection_modal.network.keepAliveSQL.readOnly": "Custom keep-alive SQL must be one SELECT or WITH statement.", "connection_modal.network.proxy": "Proxy", "connection_modal.network.proxy_disabled_hint": "Enable proxy to route this connection through SOCKS5 or HTTP CONNECT.", "connection_modal.network.proxy_panel_description": "Configure per-connection proxy settings.", diff --git a/shared/i18n/ja-JP.json b/shared/i18n/ja-JP.json index 4e23b0b8..fa5d0c0f 100644 --- a/shared/i18n/ja-JP.json +++ b/shared/i18n/ja-JP.json @@ -3527,6 +3527,10 @@ "connection_modal.network.http_tunnel_disabled_hint": "HTTP Tunnel を有効にすると Tunnel ホスト、ポート、認証情報を設定できます。", "connection_modal.network.http_tunnel_mutex_hint": "HTTP Tunnel は SSH Tunnel またはプロキシと同時に使用できません。", "connection_modal.network.http_tunnel_panel_description": "この接続の HTTP Tunnel ゲートウェイを設定します。", + "connection_modal.network.keepAliveSQL.help": "空欄の場合はドライバーの Ping を使用します。SELECT/WITH 文を 1 つだけ指定し、返すデータ量が少ない軽量なクエリと読み取り専用のデータベースアカウントを使用してください。この値は接続情報とともに平文で保存されるため、認証情報を入力しないでください。", + "connection_modal.network.keepAliveSQL.label": "カスタムキープアライブ SQL", + "connection_modal.network.keepAliveSQL.maxLength": "カスタムキープアライブ SQL は 4096 文字以内で指定してください。", + "connection_modal.network.keepAliveSQL.readOnly": "カスタムキープアライブ SQL は SELECT または WITH 文 1 つである必要があります。", "connection_modal.network.proxy": "プロキシ", "connection_modal.network.proxy_disabled_hint": "プロキシを有効にすると、この接続は SOCKS5 または HTTP CONNECT 経由で転送されます。", "connection_modal.network.proxy_panel_description": "この接続専用のプロキシ設定を行います。", diff --git a/shared/i18n/messages.ts b/shared/i18n/messages.ts index ac7a8a39..d6cebaa0 100644 --- a/shared/i18n/messages.ts +++ b/shared/i18n/messages.ts @@ -649,9 +649,16 @@ export const messages: Record> = { "仅在跳板机 token 或长连接会话需要定期续期时开启。", "connection.modal.network.keepAliveInterval.label": "探活间隔 (分钟)", "connection.modal.network.keepAliveInterval.help": - "后台会按这个间隔对已建立的缓存连接执行 Ping，默认 240 分钟。", + "后台会按这个间隔对已建立的缓存连接执行 Ping 或自定义探活 SQL，默认 240 分钟。", "connection.modal.network.keepAliveInterval.range": "探活间隔范围: 1-1440 分钟", + "connection.modal.network.keepAliveSQL.label": "自定义探活 SQL", + "connection.modal.network.keepAliveSQL.help": + "留空时使用驱动 Ping；仅允许一条 SELECT/WITH，请使用只返回少量数据的轻量查询和数据库只读账号。配置会随连接明文保存，请勿填写凭证。", + "connection.modal.network.keepAliveSQL.maxLength": + "自定义探活 SQL 不能超过 4096 个字符", + "connection.modal.network.keepAliveSQL.readOnly": + "自定义探活 SQL 仅允许一条 SELECT 或 WITH 语句", "connection.modal.appearance.title": "外观", "connection.modal.appearance.description": "自定义图标与颜色", "connection.modal.appearance.icon": "图标", @@ -1549,9 +1556,16 @@ export const messages: Record> = { "connection.modal.network.keepAliveInterval.label": "Keep-alive interval (minutes)", "connection.modal.network.keepAliveInterval.help": - "GoNavi pings established cached connections at this interval. Default is 240 minutes.", + "GoNavi runs Ping or the custom keep-alive SQL on established cached connections at this interval. Default is 240 minutes.", "connection.modal.network.keepAliveInterval.range": "Keep-alive interval must be between 1 and 1440 minutes.", + "connection.modal.network.keepAliveSQL.label": "Custom keep-alive SQL", + "connection.modal.network.keepAliveSQL.help": + "Leave blank to use the driver Ping. Only one SELECT/WITH statement is allowed; use a lightweight query that returns little data and a database account with read-only permissions. This value is stored in plain text with the connection; do not include credentials.", + "connection.modal.network.keepAliveSQL.maxLength": + "Custom keep-alive SQL cannot exceed 4096 characters.", + "connection.modal.network.keepAliveSQL.readOnly": + "Custom keep-alive SQL must be one SELECT or WITH statement.", "connection.modal.appearance.title": "Appearance", "connection.modal.appearance.description": "Custom icon and color", "connection.modal.appearance.icon": "Icon", diff --git a/shared/i18n/ru-RU.json b/shared/i18n/ru-RU.json index ff8d3094..5f740e76 100644 --- a/shared/i18n/ru-RU.json +++ b/shared/i18n/ru-RU.json @@ -3527,6 +3527,10 @@ "connection_modal.network.http_tunnel_disabled_hint": "Включите HTTP Tunnel, чтобы настроить хост, порт и учетные данные Tunnel.", "connection_modal.network.http_tunnel_mutex_hint": "HTTP Tunnel нельзя использовать вместе с SSH Tunnel или прокси.", "connection_modal.network.http_tunnel_panel_description": "Настройте шлюз HTTP Tunnel для этого подключения.", + "connection_modal.network.keepAliveSQL.help": "Оставьте поле пустым, чтобы использовать Ping драйвера. Допускается только один запрос SELECT/WITH; используйте лёгкий запрос с небольшим результатом и учётную запись базы данных только для чтения. Значение сохраняется в открытом виде вместе с подключением; не указывайте учётные данные.", + "connection_modal.network.keepAliveSQL.label": "Пользовательский SQL для проверки соединения", + "connection_modal.network.keepAliveSQL.maxLength": "Пользовательский SQL для проверки соединения не должен превышать 4096 символов.", + "connection_modal.network.keepAliveSQL.readOnly": "Пользовательский SQL должен содержать только один запрос SELECT или WITH.", "connection_modal.network.proxy": "Прокси", "connection_modal.network.proxy_disabled_hint": "Включите прокси, чтобы направить это подключение через SOCKS5 или HTTP CONNECT.", "connection_modal.network.proxy_panel_description": "Настройте параметры прокси для этого подключения.", diff --git a/shared/i18n/zh-CN.json b/shared/i18n/zh-CN.json index a0699db2..b651021d 100644 --- a/shared/i18n/zh-CN.json +++ b/shared/i18n/zh-CN.json @@ -3527,6 +3527,10 @@ "connection_modal.network.http_tunnel_disabled_hint": "启用 HTTP Tunnel 后可配置 Tunnel 主机、端口和凭据。", "connection_modal.network.http_tunnel_mutex_hint": "HTTP Tunnel 不能与 SSH Tunnel 或代理同时使用。", "connection_modal.network.http_tunnel_panel_description": "为此连接配置 HTTP Tunnel 网关。", + "connection_modal.network.keepAliveSQL.help": "留空时使用驱动 Ping；仅允许一条 SELECT/WITH，请使用只返回少量数据的轻量查询和数据库只读账号。配置会随连接明文保存，请勿填写凭证。", + "connection_modal.network.keepAliveSQL.label": "自定义探活 SQL", + "connection_modal.network.keepAliveSQL.maxLength": "自定义探活 SQL 不能超过 4096 个字符", + "connection_modal.network.keepAliveSQL.readOnly": "自定义探活 SQL 仅允许一条 SELECT 或 WITH 语句", "connection_modal.network.proxy": "代理", "connection_modal.network.proxy_disabled_hint": "启用代理后，此连接会通过 SOCKS5 或 HTTP CONNECT 转发。", "connection_modal.network.proxy_panel_description": "配置此连接的专用代理设置。", diff --git a/shared/i18n/zh-TW.json b/shared/i18n/zh-TW.json index d598cbc1..1686422e 100644 --- a/shared/i18n/zh-TW.json +++ b/shared/i18n/zh-TW.json @@ -3527,6 +3527,10 @@ "connection_modal.network.http_tunnel_disabled_hint": "啟用 HTTP Tunnel 後可設定 Tunnel 主機、連接埠和憑據。", "connection_modal.network.http_tunnel_mutex_hint": "HTTP Tunnel 不能與 SSH Tunnel 或代理同時使用。", "connection_modal.network.http_tunnel_panel_description": "為此連線設定 HTTP Tunnel 閘道。", + "connection_modal.network.keepAliveSQL.help": "留空時使用驅動程式 Ping；僅允許一條 SELECT/WITH，請使用只回傳少量資料的輕量查詢和資料庫唯讀帳號。此設定會隨連線以純文字儲存，請勿填入憑證。", + "connection_modal.network.keepAliveSQL.label": "自訂探活 SQL", + "connection_modal.network.keepAliveSQL.maxLength": "自訂探活 SQL 不得超過 4096 個字元。", + "connection_modal.network.keepAliveSQL.readOnly": "自訂探活 SQL 僅允許單一 SELECT 或 WITH 陳述式。", "connection_modal.network.proxy": "代理", "connection_modal.network.proxy_disabled_hint": "啟用代理後，此連線會透過 SOCKS5 或 HTTP CONNECT 轉發。", "connection_modal.network.proxy_panel_description": "設定此連線的專用代理設定。",