@bytebase/dbhub 0.16.1 → 0.18.0

package/dist/{chunk-IBBG4PSO.js → chunk-WCXOWHL3.js}

@@ -296,7 +296,9 @@ var SSHTunnel = class {
           privateKey,
           targetConfig.passphrase,
           previousStream,
-          `jump host ${i + 1}`
+          `jump host ${i + 1}`,
+          targetConfig.keepaliveInterval,
+          targetConfig.keepaliveCountMax
         );
         console.error(`  \u2192 Forwarding through ${jumpHost.host}:${jumpHost.port} to ${nextHost.host}:${nextHost.port}`);
         forwardStream = await this.forwardTo(client, nextHost.host, nextHost.port);
@@ -322,7 +324,9 @@ var SSHTunnel = class {
       privateKey,
       targetConfig.passphrase,
       previousStream,
-      jumpHosts.length > 0 ? "target host" : void 0
+      jumpHosts.length > 0 ? "target host" : void 0,
+      targetConfig.keepaliveInterval,
+      targetConfig.keepaliveCountMax
     );
     this.sshClients.push(finalClient);
     return finalClient;
@@ -330,7 +334,7 @@ var SSHTunnel = class {
   /**
    * Connect to a single SSH host.
    */
-  connectToHost(hostInfo, password, privateKey, passphrase, sock, label) {
+  connectToHost(hostInfo, password, privateKey, passphrase, sock, label, keepaliveInterval, keepaliveCountMax) {
     return new Promise((resolve, reject) => {
       const client = new Client();
       const sshConfig = {
@@ -350,6 +354,17 @@ var SSHTunnel = class {
       if (sock) {
         sshConfig.sock = sock;
       }
+      if (keepaliveInterval !== void 0) {
+        if (Number.isNaN(keepaliveInterval) || keepaliveInterval < 0) {
+          const desc = label || `${hostInfo.host}:${hostInfo.port}`;
+          console.warn(
+            `Invalid SSH keepaliveInterval (${keepaliveInterval}) for ${desc}; keepalive configuration will be ignored.`
+          );
+        } else if (keepaliveInterval > 0) {
+          sshConfig.keepaliveInterval = keepaliveInterval * 1e3;
+          sshConfig.keepaliveCountMax = keepaliveCountMax ?? 3;
+        }
+      }
       const onError = (err) => {
         client.removeListener("ready", onReady);
         client.destroy();
@@ -1011,6 +1026,27 @@ function resolveSSHConfig() {
     config.proxyJump = process.env.SSH_PROXY_JUMP;
     sources.push("SSH_PROXY_JUMP from environment");
   }
+  const parseNonNegativeInteger = (value, name) => {
+    const parsed = Number(value);
+    if (!Number.isInteger(parsed) || parsed < 0) {
+      throw new Error(`Invalid value for ${name}: "${value}". Expected a non-negative integer.`);
+    }
+    return parsed;
+  };
+  if (args["ssh-keepalive-interval"]) {
+    config.keepaliveInterval = parseNonNegativeInteger(args["ssh-keepalive-interval"], "ssh-keepalive-interval");
+    sources.push("ssh-keepalive-interval from command line");
+  } else if (process.env.SSH_KEEPALIVE_INTERVAL) {
+    config.keepaliveInterval = parseNonNegativeInteger(process.env.SSH_KEEPALIVE_INTERVAL, "SSH_KEEPALIVE_INTERVAL");
+    sources.push("SSH_KEEPALIVE_INTERVAL from environment");
+  }
+  if (args["ssh-keepalive-count-max"]) {
+    config.keepaliveCountMax = parseNonNegativeInteger(args["ssh-keepalive-count-max"], "ssh-keepalive-count-max");
+    sources.push("ssh-keepalive-count-max from command line");
+  } else if (process.env.SSH_KEEPALIVE_COUNT_MAX) {
+    config.keepaliveCountMax = parseNonNegativeInteger(process.env.SSH_KEEPALIVE_COUNT_MAX, "SSH_KEEPALIVE_COUNT_MAX");
+    sources.push("SSH_KEEPALIVE_COUNT_MAX from environment");
+  }
   if (!config.host || !config.username) {
     throw new Error("SSH tunnel configuration requires at least --ssh-host and --ssh-user");
   }
@@ -1090,6 +1126,8 @@ async function resolveSourceConfigs() {
       source.ssh_password = sshResult.config.password;
       source.ssh_key = sshResult.config.privateKey;
       source.ssh_passphrase = sshResult.config.passphrase;
+      source.ssh_keepalive_interval = sshResult.config.keepaliveInterval;
+      source.ssh_keepalive_count_max = sshResult.config.keepaliveCountMax;
     }
     if (dsnResult.isDemo) {
       const { getSqliteInMemorySetupSql } = await import("./demo-loader-PSMTLZ2T.js");
@@ -1270,6 +1308,31 @@ function validateSourceConfig(source, configPath) {
       );
     }
   }
+  if (source.aws_iam_auth !== void 0 && typeof source.aws_iam_auth !== "boolean") {
+    throw new Error(
+      `Configuration file ${configPath}: source '${source.id}' has invalid aws_iam_auth. Must be a boolean (true or false).`
+    );
+  }
+  if (source.aws_region !== void 0) {
+    if (typeof source.aws_region !== "string" || source.aws_region.trim().length === 0) {
+      throw new Error(
+        `Configuration file ${configPath}: source '${source.id}' has invalid aws_region. Must be a non-empty string (e.g., "eu-west-1").`
+      );
+    }
+  }
+  if (source.aws_iam_auth === true) {
+    const validIamTypes = ["postgres", "mysql", "mariadb"];
+    if (!source.type || !validIamTypes.includes(source.type)) {
+      throw new Error(
+        `Configuration file ${configPath}: source '${source.id}' has aws_iam_auth enabled, but this is only supported for postgres, mysql, and mariadb sources.`
+      );
+    }
+    if (!source.aws_region) {
+      throw new Error(
+        `Configuration file ${configPath}: source '${source.id}' has aws_iam_auth enabled but aws_region is not specified.`
+      );
+    }
+  }
   if (source.connection_timeout !== void 0) {
     if (typeof source.connection_timeout !== "number" || source.connection_timeout <= 0) {
       throw new Error(
@@ -1339,6 +1402,18 @@ function validateSourceConfig(source, configPath) {
       );
     }
   }
+  if (source.search_path !== void 0) {
+    if (source.type !== "postgres") {
+      throw new Error(
+        `Configuration file ${configPath}: source '${source.id}' has 'search_path' but it is only supported for PostgreSQL sources.`
+      );
+    }
+    if (typeof source.search_path !== "string" || source.search_path.trim().length === 0) {
+      throw new Error(
+        `Configuration file ${configPath}: source '${source.id}' has invalid search_path. Must be a non-empty string of comma-separated schema names (e.g., "myschema,public").`
+      );
+    }
+  }
   if (source.readonly !== void 0) {
     throw new Error(
       `Configuration file ${configPath}: source '${source.id}' has 'readonly' field, but readonly must be configured per-tool, not per-source. Move 'readonly' to [[tools]] configuration instead.`
@@ -1408,7 +1483,8 @@ function buildDSNFromSource(source) {
     }
     return `sqlite:///${source.database}`;
   }
-  const passwordRequired = source.authentication !== "azure-active-directory-access-token";
+  const isAwsIamPasswordless = source.aws_iam_auth === true && ["postgres", "mysql", "mariadb"].includes(source.type);
+  const passwordRequired = source.authentication !== "azure-active-directory-access-token" && !isAwsIamPasswordless;
   if (!source.host || !source.user || !source.database) {
     throw new Error(
       `Source '${source.id}': missing required connection parameters. Required: type, host, user, database`
@@ -1416,7 +1492,7 @@ function buildDSNFromSource(source) {
   }
   if (passwordRequired && !source.password) {
     throw new Error(
-      `Source '${source.id}': password is required. (Password is optional only for azure-active-directory-access-token authentication)`
+      `Source '${source.id}': password is required. (Password is optional for azure-active-directory-access-token authentication or when aws_iam_auth=true)`
     );
   }
   const port = source.port || getDefaultPortForType(source.type);
@@ -1448,8 +1524,21 @@ function buildDSNFromSource(source) {
   return dsn;
 }
 
+// src/utils/aws-rds-signer.ts
+import { Signer } from "@aws-sdk/rds-signer";
+async function generateRdsAuthToken(params) {
+  const signer = new Signer({
+    hostname: params.hostname,
+    port: params.port,
+    username: params.username,
+    region: params.region
+  });
+  return signer.getAuthToken();
+}
+
 // src/connectors/manager.ts
 var managerInstance = null;
+var AWS_IAM_TOKEN_REFRESH_MS = 14 * 60 * 1e3;
 var ConnectorManager = class {
   // Prevent race conditions
   constructor() {
@@ -1460,6 +1549,8 @@ var ConnectorManager = class {
     // Store original source configs
     this.sourceIds = [];
     // Ordered list of source IDs (first is default)
+    this.iamRefreshTimers = /* @__PURE__ */ new Map();
+    this.isDisconnecting = false;
     // Lazy connection support
     this.lazySources = /* @__PURE__ */ new Map();
     // Sources pending lazy connection
@@ -1549,7 +1640,7 @@ var ConnectorManager = class {
    */
   async connectSource(source) {
     const sourceId = source.id;
-    const dsn = buildDSNFromSource(source);
+    const dsn = await this.buildConnectionDSN(source);
     console.error(`  - ${sourceId}: ${redactDSN(dsn)}`);
     let actualDSN = dsn;
     if (source.ssh_host) {
@@ -1565,7 +1656,9 @@ var ConnectorManager = class {
         password: source.ssh_password,
         privateKey: source.ssh_key,
         passphrase: source.ssh_passphrase,
-        proxyJump: source.ssh_proxy_jump
+        proxyJump: source.ssh_proxy_jump,
+        keepaliveInterval: source.ssh_keepalive_interval,
+        keepaliveCountMax: source.ssh_keepalive_count_max
       };
       if (!sshConfig.password && !sshConfig.privateKey) {
         throw new Error(
@@ -1606,17 +1699,26 @@ var ConnectorManager = class {
     if (source.readonly !== void 0) {
       config.readonly = source.readonly;
     }
+    if (source.search_path) {
+      config.searchPath = source.search_path;
+    }
     await connector.connect(actualDSN, source.init_script, config);
     this.connectors.set(sourceId, connector);
     if (!this.sourceIds.includes(sourceId)) {
       this.sourceIds.push(sourceId);
     }
     this.sourceConfigs.set(sourceId, source);
+    this.scheduleIamRefresh(source);
   }
   /**
    * Close all database connections
    */
   async disconnect() {
+    this.isDisconnecting = true;
+    for (const timer of this.iamRefreshTimers.values()) {
+      clearTimeout(timer);
+    }
+    this.iamRefreshTimers.clear();
     for (const [sourceId, connector] of this.connectors.entries()) {
       try {
         await connector.disconnect();
@@ -1638,6 +1740,7 @@ var ConnectorManager = class {
     this.lazySources.clear();
     this.pendingConnections.clear();
     this.sourceIds = [];
+    this.isDisconnecting = false;
   }
   /**
    * Get a connector by source ID
@@ -1738,63 +1841,317 @@ var ConnectorManager = class {
     }
     return getDefaultPortForType(type) ?? 0;
   }
+  scheduleIamRefresh(source) {
+    if (this.isDisconnecting) {
+      return;
+    }
+    const sourceId = source.id;
+    const existingTimer = this.iamRefreshTimers.get(sourceId);
+    if (existingTimer) {
+      clearTimeout(existingTimer);
+      this.iamRefreshTimers.delete(sourceId);
+    }
+    if (!source.aws_iam_auth) {
+      return;
+    }
+    const timer = setTimeout(async () => {
+      if (this.isDisconnecting) {
+        return;
+      }
+      try {
+        await this.refreshIamSourceConnection(source);
+      } catch (error) {
+        console.error(
+          `Error refreshing AWS IAM auth token for source '${sourceId}':`,
+          error
+        );
+      } finally {
+        if (!this.isDisconnecting && this.sourceConfigs.has(sourceId)) {
+          this.scheduleIamRefresh(source);
+        }
+      }
+    }, AWS_IAM_TOKEN_REFRESH_MS);
+    timer.unref?.();
+    this.iamRefreshTimers.set(sourceId, timer);
+  }
+  async refreshIamSourceConnection(source) {
+    const sourceId = source.id;
+    if (this.isDisconnecting || !source.aws_iam_auth || !this.connectors.has(sourceId)) {
+      return;
+    }
+    console.error(`Refreshing AWS IAM auth connection for source '${sourceId}'...`);
+    const existingConnector = this.connectors.get(sourceId);
+    if (existingConnector) {
+      await existingConnector.disconnect();
+      this.connectors.delete(sourceId);
+    }
+    const existingTunnel = this.sshTunnels.get(sourceId);
+    if (existingTunnel) {
+      await existingTunnel.close();
+      this.sshTunnels.delete(sourceId);
+    }
+    if (this.isDisconnecting) {
+      return;
+    }
+    await this.connectSource(source);
+  }
+  /**
+   * Build a connection DSN, optionally replacing password with
+   * an AWS RDS IAM auth token when aws_iam_auth is enabled.
+   */
+  async buildConnectionDSN(source) {
+    const dsn = buildDSNFromSource(source);
+    if (!source.aws_iam_auth) {
+      return dsn;
+    }
+    const supportedIamTypes = ["postgres", "mysql", "mariadb"];
+    if (!source.type || !supportedIamTypes.includes(source.type)) {
+      throw new Error(
+        `Source '${source.id}': aws_iam_auth is only supported for postgres, mysql, and mariadb`
+      );
+    }
+    if (!source.aws_region) {
+      throw new Error(
+        `Source '${source.id}': aws_region is required when aws_iam_auth is enabled`
+      );
+    }
+    const parsed = new SafeURL(dsn);
+    const hostname = parsed.hostname;
+    const username = source.user || parsed.username;
+    const defaultPort = getDefaultPortForType(source.type);
+    const port = parsed.port ? parseInt(parsed.port) : defaultPort;
+    if (!hostname || !username || !port) {
+      throw new Error(
+        `Source '${source.id}': unable to resolve host, username, or port for AWS IAM authentication`
+      );
+    }
+    const token = await generateRdsAuthToken({
+      hostname,
+      port,
+      username,
+      region: source.aws_region
+    });
+    const queryParams = new Map(parsed.searchParams);
+    queryParams.set("sslmode", "require");
+    const protocol = parsed.protocol.endsWith(":") ? parsed.protocol.slice(0, -1) : parsed.protocol;
+    const encodedUser = encodeURIComponent(username);
+    const encodedToken = encodeURIComponent(token);
+    const path3 = parsed.pathname || "/";
+    const query = Array.from(queryParams.entries()).map(
+      ([key, value]) => `${encodeURIComponent(key)}=${encodeURIComponent(value)}`
+    ).join("&");
+    return `${protocol}://${encodedUser}:${encodedToken}@${hostname}:${port}${path3}${query ? `?${query}` : ""}`;
+  }
 };
 
 // src/utils/sql-parser.ts
-function stripCommentsAndStrings(sql) {
-  let result = "";
+var TokenType = { Plain: 0, Comment: 1, QuotedBlock: 2 };
+function plainToken(i) {
+  return { type: TokenType.Plain, end: i + 1 };
+}
+function scanSingleLineComment(sql, i) {
+  if (sql[i] !== "-" || sql[i + 1] !== "-") {
+    return null;
+  }
+  let j = i;
+  while (j < sql.length && sql[j] !== "\n") {
+    j++;
+  }
+  return { type: TokenType.Comment, end: j };
+}
+function scanMultiLineComment(sql, i) {
+  if (sql[i] !== "/" || sql[i + 1] !== "*") {
+    return null;
+  }
+  let j = i + 2;
+  while (j < sql.length && !(sql[j] === "*" && sql[j + 1] === "/")) {
+    j++;
+  }
+  if (j < sql.length) {
+    j += 2;
+  }
+  return { type: TokenType.Comment, end: j };
+}
+function scanNestedMultiLineComment(sql, i) {
+  if (sql[i] !== "/" || sql[i + 1] !== "*") {
+    return null;
+  }
+  let j = i + 2;
+  let depth = 1;
+  while (j < sql.length && depth > 0) {
+    if (sql[j] === "/" && sql[j + 1] === "*") {
+      depth++;
+      j += 2;
+    } else if (sql[j] === "*" && sql[j + 1] === "/") {
+      depth--;
+      j += 2;
+    } else {
+      j++;
+    }
+  }
+  return { type: TokenType.Comment, end: j };
+}
+function scanSingleQuotedString(sql, i) {
+  if (sql[i] !== "'") {
+    return null;
+  }
+  let j = i + 1;
+  while (j < sql.length) {
+    if (sql[j] === "'" && sql[j + 1] === "'") {
+      j += 2;
+    } else if (sql[j] === "'") {
+      j++;
+      break;
+    } else {
+      j++;
+    }
+  }
+  return { type: TokenType.QuotedBlock, end: j };
+}
+function scanDoubleQuotedString(sql, i) {
+  if (sql[i] !== '"') {
+    return null;
+  }
+  let j = i + 1;
+  while (j < sql.length) {
+    if (sql[j] === '"' && sql[j + 1] === '"') {
+      j += 2;
+    } else if (sql[j] === '"') {
+      j++;
+      break;
+    } else {
+      j++;
+    }
+  }
+  return { type: TokenType.QuotedBlock, end: j };
+}
+var dollarQuoteOpenRegex = /^\$([a-zA-Z_]\w*)?\$/;
+function scanDollarQuotedBlock(sql, i) {
+  if (sql[i] !== "$") {
+    return null;
+  }
+  const next = sql[i + 1];
+  if (next >= "0" && next <= "9") {
+    return null;
+  }
+  const remaining = sql.substring(i);
+  const m = dollarQuoteOpenRegex.exec(remaining);
+  if (!m) {
+    return null;
+  }
+  const tag = m[0];
+  const bodyStart = i + tag.length;
+  const closeIdx = sql.indexOf(tag, bodyStart);
+  const end = closeIdx !== -1 ? closeIdx + tag.length : sql.length;
+  return { type: TokenType.QuotedBlock, end };
+}
+function scanBacktickQuotedIdentifier(sql, i) {
+  if (sql[i] !== "`") {
+    return null;
+  }
+  let j = i + 1;
+  while (j < sql.length) {
+    if (sql[j] === "`" && sql[j + 1] === "`") {
+      j += 2;
+    } else if (sql[j] === "`") {
+      j++;
+      break;
+    } else {
+      j++;
+    }
+  }
+  return { type: TokenType.QuotedBlock, end: j };
+}
+function scanBracketQuotedIdentifier(sql, i) {
+  if (sql[i] !== "[") {
+    return null;
+  }
+  let j = i + 1;
+  while (j < sql.length) {
+    if (sql[j] === "]" && sql[j + 1] === "]") {
+      j += 2;
+    } else if (sql[j] === "]") {
+      j++;
+      break;
+    } else {
+      j++;
+    }
+  }
+  return { type: TokenType.QuotedBlock, end: j };
+}
+function scanTokenAnsi(sql, i) {
+  return scanSingleLineComment(sql, i) ?? scanMultiLineComment(sql, i) ?? scanSingleQuotedString(sql, i) ?? scanDoubleQuotedString(sql, i) ?? plainToken(i);
+}
+function scanTokenPostgres(sql, i) {
+  return scanSingleLineComment(sql, i) ?? scanNestedMultiLineComment(sql, i) ?? scanSingleQuotedString(sql, i) ?? scanDoubleQuotedString(sql, i) ?? scanDollarQuotedBlock(sql, i) ?? plainToken(i);
+}
+function scanTokenMySQL(sql, i) {
+  return scanSingleLineComment(sql, i) ?? scanMultiLineComment(sql, i) ?? scanSingleQuotedString(sql, i) ?? scanDoubleQuotedString(sql, i) ?? scanBacktickQuotedIdentifier(sql, i) ?? plainToken(i);
+}
+function scanTokenSQLite(sql, i) {
+  return scanSingleLineComment(sql, i) ?? scanMultiLineComment(sql, i) ?? scanSingleQuotedString(sql, i) ?? scanDoubleQuotedString(sql, i) ?? scanBacktickQuotedIdentifier(sql, i) ?? scanBracketQuotedIdentifier(sql, i) ?? plainToken(i);
+}
+function scanTokenSQLServer(sql, i) {
+  return scanSingleLineComment(sql, i) ?? scanMultiLineComment(sql, i) ?? scanSingleQuotedString(sql, i) ?? scanDoubleQuotedString(sql, i) ?? scanBracketQuotedIdentifier(sql, i) ?? plainToken(i);
+}
+var dialectScanners = {
+  postgres: scanTokenPostgres,
+  mysql: scanTokenMySQL,
+  mariadb: scanTokenMySQL,
+  sqlite: scanTokenSQLite,
+  sqlserver: scanTokenSQLServer
+};
+function getScanner(dialect) {
+  return dialect ? dialectScanners[dialect] ?? scanTokenAnsi : scanTokenAnsi;
+}
+function stripCommentsAndStrings(sql, dialect) {
+  const scanToken = getScanner(dialect);
+  const parts = [];
+  let plainStart = -1;
   let i = 0;
   while (i < sql.length) {
-    if (sql[i] === "-" && sql[i + 1] === "-") {
-      while (i < sql.length && sql[i] !== "\n") {
-        i++;
+    const token = scanToken(sql, i);
+    if (token.type === TokenType.Plain) {
+      if (plainStart === -1) {
+        plainStart = i;
       }
-      result += " ";
-      continue;
-    }
-    if (sql[i] === "/" && sql[i + 1] === "*") {
-      i += 2;
-      while (i < sql.length && !(sql[i] === "*" && sql[i + 1] === "/")) {
-        i++;
+    } else {
+      if (plainStart !== -1) {
+        parts.push(sql.substring(plainStart, i));
+        plainStart = -1;
       }
-      i += 2;
-      result += " ";
-      continue;
+      parts.push(" ");
     }
-    if (sql[i] === "'") {
-      i++;
-      while (i < sql.length) {
-        if (sql[i] === "'" && sql[i + 1] === "'") {
-          i += 2;
-        } else if (sql[i] === "'") {
-          i++;
-          break;
-        } else {
-          i++;
-        }
+    i = token.end;
+  }
+  if (plainStart !== -1) {
+    parts.push(sql.substring(plainStart));
+  }
+  return parts.join("");
+}
+function splitSQLStatements(sql, dialect) {
+  const scanToken = getScanner(dialect);
+  const statements = [];
+  let stmtStart = 0;
+  let i = 0;
+  while (i < sql.length) {
+    if (sql[i] === ";") {
+      const trimmed2 = sql.substring(stmtStart, i).trim();
+      if (trimmed2.length > 0) {
+        statements.push(trimmed2);
       }
-      result += " ";
-      continue;
-    }
-    if (sql[i] === '"') {
+      stmtStart = i + 1;
       i++;
-      while (i < sql.length) {
-        if (sql[i] === '"' && sql[i + 1] === '"') {
-          i += 2;
-        } else if (sql[i] === '"') {
-          i++;
-          break;
-        } else {
-          i++;
-        }
-      }
-      result += " ";
       continue;
     }
-    result += sql[i];
-    i++;
+    const token = scanToken(sql, i);
+    i = token.end;
+  }
+  const trimmed = sql.substring(stmtStart).trim();
+  if (trimmed.length > 0) {
+    statements.push(trimmed);
   }
-  return result;
+  return statements;
 }
 
 // src/utils/parameter-mapper.ts
@@ -2130,12 +2487,15 @@ export {
   getDatabaseTypeFromDSN,
   getDefaultPortForType,
   stripCommentsAndStrings,
+  splitSQLStatements,
   isDemoMode,
   resolveTransport,
   resolvePort,
   resolveSourceConfigs,
   BUILTIN_TOOL_EXECUTE_SQL,
   BUILTIN_TOOL_SEARCH_OBJECTS,
+  loadTomlConfig,
+  resolveTomlConfigPath,
   ConnectorManager,
   mapArgumentsToArray,
   ToolRegistry,