Merge pull request #3433 from github/aeisenberg/re-import-test-db

Import testproj databases into workspace storage

Author: aeisenberg

extensions/ql-vscode/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## [UNRELEASED]
 
+- 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.3 - 29 February 2024
 
 - Update variant analysis view to show when cancelation is in progress. [#3405](https://github.com/github/vscode-codeql/pull/3405)

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

@@ -11,6 +11,7 @@ import {
   createWriteStream,
   remove,
   readdir,
+  copy,
 } from "fs-extra";
 import { basename, join } from "path";
 import type { Octokit } from "@octokit/rest";
@@ -64,7 +65,7 @@ export async function promptImportInternetDatabase(
 
   validateUrl(databaseUrl);
 
-  const item = await databaseArchiveFetcher(
+  const item = await fetchDatabaseToWorkspaceStorage(
     databaseUrl,
     {},
     databaseManager,
@@ -258,7 +259,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",
@@ -282,14 +283,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,
@@ -298,24 +300,27 @@ export async function importArchiveDatabase(
   cli: CodeQLCliServer,
 ): Promise<DatabaseItem | undefined> {
   try {
-    const item = await databaseArchiveFetcher(
+    const origin: DatabaseOrigin = {
+      type: databaseUrl.endsWith(".testproj") ? "testproj" : "archive",
+      path: Uri.parse(databaseUrl).fsPath,
+    };
+    const item = await fetchDatabaseToWorkspaceStorage(
       databaseUrl,
       {},
       databaseManager,
       storagePath,
       undefined,
-      {
-        type: "archive",
-        path: databaseUrl,
-      },
+      origin,
       progress,
       cli,
     );
     if (item) {
       await commandManager.execute("codeQLDatabases.focus");
       void showAndLogInformationMessage(
         extLogger,
-        "Database unzipped and imported successfully.",
+        origin.type === "testproj"
+          ? "Test database imported successfully."
+          : "Database unzipped and imported successfully.",
       );
     }
     return item;
@@ -332,10 +337,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.
@@ -346,7 +351,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,
@@ -374,7 +379,11 @@ async function databaseArchiveFetcher(
   );
 
   if (isFile(databaseUrl)) {
-    await readAndUnzip(databaseUrl, unzipPath, cli, progress);
+    if (origin.type === "testproj") {
+      await copyDatabase(databaseUrl, unzipPath, progress);
+    } else {
+      await readAndUnzip(databaseUrl, unzipPath, cli, progress);
+    }
   } else {
     await fetchAndUnzip(databaseUrl, requestHeaders, unzipPath, cli, progress);
   }
@@ -438,6 +447,8 @@ async function getStorageFolder(
     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);
     }
   }
 
@@ -484,6 +495,26 @@ function validateUrl(databaseUrl: string) {
   }
 }
 
+/**
+ * Copies a database folder from the file system into the workspace storage.
+ * @param scrDirURL the original location of the database as a URL string
+ * @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(
+  srcDirURL: string,
+  destDir: string,
+  progress?: ProgressCallback,
+) {
+  progress?.({
+    maxStep: 10,
+    step: 9,
+    message: `Copying database ${basename(destDir)} into the workspace`,
+  });
+  await ensureDir(destDir);
+  await copy(Uri.parse(srcDirURL).fsPath, 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";
@@ -145,7 +145,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;
   }
 
@@ -281,6 +282,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":
@@ -712,7 +714,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,
@@ -740,6 +742,60 @@ export class DatabaseUI extends DisposableObject {
     );
   }
 
+  private async handleImportTestDatabase(uri: Uri): Promise<void> {
+    return withProgress(
+      async (progress) => {
+        try {
+          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);
+          const baseName = basename(uri.fsPath);
+          if (existingItem !== undefined) {
+            progress({
+              maxStep: 9,
+              step: 1,
+              message: `Removing existing test database ${baseName}`,
+            });
+            await this.databaseManager.removeDatabaseItem(existingItem);
+          }
+
+          await importLocalDatabase(
+            this.app.commands,
+            uri.toString(true),
+            this.databaseManager,
+            this.storagePath,
+            progress,
+            this.queryServer.cliServer,
+          );
+
+          if (existingItem !== undefined) {
+            progress({
+              maxStep: 9,
+              step: 9,
+              message: `Successfully re-imported ${baseName}`,
+            });
+          }
+        } 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> {
@@ -959,7 +1015,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,