✨ feat(connection): 支持自定义 SQL 探活

- 新增连接级自定义探活 SQL 配置与多语言表单校验 - 限制单条 SELECT/WITH 并使用可取消的超时查询 - 隔离陈旧探活策略，避免配置错误驱逐健康连接 Fixes #611
2026-08-22 08:53:46 +08:00 · 2026-07-19 12:29:21 +08:00
parent 82c62fd55c
commit 94ad5ae2f6
25 changed files with 1099 additions and 38 deletions
--- 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");
    },
  );

--- 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",
--- 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<ConnectionModalNetworkSecu
    !effectiveUseHttpTunnel &&
    (useProxy || !!form.getFieldValue("useProxy"));
  const keepAliveEnabled = !!Form.useWatch("keepAliveEnabled", form);
+  const connectionDriver = Form.useWatch("driver", form);
+  const oceanBaseProtocol = Form.useWatch("oceanBaseProtocol", form);
+  const keepAliveSQLSupported = supportsConnectionKeepAliveSQL({
+    type: dbType,
+    driver: connectionDriver,
+    oceanBaseProtocol: oceanBaseProtocol,
+  });
  const networkItems: Array<{
    key: "ssl" | "ssh" | "proxy" | "httpTunnel";
    title: string;
@@ -1031,6 +1043,51 @@ const ConnectionModalNetworkSecuritySection: React.FC<ConnectionModalNetworkSecu
            })}
          />
        </Form.Item>
+        {keepAliveSQLSupported ? (
+          <Form.Item
+            name="keepAliveSQL"
+            label={t("connection.modal.network.keepAliveSQL.label")}
+            extra={t("connection.modal.network.keepAliveSQL.help")}
+            rules={[
+              {
+                max: MAX_CONNECTION_KEEPALIVE_SQL_LENGTH,
+                message: t("connection.modal.network.keepAliveSQL.maxLength"),
+              },
+              {
+                validator: (_, value) => {
+                  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 }}
+          >
+            <Input.TextArea
+              {...noAutoCapInputProps}
+              autoSize={{ minRows: 2, maxRows: 4 }}
+              disabled={!keepAliveEnabled}
+              maxLength={MAX_CONNECTION_KEEPALIVE_SQL_LENGTH}
+              placeholder="SELECT 1"
+              showCount
+            />
+          </Form.Item>
+        ) : null}
      </div>
    </div>
  );
--- a/frontend/src/components/connectionModal/ConnectionModalStep2.tsx
+++ b/frontend/src/components/connectionModal/ConnectionModalStep2.tsx
@@ -2590,6 +2590,7 @@ const ConnectionModalStep2: React.FC<ConnectionModalStep2Props> = (props) => {
        timeout: 30,
        keepAliveEnabled: false,
        keepAliveIntervalMinutes: 240,
+        keepAliveSQL: "",
        uri: "",
        connectionParams: "",
        restrictDataEdit: false,
--- 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 () => {
--- 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,
--- 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 () => {
--- 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);
--- 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
--- 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',
--- 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 => {
--- 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', () => {