fix(cli): close direct query resources after command execution

Fix #745

packages/cli/src/cli.ts

@@ -24,6 +24,7 @@ export interface CliResult {
 
 interface RunOptions {
   backend?: Backend;
+  backendFactory?: (env?: NodeJS.ProcessEnv) => Promise<Backend>;
   env?: NodeJS.ProcessEnv;
   diagnostics?: () => Promise<DbxDiagnostics>;
 }
@@ -56,6 +57,7 @@ class CliError extends Error {
 
 export async function runCli(argv: string[], options: RunOptions = {}): Promise<CliResult> {
   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(() => {});
   }
 }
 

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(),

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<TableInfo[]>;
   describeTable(config: ConnectionConfig, table: string, schema?: string): Promise<ColumnInfo[]>;
   executeQuery(config: ConnectionConfig, sql: string, options?: QueryOptions): Promise<QueryResult>;
+  close?(): Promise<void>;
 }
 
 export async function createBackend(env: NodeJS.ProcessEnv = process.env): Promise<Backend> {
@@ -35,5 +37,6 @@ export async function createBackend(env: NodeJS.ProcessEnv = process.env): Promi
     listTables: desktopListTables,
     describeTable: desktopDescribeTable,
     executeQuery: desktopExecuteQuery,
+    close: desktopCloseDatabaseResources,
   };
 }

packages/node-core/src/database.ts

@@ -54,7 +54,7 @@ interface RqliteResponse {
 }
 
 const pools = new Map<string, PoolEntry>();
-const proxyTunnels = new Map<string, { server: Server; port: number }>();
+const proxyTunnels = new Map<string, { server: Server; port: number; sockets: Set<Socket> }>();
 
 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<void> {
+  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<void>((resolve) => {
+          for (const socket of sockets) socket.destroy();
+          server.close(() => resolve());
+        }),
+    ),
+  );
+}
+
 async function getPgPool(config: ConnectionConfig): Promise<import("pg").Pool> {
   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<Socket>();
   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 };
 }