Merge pull request #2328 from github/yer-a-single-rooted-workspace-query

Ensure we're selecting database in single rooted workspace

Author: elenatanasoiu

extensions/ql-vscode/src/databaseFetcher.ts

@@ -317,13 +317,15 @@ async function databaseArchiveFetcher(
     });
     await ensureZippedSourceLocation(dbPath);
 
+    const makeSelected = true;
+
     const item = await databaseManager.openDatabase(
       progress,
       token,
       Uri.file(dbPath),
+      makeSelected,
       nameOverride,
     );
-    await databaseManager.setCurrentDatabaseItem(item);
     return item;
   } else {
     throw new Error("Database not found in archive.");

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

@@ -306,18 +306,21 @@ export class DatabaseUI extends DisposableObject {
               `${workspace.workspaceFolders[0].uri}/.tours/codeql-tutorial-database`,
             );
 
-            let databaseItem = this.databaseManager.findDatabaseItem(uri);
-            const isTutorialDatabase = true;
+            const databaseItem = this.databaseManager.findDatabaseItem(uri);
             if (databaseItem === undefined) {
-              databaseItem = await this.databaseManager.openDatabase(
+              const makeSelected = true;
+              const nameOverride = "CodeQL Tutorial Database";
+              const isTutorialDatabase = true;
+
+              await this.databaseManager.openDatabase(
                 progress,
                 token,
                 uri,
-                "CodeQL Tutorial Database",
+                makeSelected,
+                nameOverride,
                 isTutorialDatabase,
               );
             }
-            await this.databaseManager.setCurrentDatabaseItem(databaseItem);
             await this.handleTourDependencies();
           }
         } catch (e) {
@@ -630,7 +633,7 @@ export class DatabaseUI extends DisposableObject {
               this.queryServer?.cliServer,
             );
           } else {
-            await this.setCurrentDatabase(progress, token, uri);
+            await this.databaseManager.openDatabase(progress, token, uri);
           }
         } catch (e) {
           // rethrow and let this be handled by default error handling.
@@ -752,24 +755,6 @@ export class DatabaseUI extends DisposableObject {
     return this.databaseManager.currentDatabaseItem;
   }
 
-  private async setCurrentDatabase(
-    progress: ProgressCallback,
-    token: CancellationToken,
-    uri: Uri,
-  ): Promise<DatabaseItem | undefined> {
-    let databaseItem = this.databaseManager.findDatabaseItem(uri);
-    if (databaseItem === undefined) {
-      databaseItem = await this.databaseManager.openDatabase(
-        progress,
-        token,
-        uri,
-      );
-    }
-    await this.databaseManager.setCurrentDatabaseItem(databaseItem);
-
-    return databaseItem;
-  }
-
   /**
    * Ask the user for a database directory. Returns the chosen database, or `undefined` if the
    * operation was canceled.
@@ -789,7 +774,11 @@ export class DatabaseUI extends DisposableObject {
         if (byFolder) {
           const fixedUri = await this.fixDbUri(uri);
           // we are selecting a database folder
-          return await this.setCurrentDatabase(progress, token, fixedUri);
+          return await this.databaseManager.openDatabase(
+            progress,
+            token,
+            fixedUri,
+          );
         } else {
           // we are selecting a database archive. Must unzip into a workspace-controlled area
           // before importing.

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

@@ -621,6 +621,7 @@ export class DatabaseManager extends DisposableObject {
     progress: ProgressCallback,
     token: vscode.CancellationToken,
     uri: vscode.Uri,
+    makeSelected = false,
     displayName?: string,
     isTutorialDatabase?: boolean,
   ): Promise<DatabaseItem> {
@@ -629,6 +630,7 @@ export class DatabaseManager extends DisposableObject {
     return await this.addExistingDatabaseItem(
       databaseItem,
       progress,
+      makeSelected,
       token,
       isTutorialDatabase,
     );
@@ -643,6 +645,7 @@ export class DatabaseManager extends DisposableObject {
   public async addExistingDatabaseItem(
     databaseItem: DatabaseItem,
     progress: ProgressCallback,
+    makeSelected = true,
     token: vscode.CancellationToken,
     isTutorialDatabase?: boolean,
   ): Promise<DatabaseItem> {
@@ -652,6 +655,9 @@ export class DatabaseManager extends DisposableObject {
     }
 
     await this.addDatabaseItem(progress, token, databaseItem);
+    if (makeSelected) {
+      await this.setCurrentDatabaseItem(databaseItem);
+    }
     await this.addDatabaseSourceArchiveFolder(databaseItem);
 
     if (isCodespacesTemplate() && !isTutorialDatabase) {

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

@@ -64,15 +64,14 @@ export class SkeletonQueryWizard {
     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();
     }
 
+    // select existing database for language or download a new one
+    await this.selectOrDownloadDatabase();
+
     // open a query file
 
     try {
@@ -216,41 +215,19 @@ export class SkeletonQueryWizard {
     );
   }
 
-  private async selectExistingDatabase() {
-    if (this.language === undefined) {
-      throw new Error("Language is undefined");
-    }
-
+  private async selectOrDownloadDatabase() {
     if (this.qlPackStoragePath === undefined) {
       throw new Error("QL Pack storage path is undefined");
     }
 
-    const databaseNwo = QUERY_LANGUAGE_TO_DATABASE_REPO[this.language];
-
-    const existingDatabaseItem = await this.findDatabaseItemByNwo(
-      this.language,
-      databaseNwo,
-      this.databaseManager.databaseItems,
-    );
+    const existingDatabaseItem = await this.findExistingDatabaseItem();
 
     if (existingDatabaseItem) {
       // select the found database
       await this.databaseManager.setCurrentDatabaseItem(existingDatabaseItem);
     } else {
-      const sameLanguageDatabaseItem = await this.findDatabaseItemByLanguage(
-        this.language,
-        this.databaseManager.databaseItems,
-      );
-
-      if (sameLanguageDatabaseItem) {
-        // select the found database
-        await this.databaseManager.setCurrentDatabaseItem(
-          sameLanguageDatabaseItem,
-        );
-      } else {
-        // download new database and select it
-        await this.downloadDatabase();
-      }
+      // download new database and select it
+      await this.downloadDatabase();
     }
   }
 
@@ -286,4 +263,27 @@ export class SkeletonQueryWizard {
     }
     return dbs[0];
   }
+
+  private async findExistingDatabaseItem() {
+    if (this.language === undefined) {
+      throw new Error("Language is undefined");
+    }
+
+    const defaultDatabaseNwo = QUERY_LANGUAGE_TO_DATABASE_REPO[this.language];
+
+    const defaultDatabaseItem = await this.findDatabaseItemByNwo(
+      this.language,
+      defaultDatabaseNwo,
+      this.databaseManager.databaseItems,
+    );
+
+    if (defaultDatabaseItem !== undefined) {
+      return defaultDatabaseItem;
+    }
+
+    return await this.findDatabaseItemByLanguage(
+      this.language,
+      this.databaseManager.databaseItems,
+    );
+  }
 }

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

@@ -708,6 +708,7 @@ describe("local databases", () => {
   describe("openDatabase", () => {
     let createSkeletonPacksSpy: jest.SpyInstance;
     let resolveDatabaseContentsSpy: jest.SpyInstance;
+    let setCurrentDatabaseItemSpy: jest.SpyInstance;
     let addDatabaseSourceArchiveFolderSpy: jest.SpyInstance;
     let mockDbItem: DatabaseItemImpl;
 
@@ -722,6 +723,11 @@ describe("local databases", () => {
         .spyOn(DatabaseResolver, "resolveDatabaseContents")
         .mockResolvedValue({} as DatabaseContentsWithDbScheme);
 
+      setCurrentDatabaseItemSpy = jest.spyOn(
+        databaseManager,
+        "setCurrentDatabaseItem",
+      );
+
       addDatabaseSourceArchiveFolderSpy = jest.spyOn(
         databaseManager,
         "addDatabaseSourceArchiveFolder",
@@ -746,6 +752,19 @@ describe("local databases", () => {
       expect(resolveDatabaseContentsSpy).toBeCalledTimes(1);
     });
 
+    it("should set the database as the currently selected one", async () => {
+      const makeSelected = true;
+
+      await databaseManager.openDatabase(
+        {} as ProgressCallback,
+        {} as CancellationToken,
+        mockDbItem.databaseUri,
+        makeSelected,
+      );
+
+      expect(setCurrentDatabaseItemSpy).toBeCalledTimes(1);
+    });
+
     it("should add database source archive folder", async () => {
       await databaseManager.openDatabase(
         {} as ProgressCallback,
@@ -762,12 +781,15 @@ describe("local databases", () => {
           jest.spyOn(Setting.prototype, "getValue").mockReturnValue(true);
 
           const isTutorialDatabase = true;
+          const makeSelected = true;
+          const nameOverride = "CodeQL Tutorial Database";
 
           await databaseManager.openDatabase(
             {} as ProgressCallback,
             {} as CancellationToken,
             mockDbItem.databaseUri,
-            "CodeQL Tutorial Database",
+            makeSelected,
+            nameOverride,
             isTutorialDatabase,
           );