diff --git a/packages/cli/src/cli.ts b/packages/cli/src/cli.ts index 0e24ec282..e7bd7c426 100644 --- a/packages/cli/src/cli.ts +++ b/packages/cli/src/cli.ts @@ -24,6 +24,7 @@ export interface CliResult { interface RunOptions { backend?: Backend; + backendFactory?: (env?: NodeJS.ProcessEnv) => Promise; env?: NodeJS.ProcessEnv; diagnostics?: () => Promise; } @@ -56,6 +57,7 @@ class CliError extends Error { export async function runCli(argv: string[], options: RunOptions = {}): Promise { const env = options.env ?? process.env; + let ownedBackend: Backend | undefined; try { const flags = parseFlags(argv); @@ -69,7 +71,8 @@ export async function runCli(argv: string[], options: RunOptions = {}): Promise< return ok(`${usage()}\n`); } - const backend = options.backend ?? (await createBackend(env)); + const backendFactory = options.backendFactory ?? createBackend; + const backend = options.backend ?? (ownedBackend = await backendFactory(env)); if (args[0] === "doctor") { ensureArgCount(args, 1, "dbx doctor"); @@ -242,6 +245,8 @@ export async function runCli(argv: string[], options: RunOptions = {}): Promise< : "ERROR"; const wantsJson = argv.includes("--json"); return fail(code, message, wantsJson); + } finally { + await ownedBackend?.close?.().catch(() => {}); } } diff --git a/packages/cli/tests/cli.test.ts b/packages/cli/tests/cli.test.ts index 43742b760..220448682 100644 --- a/packages/cli/tests/cli.test.ts +++ b/packages/cli/tests/cli.test.ts @@ -87,6 +87,54 @@ test("runs read query as json", async () => { }); }); +test("closes backend resources created by the CLI", async () => { + let closed = false; + const result = await runCli(["query", "local", "select count(*) as total from users", "--json"], { + backendFactory: async () => + fakeBackend({ + close: async () => { + closed = true; + }, + }), + }); + + assert.equal(result.exitCode, 0); + assert.equal(closed, true); +}); + +test("closes backend resources created by the CLI after failures", async () => { + let closed = false; + const result = await runCli(["query", "local", "select count(*) as total from users", "--json"], { + backendFactory: async () => + fakeBackend({ + executeQuery: async () => { + throw new Error("boom"); + }, + close: async () => { + closed = true; + }, + }), + }); + + assert.equal(result.exitCode, 1); + assert.equal(JSON.parse(result.stderr).error.message, "boom"); + assert.equal(closed, true); +}); + +test("does not close caller-provided backend resources", async () => { + let closed = false; + const result = await runCli(["query", "local", "select count(*) as total from users", "--json"], { + backend: fakeBackend({ + close: async () => { + closed = true; + }, + }), + }); + + assert.equal(result.exitCode, 0); + assert.equal(closed, false); +}); + test("runs query as csv", async () => { const result = await runCli(["query", "local", "select count(*) as total from users", "--format", "csv"], { backend: fakeBackend(), diff --git a/packages/node-core/src/backend.ts b/packages/node-core/src/backend.ts index 6c053f175..d8a4f6685 100644 --- a/packages/node-core/src/backend.ts +++ b/packages/node-core/src/backend.ts @@ -5,6 +5,7 @@ import { removeConnection as desktopRemoveConnection, } from "./connections.js"; import { + closeDatabaseResources as desktopCloseDatabaseResources, describeTable as desktopDescribeTable, executeQuery as desktopExecuteQuery, listTables as desktopListTables, @@ -20,6 +21,7 @@ export interface Backend { listTables(config: ConnectionConfig, schema?: string): Promise; describeTable(config: ConnectionConfig, table: string, schema?: string): Promise; executeQuery(config: ConnectionConfig, sql: string, options?: QueryOptions): Promise; + close?(): Promise; } export async function createBackend(env: NodeJS.ProcessEnv = process.env): Promise { @@ -35,5 +37,6 @@ export async function createBackend(env: NodeJS.ProcessEnv = process.env): Promi listTables: desktopListTables, describeTable: desktopDescribeTable, executeQuery: desktopExecuteQuery, + close: desktopCloseDatabaseResources, }; } diff --git a/packages/node-core/src/database.ts b/packages/node-core/src/database.ts index 07db75ba6..f307b71d9 100644 --- a/packages/node-core/src/database.ts +++ b/packages/node-core/src/database.ts @@ -54,7 +54,7 @@ interface RqliteResponse { } const pools = new Map(); -const proxyTunnels = new Map(); +const proxyTunnels = new Map }>(); function poolKey(config: ConnectionConfig): string { return `${config.id}:${config.database || ""}`; @@ -62,6 +62,7 @@ function poolKey(config: ConnectionConfig): string { function evictPool(key: string, entry: PoolEntry) { pools.delete(key); + clearTimeout(entry.timer); if (entry.type === "pg") { (entry.pool as import("pg").Pool).end().catch(() => {}); } else { @@ -74,6 +75,33 @@ function resetIdleTimer(key: string, entry: PoolEntry) { entry.timer = setTimeout(() => evictPool(key, entry), IDLE_TIMEOUT_MS); } +export async function closeDatabaseResources(): Promise { + const poolEntries = [...pools.entries()]; + pools.clear(); + await Promise.all( + poolEntries.map(async ([, entry]) => { + clearTimeout(entry.timer); + if (entry.type === "pg") { + await (entry.pool as import("pg").Pool).end().catch(() => {}); + } else { + await (entry.pool as import("mysql2/promise").Pool).end().catch(() => {}); + } + }), + ); + + const tunnels = [...proxyTunnels.values()]; + proxyTunnels.clear(); + await Promise.all( + tunnels.map( + ({ server, sockets }) => + new Promise((resolve) => { + for (const socket of sockets) socket.destroy(); + server.close(() => resolve()); + }), + ), + ); +} + async function getPgPool(config: ConnectionConfig): Promise { const key = poolKey(config); const existing = pools.get(key); @@ -124,9 +152,14 @@ async function connectionEndpoint(config: ConnectionConfig): Promise<{ host: str const existing = proxyTunnels.get(config.id); if (existing) return { host: "127.0.0.1", port: existing.port }; + const sockets = new Set(); const server = createServer((inbound) => { + sockets.add(inbound); + inbound.once("close", () => sockets.delete(inbound)); connectViaProxy(config) .then((outbound) => { + sockets.add(outbound); + outbound.once("close", () => sockets.delete(outbound)); inbound.pipe(outbound); outbound.pipe(inbound); }) @@ -140,7 +173,7 @@ async function connectionEndpoint(config: ConnectionConfig): Promise<{ host: str else reject(new Error("Failed to bind proxy tunnel")); }); }); - proxyTunnels.set(config.id, { server, port }); + proxyTunnels.set(config.id, { server, port, sockets }); return { host: "127.0.0.1", port }; }