Import testproj databases into workspace storage

extensions/ql-vscode/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 - Remove support for CodeQL CLI versions older than 2.13.5. [#3371](https://github.com/github/vscode-codeql/pull/3371)
 - Add a timeout to downloading databases and the CodeQL CLI. These can be changed using the `codeQL.addingDatabases.downloadTimeout` and `codeQL.cli.downloadTimeout` settings respectively. [#3373](https://github.com/github/vscode-codeql/pull/3373)
+- Databases created from [CodeQL test cases](https://docs.github.com/en/code-security/codeql-cli/using-the-advanced-functionality-of-the-codeql-cli/testing-custom-queries) are now copied into a shared VS Code storage location. This avoids a bug where re-running test cases would fail if the test's database is already imported into the workspace. [#3433](https://github.com/github/vscode-codeql/pull/3433)
 
 ## 1.12.2 - 14 February 2024
 

extensions/ql-vscode/package.json

@@ -738,6 +738,10 @@
         "command": "codeQL.setCurrentDatabase",
         "title": "CodeQL: Set Current Database"
       },
+      {
+        "command": "codeQL.importTestDatabase",
+        "title": "CodeQL: (Re-)Import Test Database"
+      },
       {
         "command": "codeQL.getCurrentDatabase",
         "title": "CodeQL: Get Current Database"
@@ -1322,7 +1326,12 @@
         {
           "command": "codeQL.setCurrentDatabase",
           "group": "9_qlCommands",
-          "when": "resourceScheme == codeql-zip-archive || explorerResourceIsFolder || resourceExtname == .zip"
+          "when": "resourceExtname != .testproj && (resourceScheme == codeql-zip-archive || explorerResourceIsFolder || resourceExtname == .zipz)"
+        },
+        {
+          "command": "codeQL.importTestDatabase",
+          "group": "9_qlCommands",
+          "when": "explorerResourceIsFolder && resourceExtname == .testproj"
         },
         {
           "command": "codeQL.viewAstContextExplorer",
@@ -1476,6 +1485,10 @@
           "command": "codeQL.setCurrentDatabase",
           "when": "false"
         },
+        {
+          "command": "codeQL.importTestDatabase",
+          "when": "false"
+        },
         {
           "command": "codeQL.getCurrentDatabase",
           "when": "false"
@@ -2018,7 +2031,6 @@
     "@types/tar-stream": "^3.1.3",
     "@types/through2": "^2.0.36",
     "@types/tmp": "^0.2.6",
-    "@types/unzipper": "^0.10.1",
     "@types/vscode": "^1.82.0",
     "@types/yauzl": "^2.10.3",
     "@typescript-eslint/eslint-plugin": "^6.19.0",

extensions/ql-vscode/src/common/commands.ts

@@ -220,6 +220,7 @@ export type LocalDatabasesCommands = {
 
   // Explorer context menu
   "codeQL.setCurrentDatabase": (uri: Uri) => Promise<void>;
+  "codeQL.importTestDatabase": (uri: Uri) => Promise<void>;
 
   // Database panel view title commands
   "codeQLDatabases.chooseDatabaseFolder": () => Promise<void>;

extensions/ql-vscode/src/databases/database-fetcher.ts

@@ -12,6 +12,7 @@ import {
   remove,
   stat,
   readdir,
+  copy,
 } from "fs-extra";
 import { basename, join } from "path";
 import type { Octokit } from "@octokit/rest";
@@ -61,7 +62,7 @@ export async function promptImportInternetDatabase(
 
   validateUrl(databaseUrl);
 
-  const item = await databaseArchiveFetcher(
+  const item = await fetchDatabaseToWorkspaceStorage(
     databaseUrl,
     {},
     databaseManager,
@@ -254,7 +255,7 @@ export async function downloadGitHubDatabaseFromUrl(
    * We only need the actual token string.
    */
   const octokitToken = ((await octokit.auth()) as { token: string })?.token;
-  return await databaseArchiveFetcher(
+  return await fetchDatabaseToWorkspaceStorage(
     databaseUrl,
     {
       Accept: "application/zip",
@@ -278,14 +279,15 @@ export async function downloadGitHubDatabaseFromUrl(
 }
 
 /**
- * Imports a database from a local archive.
+ * Imports a database from a local archive or a test database that is in a folder
+ * ending with `.testproj`.
  *
- * @param databaseUrl the file url of the archive to import
+ * @param databaseUrl the file url of the archive or directory to import
  * @param databaseManager the DatabaseManager
  * @param storagePath where to store the unzipped database.
  * @param cli the CodeQL CLI server
  */
-export async function importArchiveDatabase(
+export async function importLocalDatabase(
   commandManager: AppCommandManager,
   databaseUrl: string,
   databaseManager: DatabaseManager,
@@ -294,16 +296,18 @@ export async function importArchiveDatabase(
   cli: CodeQLCliServer,
 ): Promise<DatabaseItem | undefined> {
   try {
-    const item = await databaseArchiveFetcher(
+    const origin: DatabaseOrigin = {
+      type: databaseUrl.endsWith(".testproj") ? "testproj" : "archive",
+      // TODO validate that archive origins can use a file path, not a URI
+      path: Uri.parse(databaseUrl).fsPath,
+    };
+    const item = await fetchDatabaseToWorkspaceStorage(
       databaseUrl,
       {},
       databaseManager,
       storagePath,
       undefined,
-      {
-        type: "archive",
-        path: databaseUrl,
-      },
+      origin,
       progress,
       cli,
     );
@@ -328,10 +332,10 @@ export async function importArchiveDatabase(
 }
 
 /**
- * Fetches an archive database. The database might be on the internet
+ * Fetches a database into workspace storage. The database might be on the internet
  * or in the local filesystem.
  *
- * @param databaseUrl URL from which to grab the database
+ * @param databaseUrl URL from which to grab the database. This could be a local archive file, a local directory, or a remote URL.
  * @param requestHeaders Headers to send with the request
  * @param databaseManager the DatabaseManager
  * @param storagePath where to store the unzipped database.
@@ -342,7 +346,7 @@ export async function importArchiveDatabase(
  * @param makeSelected make the new database selected in the databases panel (default: true)
  * @param addSourceArchiveFolder whether to add a workspace folder containing the source archive to the workspace
  */
-async function databaseArchiveFetcher(
+async function fetchDatabaseToWorkspaceStorage(
   databaseUrl: string,
   requestHeaders: { [key: string]: string },
   databaseManager: DatabaseManager,
@@ -366,7 +370,11 @@ async function databaseArchiveFetcher(
   const unzipPath = await getStorageFolder(storagePath, databaseUrl);
 
   if (isFile(databaseUrl)) {
-    await readAndUnzip(databaseUrl, unzipPath, cli, progress);
+    if (origin.type === "testproj") {
+      await copyDatabase(origin.path, unzipPath, progress);
+    } else {
+      await readAndUnzip(databaseUrl, unzipPath, cli, progress);
+    }
   } else {
     await fetchAndUnzip(databaseUrl, requestHeaders, unzipPath, cli, progress);
   }
@@ -416,6 +424,8 @@ async function getStorageFolder(storagePath: string, urlStr: string) {
   let lastName = basename(url.path).substring(0, 250);
   if (lastName.endsWith(".zip")) {
     lastName = lastName.substring(0, lastName.length - 4);
+  } else if (lastName.endsWith(".testproj")) {
+    lastName = lastName.substring(0, lastName.length - 9);
   }
 
   const realpath = await fs_realpath(storagePath);
@@ -446,6 +456,26 @@ function validateUrl(databaseUrl: string) {
   }
 }
 
+/**
+ * Copies a database folder from the file system into the workspace storage.
+ * @param scrDir the original location of the database
+ * @param destDir the location to copy the database to. This should be a folder in the workspace storage.
+ * @param progress callback to send progress messages to
+ */
+async function copyDatabase(
+  scrDir: string,
+  destDir: string,
+  progress?: ProgressCallback,
+) {
+  progress?.({
+    maxStep: 10,
+    step: 9,
+    message: `Copying database ${basename(destDir)} into the workspace`,
+  });
+  await ensureDir(destDir);
+  await copy(scrDir, destDir);
+}
+
 async function readAndUnzip(
   zipUrl: string,
   unzipPath: string,

extensions/ql-vscode/src/databases/local-databases-ui.ts

@@ -43,7 +43,7 @@ import {
   showAndLogErrorMessage,
 } from "../common/logging";
 import {
-  importArchiveDatabase,
+  importLocalDatabase,
   promptImportGithubDatabase,
   promptImportInternetDatabase,
 } from "./database-fetcher";
@@ -146,7 +146,8 @@ class DatabaseTreeDataProvider
       item.iconPath = new ThemeIcon("error", new ThemeColor("errorForeground"));
     }
     item.tooltip = element.databaseUri.fsPath;
-    item.description = element.language;
+    item.description =
+      element.language + (element.origin?.type === "testproj" ? " (test)" : "");
     return item;
   }
 
@@ -282,6 +283,7 @@ export class DatabaseUI extends DisposableObject {
         this.handleChooseDatabaseInternet.bind(this),
       "codeQL.chooseDatabaseGithub": this.handleChooseDatabaseGithub.bind(this),
       "codeQL.setCurrentDatabase": this.handleSetCurrentDatabase.bind(this),
+      "codeQL.importTestDatabase": this.handleImportTestDatabase.bind(this),
       "codeQL.setDefaultTourDatabase":
         this.handleSetDefaultTourDatabase.bind(this),
       "codeQL.upgradeCurrentDatabase":
@@ -716,7 +718,7 @@ export class DatabaseUI extends DisposableObject {
         try {
           // Assume user has selected an archive if the file has a .zip extension
           if (uri.path.endsWith(".zip")) {
-            await importArchiveDatabase(
+            await importLocalDatabase(
               this.app.commands,
               uri.toString(true),
               this.databaseManager,
@@ -744,6 +746,54 @@ export class DatabaseUI extends DisposableObject {
     );
   }
 
+  private async handleImportTestDatabase(uri: Uri): Promise<void> {
+    return withProgress(
+      async (progress) => {
+        try {
+          // Assume user has selected an archive if the file has a .zip extension
+          if (!uri.path.endsWith(".testproj")) {
+            throw new Error(
+              "Please select a valid test database to import. Test databases end with `.testproj`.",
+            );
+          }
+
+          // Check if the database is already in the workspace. If
+          // so, delete it first before importing the new one.
+          const existingItem = this.databaseManager.findTestDatabase(uri);
+          if (existingItem !== undefined) {
+            progress({
+              maxStep: 9,
+              step: 1,
+              message: `Removing existing test database ${basename(
+                uri.fsPath,
+              )}`,
+            });
+            await this.databaseManager.removeDatabaseItem(existingItem);
+          }
+
+          await importLocalDatabase(
+            this.app.commands,
+            uri.toString(true),
+            this.databaseManager,
+            this.storagePath,
+            progress,
+            this.queryServer.cliServer,
+          );
+        } catch (e) {
+          // rethrow and let this be handled by default error handling.
+          throw new Error(
+            `Could not set database to ${basename(
+              uri.fsPath,
+            )}. Reason: ${getErrorMessage(e)}`,
+          );
+        }
+      },
+      {
+        title: "(Re-)importing test database from directory",
+      },
+    );
+  }
+
   private async handleRemoveDatabase(
     databaseItems: DatabaseItem[],
   ): Promise<void> {
@@ -963,7 +1013,7 @@ export class DatabaseUI extends DisposableObject {
         } else {
           // we are selecting a database archive. Must unzip into a workspace-controlled area
           // before importing.
-          return await importArchiveDatabase(
+          return await importLocalDatabase(
             this.app.commands,
             uri.toString(true),
             this.databaseManager,