From 4d4a72bddb26d2a32c7357c7acd6deca8441f83d Mon Sep 17 00:00:00 2001
From: Koen Vlaswinkel <koesie10@github.com>
Date: Fri, 15 Sep 2023 11:45:23 +0200
Subject: [PATCH] Cleanup automodel temporary pack after use

---
 .../model-editor/auto-model-codeml-queries.ts | 22 +++++++++++++------
 .../auto-model-codeml-queries.test.ts         | 14 ++++++++----
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/extensions/ql-vscode/src/model-editor/auto-model-codeml-queries.ts b/extensions/ql-vscode/src/model-editor/auto-model-codeml-queries.ts
index eed7827ea..8773bfc13 100644
--- a/extensions/ql-vscode/src/model-editor/auto-model-codeml-queries.ts
+++ b/extensions/ql-vscode/src/model-editor/auto-model-codeml-queries.ts
@@ -63,10 +63,8 @@ export async function runAutoModelQueries({
   );
 
   // Generate a pack containing the candidate filters
-  const filterPackDir = await generateCandidateFilterPack(
-    databaseItem.language,
-    candidateMethods,
-  );
+  const { packDir: filterPackDir, cleanup: cleanupFilterPack } =
+    await generateCandidateFilterPack(databaseItem.language, candidateMethods);
 
   const additionalPacks = [...getOnDiskWorkspaceFolders(), filterPackDir];
   const extensionPacks = Object.keys(
@@ -85,6 +83,8 @@ export async function runAutoModelQueries({
     token: cancellationTokenSource.token,
   });
 
+  await cleanupFilterPack();
+
   if (!completedQuery) {
     return undefined;
   }
@@ -155,6 +155,11 @@ async function resolveAutomodelQuery(
   return queries[0];
 }
 
+type CandidateFilterPackResult = {
+  packDir: string;
+  cleanup: () => Promise<void>;
+};
+
 /**
  * generateCandidateFilterPack will create a temporary extension pack.
  * This pack will contain a filter that will restrict the automodel queries
@@ -167,9 +172,9 @@ async function resolveAutomodelQuery(
 export async function generateCandidateFilterPack(
   language: string,
   candidateMethods: MethodSignature[],
-): Promise<string> {
+): Promise<CandidateFilterPackResult> {
   // Pack resides in a temporary directory, to not pollute the workspace.
-  const packDir = (await dir({ unsafeCleanup: true })).path;
+  const { path: packDir, cleanup } = await dir({ unsafeCleanup: true });
 
   const syntheticConfigPack = {
     name: "codeql/automodel-filter",
@@ -208,7 +213,10 @@ export async function generateCandidateFilterPack(
   const filterFile = join(packDir, "filter.yml");
   await writeFile(filterFile, dumpYaml(filter), "utf8");
 
-  return packDir;
+  return {
+    packDir,
+    cleanup,
+  };
 }
 
 async function interpretAutomodelResults(
diff --git a/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/auto-model-codeml-queries.test.ts b/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/auto-model-codeml-queries.test.ts
index f447d35a6..465db8ebd 100644
--- a/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/auto-model-codeml-queries.test.ts
+++ b/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/auto-model-codeml-queries.test.ts
@@ -16,7 +16,7 @@ import { QueryRunner } from "../../../../src/query-server";
 import * as queryResolver from "../../../../src/local-queries/query-resolver";
 import { MethodSignature } from "../../../../src/model-editor/method";
 import { join } from "path";
-import { exists, readFile } from "fs-extra";
+import { pathExists, readFile } from "fs-extra";
 import { load as loadYaml } from "js-yaml";
 import { CancellationTokenSource } from "vscode-jsonrpc";
 import { QueryOutputDir } from "../../../../src/run-queries-shared";
@@ -176,12 +176,15 @@ describe("generateCandidateFilterPack", () => {
         methodParameters: "()",
       },
     ];
-    const packDir = await generateCandidateFilterPack("java", candidateMethods);
+    const { packDir, cleanup } = await generateCandidateFilterPack(
+      "java",
+      candidateMethods,
+    );
     expect(packDir).not.toBeUndefined();
     const qlpackFile = join(packDir, "codeql-pack.yml");
-    expect(await exists(qlpackFile)).toBe(true);
+    expect(await pathExists(qlpackFile)).toBe(true);
     const filterFile = join(packDir, "filter.yml");
-    expect(await exists(filterFile)).toBe(true);
+    expect(await pathExists(filterFile)).toBe(true);
     // Read the contents of filterFile and parse as yaml
     const yaml = await loadYaml(await readFile(filterFile, "utf8"));
     const extensions = yaml.extensions;
@@ -193,5 +196,8 @@ describe("generateCandidateFilterPack", () => {
     expect(extension.data).toBeInstanceOf(Array);
     expect(extension.data).toHaveLength(1);
     expect(extension.data[0]).toEqual(["org.my", "A", "x", "()"]);
+
+    await cleanup();
+    expect(await pathExists(packDir)).toBe(false);
   });
 });