feat: add manager with hash cache, tunnel rebuild

2026-06-12 00:01:50 +05:00
parent fe6e0bcbfa
commit 52baa59f97
2 changed files with 302 additions and 0 deletions
@@ -0,0 +1,194 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import type { Registry } from '../../../src/config/registry.js'
+import type { ConnectionConfig, ResolvedConnection, SshConfig } from '../../../src/config/types.js'
+import type { Driver, DriverTarget } from '../../../src/db/driver.js'
+import { createManager } from '../../../src/db/manager.js'
+import type { Tunnel } from '../../../src/net/tunnel.js'
+
+const config = (extra: Partial<ConnectionConfig> = {}): ConnectionConfig => ({
+    name: 'c',
+    type: 'postgres',
+    host: 'db-host',
+    user: 'u',
+    readonly: false,
+    ...extra
+})
+
+const fakeDriver = (): Driver & { disposed: boolean } => {
+    const driver = {
+        disposed: false,
+        query: vi.fn(),
+        listDatabases: vi.fn(),
+        listTables: vi.fn(),
+        describeTable: vi.fn(),
+        serverVersion: vi.fn(),
+        dispose: vi.fn(async () => {
+            driver.disposed = true
+        })
+    }
+    return driver
+}
+
+type FakeTunnel = Tunnel & { closed: boolean }
+
+describe('createManager', () => {
+    let resolved: ResolvedConnection
+    let registry: Registry
+    let drivers: Array<Driver & { disposed: boolean }>
+    let targets: DriverTarget[]
+    let tunnels: FakeTunnel[]
+    let createDriver: (target: DriverTarget) => Driver
+    let createTunnel: ReturnType<typeof vi.fn>
+
+    beforeEach(() => {
+        resolved = { config: config(), source: 'store', hash: 'h1' }
+        registry = {
+            list: vi.fn(() => [resolved]),
+            resolve: vi.fn(() => resolved),
+            add: vi.fn(),
+            update: vi.fn(),
+            remove: vi.fn()
+        }
+        drivers = []
+        targets = []
+        tunnels = []
+        createDriver = (target: DriverTarget) => {
+            targets.push(target)
+            const driver = fakeDriver()
+            drivers.push(driver)
+            return driver
+        }
+        createTunnel = vi.fn(async () => {
+            const tunnel: FakeTunnel = {
+                localHost: '127.0.0.1',
+                localPort: 54321,
+                closed: false,
+                isClosed: () => tunnel.closed,
+                close: async () => {
+                    tunnel.closed = true
+                }
+            }
+            tunnels.push(tunnel)
+            return tunnel
+        })
+    })
+
+    const sshExtra = (): { ssh: SshConfig } => ({
+        ssh: { host: 'bastion', port: 22, user: 'root', password: 'x' }
+    })
+
+    it('builds lazily and caches by hash', async () => {
+        const manager = createManager(registry, { createDriver, createTunnel })
+        const first = await manager.get('c')
+        const second = await manager.get('c')
+        expect(first.driver).toBe(second.driver)
+        expect(drivers).toHaveLength(1)
+        expect(targets[0]).toMatchObject({ host: 'db-host', port: 5432 })
+        expect(createTunnel).not.toHaveBeenCalled()
+    })
+
+    it('rebuilds and disposes the old driver when the hash changes', async () => {
+        const manager = createManager(registry, { createDriver, createTunnel })
+        const first = await manager.get('c')
+        resolved = { config: config({ port: 9999 }), source: 'store', hash: 'h2' }
+        const second = await manager.get('c')
+        expect(second.driver).not.toBe(first.driver)
+        expect(drivers[0].disposed).toBe(true)
+        expect(targets[1]).toMatchObject({ port: 9999 })
+    })
+
+    it('routes through the tunnel endpoint when ssh is configured', async () => {
+        resolved = { config: config(sshExtra()), source: 'store', hash: 'ssh1' }
+        const manager = createManager(registry, { createDriver, createTunnel })
+        await manager.get('c')
+        expect(createTunnel).toHaveBeenCalledWith(
+            expect.objectContaining({ host: 'bastion' }),
+            'db-host',
+            5432
+        )
+        expect(targets[0]).toMatchObject({ host: '127.0.0.1', port: 54321 })
+    })
+
+    it('uses the engine default port for mysql', async () => {
+        resolved = { config: config({ type: 'mysql' }), source: 'store', hash: 'm1' }
+        const manager = createManager(registry, { createDriver, createTunnel })
+        await manager.get('c')
+        expect(targets[0]).toMatchObject({ port: 3306 })
+    })
+
+    it('rebuilds when the tunnel reports closed', async () => {
+        resolved = { config: config(sshExtra()), source: 'store', hash: 'ssh1' }
+        const manager = createManager(registry, { createDriver, createTunnel })
+        const first = await manager.get('c')
+        tunnels[0].closed = true
+        const second = await manager.get('c')
+        expect(second.driver).not.toBe(first.driver)
+        expect(drivers[0].disposed).toBe(true)
+        expect(tunnels).toHaveLength(2)
+        expect(second.driver).toBe(drivers[1])
+    })
+
+    it('deduplicates concurrent builds', async () => {
+        const manager = createManager(registry, { createDriver, createTunnel })
+        const [a, b] = await Promise.all([manager.get('c'), manager.get('c')])
+        expect(a.driver).toBe(b.driver)
+        expect(drivers).toHaveLength(1)
+    })
+
+    it('clears the cache when a build fails, allowing retry', async () => {
+        const failing = vi
+            .fn()
+            .mockRejectedValueOnce(new Error('tunnel down'))
+            .mockImplementation(createTunnel)
+        resolved = { config: config(sshExtra()), source: 'store', hash: 'ssh1' }
+        const manager = createManager(registry, { createDriver, createTunnel: failing })
+        await expect(manager.get('c')).rejects.toThrow('tunnel down')
+        const retried = await manager.get('c')
+        expect(retried.driver).toBe(drivers[0])
+    })
+
+    it('invalidate disposes driver and tunnel', async () => {
+        resolved = { config: config(sshExtra()), source: 'store', hash: 'ssh1' }
+        const manager = createManager(registry, { createDriver, createTunnel })
+        await manager.get('c')
+        await manager.invalidate('c')
+        expect(drivers[0].disposed).toBe(true)
+        expect(tunnels[0].closed).toBe(true)
+    })
+
+    it('invalidate of an unknown name is a no-op', async () => {
+        const manager = createManager(registry, { createDriver, createTunnel })
+        await expect(manager.invalidate('ghost')).resolves.toBeUndefined()
+    })
+
+    it('invalidate during an in-flight build disposes the entry once built', async () => {
+        let release: () => void = () => {}
+        const gate = new Promise<void>((resolve) => {
+            release = resolve
+        })
+        const gatedTunnel = vi.fn(async (ssh: SshConfig, host: string, port: number) => {
+            await gate
+            return createTunnel(ssh, host, port)
+        })
+        resolved = { config: config(sshExtra()), source: 'store', hash: 'ssh1' }
+        const manager = createManager(registry, { createDriver, createTunnel: gatedTunnel })
+        const pending = manager.get('c')
+        const invalidated = manager.invalidate('c')
+        release()
+        await pending
+        await invalidated
+        expect(drivers[0].disposed).toBe(true)
+        expect(tunnels[0].closed).toBe(true)
+        const rebuilt = await manager.get('c')
+        expect(rebuilt.driver).toBe(drivers[1])
+    })
+
+    it('disposeAll disposes everything', async () => {
+        const manager = createManager(registry, { createDriver, createTunnel })
+        await manager.get('c')
+        await manager.disposeAll()
+        expect(drivers[0].disposed).toBe(true)
+        const again = await manager.get('c')
+        expect(again.driver).toBe(drivers[1])
+    })
+})