Merge pull request #3012 from github/koesie10/use-selected-queries-item

Create new queries in selected folder of queries panel

extensions/ql-vscode/src/extension.ts

@@ -802,7 +802,11 @@ async function activateWithInstalledDistribution(
   );
   ctx.subscriptions.push(databaseUI);
 
-  QueriesModule.initialize(app, languageContext, cliServer);
+  const queriesModule = QueriesModule.initialize(
+    app,
+    languageContext,
+    cliServer,
+  );
 
   void extLogger.log("Initializing evaluator log viewer.");
   const evalLogViewer = new EvalLogViewer();
@@ -941,6 +945,10 @@ async function activateWithInstalledDistribution(
   );
   ctx.subscriptions.push(localQueries);
 
+  queriesModule.onDidChangeSelection((event) =>
+    localQueries.setSelectedQueryTreeViewItems(event.selection),
+  );
+
   void extLogger.log("Initializing debugger factory.");
   ctx.subscriptions.push(
     new QLDebugAdapterDescriptorFactory(queryStorageDir, qs, localQueries),

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

@@ -63,6 +63,8 @@ export enum QuickEvalType {
 }
 
 export class LocalQueries extends DisposableObject {
+  private selectedQueryTreeViewItems: readonly QueryTreeViewItem[] = [];
+
   public constructor(
     private readonly app: App,
     private readonly queryRunner: QueryRunner,
@@ -77,6 +79,12 @@ export class LocalQueries extends DisposableObject {
     super();
   }
 
+  public setSelectedQueryTreeViewItems(
+    selection: readonly QueryTreeViewItem[],
+  ) {
+    this.selectedQueryTreeViewItems = selection;
+  }
+
   public getCommands(): LocalQueryCommands {
     return {
       "codeQL.runQuery": this.runQuery.bind(this),
@@ -333,6 +341,7 @@ export class LocalQueries extends DisposableObject {
           this.app.logger,
           this.databaseManager,
           contextStoragePath,
+          this.selectedQueryTreeViewItems,
           language,
         );
         await skeletonQueryWizard.execute();

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

@@ -1,13 +1,10 @@
-import { join } from "path";
+import { basename, dirname, join } from "path";
-import { Uri, workspace, window as Window } from "vscode";
+import { Uri, window as Window, workspace } from "vscode";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { BaseLogger } from "../common/logging";
 import { Credentials } from "../common/authentication";
 import { QueryLanguage } from "../common/query-language";
-import {
+import { getFirstWorkspaceFolder } from "../common/vscode/workspace-folders";
-  getFirstWorkspaceFolder,
-  isFolderAlreadyInWorkspace,
-} from "../common/vscode/workspace-folders";
 import { getErrorMessage } from "../common/helpers-pure";
 import { QlPackGenerator } from "./qlpack-generator";
 import { DatabaseItem, DatabaseManager } from "../databases/local-databases";
@@ -24,8 +21,9 @@ import {
   isCodespacesTemplate,
   setQlPackLocation,
 } from "../config";
-import { existsSync } from "fs-extra";
+import { lstat, pathExists } from "fs-extra";
 import { askForLanguage } from "../codeql-cli/query-language";
+import { QueryTreeViewItem } from "../queries-panel/query-tree-view-item";
 
 type QueryLanguagesToDatabaseMap = Record<string, string>;
 
@@ -51,6 +49,7 @@ export class SkeletonQueryWizard {
     private readonly logger: BaseLogger,
     private readonly databaseManager: DatabaseManager,
     private readonly databaseStoragePath: string | undefined,
+    private readonly selectedItems: readonly QueryTreeViewItem[],
     private language: QueryLanguage | undefined = undefined,
   ) {}
 
@@ -70,9 +69,9 @@ export class SkeletonQueryWizard {
 
     this.qlPackStoragePath = await this.determineStoragePath();
 
-    const skeletonPackAlreadyExists =
+    const skeletonPackAlreadyExists = await pathExists(
-      existsSync(join(this.qlPackStoragePath, this.folderName)) ||
+      join(this.qlPackStoragePath, this.folderName),
-      isFolderAlreadyInWorkspace(this.folderName);
+    );
 
     if (skeletonPackAlreadyExists) {
       // just create a new example query file in skeleton QL pack
@@ -109,7 +108,41 @@ export class SkeletonQueryWizard {
     });
   }
 
-  public async determineStoragePath() {
+  public async determineStoragePath(): Promise<string> {
+    if (this.selectedItems.length === 0) {
+      return this.determineRootStoragePath();
+    }
+
+    const storagePath = await this.determineStoragePathFromSelection();
+
+    // If the user has selected a folder or file within a folder that matches the current
+    // folder name, we should create a query rather than a query pack
+    if (basename(storagePath) === this.folderName) {
+      return dirname(storagePath);
+    }
+
+    return storagePath;
+  }
+
+  private async determineStoragePathFromSelection(): Promise<string> {
+    // Just like VS Code's "New File" command, if the user has selected multiple files/folders in the queries panel,
+    // we will create the new file in the same folder as the first selected item.
+    // See https://github.com/microsoft/vscode/blob/a8b7239d0311d4915b57c837972baf4b01394491/src/vs/workbench/contrib/files/browser/fileActions.ts#L893-L900
+    const selectedItem = this.selectedItems[0];
+
+    const path = selectedItem.path;
+
+    // We use stat to protect against outdated query tree items
+    const fileStat = await lstat(path);
+
+    if (fileStat.isDirectory()) {
+      return path;
+    }
+
+    return dirname(path);
+  }
+
+  public async determineRootStoragePath() {
     const firstStorageFolder = getFirstWorkspaceFolder();
 
     if (isCodespacesTemplate()) {
@@ -118,7 +151,7 @@ export class SkeletonQueryWizard {
 
     let storageFolder = getQlPackLocation();
 
-    if (storageFolder === undefined || !existsSync(storageFolder)) {
+    if (storageFolder === undefined || !(await pathExists(storageFolder))) {
       storageFolder = await Window.showInputBox({
         title:
           "Please choose a folder in which to create your new query pack. You can change this in the extension settings.",
@@ -131,7 +164,7 @@ export class SkeletonQueryWizard {
       throw new UserCancellationException("No storage folder entered.");
     }
 
-    if (!existsSync(storageFolder)) {
+    if (!(await pathExists(storageFolder))) {
       throw new UserCancellationException(
         "Invalid folder. Must be a folder that already exists.",
       );
@@ -208,7 +241,7 @@ export class SkeletonQueryWizard {
       await qlPackGenerator.createExampleQlFile(this.fileName);
     } catch (e: unknown) {
       void this.logger.log(
-        `Could not create skeleton QL pack: ${getErrorMessage(e)}`,
+        `Could not create query example file: ${getErrorMessage(e)}`,
       );
     }
   }

extensions/ql-vscode/src/queries-panel/queries-module.ts

@@ -7,9 +7,18 @@ import { QueriesPanel } from "./queries-panel";
 import { QueryDiscovery } from "./query-discovery";
 import { QueryPackDiscovery } from "./query-pack-discovery";
 import { LanguageContextStore } from "../language-context-store";
+import { TreeViewSelectionChangeEvent } from "vscode";
+import { QueryTreeViewItem } from "./query-tree-view-item";
 
 export class QueriesModule extends DisposableObject {
   private queriesPanel: QueriesPanel | undefined;
+  private readonly onDidChangeSelectionEmitter = this.push(
+    this.app.createEventEmitter<
+      TreeViewSelectionChangeEvent<QueryTreeViewItem>
+    >(),
+  );
+
+  public readonly onDidChangeSelection = this.onDidChangeSelectionEmitter.event;
 
   private constructor(readonly app: App) {
     super();
@@ -52,6 +61,9 @@ export class QueriesModule extends DisposableObject {
     void queryDiscovery.initialRefresh();
 
     this.queriesPanel = new QueriesPanel(queryDiscovery, app);
+    this.queriesPanel.onDidChangeSelection((event) =>
+      this.onDidChangeSelectionEmitter.fire(event),
+    );
     this.push(this.queriesPanel);
   }
 }

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

@@ -1,7 +1,13 @@
 import { DisposableObject } from "../common/disposable-object";
 import { QueryTreeDataProvider } from "./query-tree-data-provider";
 import { QueryDiscovery } from "./query-discovery";
-import { TextEditor, TreeView, window } from "vscode";
+import {
+  Event,
+  TextEditor,
+  TreeView,
+  TreeViewSelectionChangeEvent,
+  window,
+} from "vscode";
 import { App } from "../common/app";
 import { QueryTreeViewItem } from "./query-tree-view-item";
 
@@ -16,6 +22,7 @@ export class QueriesPanel extends DisposableObject {
     super();
 
     this.dataProvider = new QueryTreeDataProvider(queryDiscovery, app);
+    this.push(this.dataProvider);
 
     this.treeView = window.createTreeView("codeQLQueries", {
       treeDataProvider: this.dataProvider,
@@ -25,6 +32,12 @@ export class QueriesPanel extends DisposableObject {
     this.subscribeToTreeSelectionEvents();
   }
 
+  public get onDidChangeSelection(): Event<
+    TreeViewSelectionChangeEvent<QueryTreeViewItem>
+  > {
+    return this.treeView.onDidChangeSelection;
+  }
+
   private subscribeToTreeSelectionEvents(): void {
     // Keep track of whether the user has changed their text editor while
     // the tree view was not visible. If so, we will focus the text editor

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

@@ -8,9 +8,14 @@ import * as tmp from "tmp";
 import { TextDocument, window, workspace, WorkspaceFolder } from "vscode";
 import { extLogger } from "../../../../src/common/logging/vscode";
 import { QlPackGenerator } from "../../../../src/local-queries/qlpack-generator";
-import * as workspaceFolders from "../../../../src/common/vscode/workspace-folders";
+import {
-import { createFileSync, ensureDirSync, removeSync } from "fs-extra";
+  createFileSync,
-import { join } from "path";
+  ensureDir,
+  ensureDirSync,
+  ensureFile,
+  removeSync,
+} from "fs-extra";
+import { dirname, join } from "path";
 import { testCredentialsWithStub } from "../../../factories/authentication";
 import {
   DatabaseItem,
@@ -22,6 +27,11 @@ import { createMockDB } from "../../../factories/databases/databases";
 import { asError } from "../../../../src/common/helpers-pure";
 import { Setting } from "../../../../src/config";
 import { QueryLanguage } from "../../../../src/common/query-language";
+import {
+  createQueryTreeFileItem,
+  createQueryTreeFolderItem,
+  QueryTreeViewItem,
+} from "../../../../src/queries-panel/query-tree-view-item";
 
 describe("SkeletonQueryWizard", () => {
   let mockCli: CodeQLCliServer;
@@ -49,6 +59,7 @@ describe("SkeletonQueryWizard", () => {
 
   const credentials = testCredentialsWithStub();
   const chosenLanguage = "ruby";
+  const selectedItems: QueryTreeViewItem[] = [];
 
   jest.spyOn(extLogger, "log").mockResolvedValue(undefined);
 
@@ -117,6 +128,7 @@ describe("SkeletonQueryWizard", () => {
       extLogger,
       mockDatabaseManager,
       storagePath,
+      selectedItems,
     );
 
     askForGitHubRepoSpy = jest
@@ -144,6 +156,7 @@ describe("SkeletonQueryWizard", () => {
         extLogger,
         mockDatabaseManager,
         storagePath,
+        selectedItems,
         QueryLanguage.Swift,
       );
     });
@@ -157,11 +170,6 @@ describe("SkeletonQueryWizard", () => {
   });
 
   describe("if QL pack doesn't exist", () => {
-    beforeEach(() => {
-      jest
-        .spyOn(workspaceFolders, "isFolderAlreadyInWorkspace")
-        .mockReturnValue(false);
-    });
     it("should try to create a new QL pack based on the language", async () => {
       await wizard.execute();
 
@@ -187,10 +195,6 @@ describe("SkeletonQueryWizard", () => {
 
   describe("if QL pack exists", () => {
     beforeEach(async () => {
-      jest
-        .spyOn(workspaceFolders, "isFolderAlreadyInWorkspace")
-        .mockReturnValue(true);
-
       // create a skeleton codeql-custom-queries-${language} folder
       // with an example QL file inside
       ensureDirSync(
@@ -272,6 +276,7 @@ describe("SkeletonQueryWizard", () => {
           extLogger,
           mockDatabaseManagerWithItems,
           storagePath,
+          selectedItems,
         );
       });
 
@@ -407,7 +412,7 @@ describe("SkeletonQueryWizard", () => {
   });
 
   describe("determineStoragePath", () => {
-    it("should prompt the user to provide a storage path", async () => {
+    it("should prompt the user to provide a storage path when no items are selected", async () => {
       const chosenPath = await wizard.determineStoragePath();
 
       expect(showInputBoxSpy).toHaveBeenCalledWith(
@@ -416,10 +421,180 @@ describe("SkeletonQueryWizard", () => {
       expect(chosenPath).toEqual(storagePath);
     });
 
+    describe("with folders and files", () => {
+      let queriesDir: tmp.DirResult;
+
+      beforeEach(async () => {
+        queriesDir = tmp.dirSync({
+          prefix: "queries_",
+          unsafeCleanup: true,
+        });
+
+        await ensureDir(join(queriesDir.name, "folder"));
+        await ensureFile(join(queriesDir.name, "queries-java", "example.ql"));
+        await ensureFile(
+          join(queriesDir.name, "codeql-custom-queries-swift", "example.ql"),
+        );
+      });
+
+      describe("with selected folder", () => {
+        let selectedItems: QueryTreeViewItem[];
+
+        beforeEach(async () => {
+          selectedItems = [
+            createQueryTreeFolderItem(
+              "folder",
+              join(queriesDir.name, "folder"),
+              [
+                createQueryTreeFileItem(
+                  "example.ql",
+                  join(queriesDir.name, "folder", "example.ql"),
+                  "java",
+                ),
+              ],
+            ),
+          ];
+
+          wizard = new SkeletonQueryWizard(
+            mockCli,
+            jest.fn(),
+            credentials,
+            extLogger,
+            mockDatabaseManager,
+            storagePath,
+            selectedItems,
+          );
+        });
+
+        it("returns the selected folder path", async () => {
+          const chosenPath = await wizard.determineStoragePath();
+
+          expect(chosenPath).toEqual(selectedItems[0].path);
+        });
+      });
+
+      describe("with selected file", () => {
+        let selectedItems: QueryTreeViewItem[];
+
+        beforeEach(async () => {
+          selectedItems = [
+            createQueryTreeFileItem(
+              "example.ql",
+              join(queriesDir.name, "queries-java", "example.ql"),
+              "java",
+            ),
+          ];
+
+          wizard = new SkeletonQueryWizard(
+            mockCli,
+            jest.fn(),
+            credentials,
+            extLogger,
+            mockDatabaseManager,
+            storagePath,
+            selectedItems,
+          );
+        });
+
+        it("returns the selected file path", async () => {
+          const chosenPath = await wizard.determineStoragePath();
+
+          expect(chosenPath).toEqual(dirname(selectedItems[0].path));
+        });
+      });
+
+      describe("with selected file with same name", () => {
+        let selectedItems: QueryTreeViewItem[];
+
+        beforeEach(async () => {
+          selectedItems = [
+            createQueryTreeFileItem(
+              "example.ql",
+              join(
+                queriesDir.name,
+                "codeql-custom-queries-swift",
+                "example.ql",
+              ),
+              "java",
+            ),
+          ];
+
+          wizard = new SkeletonQueryWizard(
+            mockCli,
+            jest.fn(),
+            credentials,
+            extLogger,
+            mockDatabaseManager,
+            storagePath,
+            selectedItems,
+            QueryLanguage.Swift,
+          );
+        });
+
+        it("returns the parent path", async () => {
+          const chosenPath = await wizard.determineStoragePath();
+
+          expect(chosenPath).toEqual(queriesDir.name);
+        });
+      });
+
+      describe("with multiple selected items", () => {
+        let selectedItems: QueryTreeViewItem[];
+
+        beforeEach(async () => {
+          selectedItems = [
+            createQueryTreeFileItem(
+              "example.ql",
+              join(queriesDir.name, "queries-java", "example.ql"),
+              "java",
+            ),
+            createQueryTreeFolderItem(
+              "folder",
+              join(queriesDir.name, "folder"),
+              [
+                createQueryTreeFileItem(
+                  "example.ql",
+                  join(queriesDir.name, "folder", "example.ql"),
+                  "java",
+                ),
+              ],
+            ),
+          ];
+
+          wizard = new SkeletonQueryWizard(
+            mockCli,
+            jest.fn(),
+            credentials,
+            extLogger,
+            mockDatabaseManager,
+            storagePath,
+            selectedItems,
+          );
+        });
+
+        it("returns the first selected item path", async () => {
+          const chosenPath = await wizard.determineStoragePath();
+
+          expect(chosenPath).toEqual(dirname(selectedItems[0].path));
+        });
+      });
+    });
+  });
+
+  describe("determineRootStoragePath", () => {
+    it("should prompt the user to provide a storage path", async () => {
+      const chosenPath = await wizard.determineRootStoragePath();
+
+      expect(showInputBoxSpy).toHaveBeenCalledWith(
+        expect.objectContaining({ value: storagePath }),
+      );
+      expect(chosenPath).toEqual(storagePath);
+    });
+
     it("should write the chosen folder to settings", async () => {
       const updateValueSpy = jest.spyOn(Setting.prototype, "updateValue");
 
-      await wizard.determineStoragePath();
+      await wizard.determineRootStoragePath();
 
       expect(updateValueSpy).toHaveBeenCalledWith(storagePath, 2);
     });
@@ -453,7 +628,7 @@ describe("SkeletonQueryWizard", () => {
       });
 
       it("should not prompt the user", async () => {
-        const chosenPath = await wizard.determineStoragePath();
+        const chosenPath = await wizard.determineRootStoragePath();
 
         expect(showInputBoxSpy).not.toHaveBeenCalled();
         expect(chosenPath).toEqual(storagePath);
@@ -484,7 +659,7 @@ describe("SkeletonQueryWizard", () => {
         });
 
         it("should return it and not prompt the user", async () => {
-          const chosenPath = await wizard.determineStoragePath();
+          const chosenPath = await wizard.determineRootStoragePath();
 
           expect(showInputBoxSpy).not.toHaveBeenCalled();
           expect(chosenPath).toEqual(storedPath);
@@ -513,7 +688,7 @@ describe("SkeletonQueryWizard", () => {
         });
 
         it("should prompt the user for to provide a new folder name", async () => {
-          const chosenPath = await wizard.determineStoragePath();
+          const chosenPath = await wizard.determineRootStoragePath();
 
           expect(showInputBoxSpy).toHaveBeenCalled();
           expect(chosenPath).toEqual(storagePath);