Skip to content

Commit 9027fe3

Browse files
tianzhouclaude
andauthored
fix: reload TOML config per source instead of reconnecting every source (#461)
* fix: reload TOML config per source instead of reconnecting every source The config watcher used to call `disconnect()` on every source and then reconnect them all whenever `dbhub.toml` changed. In-flight queries finished, but until the slowest one did, every new call on every source failed with "Cannot use a pool after calling end on the pool", including sources the edit never touched (#454, problem 1). The reload is now a diff against the last applied source list: - Unchanged sources are left alone and keep their pools. - New sources are connected (or registered lazily), removed ones are disconnected, and edited ones are reconnected, one at a time. - Rollback is per source: a changed source whose new config fails to connect is restored to its previous config; a new source that fails is skipped. Tools referencing a source that is not live are dropped from the registry so one bad source cannot take the rest down. - Source order follows the file so the first entry stays the default. `ConnectorManager` gains `addSource()`, `removeSource()` and `reorderSources()`; `disconnect()` is unchanged and still used at shutdown. Refs #454 Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HPXjjvWtnhfENJ7FWHoyfH * fix: close IAM refresh race and keep previous tools for rolled-back sources Two review findings on the per-source reload: - An IAM refresh already running when removeSource() was called could finish its reconnect after the removal and put the source back (or clobber a same-id replacement), then re-arm its timer. Track the in-flight refresh per source and have removeSource() await it, and guard the refresh's reconnect and re-arm on the source still being the registered one. - A source rolled back to its previous config was still given the new file's tool definitions, which were written against the config that did not apply. The watcher now keeps the previous tools for rolled-back ids and uses the new tools only for sources that are live under the new config. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HPXjjvWtnhfENJ7FWHoyfH * fix: validate the registry before applying a reload and close tunnels on every connect failure Review findings on the per-source reload: - The tool registry validates more than the TOML loader (custom tool parameters, duplicate tool names), so it could throw after the source diff had already been applied, leaving the manager on the new sources with the old registry and a stale diff baseline. The watcher now builds a registry from the new file before touching any connection, so a bad file is rejected whole. The one remaining conflict, a rolled-back source's previous tools clashing with the new file's tools, takes that source offline instead of serving it with the wrong tools. The diff baseline is updated as soon as the sources are committed. - connectSource() closed the SSH tunnel only when connector.connect() rejected. A failure between tunnel establishment and that call left the tunnel in the map, and a later retry would orphan it. The cleanup now covers everything after the tunnel is up. - The reload log reported success even when a changed source had been rolled back; the degraded message now counts rolled-back sources too. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HPXjjvWtnhfENJ7FWHoyfH --------- Co-authored-by: Claude <noreply@anthropic.com>
1 parent 28f99f0 commit 9027fe3

5 files changed

Lines changed: 811 additions & 153 deletions

File tree

‎docs/config/toml.mdx‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,8 @@ When using TOML configuration, DBHub automatically watches `dbhub.toml` for chan
3232

3333
1. DBHub detects file changes and waits 500ms (debounce) to handle editors that write in multiple steps
3434
2. The new configuration is parsed and validated — if invalid, existing connections are preserved
35-
3. All database connections are disconnected and reconnected with the new configuration
36-
4. If reconnection fails, DBHub automatically rolls back to the previous working configuration
35+
3. The new source list is compared with the running one, and only what changed is applied: new `[[sources]]` entries are connected, removed ones are disconnected, and edited ones are reconnected. Sources that did not change keep serving queries throughout
36+
4. If a changed source fails to connect, that source alone is rolled back to its previous configuration; a new source that fails to connect is skipped. The other sources are unaffected
3737

3838
```
3939
# Edit dbhub.toml while DBHub is running — changes apply automatically

‎src/connectors/__tests__/manager.test.ts‎

Lines changed: 223 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -390,6 +390,31 @@ describe("ConnectorManager IAM refresh recovery", () => {
390390
expect(instances).toHaveLength(2);
391391
});
392392

393+
it("should close the SSH tunnel when no connector accepts the DSN after the tunnel is up", async () => {
394+
mocks.looksLikeSSHAlias.mockReturnValue(false);
395+
const establishSpy = vi
396+
.spyOn(SSHTunnel.prototype, "establish")
397+
.mockResolvedValue({ localPort: 55555, targetHost: "db.internal", targetPort: 5432 });
398+
const closeSpy = vi.spyOn(SSHTunnel.prototype, "close").mockResolvedValue(undefined);
399+
vi.spyOn(ConnectorRegistry, "getConnectorForDSN").mockReturnValue(null);
400+
401+
const manager = new ConnectorManager();
402+
const source: SourceConfig = {
403+
id: "pg_ssh",
404+
type: "postgres",
405+
dsn: "postgres://user:pass@db.internal:5432/mydb",
406+
ssh_host: "bastion.example.com",
407+
ssh_user: "ubuntu",
408+
ssh_password: "secret",
409+
};
410+
411+
await expect(manager.addSource(source)).rejects.toThrow("No connector found");
412+
expect(establishSpy).toHaveBeenCalledTimes(1);
413+
expect(closeSpy).toHaveBeenCalledTimes(1);
414+
expect((manager as any).sshTunnels.size).toBe(0);
415+
expect(manager.getSourceIds()).toEqual([]);
416+
});
417+
393418
it("should not list a source as available when it can neither serve nor reconnect", () => {
394419
const manager = new ConnectorManager();
395420
(manager as any).sourceIds = ["alive", "dead"];
@@ -510,3 +535,201 @@ describe("PostgreSQL IAM authentication on demand", () => {
510535
expect(manager.getConnector(source.id)).toBeDefined();
511536
});
512537
});
538+
539+
describe("ConnectorManager per-source add/remove", () => {
540+
function stubConnectorRegistry() {
541+
const instances: Array<{ connect: ReturnType<typeof vi.fn>; disconnect: ReturnType<typeof vi.fn> }> = [];
542+
const prototype = {
543+
id: "postgres",
544+
clone: () => {
545+
const instance = {
546+
id: "postgres",
547+
connect: vi.fn().mockResolvedValue(undefined),
548+
disconnect: vi.fn().mockResolvedValue(undefined),
549+
};
550+
instances.push(instance);
551+
return instance;
552+
},
553+
};
554+
vi.spyOn(ConnectorRegistry, "getConnectorForDSN").mockReturnValue(prototype as any);
555+
return instances;
556+
}
557+
558+
const srcA: SourceConfig = { id: "a", type: "postgres", dsn: "postgres://u:p@h/a" };
559+
const srcB: SourceConfig = { id: "b", type: "postgres", dsn: "postgres://u:p@h/b" };
560+
561+
afterEach(() => {
562+
vi.restoreAllMocks();
563+
});
564+
565+
it("adds an eager source without touching the existing one", async () => {
566+
const instances = stubConnectorRegistry();
567+
const manager = new ConnectorManager();
568+
await manager.connectWithSources([srcA]);
569+
570+
await manager.addSource(srcB);
571+
572+
expect(instances).toHaveLength(2);
573+
expect(instances[0].disconnect).not.toHaveBeenCalled();
574+
expect(manager.getConnector("a")).toBe(instances[0]);
575+
expect(manager.getConnector("b")).toBe(instances[1]);
576+
expect(manager.getSourceIds()).toEqual(["a", "b"]);
577+
expect(manager.getAllSourceConfigs()).toEqual([srcA, srcB]);
578+
});
579+
580+
it("registers a lazy source and connects it on first use", async () => {
581+
const instances = stubConnectorRegistry();
582+
const manager = new ConnectorManager();
583+
await manager.connectWithSources([srcA]);
584+
585+
await manager.addSource({ ...srcB, lazy: true });
586+
expect(instances).toHaveLength(1);
587+
expect(manager.getSourceIds()).toEqual(["a", "b"]);
588+
589+
await manager.ensureConnected("b");
590+
expect(instances).toHaveLength(2);
591+
expect(manager.getConnector("b")).toBe(instances[1]);
592+
});
593+
594+
it("rejects a duplicate source id", async () => {
595+
stubConnectorRegistry();
596+
const manager = new ConnectorManager();
597+
await manager.connectWithSources([srcA]);
598+
599+
await expect(manager.addSource(srcA)).rejects.toThrow("already exists");
600+
});
601+
602+
it("removes one source and leaves the other connected", async () => {
603+
const instances = stubConnectorRegistry();
604+
const manager = new ConnectorManager();
605+
await manager.connectWithSources([srcA, srcB]);
606+
607+
await manager.removeSource("a");
608+
609+
expect(instances[0].disconnect).toHaveBeenCalledTimes(1);
610+
expect(instances[1].disconnect).not.toHaveBeenCalled();
611+
expect(manager.getSourceIds()).toEqual(["b"]);
612+
expect(manager.getConnector()).toBe(instances[1]);
613+
expect(() => manager.getConnector("a")).toThrow("Source 'a' not found");
614+
});
615+
616+
it("removing an unknown source is a no-op", async () => {
617+
stubConnectorRegistry();
618+
const manager = new ConnectorManager();
619+
await manager.connectWithSources([srcA]);
620+
621+
await expect(manager.removeSource("nope")).resolves.toBeUndefined();
622+
expect(manager.getSourceIds()).toEqual(["a"]);
623+
});
624+
625+
it("reorders sources so the requested first id becomes the default", async () => {
626+
const instances = stubConnectorRegistry();
627+
const manager = new ConnectorManager();
628+
await manager.connectWithSources([srcA, srcB]);
629+
630+
manager.reorderSources(["b", "unknown", "a"]);
631+
632+
expect(manager.getSourceIds()).toEqual(["b", "a"]);
633+
expect(manager.getConnector()).toBe(instances[1]);
634+
});
635+
});
636+
637+
describe("ConnectorManager IAM refresh racing with removeSource", () => {
638+
const AWS_IAM_TOKEN_REFRESH_MS = 14 * 60 * 1000;
639+
640+
function makeIamSource(dsn = "mysql://dbuser:ignored@mydb.abc123.eu-west-1.rds.amazonaws.com:3306/mydb"): SourceConfig {
641+
return { id: "mysql_iam", type: "mysql", dsn, aws_iam_auth: true, aws_region: "eu-west-1" };
642+
}
643+
644+
/** Connector stubs; the first one's disconnect() blocks until the test releases it. */
645+
function stubSlowDisconnectRegistry() {
646+
const instances: Array<{ connect: ReturnType<typeof vi.fn>; disconnect: ReturnType<typeof vi.fn>; release: () => void }> = [];
647+
const prototype = {
648+
id: "mysql",
649+
clone: () => {
650+
let release: () => void = () => {};
651+
const disconnect = instances.length === 0
652+
? vi.fn().mockImplementation(() => new Promise<void>(resolve => { release = resolve; }))
653+
: vi.fn().mockResolvedValue(undefined);
654+
const instance = {
655+
id: "mysql",
656+
connect: vi.fn().mockResolvedValue(undefined),
657+
disconnect,
658+
release: () => release(),
659+
};
660+
instances.push(instance);
661+
return instance;
662+
},
663+
};
664+
vi.spyOn(ConnectorRegistry, "getConnectorForDSN").mockReturnValue(prototype as any);
665+
return instances;
666+
}
667+
668+
beforeEach(() => {
669+
vi.clearAllMocks();
670+
vi.useFakeTimers();
671+
mocks.generateRdsAuthToken.mockResolvedValue("token");
672+
});
673+
674+
afterEach(() => {
675+
vi.useRealTimers();
676+
vi.restoreAllMocks();
677+
});
678+
679+
it("does not resurrect a source removed while its IAM refresh was in flight", async () => {
680+
const instances = stubSlowDisconnectRegistry();
681+
const manager = new ConnectorManager();
682+
await manager.connectWithSources([makeIamSource()]);
683+
684+
// Refresh fires and blocks inside the old connector's disconnect().
685+
await vi.advanceTimersByTimeAsync(AWS_IAM_TOKEN_REFRESH_MS);
686+
expect(instances[0].disconnect).toHaveBeenCalledTimes(1);
687+
688+
// Config reload drops the source while the refresh is parked.
689+
const removal = manager.removeSource("mysql_iam");
690+
instances[0].release();
691+
await removal;
692+
await vi.advanceTimersByTimeAsync(0);
693+
694+
// The refresh finished its reconnect, and removal then tore that connector down too.
695+
expect(instances).toHaveLength(2);
696+
expect(instances[1].disconnect).toHaveBeenCalledTimes(1);
697+
expect(manager.getSourceIds()).toEqual([]);
698+
expect(() => manager.getConnector("mysql_iam")).toThrow();
699+
700+
// And no refresh timer was re-armed for the dead source.
701+
await vi.advanceTimersByTimeAsync(AWS_IAM_TOKEN_REFRESH_MS);
702+
expect(instances).toHaveLength(2);
703+
});
704+
705+
it("does not clobber a same-id replacement added while the old refresh was in flight", async () => {
706+
const instances = stubSlowDisconnectRegistry();
707+
const manager = new ConnectorManager();
708+
await manager.connectWithSources([makeIamSource()]);
709+
710+
await vi.advanceTimersByTimeAsync(AWS_IAM_TOKEN_REFRESH_MS);
711+
expect(instances[0].disconnect).toHaveBeenCalledTimes(1);
712+
713+
const replacement = makeIamSource("mysql://dbuser:ignored@other.abc123.eu-west-1.rds.amazonaws.com:3306/mydb");
714+
const swap = (async () => {
715+
await manager.removeSource("mysql_iam");
716+
await manager.addSource(replacement);
717+
})();
718+
instances[0].release();
719+
await swap;
720+
await vi.advanceTimersByTimeAsync(0);
721+
722+
// instances[1] is the refresh's reconnect of the old config, torn down by removeSource;
723+
// instances[2] is the replacement and is what the manager serves.
724+
expect(instances).toHaveLength(3);
725+
expect(instances[1].disconnect).toHaveBeenCalledTimes(1);
726+
expect(manager.getConnector("mysql_iam")).toBe(instances[2]);
727+
expect(manager.getSourceConfig("mysql_iam")).toBe(replacement);
728+
expect(instances[2].connect.mock.calls[0][0]).toContain("other.abc123");
729+
730+
// The next refresh tick belongs to the replacement, not the removed config.
731+
await vi.advanceTimersByTimeAsync(AWS_IAM_TOKEN_REFRESH_MS);
732+
expect(instances).toHaveLength(4);
733+
expect(instances[3].connect.mock.calls[0][0]).toContain("other.abc123");
734+
});
735+
});

0 commit comments

Comments
 (0)