Merge pull request #2250 from github/elena/yer-a-wizard-query

Adapt skeleton pack generation to work with new queries panel

extensions/ql-vscode/package.json

@@ -361,6 +361,10 @@
         "command": "codeQL.quickQuery",
         "title": "CodeQL: Quick Query"
       },
+      {
+        "command": "codeQL.createSkeletonQuery",
+        "title": "CodeQL: Create Query"
+      },
       {
         "command": "codeQL.openDocumentation",
         "title": "CodeQL: Open Documentation"

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

@@ -99,6 +99,7 @@ export type LocalQueryCommands = {
   "codeQL.quickEvalContextEditor": (uri: Uri) => Promise<void>;
   "codeQL.codeLensQuickEval": (uri: Uri, range: Range) => Promise<void>;
   "codeQL.quickQuery": () => Promise<void>;
+  "codeQL.createSkeletonQuery": () => Promise<void>;
 };
 
 export type ResultsViewCommands = {

extensions/ql-vscode/src/databaseFetcher.ts

@@ -78,6 +78,10 @@ export async function promptImportInternetDatabase(
  *
  * @param databaseManager the DatabaseManager
  * @param storagePath where to store the unzipped database.
+ * @param credentials the credentials to use to authenticate with GitHub
+ * @param progress the progress callback
+ * @param token the cancellation token
+ * @param cli the CodeQL CLI server
  */
 export async function promptImportGithubDatabase(
   commandManager: AppCommandManager,
@@ -103,6 +107,49 @@ export async function promptImportGithubDatabase(
     return;
   }
 
+  const databaseItem = await downloadGitHubDatabase(
+    githubRepo,
+    databaseManager,
+    storagePath,
+    credentials,
+    progress,
+    token,
+    cli,
+  );
+
+  if (databaseItem) {
+    await commandManager.execute("codeQLDatabases.focus");
+    void showAndLogInformationMessage(
+      "Database downloaded and imported successfully.",
+    );
+    return databaseItem;
+  }
+
+  return;
+}
+
+/**
+ * Downloads a database from GitHub
+ *
+ * @param githubRepo the GitHub repository to download the database from
+ * @param databaseManager the DatabaseManager
+ * @param storagePath where to store the unzipped database.
+ * @param credentials the credentials to use to authenticate with GitHub
+ * @param progress the progress callback
+ * @param token the cancellation token
+ * @param cli the CodeQL CLI server
+ * @param language the language to download. If undefined, the user will be prompted to choose a language.
+ **/
+export async function downloadGitHubDatabase(
+  githubRepo: string,
+  databaseManager: DatabaseManager,
+  storagePath: string,
+  credentials: Credentials | undefined,
+  progress: ProgressCallback,
+  token: CancellationToken,
+  cli?: CodeQLCliServer,
+  language?: string,
+): Promise<DatabaseItem | undefined> {
   const nwo = getNwoFromGitHubUrl(githubRepo) || githubRepo;
   if (!isValidGitHubNwo(nwo)) {
     throw new Error(`Invalid GitHub repository: ${githubRepo}`);
@@ -112,7 +159,12 @@ export async function promptImportGithubDatabase(
     ? await credentials.getOctokit()
     : new Octokit.Octokit({ retry });
 
-  const result = await convertGithubNwoToDatabaseUrl(nwo, octokit, progress);
+  const result = await convertGithubNwoToDatabaseUrl(
+    nwo,
+    octokit,
+    progress,
+    language,
+  );
   if (!result) {
     return;
   }
@@ -130,7 +182,7 @@ export async function promptImportGithubDatabase(
    * We only need the actual token string.
    */
   const octokitToken = ((await octokit.auth()) as { token: string })?.token;
-  const item = await databaseArchiveFetcher(
+  return await databaseArchiveFetcher(
     databaseUrl,
     {
       Accept: "application/zip",
@@ -143,14 +195,6 @@ export async function promptImportGithubDatabase(
     token,
     cli,
   );
-  if (item) {
-    await commandManager.execute("codeQLDatabases.focus");
-    void showAndLogInformationMessage(
-      "Database downloaded and imported successfully.",
-    );
-    return item;
-  }
-  return;
 }
 
 /**
@@ -450,6 +494,7 @@ export async function convertGithubNwoToDatabaseUrl(
   nwo: string,
   octokit: Octokit.Octokit,
   progress: ProgressCallback,
+  language?: string,
 ): Promise<
   | {
       databaseUrl: string;
@@ -468,9 +513,11 @@ export async function convertGithubNwoToDatabaseUrl(
 
     const languages = response.data.map((db: any) => db.language);
 
-    const language = await promptForLanguage(languages, progress);
-    if (!language) {
-      return;
+    if (!language || !languages.includes(language)) {
+      language = await promptForLanguage(languages, progress);
+      if (!language) {
+        return;
+      }
     }
 
     return {
@@ -484,7 +531,7 @@ export async function convertGithubNwoToDatabaseUrl(
   }
 }
 
-async function promptForLanguage(
+export async function promptForLanguage(
   languages: string[],
   progress: ProgressCallback,
 ): Promise<string | undefined> {

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

@@ -897,6 +897,20 @@ export class DatabaseManager extends DisposableObject {
     }
   }
 
+  public async digForDatabaseItem(
+    language: string,
+    databaseNwo: string,
+  ): Promise<DatabaseItem | undefined> {
+    const dbItems = this.databaseItems || [];
+    const dbs = dbItems.filter(
+      (db) => db.language === language && db.name === databaseNwo,
+    );
+    if (dbs.length === 0) {
+      return undefined;
+    }
+    return dbs[0];
+  }
+
   /**
    * Returns the index of the workspace folder that corresponds to the source archive of `item`
    * if there is one, and -1 otherwise.

extensions/ql-vscode/src/local-queries.ts

@@ -8,7 +8,7 @@ import {
   window,
 } from "vscode";
 import { BaseLogger, extLogger, Logger, TeeLogger } from "./common";
-import { MAX_QUERIES } from "./config";
+import { isCanary, MAX_QUERIES } from "./config";
 import { gatherQlFiles } from "./pure/files";
 import { basename } from "path";
 import {
@@ -51,6 +51,7 @@ import { App } from "./common/app";
 import { DisposableObject } from "./pure/disposable-object";
 import { QueryResultType } from "./pure/new-messages";
 import { redactableError } from "./pure/errors";
+import { SkeletonQueryWizard } from "./skeleton-query-wizard";
 
 interface DatabaseQuickPickItem extends QuickPickItem {
   databaseItem: DatabaseItem;
@@ -237,6 +238,7 @@ export class LocalQueries extends DisposableObject {
       "codeQL.quickEvalContextEditor": this.quickEval.bind(this),
       "codeQL.codeLensQuickEval": this.codeLensQuickEval.bind(this),
       "codeQL.quickQuery": this.quickQuery.bind(this),
+      "codeQL.createSkeletonQuery": this.createSkeletonQuery.bind(this),
     };
   }
 
@@ -375,6 +377,26 @@ export class LocalQueries extends DisposableObject {
     );
   }
 
+  private async createSkeletonQuery(): Promise<void> {
+    await withProgress(
+      async (progress: ProgressCallback, token: CancellationToken) => {
+        const credentials = isCanary() ? this.app.credentials : undefined;
+        const skeletonQueryWizard = new SkeletonQueryWizard(
+          this.cliServer,
+          progress,
+          credentials,
+          extLogger,
+          this.databaseManager,
+          token,
+        );
+        await skeletonQueryWizard.execute();
+      },
+      {
+        title: "Create Query",
+      },
+    );
+  }
+
   /**
    * Creates a new `LocalQueryRun` object to track a query evaluation. This creates a timestamp
    * file in the query's output directory, creates a `LocalQueryInfo` object, and registers that

extensions/ql-vscode/src/qlpack-generator.ts

@@ -66,8 +66,8 @@ export class QlPackGenerator {
     await writeFile(qlPackFilePath, this.header + dump(qlPackYml), "utf8");
   }
 
-  private async createExampleQlFile() {
-    const exampleQlFilePath = join(this.folderUri.fsPath, "example.ql");
+  public async createExampleQlFile(fileName = "example.ql") {
+    const exampleQlFilePath = join(this.folderUri.fsPath, fileName);
 
     const exampleQl = `
 /**

extensions/ql-vscode/src/skeleton-query-wizard.ts

@@ -0,0 +1,247 @@
+import { join } from "path";
+import { CancellationToken, Uri, workspace, window as Window } from "vscode";
+import { CodeQLCliServer } from "./cli";
+import { OutputChannelLogger } from "./common";
+import { Credentials } from "./common/authentication";
+import { QueryLanguage } from "./common/query-language";
+import { askForLanguage, isFolderAlreadyInWorkspace } from "./helpers";
+import { getErrorMessage } from "./pure/helpers-pure";
+import { QlPackGenerator } from "./qlpack-generator";
+import { DatabaseManager } from "./local-databases";
+import * as databaseFetcher from "./databaseFetcher";
+import { ProgressCallback } from "./progress";
+
+type QueryLanguagesToDatabaseMap = Record<string, string>;
+
+export const QUERY_LANGUAGE_TO_DATABASE_REPO: QueryLanguagesToDatabaseMap = {
+  cpp: "protocolbuffers/protobuf",
+  csharp: "dotnet/efcore",
+  go: "evanw/esbuild",
+  java: "google/guava",
+  javascript: "facebook/react",
+  python: "pallets/flask",
+  ruby: "rails/rails",
+  swift: "Alamofire/Alamofire",
+};
+
+export class SkeletonQueryWizard {
+  private language: string | undefined;
+  private fileName = "example.ql";
+  private storagePath: string | undefined;
+
+  constructor(
+    private readonly cliServer: CodeQLCliServer,
+    private readonly progress: ProgressCallback,
+    private readonly credentials: Credentials | undefined,
+    private readonly extLogger: OutputChannelLogger,
+    private readonly databaseManager: DatabaseManager,
+    private readonly token: CancellationToken,
+  ) {}
+
+  private get folderName() {
+    return `codeql-custom-queries-${this.language}`;
+  }
+
+  public async execute() {
+    // show quick pick to choose language
+    this.language = await this.chooseLanguage();
+    if (!this.language) {
+      return;
+    }
+
+    this.storagePath = this.getFirstStoragePath();
+
+    const skeletonPackAlreadyExists = isFolderAlreadyInWorkspace(
+      this.folderName,
+    );
+
+    if (skeletonPackAlreadyExists) {
+      // just create a new example query file in skeleton QL pack
+      await this.createExampleFile();
+      // select existing database for language
+      await this.selectExistingDatabase();
+    } else {
+      // generate a new skeleton QL pack with query file
+      await this.createQlPack();
+      // download database based on language and select it
+      await this.downloadDatabase();
+    }
+
+    // open a query file
+    await this.openExampleFile();
+  }
+
+  private async openExampleFile() {
+    if (this.folderName === undefined || this.storagePath === undefined) {
+      throw new Error("Path to folder is undefined");
+    }
+
+    const queryFileUri = Uri.file(
+      join(this.storagePath, this.folderName, this.fileName),
+    );
+
+    try {
+      void workspace.openTextDocument(queryFileUri).then((doc) => {
+        void Window.showTextDocument(doc);
+      });
+    } catch (e: unknown) {
+      void this.extLogger.log(
+        `Could not open example query file: ${getErrorMessage(e)}`,
+      );
+    }
+  }
+
+  public getFirstStoragePath() {
+    const workspaceFolders = workspace.workspaceFolders;
+
+    if (!workspaceFolders || workspaceFolders.length === 0) {
+      throw new Error("No workspace folders found");
+    }
+
+    const firstFolder = workspaceFolders[0];
+
+    // For the vscode-codeql-starter repo, the first folder will be a ql pack
+    // so we need to get the parent folder
+    if (firstFolder.uri.path.includes("codeql-custom-queries")) {
+      // slice off the last part of the path and return the parent folder
+      return firstFolder.uri.path.split("/").slice(0, -1).join("/");
+    } else {
+      // if the first folder is not a ql pack, then we are in a normal workspace
+      return firstFolder.uri.path;
+    }
+  }
+
+  private async chooseLanguage() {
+    this.progress({
+      message: "Choose language",
+      step: 1,
+      maxStep: 1,
+    });
+
+    return await askForLanguage(this.cliServer, false);
+  }
+
+  private async createQlPack() {
+    if (this.folderName === undefined) {
+      throw new Error("Folder name is undefined");
+    }
+
+    this.progress({
+      message: "Creating skeleton QL pack around query",
+      step: 2,
+      maxStep: 2,
+    });
+
+    try {
+      const qlPackGenerator = new QlPackGenerator(
+        this.folderName,
+        this.language as QueryLanguage,
+        this.cliServer,
+        this.storagePath,
+      );
+
+      await qlPackGenerator.generate();
+    } catch (e: unknown) {
+      void this.extLogger.log(
+        `Could not create skeleton QL pack: ${getErrorMessage(e)}`,
+      );
+    }
+  }
+
+  private async createExampleFile() {
+    if (this.folderName === undefined) {
+      throw new Error("Folder name is undefined");
+    }
+
+    this.progress({
+      message:
+        "Skeleton query pack already exists. Creating additional query example file.",
+      step: 2,
+      maxStep: 2,
+    });
+
+    try {
+      const qlPackGenerator = new QlPackGenerator(
+        this.folderName,
+        this.language as QueryLanguage,
+        this.cliServer,
+        this.storagePath,
+      );
+
+      this.fileName = await this.determineNextFileName(this.folderName);
+      await qlPackGenerator.createExampleQlFile(this.fileName);
+    } catch (e: unknown) {
+      void this.extLogger.log(
+        `Could not create skeleton QL pack: ${getErrorMessage(e)}`,
+      );
+    }
+  }
+
+  private async determineNextFileName(folderName: string): Promise<string> {
+    if (this.storagePath === undefined) {
+      throw new Error("Workspace storage path is undefined");
+    }
+
+    const folderUri = Uri.file(join(this.storagePath, folderName));
+    const files = await workspace.fs.readDirectory(folderUri);
+    const qlFiles = files.filter(([filename, _fileType]) =>
+      filename.match(/example[0-9]*.ql/),
+    );
+
+    return `example${qlFiles.length + 1}.ql`;
+  }
+
+  private async downloadDatabase() {
+    if (this.storagePath === undefined) {
+      throw new Error("Workspace storage path is undefined");
+    }
+
+    if (this.language === undefined) {
+      throw new Error("Language is undefined");
+    }
+
+    this.progress({
+      message: "Downloading database",
+      step: 3,
+      maxStep: 3,
+    });
+
+    const githubRepoNwo = QUERY_LANGUAGE_TO_DATABASE_REPO[this.language];
+
+    await databaseFetcher.downloadGitHubDatabase(
+      githubRepoNwo,
+      this.databaseManager,
+      this.storagePath,
+      this.credentials,
+      this.progress,
+      this.token,
+      this.cliServer,
+      this.language,
+    );
+  }
+
+  private async selectExistingDatabase() {
+    if (this.language === undefined) {
+      throw new Error("Language is undefined");
+    }
+
+    if (this.storagePath === undefined) {
+      throw new Error("Workspace storage path is undefined");
+    }
+
+    const databaseNwo = QUERY_LANGUAGE_TO_DATABASE_REPO[this.language];
+
+    const databaseItem = await this.databaseManager.digForDatabaseItem(
+      this.language,
+      databaseNwo,
+    );
+
+    if (databaseItem) {
+      // select the found database
+      await this.databaseManager.setCurrentDatabaseItem(databaseItem);
+    } else {
+      // download new database and select it
+      await this.downloadDatabase();
+    }
+  }
+}

extensions/ql-vscode/test/vscode-tests/cli-integration/skeleton-query-wizard.test.ts

@@ -0,0 +1,308 @@
+import { CodeQLCliServer } from "../../../src/cli";
+import {
+  QUERY_LANGUAGE_TO_DATABASE_REPO,
+  SkeletonQueryWizard,
+} from "../../../src/skeleton-query-wizard";
+import { mockedObject, mockedQuickPickItem } from "../utils/mocking.helpers";
+import * as tmp from "tmp";
+import { TextDocument, window, workspace, WorkspaceFolder } from "vscode";
+import { extLogger } from "../../../src/common";
+import { QlPackGenerator } from "../../../src/qlpack-generator";
+import * as helpers from "../../../src/helpers";
+import { createFileSync, ensureDirSync, removeSync } from "fs-extra";
+import { join } from "path";
+import { CancellationTokenSource } from "vscode-jsonrpc";
+import { testCredentialsWithStub } from "../../factories/authentication";
+import { DatabaseItem, DatabaseManager } from "../../../src/local-databases";
+import * as databaseFetcher from "../../../src/databaseFetcher";
+
+jest.setTimeout(40_000);
+
+describe("SkeletonQueryWizard", () => {
+  let wizard: SkeletonQueryWizard;
+  let dir: tmp.DirResult;
+  let storagePath: string;
+  let quickPickSpy: jest.SpiedFunction<typeof window.showQuickPick>;
+  let generateSpy: jest.SpiedFunction<
+    typeof QlPackGenerator.prototype.generate
+  >;
+  let createExampleQlFileSpy: jest.SpiedFunction<
+    typeof QlPackGenerator.prototype.createExampleQlFile
+  >;
+  let downloadGitHubDatabaseSpy: jest.SpiedFunction<
+    typeof databaseFetcher.downloadGitHubDatabase
+  >;
+  let openTextDocumentSpy: jest.SpiedFunction<
+    typeof workspace.openTextDocument
+  >;
+
+  const token = new CancellationTokenSource().token;
+  const credentials = testCredentialsWithStub();
+  const chosenLanguage = "ruby";
+  const mockDatabaseManager = mockedObject<DatabaseManager>({
+    setCurrentDatabaseItem: jest.fn(),
+    digForDatabaseItem: jest.fn(),
+  });
+  const mockCli = mockedObject<CodeQLCliServer>({
+    resolveLanguages: jest
+      .fn()
+      .mockResolvedValue([
+        "ruby",
+        "javascript",
+        "go",
+        "java",
+        "python",
+        "csharp",
+        "cpp",
+      ]),
+    getSupportedLanguages: jest.fn(),
+  });
+
+  beforeEach(async () => {
+    dir = tmp.dirSync({
+      prefix: "skeleton_query_wizard_",
+      unsafeCleanup: true,
+    });
+
+    storagePath = dir.name;
+
+    jest.spyOn(workspace, "workspaceFolders", "get").mockReturnValue([
+      {
+        name: `codespaces-codeql`,
+        uri: { path: storagePath },
+      },
+      {
+        name: "/second/folder/path",
+        uri: { path: storagePath },
+      },
+    ] as WorkspaceFolder[]);
+
+    quickPickSpy = jest
+      .spyOn(window, "showQuickPick")
+      .mockResolvedValueOnce(mockedQuickPickItem(chosenLanguage));
+    generateSpy = jest
+      .spyOn(QlPackGenerator.prototype, "generate")
+      .mockResolvedValue(undefined);
+    createExampleQlFileSpy = jest
+      .spyOn(QlPackGenerator.prototype, "createExampleQlFile")
+      .mockResolvedValue(undefined);
+    downloadGitHubDatabaseSpy = jest
+      .spyOn(databaseFetcher, "downloadGitHubDatabase")
+      .mockResolvedValue(undefined);
+    openTextDocumentSpy = jest
+      .spyOn(workspace, "openTextDocument")
+      .mockResolvedValue({} as TextDocument);
+
+    wizard = new SkeletonQueryWizard(
+      mockCli,
+      jest.fn(),
+      credentials,
+      extLogger,
+      mockDatabaseManager,
+      token,
+    );
+  });
+
+  afterEach(async () => {
+    dir.removeCallback();
+  });
+
+  it("should prompt for language", async () => {
+    await wizard.execute();
+
+    expect(mockCli.getSupportedLanguages).toHaveBeenCalled();
+    expect(quickPickSpy).toHaveBeenCalled();
+  });
+
+  describe("if QL pack doesn't exist", () => {
+    beforeEach(() => {
+      jest.spyOn(helpers, "isFolderAlreadyInWorkspace").mockReturnValue(false);
+    });
+    it("should try to create a new QL pack based on the language", async () => {
+      await wizard.execute();
+
+      expect(generateSpy).toHaveBeenCalled();
+    });
+
+    it("should download database for selected language", async () => {
+      await wizard.execute();
+
+      expect(downloadGitHubDatabaseSpy).toHaveBeenCalled();
+    });
+
+    it("should open the query file", async () => {
+      await wizard.execute();
+
+      expect(openTextDocumentSpy).toHaveBeenCalledWith(
+        expect.objectContaining({
+          path: expect.stringMatching("example.ql"),
+        }),
+      );
+    });
+  });
+
+  describe("if QL pack exists", () => {
+    beforeEach(async () => {
+      jest.spyOn(helpers, "isFolderAlreadyInWorkspace").mockReturnValue(true);
+
+      // create a skeleton codeql-custom-queries-${language} folder
+      // with an example QL file inside
+      ensureDirSync(
+        join(dir.name, `codeql-custom-queries-${chosenLanguage}`, "example.ql"),
+      );
+    });
+
+    it("should create new query file in the same QL pack folder", async () => {
+      await wizard.execute();
+
+      expect(createExampleQlFileSpy).toHaveBeenCalledWith("example2.ql");
+    });
+
+    it("should only take into account example QL files", async () => {
+      createFileSync(
+        join(dir.name, `codeql-custom-queries-${chosenLanguage}`, "MyQuery.ql"),
+      );
+
+      await wizard.execute();
+
+      expect(createExampleQlFileSpy).toHaveBeenCalledWith("example2.ql");
+    });
+
+    describe("if QL pack has no query file", () => {
+      it("should create new query file in the same QL pack folder", async () => {
+        removeSync(
+          join(
+            dir.name,
+            `codeql-custom-queries-${chosenLanguage}`,
+            "example.ql",
+          ),
+        );
+        await wizard.execute();
+
+        expect(createExampleQlFileSpy).toHaveBeenCalledWith("example1.ql");
+      });
+
+      it("should open the query file", async () => {
+        removeSync(
+          join(
+            dir.name,
+            `codeql-custom-queries-${chosenLanguage}`,
+            "example.ql",
+          ),
+        );
+
+        await wizard.execute();
+
+        expect(openTextDocumentSpy).toHaveBeenCalledWith(
+          expect.objectContaining({
+            path: expect.stringMatching("example1.ql"),
+          }),
+        );
+      });
+    });
+
+    describe("if database is also already downloaded", () => {
+      let databaseNwo: string;
+      let databaseItem: DatabaseItem;
+
+      beforeEach(async () => {
+        databaseNwo = QUERY_LANGUAGE_TO_DATABASE_REPO[chosenLanguage];
+
+        databaseItem = {
+          name: databaseNwo,
+          language: chosenLanguage,
+        } as DatabaseItem;
+
+        jest
+          .spyOn(mockDatabaseManager, "digForDatabaseItem")
+          .mockResolvedValue([databaseItem] as any);
+      });
+
+      it("should not download a new database for language", async () => {
+        await wizard.execute();
+
+        expect(downloadGitHubDatabaseSpy).not.toHaveBeenCalled();
+      });
+
+      it("should select an existing database", async () => {
+        await wizard.execute();
+
+        expect(mockDatabaseManager.setCurrentDatabaseItem).toHaveBeenCalledWith(
+          [databaseItem],
+        );
+      });
+
+      it("should open the new query file", async () => {
+        await wizard.execute();
+
+        expect(openTextDocumentSpy).toHaveBeenCalledWith(
+          expect.objectContaining({
+            path: expect.stringMatching("example2.ql"),
+          }),
+        );
+      });
+    });
+
+    describe("if database is missing", () => {
+      beforeEach(async () => {
+        jest
+          .spyOn(mockDatabaseManager, "digForDatabaseItem")
+          .mockResolvedValue(undefined);
+      });
+
+      it("should download a new database for language", async () => {
+        await wizard.execute();
+
+        expect(downloadGitHubDatabaseSpy).toHaveBeenCalled();
+      });
+    });
+  });
+
+  describe("getFirstStoragePath", () => {
+    it("should return the first workspace folder", async () => {
+      jest.spyOn(workspace, "workspaceFolders", "get").mockReturnValue([
+        {
+          name: "codeql-custom-queries-cpp",
+          uri: { path: "codespaces-codeql" },
+        },
+      ] as WorkspaceFolder[]);
+
+      wizard = new SkeletonQueryWizard(
+        mockCli,
+        jest.fn(),
+        credentials,
+        extLogger,
+        mockDatabaseManager,
+        token,
+      );
+
+      expect(wizard.getFirstStoragePath()).toEqual("codespaces-codeql");
+    });
+
+    describe("if user is in vscode-codeql-starter workspace", () => {
+      it("should set storage path to parent folder", async () => {
+        jest.spyOn(workspace, "workspaceFolders", "get").mockReturnValue([
+          {
+            name: "codeql-custom-queries-cpp",
+            uri: { path: "vscode-codeql-starter/codeql-custom-queries-cpp" },
+          },
+          {
+            name: "codeql-custom-queries-csharp",
+            uri: { path: "vscode-codeql-starter/codeql-custom-queries-csharp" },
+          },
+        ] as WorkspaceFolder[]);
+
+        wizard = new SkeletonQueryWizard(
+          mockCli,
+          jest.fn(),
+          credentials,
+          extLogger,
+          mockDatabaseManager,
+          token,
+        );
+
+        expect(wizard.getFirstStoragePath()).toEqual("vscode-codeql-starter");
+      });
+    });
+  });
+});

extensions/ql-vscode/test/vscode-tests/minimal-workspace/local-databases.test.ts

@@ -785,6 +785,39 @@ describe("local databases", () => {
     });
   });
 
+  describe("digForDatabaseItem", () => {
+    describe("when the item exists", () => {
+      it("should return the database item", async () => {
+        const mockDbItem = createMockDB();
+
+        await (databaseManager as any).addDatabaseItem(
+          {} as ProgressCallback,
+          {} as CancellationToken,
+          mockDbItem,
+        );
+
+        const databaseItem = await databaseManager.digForDatabaseItem(
+          mockDbItem.language,
+          mockDbItem.name,
+        );
+
+        expect(databaseItem!.language).toEqual(mockDbItem.language);
+        expect(databaseItem!.name).toEqual(mockDbItem.name);
+      });
+    });
+
+    describe("when the item doesn't exist", () => {
+      it("should return nothing", async () => {
+        const databaseItem = await databaseManager.digForDatabaseItem(
+          "ruby",
+          "mock-database-name",
+        );
+
+        expect(databaseItem).toBeUndefined();
+      });
+    });
+  });
+
   function createMockDB(
     mockDbOptions = MOCK_DB_OPTIONS,
     // source archive location must be a real(-ish) location since

extensions/ql-vscode/test/vscode-tests/no-workspace/databaseFetcher.test.ts

@@ -23,6 +23,48 @@ describe("databaseFetcher", () => {
       request: mockRequest,
     } as unknown as Octokit.Octokit;
 
+    // We can't make the real octokit request (since we need credentials), so we mock the response.
+    const successfullMockApiResponse = {
+      data: [
+        {
+          id: 1495869,
+          name: "csharp-database",
+          language: "csharp",
+          uploader: {},
+          content_type: "application/zip",
+          state: "uploaded",
+          size: 55599715,
+          created_at: "2022-03-24T10:46:24Z",
+          updated_at: "2022-03-24T10:46:27Z",
+          url: "https://api.github.com/repositories/143040428/code-scanning/codeql/databases/csharp",
+        },
+        {
+          id: 1100671,
+          name: "database.zip",
+          language: "javascript",
+          uploader: {},
+          content_type: "application/zip",
+          state: "uploaded",
+          size: 29294434,
+          created_at: "2022-03-01T16:00:04Z",
+          updated_at: "2022-03-01T16:00:06Z",
+          url: "https://api.github.com/repositories/143040428/code-scanning/codeql/databases/javascript",
+        },
+        {
+          id: 648738,
+          name: "ql-database",
+          language: "ql",
+          uploader: {},
+          content_type: "application/json; charset=utf-8",
+          state: "uploaded",
+          size: 39735500,
+          created_at: "2022-02-02T09:38:50Z",
+          updated_at: "2022-02-02T09:38:51Z",
+          url: "https://api.github.com/repositories/143040428/code-scanning/codeql/databases/ql",
+        },
+      ],
+    };
+
     beforeEach(() => {
       quickPickSpy = jest
         .spyOn(window, "showQuickPick")
@@ -30,48 +72,7 @@ describe("databaseFetcher", () => {
     });
 
     it("should convert a GitHub nwo to a database url", async () => {
-      // We can't make the real octokit request (since we need credentials), so we mock the response.
-      const mockApiResponse = {
-        data: [
-          {
-            id: 1495869,
-            name: "csharp-database",
-            language: "csharp",
-            uploader: {},
-            content_type: "application/zip",
-            state: "uploaded",
-            size: 55599715,
-            created_at: "2022-03-24T10:46:24Z",
-            updated_at: "2022-03-24T10:46:27Z",
-            url: "https://api.github.com/repositories/143040428/code-scanning/codeql/databases/csharp",
-          },
-          {
-            id: 1100671,
-            name: "database.zip",
-            language: "javascript",
-            uploader: {},
-            content_type: "application/zip",
-            state: "uploaded",
-            size: 29294434,
-            created_at: "2022-03-01T16:00:04Z",
-            updated_at: "2022-03-01T16:00:06Z",
-            url: "https://api.github.com/repositories/143040428/code-scanning/codeql/databases/javascript",
-          },
-          {
-            id: 648738,
-            name: "ql-database",
-            language: "ql",
-            uploader: {},
-            content_type: "application/json; charset=utf-8",
-            state: "uploaded",
-            size: 39735500,
-            created_at: "2022-02-02T09:38:50Z",
-            updated_at: "2022-02-02T09:38:51Z",
-            url: "https://api.github.com/repositories/143040428/code-scanning/codeql/databases/ql",
-          },
-        ],
-      };
-      mockRequest.mockResolvedValue(mockApiResponse);
+      mockRequest.mockResolvedValue(successfullMockApiResponse);
       quickPickSpy.mockResolvedValue(mockedQuickPickItem("javascript"));
       const githubRepo = "github/codeql";
       const result = await convertGithubNwoToDatabaseUrl(
@@ -127,6 +128,45 @@ describe("databaseFetcher", () => {
       ).rejects.toThrow(/Unable to get database/);
       expect(progressSpy).toBeCalledTimes(1);
     });
+
+    describe("when language is already provided", () => {
+      describe("when language is valid", () => {
+        it("should not prompt the user", async () => {
+          mockRequest.mockResolvedValue(successfullMockApiResponse);
+          const githubRepo = "github/codeql";
+          await convertGithubNwoToDatabaseUrl(
+            githubRepo,
+            octokit,
+            progressSpy,
+            "javascript",
+          );
+          expect(quickPickSpy).not.toHaveBeenCalled();
+        });
+      });
+
+      describe("when language is invalid", () => {
+        it("should prompt for language", async () => {
+          mockRequest.mockResolvedValue(successfullMockApiResponse);
+          const githubRepo = "github/codeql";
+          await convertGithubNwoToDatabaseUrl(
+            githubRepo,
+            octokit,
+            progressSpy,
+            "invalid-language",
+          );
+          expect(quickPickSpy).toHaveBeenCalled();
+        });
+      });
+    });
+
+    describe("when language is not provided", () => {
+      it("should prompt for language", async () => {
+        mockRequest.mockResolvedValue(successfullMockApiResponse);
+        const githubRepo = "github/codeql";
+        await convertGithubNwoToDatabaseUrl(githubRepo, octokit, progressSpy);
+        expect(quickPickSpy).toHaveBeenCalled();
+      });
+    });
   });
 
   describe("findDirWithFile", () => {