@bytebase/dbhub 1.4.0 → 1.5.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,9 @@
1
1
  import {
2
2
  isDriverNotInstalled
3
3
  } from "./chunk-WVVMH6FJ.js";
4
+ import {
5
+ MAX_QUERY_TIMEOUT_SECONDS
6
+ } from "./chunk-5XSBOK3W.js";
4
7
  import {
5
8
  ConnectorRegistry,
6
9
  SafeURL,
@@ -27,8 +30,8 @@ var ALL_BUILTIN_TOOL_NAMES = [
27
30
  ];
28
31
 
29
32
  // src/utils/ssh-tunnel.ts
30
- import { Client } from "ssh2";
31
- import { readFileSync as readFileSync2 } from "fs";
33
+ import ssh2, { Client } from "ssh2";
34
+ import { existsSync, readFileSync as readFileSync2 } from "fs";
32
35
  import { createServer } from "net";
33
36
 
34
37
  // src/utils/ssh-config-parser.ts
@@ -114,6 +117,9 @@ function parseSSHConfig(hostAlias, configPath, options = {}) {
114
117
  sshConfig.privateKey = defaultKey;
115
118
  }
116
119
  }
120
+ if (sshConfig.privateKey) {
121
+ sshConfig.privateKeyDiscovered = true;
122
+ }
117
123
  if (hostConfig.ProxyJump) {
118
124
  sshConfig.proxyJump = hostConfig.ProxyJump;
119
125
  }
@@ -265,16 +271,61 @@ var SSHTunnel = class {
265
271
  try {
266
272
  const jumpHosts = config.resolvedJumpHosts ?? (config.proxyJump ? parseJumpHosts(config.proxyJump) : []);
267
273
  const privateKeyBuffer = config.privateKey ? this.loadPrivateKey(config.privateKey) : void 0;
268
- if (!config.password && !privateKeyBuffer) {
269
- throw new Error("Either password or privateKey must be provided for SSH authentication");
274
+ const agent = this.resolveAgent(config.agent);
275
+ if (!config.password && !privateKeyBuffer && !agent) {
276
+ throw new Error("Either password, privateKey, or an SSH agent (agent or SSH_AUTH_SOCK) must be provided for SSH authentication");
270
277
  }
271
- const finalClient = await this.establishChain(jumpHosts, config, privateKeyBuffer);
278
+ const finalClient = await this.establishChain(jumpHosts, config, privateKeyBuffer, agent);
272
279
  return await this.createLocalTunnel(finalClient, options);
273
280
  } catch (error) {
274
281
  this.cleanup();
275
282
  throw error;
276
283
  }
277
284
  }
285
+ /**
286
+ * Pick the SSH agent socket: an explicitly configured one wins over the ambient
287
+ * SSH_AUTH_SOCK. A configured socket that does not exist is an error; a stale
288
+ * SSH_AUTH_SOCK is ignored, since the user never asked DBHub to use it.
289
+ */
290
+ resolveAgent(configuredAgent) {
291
+ if (configuredAgent) {
292
+ if (!this.agentSocketExists(configuredAgent)) {
293
+ throw new Error(`SSH agent socket not found: ${configuredAgent}`);
294
+ }
295
+ return configuredAgent;
296
+ }
297
+ const ambientAgent = process.env.SSH_AUTH_SOCK;
298
+ if (!ambientAgent) {
299
+ return void 0;
300
+ }
301
+ if (!this.agentSocketExists(ambientAgent)) {
302
+ console.warn(`Ignoring SSH_AUTH_SOCK: no SSH agent socket at ${ambientAgent}`);
303
+ return void 0;
304
+ }
305
+ return ambientAgent;
306
+ }
307
+ /**
308
+ * Windows agents are named pipes or Pageant rather than files, so only Unix
309
+ * socket paths are checked.
310
+ */
311
+ agentSocketExists(agent) {
312
+ return process.platform === "win32" || existsSync(agent);
313
+ }
314
+ /**
315
+ * Why ssh2 would refuse this key, mirroring the checks its connect() runs before
316
+ * trying any auth method, or undefined when the key is usable.
317
+ */
318
+ unusableKeyReason(privateKey, passphrase) {
319
+ const parsed = ssh2.utils.parseKey(privateKey, passphrase);
320
+ if (parsed instanceof Error) {
321
+ return parsed.message;
322
+ }
323
+ const key = Array.isArray(parsed) ? parsed[0] : parsed;
324
+ if (key.getPrivatePEM() === null) {
325
+ return "not a private key";
326
+ }
327
+ return void 0;
328
+ }
278
329
  /**
279
330
  * Load an SSH private key, supporting both a file path (with symlink resolution)
280
331
  * and base64-encoded key content.
@@ -303,14 +354,18 @@ var SSHTunnel = class {
303
354
  * Establish a chain of SSH connections through jump hosts.
304
355
  * @returns The final SSH client connected to the target host
305
356
  */
306
- async establishChain(jumpHosts, targetConfig, privateKey) {
357
+ async establishChain(jumpHosts, targetConfig, privateKey, agent) {
307
358
  let previousStream;
308
359
  for (let i = 0; i < jumpHosts.length; i++) {
309
360
  const jumpHost = jumpHosts[i];
310
361
  const nextHost = i + 1 < jumpHosts.length ? jumpHosts[i + 1] : { host: targetConfig.host, port: targetConfig.port || 22 };
311
- const hopPrivateKey = jumpHost.privateKey ? this.loadPrivateKey(jumpHost.privateKey) : privateKey;
312
- const hopPassword = targetConfig.password;
313
- const hopPassphrase = jumpHost.passphrase ?? targetConfig.passphrase;
362
+ const hopAuth = {
363
+ password: targetConfig.password,
364
+ privateKey: jumpHost.privateKey ? this.loadPrivateKey(jumpHost.privateKey) : privateKey,
365
+ passphrase: jumpHost.passphrase ?? targetConfig.passphrase,
366
+ privateKeyDiscovered: jumpHost.privateKey ? true : targetConfig.privateKeyDiscovered,
367
+ agent
368
+ };
314
369
  let client = null;
315
370
  let forwardStream;
316
371
  try {
@@ -320,9 +375,7 @@ var SSHTunnel = class {
320
375
  port: jumpHost.port,
321
376
  username: jumpHost.username || targetConfig.username
322
377
  },
323
- hopPassword,
324
- hopPrivateKey,
325
- hopPassphrase,
378
+ hopAuth,
326
379
  previousStream,
327
380
  `jump host ${i + 1}`,
328
381
  targetConfig.keepaliveInterval,
@@ -348,9 +401,13 @@ var SSHTunnel = class {
348
401
  port: targetConfig.port || 22,
349
402
  username: targetConfig.username
350
403
  },
351
- targetConfig.password,
352
- privateKey,
353
- targetConfig.passphrase,
404
+ {
405
+ password: targetConfig.password,
406
+ privateKey,
407
+ passphrase: targetConfig.passphrase,
408
+ privateKeyDiscovered: targetConfig.privateKeyDiscovered,
409
+ agent
410
+ },
354
411
  previousStream,
355
412
  jumpHosts.length > 0 ? "target host" : void 0,
356
413
  targetConfig.keepaliveInterval,
@@ -362,7 +419,7 @@ var SSHTunnel = class {
362
419
  /**
363
420
  * Connect to a single SSH host.
364
421
  */
365
- connectToHost(hostInfo, password, privateKey, passphrase, sock, label, keepaliveInterval, keepaliveCountMax) {
422
+ connectToHost(hostInfo, auth, sock, label, keepaliveInterval, keepaliveCountMax) {
366
423
  return new Promise((resolve, reject) => {
367
424
  const client = new Client();
368
425
  const sshConfig = {
@@ -370,15 +427,28 @@ var SSHTunnel = class {
370
427
  port: hostInfo.port,
371
428
  username: hostInfo.username
372
429
  };
430
+ const { password, privateKey, passphrase, agent } = auth;
373
431
  if (password) {
374
432
  sshConfig.password = password;
375
433
  }
376
434
  if (privateKey) {
377
- sshConfig.privateKey = privateKey;
378
- if (passphrase) {
379
- sshConfig.passphrase = passphrase;
435
+ const canSkip = auth.privateKeyDiscovered && Boolean(password || agent);
436
+ const unusableReason = canSkip ? this.unusableKeyReason(privateKey, passphrase) : void 0;
437
+ if (unusableReason) {
438
+ const desc = label || `${hostInfo.host}:${hostInfo.port}`;
439
+ console.warn(
440
+ `Skipping unusable SSH private key from ~/.ssh/config for ${desc} (${unusableReason}).`
441
+ );
442
+ } else {
443
+ sshConfig.privateKey = privateKey;
444
+ if (passphrase) {
445
+ sshConfig.passphrase = passphrase;
446
+ }
380
447
  }
381
448
  }
449
+ if (agent) {
450
+ sshConfig.agent = agent;
451
+ }
382
452
  if (sock) {
383
453
  sshConfig.sock = sock;
384
454
  }
@@ -887,12 +957,14 @@ function resolveSSHConfig() {
887
957
  }
888
958
  if (args["ssh-key"]) {
889
959
  config.privateKey = args["ssh-key"];
960
+ config.privateKeyDiscovered = false;
890
961
  if (config.privateKey.startsWith("~/")) {
891
962
  config.privateKey = path.join(process.env.HOME || "", config.privateKey.substring(2));
892
963
  }
893
964
  sources.push("ssh-key from command line");
894
965
  } else if (process.env.SSH_KEY) {
895
966
  config.privateKey = process.env.SSH_KEY;
967
+ config.privateKeyDiscovered = false;
896
968
  if (config.privateKey.startsWith("~/")) {
897
969
  config.privateKey = path.join(process.env.HOME || "", config.privateKey.substring(2));
898
970
  }
@@ -905,6 +977,14 @@ function resolveSSHConfig() {
905
977
  config.passphrase = process.env.SSH_PASSPHRASE;
906
978
  sources.push("SSH_PASSPHRASE from environment");
907
979
  }
980
+ const cliAgent = requireFlagValue("ssh-agent", args, "~/.ssh/agent.sock");
981
+ if (cliAgent) {
982
+ config.agent = cliAgent;
983
+ if (config.agent.startsWith("~/")) {
984
+ config.agent = path.join(process.env.HOME || "", config.agent.substring(2));
985
+ }
986
+ sources.push("ssh-agent from command line");
987
+ }
908
988
  if (args["ssh-proxy-jump"]) {
909
989
  config.proxyJump = args["ssh-proxy-jump"];
910
990
  sources.push("ssh-proxy-jump from command line");
@@ -936,8 +1016,8 @@ function resolveSSHConfig() {
936
1016
  if (!config.host || !config.username) {
937
1017
  throw new Error("SSH tunnel configuration requires at least --ssh-host and --ssh-user");
938
1018
  }
939
- if (!config.password && !config.privateKey) {
940
- throw new Error("SSH tunnel configuration requires either --ssh-password or --ssh-key for authentication");
1019
+ if (!config.password && !config.privateKey && !config.agent && !process.env.SSH_AUTH_SOCK) {
1020
+ throw new Error("SSH tunnel configuration requires either --ssh-password or --ssh-key (or an SSH agent via --ssh-agent or SSH_AUTH_SOCK) for authentication");
941
1021
  }
942
1022
  return {
943
1023
  config,
@@ -1021,7 +1101,9 @@ async function resolveSourceConfigs() {
1021
1101
  source.ssh_user = sshResult.config.username;
1022
1102
  source.ssh_password = sshResult.config.password;
1023
1103
  source.ssh_key = sshResult.config.privateKey;
1104
+ source.ssh_key_discovered = sshResult.config.privateKeyDiscovered;
1024
1105
  source.ssh_passphrase = sshResult.config.passphrase;
1106
+ source.ssh_agent = sshResult.config.agent;
1025
1107
  source.ssh_keepalive_interval = sshResult.config.keepaliveInterval;
1026
1108
  source.ssh_keepalive_count_max = sshResult.config.keepaliveCountMax;
1027
1109
  }
@@ -1189,6 +1271,9 @@ function getRawDSNQueryParam(dsn, key) {
1189
1271
  if (queryStart === -1) {
1190
1272
  return null;
1191
1273
  }
1274
+ if (key === "sslmode" && dsn.startsWith("sqlserver://")) {
1275
+ return new URLSearchParams(dsn.substring(queryStart + 1)).get(key);
1276
+ }
1192
1277
  for (const pair of dsn.substring(queryStart + 1).split("&")) {
1193
1278
  if (pair === "") {
1194
1279
  continue;
@@ -1280,6 +1365,9 @@ function validateDSNFieldConflicts(source, configPath) {
1280
1365
  }
1281
1366
  const dsnSslmode = getRawDSNQueryParam(source.dsn, "sslmode");
1282
1367
  if (source.sslmode && dsnSslmode !== null && dsnSslmode !== source.sslmode) {
1368
+ if (source.type === "sqlserver") {
1369
+ throw new Error("Conflicting SQL Server sslmode. Set sslmode in only one place, or make the two values match.");
1370
+ }
1283
1371
  conflict("sslmode", source.sslmode, dsnSslmode);
1284
1372
  }
1285
1373
  const dsnSslrootcert = getRawDSNQueryParam(source.dsn, "sslrootcert");
@@ -1369,9 +1457,9 @@ function validateSourceConfig(source, configPath) {
1369
1457
  }
1370
1458
  }
1371
1459
  if (source.query_timeout !== void 0) {
1372
- if (typeof source.query_timeout !== "number" || source.query_timeout <= 0) {
1460
+ if (!Number.isFinite(source.query_timeout) || source.query_timeout <= 0 || source.query_timeout > MAX_QUERY_TIMEOUT_SECONDS) {
1373
1461
  throw new Error(
1374
- `Configuration file ${configPath}: source '${source.id}' has invalid query_timeout. Must be a positive number (in seconds).`
1462
+ `Configuration file ${configPath}: source '${source.id}' has invalid query_timeout. Must be a positive number of seconds, at most ${MAX_QUERY_TIMEOUT_SECONDS}.`
1375
1463
  );
1376
1464
  }
1377
1465
  }
@@ -1394,6 +1482,13 @@ function validateSourceConfig(source, configPath) {
1394
1482
  );
1395
1483
  }
1396
1484
  }
1485
+ if (source.ssh_agent !== void 0) {
1486
+ if (typeof source.ssh_agent !== "string" || source.ssh_agent.trim() === "") {
1487
+ throw new Error(
1488
+ `Configuration file ${configPath}: source '${source.id}' has invalid ssh_agent. Must be a path to an SSH agent socket.`
1489
+ );
1490
+ }
1491
+ }
1397
1492
  if (source.sslmode !== void 0) {
1398
1493
  if (source.type === "sqlite") {
1399
1494
  throw new Error(
@@ -1408,6 +1503,7 @@ function validateSourceConfig(source, configPath) {
1408
1503
  }
1409
1504
  const verifyModesByType = {
1410
1505
  postgres: ["verify-ca", "verify-full"],
1506
+ sqlserver: ["verify-full"],
1411
1507
  oracle: ["verify-full"]
1412
1508
  };
1413
1509
  if ((source.sslmode === "verify-ca" || source.sslmode === "verify-full") && !(verifyModesByType[source.type] ?? []).includes(source.sslmode)) {
@@ -1566,6 +1662,10 @@ function processSourceConfigs(sources, configPath) {
1566
1662
  if (processed.ssh_key) {
1567
1663
  processed.ssh_key = expandHomeDir(processed.ssh_key);
1568
1664
  }
1665
+ delete processed.ssh_key_discovered;
1666
+ if (typeof processed.ssh_agent === "string" && processed.ssh_agent) {
1667
+ processed.ssh_agent = expandHomeDir(processed.ssh_agent);
1668
+ }
1569
1669
  for (const field of ["sslrootcert", "sslcert", "sslkey"]) {
1570
1670
  if (processed[field]) {
1571
1671
  processed[field] = expandHomeDir(processed[field]);
@@ -1599,7 +1699,7 @@ function processSourceConfigs(sources, configPath) {
1599
1699
  try {
1600
1700
  const url = new SafeURL(processed.dsn);
1601
1701
  const dsnSslmode = url.getSearchParam("sslmode");
1602
- if (!processed.sslmode && dsnSslmode) {
1702
+ if (processed.sslmode === void 0 && dsnSslmode) {
1603
1703
  processed.sslmode = dsnSslmode;
1604
1704
  }
1605
1705
  for (const field of ["sslrootcert", "sslcert", "sslkey"]) {
@@ -1837,7 +1937,6 @@ function classifyConnectionError(error, connectorType, sourceId) {
1837
1937
  var managerInstance = null;
1838
1938
  var AWS_IAM_TOKEN_REFRESH_MS = 14 * 60 * 1e3;
1839
1939
  var ConnectorManager = class {
1840
- // Prevent race conditions
1841
1940
  constructor() {
1842
1941
  // Maps for multi-source support
1843
1942
  this.connectors = /* @__PURE__ */ new Map();
@@ -1847,11 +1946,17 @@ var ConnectorManager = class {
1847
1946
  this.sourceIds = [];
1848
1947
  // Ordered list of source IDs (first is default)
1849
1948
  this.iamRefreshTimers = /* @__PURE__ */ new Map();
1949
+ this.pendingIamRefreshes = /* @__PURE__ */ new Map();
1950
+ // In-flight refresh per source
1850
1951
  this.isDisconnecting = false;
1851
1952
  // Lazy connection support
1852
1953
  this.lazySources = /* @__PURE__ */ new Map();
1853
1954
  // Sources pending lazy connection
1854
1955
  this.pendingConnections = /* @__PURE__ */ new Map();
1956
+ // Prevent race conditions
1957
+ // A socket timeout does not cancel a credential helper. Share its in-flight
1958
+ // attempt across retries of the same config, but never across changed profiles.
1959
+ this.pendingIamTokens = /* @__PURE__ */ new WeakMap();
1855
1960
  if (!managerInstance) {
1856
1961
  managerInstance = this;
1857
1962
  }
@@ -1937,7 +2042,8 @@ var ConnectorManager = class {
1937
2042
  */
1938
2043
  async connectSource(source) {
1939
2044
  const sourceId = source.id;
1940
- const dsn = await this.buildConnectionDSN(source);
2045
+ const config = {};
2046
+ const dsn = await this.buildConnectionDSN(source, config);
1941
2047
  console.error(` - ${sourceId}: ${redactDSN(dsn)}`);
1942
2048
  let actualDSN = dsn;
1943
2049
  let tunnel;
@@ -1965,7 +2071,9 @@ var ConnectorManager = class {
1965
2071
  username: username || "",
1966
2072
  password: source.ssh_password,
1967
2073
  privateKey: source.ssh_key || resolvedSSHConfig?.privateKey,
2074
+ privateKeyDiscovered: source.ssh_key ? source.ssh_key_discovered : resolvedSSHConfig?.privateKeyDiscovered,
1968
2075
  passphrase: source.ssh_passphrase,
2076
+ agent: source.ssh_agent,
1969
2077
  proxyJump,
1970
2078
  resolvedJumpHosts,
1971
2079
  keepaliveInterval: source.ssh_keepalive_interval,
@@ -1976,9 +2084,9 @@ var ConnectorManager = class {
1976
2084
  `Source '${sourceId}': SSH tunnel requires ssh_user (or a matching Host entry in ~/.ssh/config with User)`
1977
2085
  );
1978
2086
  }
1979
- if (!sshConfig.password && !sshConfig.privateKey) {
2087
+ if (!sshConfig.password && !sshConfig.privateKey && !sshConfig.agent && !process.env.SSH_AUTH_SOCK) {
1980
2088
  throw new Error(
1981
- `Source '${sourceId}': SSH tunnel requires either ssh_password or ssh_key (or a matching Host entry in ~/.ssh/config with IdentityFile)`
2089
+ `Source '${sourceId}': SSH tunnel requires either ssh_password or ssh_key (or a matching Host entry in ~/.ssh/config with IdentityFile, or an SSH agent via ssh_agent or SSH_AUTH_SOCK)`
1982
2090
  );
1983
2091
  }
1984
2092
  const url = new URL(dsn);
@@ -2005,37 +2113,37 @@ var ConnectorManager = class {
2005
2113
  ` SSH tunnel established through localhost:${tunnelInfo.localPort}`
2006
2114
  );
2007
2115
  }
2008
- const connectorPrototype = ConnectorRegistry.getConnectorForDSN(actualDSN);
2009
- if (!connectorPrototype) {
2010
- throw new Error(
2011
- `Source '${sourceId}': No connector found for DSN: ${actualDSN}`
2012
- );
2013
- }
2014
- const connector = connectorPrototype.clone();
2015
- connector.sourceId = sourceId;
2016
- const config = {};
2017
- if (source.connection_timeout !== void 0) {
2018
- config.connectionTimeoutSeconds = source.connection_timeout;
2019
- }
2020
- if (source.query_timeout !== void 0 && connector.id !== "sqlite") {
2021
- config.queryTimeoutSeconds = source.query_timeout;
2022
- }
2023
- if (source.pool_max_connections !== void 0) {
2024
- config.poolMaxConnections = source.pool_max_connections;
2025
- }
2026
- if (source.search_path) {
2027
- config.searchPath = source.search_path;
2028
- }
2029
- if (source.timezone) {
2030
- config.timezone = source.timezone;
2031
- }
2032
- if (source.charset) {
2033
- config.charset = source.charset;
2034
- }
2035
- if (source.collation) {
2036
- config.collation = source.collation;
2037
- }
2116
+ let connector;
2038
2117
  try {
2118
+ const connectorPrototype = ConnectorRegistry.getConnectorForDSN(actualDSN);
2119
+ if (!connectorPrototype) {
2120
+ throw new Error(
2121
+ `Source '${sourceId}': No connector found for DSN: ${actualDSN}`
2122
+ );
2123
+ }
2124
+ connector = connectorPrototype.clone();
2125
+ connector.sourceId = sourceId;
2126
+ if (source.connection_timeout !== void 0) {
2127
+ config.connectionTimeoutSeconds = source.connection_timeout;
2128
+ }
2129
+ if (source.query_timeout !== void 0 && connector.id !== "sqlite") {
2130
+ config.queryTimeoutSeconds = source.query_timeout;
2131
+ }
2132
+ if (source.pool_max_connections !== void 0) {
2133
+ config.poolMaxConnections = source.pool_max_connections;
2134
+ }
2135
+ if (source.search_path) {
2136
+ config.searchPath = source.search_path;
2137
+ }
2138
+ if (source.timezone) {
2139
+ config.timezone = source.timezone;
2140
+ }
2141
+ if (source.charset) {
2142
+ config.charset = source.charset;
2143
+ }
2144
+ if (source.collation) {
2145
+ config.collation = source.collation;
2146
+ }
2039
2147
  await connector.connect(actualDSN, source.init_script, config);
2040
2148
  } catch (error) {
2041
2149
  if (tunnel) {
@@ -2055,6 +2163,77 @@ var ConnectorManager = class {
2055
2163
  this.sourceConfigs.set(sourceId, source);
2056
2164
  this.scheduleIamRefresh(source);
2057
2165
  }
2166
+ /**
2167
+ * Add a single source without touching the others. Eager sources connect now;
2168
+ * lazy ones are registered and connect on first use. Used by the TOML hot reload
2169
+ * to apply only the entries that changed.
2170
+ */
2171
+ async addSource(source) {
2172
+ if (this.sourceIds.includes(source.id)) {
2173
+ throw new Error(`Source '${source.id}' already exists`);
2174
+ }
2175
+ if (source.lazy) {
2176
+ this.registerLazySource(source);
2177
+ } else {
2178
+ await this.connectSource(source);
2179
+ }
2180
+ }
2181
+ /**
2182
+ * Disconnect and forget a single source, leaving every other source's pool and
2183
+ * tunnel untouched. Resolves silently for an unknown id.
2184
+ */
2185
+ async removeSource(sourceId) {
2186
+ const pending = this.pendingConnections.get(sourceId);
2187
+ if (pending) {
2188
+ try {
2189
+ await pending;
2190
+ } catch {
2191
+ }
2192
+ }
2193
+ const refresh = this.pendingIamRefreshes.get(sourceId);
2194
+ if (refresh) {
2195
+ await refresh;
2196
+ }
2197
+ const timer = this.iamRefreshTimers.get(sourceId);
2198
+ if (timer) {
2199
+ clearTimeout(timer);
2200
+ this.iamRefreshTimers.delete(sourceId);
2201
+ }
2202
+ const connector = this.connectors.get(sourceId);
2203
+ this.connectors.delete(sourceId);
2204
+ if (connector) {
2205
+ try {
2206
+ await connector.disconnect();
2207
+ console.error(`Disconnected from source '${sourceId}'`);
2208
+ } catch (error) {
2209
+ console.error(`Error disconnecting from source '${sourceId}':`, error);
2210
+ }
2211
+ }
2212
+ const tunnel = this.sshTunnels.get(sourceId);
2213
+ this.sshTunnels.delete(sourceId);
2214
+ if (tunnel) {
2215
+ try {
2216
+ await tunnel.close();
2217
+ } catch (error) {
2218
+ console.error(`Error closing SSH tunnel for source '${sourceId}':`, error);
2219
+ }
2220
+ }
2221
+ this.sourceConfigs.delete(sourceId);
2222
+ this.lazySources.delete(sourceId);
2223
+ this.pendingConnections.delete(sourceId);
2224
+ this.sourceIds = this.sourceIds.filter((id) => id !== sourceId);
2225
+ }
2226
+ /**
2227
+ * Reorder known sources to match `orderedIds` (the first entry is the default
2228
+ * source). Unknown ids are ignored; known ids missing from the list keep their
2229
+ * relative order after the listed ones.
2230
+ */
2231
+ reorderSources(orderedIds) {
2232
+ const known = new Set(this.sourceIds);
2233
+ const ordered = orderedIds.filter((id) => known.has(id));
2234
+ const listed = new Set(ordered);
2235
+ this.sourceIds = [...ordered, ...this.sourceIds.filter((id) => !listed.has(id))];
2236
+ }
2058
2237
  /**
2059
2238
  * Close all database connections
2060
2239
  */
@@ -2084,6 +2263,7 @@ var ConnectorManager = class {
2084
2263
  this.sourceConfigs.clear();
2085
2264
  this.lazySources.clear();
2086
2265
  this.pendingConnections.clear();
2266
+ this.pendingIamRefreshes.clear();
2087
2267
  this.sourceIds = [];
2088
2268
  this.isDisconnecting = false;
2089
2269
  }
@@ -2204,25 +2384,29 @@ var ConnectorManager = class {
2204
2384
  clearTimeout(existingTimer);
2205
2385
  this.iamRefreshTimers.delete(sourceId);
2206
2386
  }
2207
- if (!source.aws_iam_auth) {
2387
+ if (!source.aws_iam_auth || source.type === "postgres") {
2208
2388
  return;
2209
2389
  }
2210
- const timer = setTimeout(async () => {
2390
+ const timer = setTimeout(() => {
2211
2391
  if (this.isDisconnecting) {
2212
2392
  return;
2213
2393
  }
2214
- try {
2215
- await this.refreshIamSourceConnection(source);
2216
- } catch (error) {
2217
- console.error(
2218
- `Error refreshing AWS IAM auth token for source '${sourceId}':`,
2219
- error
2220
- );
2221
- } finally {
2222
- if (!this.isDisconnecting && this.connectors.has(sourceId)) {
2223
- this.scheduleIamRefresh(source);
2394
+ const run = (async () => {
2395
+ try {
2396
+ await this.refreshIamSourceConnection(source);
2397
+ } catch (error) {
2398
+ console.error(
2399
+ `Error refreshing AWS IAM auth token for source '${sourceId}':`,
2400
+ error
2401
+ );
2402
+ } finally {
2403
+ this.pendingIamRefreshes.delete(sourceId);
2404
+ if (!this.isDisconnecting && this.ownsSource(source) && this.connectors.has(sourceId)) {
2405
+ this.scheduleIamRefresh(source);
2406
+ }
2224
2407
  }
2225
- }
2408
+ })();
2409
+ this.pendingIamRefreshes.set(sourceId, run);
2226
2410
  }, AWS_IAM_TOKEN_REFRESH_MS);
2227
2411
  timer.unref?.();
2228
2412
  this.iamRefreshTimers.set(sourceId, timer);
@@ -2243,23 +2427,31 @@ var ConnectorManager = class {
2243
2427
  await existingTunnel.close();
2244
2428
  this.sshTunnels.delete(sourceId);
2245
2429
  }
2246
- if (this.isDisconnecting) {
2430
+ if (this.isDisconnecting || !this.ownsSource(source)) {
2247
2431
  return;
2248
2432
  }
2249
2433
  try {
2250
2434
  await this.connectSource(source);
2251
2435
  } catch (error) {
2252
- if (!this.isDisconnecting && this.sourceConfigs.has(sourceId)) {
2436
+ if (!this.isDisconnecting && this.ownsSource(source)) {
2253
2437
  this.lazySources.set(sourceId, source);
2254
2438
  }
2255
2439
  throw error;
2256
2440
  }
2257
2441
  }
2442
+ /**
2443
+ * True while `source` is the config object registered under its id. Every
2444
+ * registration path stores the same object, so identity tells an in-flight
2445
+ * operation whether its source was removed or replaced underneath it.
2446
+ */
2447
+ ownsSource(source) {
2448
+ return this.sourceConfigs.get(source.id) === source;
2449
+ }
2258
2450
  /**
2259
2451
  * Build a connection DSN, optionally replacing password with
2260
2452
  * an AWS RDS IAM auth token when aws_iam_auth is enabled.
2261
2453
  */
2262
- async buildConnectionDSN(source) {
2454
+ async buildConnectionDSN(source, config = {}) {
2263
2455
  const dsn = buildDSNFromSource(source);
2264
2456
  if (!source.aws_iam_auth) {
2265
2457
  return dsn;
@@ -2285,13 +2477,26 @@ var ConnectorManager = class {
2285
2477
  `Source '${source.id}': unable to resolve host, username, or port for AWS IAM authentication`
2286
2478
  );
2287
2479
  }
2288
- const token = await generateRdsAuthToken({
2289
- hostname,
2290
- port,
2291
- username,
2292
- region: source.aws_region,
2293
- profile: source.aws_profile
2294
- });
2480
+ const password = () => {
2481
+ let pending = this.pendingIamTokens.get(source);
2482
+ if (!pending) {
2483
+ pending = generateRdsAuthToken({
2484
+ hostname,
2485
+ port,
2486
+ username,
2487
+ region: source.aws_region,
2488
+ profile: source.aws_profile
2489
+ }).finally(() => {
2490
+ this.pendingIamTokens.delete(source);
2491
+ });
2492
+ this.pendingIamTokens.set(source, pending);
2493
+ }
2494
+ return pending;
2495
+ };
2496
+ if (source.type === "postgres") {
2497
+ config.password = password;
2498
+ }
2499
+ const token = source.type === "postgres" ? "" : await password();
2295
2500
  const queryParams = new Map(parsed.searchParams);
2296
2501
  const currentSslMode = queryParams.get("sslmode");
2297
2502
  if (currentSslMode !== "verify-ca" && currentSslMode !== "verify-full") {