Merge pull request #3458 from github/robertbrignull/database-prompting

Refactor database-fetcher.ts into a class and simplify

Author: robertbrignull

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

@@ -2,10 +2,8 @@ import { window } from "vscode";
 import type { Octokit } from "@octokit/rest";
 import { showNeverAskAgainDialog } from "../../common/vscode/dialog";
 import { getLanguageDisplayName } from "../../common/query-language";
-import { downloadGitHubDatabaseFromUrl } from "../database-fetcher";
+import type { DatabaseFetcher } from "../database-fetcher";
 import { withProgress } from "../../common/vscode/progress";
-import type { DatabaseManager } from "../local-databases";
-import type { CodeQLCliServer } from "../../codeql-cli/cli";
 import type { AppCommandManager } from "../../common/commands";
 import type { GitHubDatabaseConfig } from "../../config";
 import type { CodeqlDatabase } from "./api";
@@ -58,9 +56,7 @@ export async function downloadDatabaseFromGitHub(
   owner: string,
   repo: string,
   databases: CodeqlDatabase[],
-  databaseManager: DatabaseManager,
-  storagePath: string,
-  cliServer: CodeQLCliServer,
+  databaseFetcher: DatabaseFetcher,
   commandManager: AppCommandManager,
 ): Promise<void> {
   const selectedDatabases = await promptForDatabases(databases);
@@ -72,7 +68,7 @@ export async function downloadDatabaseFromGitHub(
     selectedDatabases.map((database) =>
       withProgress(
         async (progress) => {
-          await downloadGitHubDatabaseFromUrl(
+          await databaseFetcher.downloadGitHubDatabaseFromUrl(
             database.url,
             database.id,
             database.created_at,
@@ -81,9 +77,6 @@ export async function downloadDatabaseFromGitHub(
             repo,
             octokit,
             progress,
-            databaseManager,
-            storagePath,
-            cliServer,
             true,
             false,
           );

extensions/ql-vscode/src/databases/github-databases/download.ts

@@ -14,7 +14,6 @@ import {
 } from "./download";
 import type { GitHubDatabaseConfig } from "../../config";
 import type { DatabaseManager } from "../local-databases";
-import type { CodeQLCliServer } from "../../codeql-cli/cli";
 import type { CodeqlDatabase, ListDatabasesResult } from "./api";
 import { listDatabases } from "./api";
 import type { DatabaseUpdate } from "./updates";
@@ -24,6 +23,7 @@ import {
   isNewerDatabaseAvailable,
 } from "./updates";
 import type { Octokit } from "@octokit/rest";
+import type { DatabaseFetcher } from "../database-fetcher";
 
 export class GitHubDatabasesModule extends DisposableObject {
   /**
@@ -33,8 +33,7 @@ export class GitHubDatabasesModule extends DisposableObject {
   constructor(
     private readonly app: App,
     private readonly databaseManager: DatabaseManager,
-    private readonly databaseStoragePath: string,
-    private readonly cliServer: CodeQLCliServer,
+    private readonly databaseFetcher: DatabaseFetcher,
     private readonly config: GitHubDatabaseConfig,
   ) {
     super();
@@ -43,15 +42,13 @@ export class GitHubDatabasesModule extends DisposableObject {
   public static async initialize(
     app: App,
     databaseManager: DatabaseManager,
-    databaseStoragePath: string,
-    cliServer: CodeQLCliServer,
+    databaseFetcher: DatabaseFetcher,
     config: GitHubDatabaseConfig,
   ): Promise<GitHubDatabasesModule> {
     const githubDatabasesModule = new GitHubDatabasesModule(
       app,
       databaseManager,
-      databaseStoragePath,
-      cliServer,
+      databaseFetcher,
       config,
     );
     app.subscriptions.push(githubDatabasesModule);
@@ -185,9 +182,7 @@ export class GitHubDatabasesModule extends DisposableObject {
       owner,
       repo,
       databases,
-      this.databaseManager,
-      this.databaseStoragePath,
-      this.cliServer,
+      this.databaseFetcher,
       this.app.commands,
     );
   }
@@ -212,8 +207,7 @@ export class GitHubDatabasesModule extends DisposableObject {
       repo,
       databaseUpdates,
       this.databaseManager,
-      this.databaseStoragePath,
-      this.cliServer,
+      this.databaseFetcher,
       this.app.commands,
     );
   }

extensions/ql-vscode/src/databases/github-databases/github-databases-module.ts

@@ -1,11 +1,10 @@
 import type { CodeqlDatabase } from "./api";
 import type { DatabaseItem, DatabaseManager } from "../local-databases";
 import type { Octokit } from "@octokit/rest";
-import type { CodeQLCliServer } from "../../codeql-cli/cli";
 import type { AppCommandManager } from "../../common/commands";
 import { getLanguageDisplayName } from "../../common/query-language";
 import { showNeverAskAgainDialog } from "../../common/vscode/dialog";
-import { downloadGitHubDatabaseFromUrl } from "../database-fetcher";
+import type { DatabaseFetcher } from "../database-fetcher";
 import { withProgress } from "../../common/vscode/progress";
 import { window } from "vscode";
 import type { GitHubDatabaseConfig } from "../../config";
@@ -156,8 +155,7 @@ export async function downloadDatabaseUpdateFromGitHub(
   repo: string,
   updates: DatabaseUpdate[],
   databaseManager: DatabaseManager,
-  storagePath: string,
-  cliServer: CodeQLCliServer,
+  databaseFetcher: DatabaseFetcher,
   commandManager: AppCommandManager,
 ): Promise<void> {
   const selectedDatabases = await promptForDatabases(
@@ -179,21 +177,19 @@ export async function downloadDatabaseUpdateFromGitHub(
 
       return withProgress(
         async (progress) => {
-          const newDatabase = await downloadGitHubDatabaseFromUrl(
-            database.url,
-            database.id,
-            database.created_at,
-            database.commit_oid ?? null,
-            owner,
-            repo,
-            octokit,
-            progress,
-            databaseManager,
-            storagePath,
-            cliServer,
-            databaseManager.currentDatabaseItem === update.databaseItem,
-            update.databaseItem.hasSourceArchiveInExplorer(),
-          );
+          const newDatabase =
+            await databaseFetcher.downloadGitHubDatabaseFromUrl(
+              database.url,
+              database.id,
+              database.created_at,
+              database.commit_oid ?? null,
+              owner,
+              repo,
+              octokit,
+              progress,
+              databaseManager.currentDatabaseItem === update.databaseItem,
+              update.databaseItem.hasSourceArchiveInExplorer(),
+            );
           if (newDatabase === undefined) {
             return;
           }

extensions/ql-vscode/src/databases/github-databases/updates.ts

@@ -37,11 +37,7 @@ import {
   showAndLogExceptionWithTelemetry,
   showAndLogErrorMessage,
 } from "../common/logging";
-import {
-  importLocalDatabase,
-  promptImportGithubDatabase,
-  promptImportInternetDatabase,
-} from "./database-fetcher";
+import type { DatabaseFetcher } from "./database-fetcher";
 import { asError, asyncFilter, getErrorMessage } from "../common/helpers-pure";
 import type { QueryRunner } from "../query-server";
 import type { App } from "../common/app";
@@ -248,6 +244,7 @@ export class DatabaseUI extends DisposableObject {
   public constructor(
     private app: App,
     private databaseManager: DatabaseManager,
+    private readonly databaseFetcher: DatabaseFetcher,
     languageContext: LanguageContextStore,
     private readonly queryServer: QueryRunner,
     private readonly storagePath: string,
@@ -535,13 +532,7 @@ export class DatabaseUI extends DisposableObject {
   private async handleChooseDatabaseInternet(): Promise<void> {
     return withProgress(
       async (progress) => {
-        await promptImportInternetDatabase(
-          this.app.commands,
-          this.databaseManager,
-          this.storagePath,
-          progress,
-          this.queryServer.cliServer,
-        );
+        await this.databaseFetcher.promptImportInternetDatabase(progress);
       },
       {
         title: "Adding database from URL",
@@ -552,13 +543,7 @@ export class DatabaseUI extends DisposableObject {
   private async handleChooseDatabaseGithub(): Promise<void> {
     return withProgress(
       async (progress) => {
-        await promptImportGithubDatabase(
-          this.app,
-          this.databaseManager,
-          this.storagePath,
-          progress,
-          this.queryServer.cliServer,
-        );
+        await this.databaseFetcher.promptImportGithubDatabase(progress);
       },
       {
         title: "Adding database from GitHub",
@@ -707,13 +692,9 @@ 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 importLocalDatabase(
-              this.app.commands,
+            await this.databaseFetcher.importLocalDatabase(
               uri.toString(true),
-              this.databaseManager,
-              this.storagePath,
               progress,
-              this.queryServer.cliServer,
             );
           } else {
             await this.databaseManager.openDatabase(uri, {
@@ -758,13 +739,9 @@ export class DatabaseUI extends DisposableObject {
             await this.databaseManager.removeDatabaseItem(existingItem);
           }
 
-          await importLocalDatabase(
-            this.app.commands,
+          await this.databaseFetcher.importLocalDatabase(
             uri.toString(true),
-            this.databaseManager,
-            this.storagePath,
             progress,
-            this.queryServer.cliServer,
           );
 
           if (existingItem !== undefined) {
@@ -1005,13 +982,9 @@ export class DatabaseUI extends DisposableObject {
       // we are selecting a database archive or a testproj.
       // Unzip archives (if an archive) and copy into a workspace-controlled area
       // before importing.
-      return await importLocalDatabase(
-        this.app.commands,
+      return await this.databaseFetcher.importLocalDatabase(
         uri.toString(true),
-        this.databaseManager,
-        this.storagePath,
         progress,
-        this.queryServer.cliServer,
       );
     }
   }

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

@@ -1,4 +1,7 @@
+import { pathExists, remove } from "fs-extra";
+import { join } from "path";
 import type { Uri } from "vscode";
+import { zip } from "zip-a-folder";
 
 /**
  * The layout of the database.
@@ -28,3 +31,26 @@ export interface DatabaseContents {
 export interface DatabaseContentsWithDbScheme extends DatabaseContents {
   dbSchemeUri: Uri; // Always present
 }
+
+/**
+ * Databases created by the old odasa tool will not have a zipped
+ * source location. However, this extension works better if sources
+ * are zipped.
+ *
+ * This function ensures that the source location is zipped. If the
+ * `src` folder exists and the `src.zip` file does not, the `src`
+ * folder will be zipped and then deleted.
+ *
+ * @param databasePath The full path to the unzipped database
+ */
+export async function ensureZippedSourceLocation(
+  databasePath: string,
+): Promise<void> {
+  const srcFolderPath = join(databasePath, "src");
+  const srcZipPath = `${srcFolderPath}.zip`;
+
+  if ((await pathExists(srcFolderPath)) && !(await pathExists(srcZipPath))) {
+    await zip(srcFolderPath, srcZipPath);
+    await remove(srcFolderPath);
+  }
+}

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

@@ -43,7 +43,7 @@ import { DatabaseResolver } from "./database-resolver";
 import { telemetryListener } from "../../common/vscode/telemetry";
 import type { LanguageContextStore } from "../../language-context-store";
 import type { DatabaseOrigin } from "./database-origin";
-import { ensureZippedSourceLocation } from "../database-fetcher";
+import { ensureZippedSourceLocation } from "./database-contents";
 
 /**
  * The name of the key in the workspaceState dictionary in which we

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

@@ -133,6 +133,7 @@ import { OpenReferencedFileCodeLensProvider } from "./local-queries/open-referen
 import { LanguageContextStore } from "./language-context-store";
 import { LanguageSelectionPanel } from "./language-selection-panel/language-selection-panel";
 import { GitHubDatabasesModule } from "./databases/github-databases";
+import { DatabaseFetcher } from "./databases/database-fetcher";
 
 /**
  * extension.ts
@@ -799,12 +800,20 @@ async function activateWithInstalledDistribution(
   // Let this run async.
   void dbm.loadPersistedState();
 
+  const databaseFetcher = new DatabaseFetcher(
+    app,
+    dbm,
+    getContextStoragePath(ctx),
+    cliServer,
+  );
+
   ctx.subscriptions.push(dbm);
 
   void extLogger.log("Initializing database panel.");
   const databaseUI = new DatabaseUI(
     app,
     dbm,
+    databaseFetcher,
     languageContext,
     qs,
     getContextStoragePath(ctx),
@@ -881,8 +890,7 @@ async function activateWithInstalledDistribution(
   await GitHubDatabasesModule.initialize(
     app,
     dbm,
-    getContextStoragePath(ctx),
-    cliServer,
+    databaseFetcher,
     githubDatabaseConfigListener,
   );
 
@@ -953,6 +961,7 @@ async function activateWithInstalledDistribution(
     qs,
     qhm,
     dbm,
+    databaseFetcher,
     cliServer,
     databaseUI,
     localQueryResultsView,
@@ -977,6 +986,7 @@ async function activateWithInstalledDistribution(
   const modelEditorModule = await ModelEditorModule.initialize(
     app,
     dbm,
+    databaseFetcher,
     variantAnalysisManager,
     cliServer,
     qs,

extensions/ql-vscode/src/extension.ts

@@ -54,6 +54,7 @@ import type { QueryTreeViewItem } from "../queries-panel/query-tree-view-item";
 import { tryGetQueryLanguage } from "../common/query-language";
 import type { LanguageContextStore } from "../language-context-store";
 import type { ExtensionApp } from "../common/vscode/vscode-app";
+import type { DatabaseFetcher } from "../databases/database-fetcher";
 
 export enum QuickEvalType {
   None,
@@ -69,6 +70,7 @@ export class LocalQueries extends DisposableObject {
     private readonly queryRunner: QueryRunner,
     private readonly queryHistoryManager: QueryHistoryManager,
     private readonly databaseManager: DatabaseManager,
+    private readonly databaseFetcher: DatabaseFetcher,
     private readonly cliServer: CodeQLCliServer,
     private readonly databaseUI: DatabaseUI,
     private readonly localQueryResultsView: ResultsView,
@@ -319,15 +321,13 @@ export class LocalQueries extends DisposableObject {
   private async createSkeletonQuery(): Promise<void> {
     await withProgress(
       async (progress: ProgressCallback) => {
-        const contextStoragePath =
-          this.app.workspaceStoragePath || this.app.globalStoragePath;
         const language = this.languageContextStore.selectedLanguage;
         const skeletonQueryWizard = new SkeletonQueryWizard(
           this.cliServer,
           progress,
           this.app,
           this.databaseManager,
-          contextStoragePath,
+          this.databaseFetcher,
           this.selectedQueryTreeViewItems,
           language,
         );