Pull out common code for importing and resetting the current database

Author: robertbrignull

extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts

@@ -14,7 +14,11 @@ import {
   FromDataExtensionsEditorMessage,
   ToDataExtensionsEditorMessage,
 } from "../common/interface-types";
-import { ProgressUpdate, withProgress } from "../common/vscode/progress";
+import {
+  ProgressCallback,
+  ProgressUpdate,
+  withProgress,
+} from "../common/vscode/progress";
 import { QueryRunner } from "../query-server";
 import {
   showAndLogExceptionWithTelemetry,
@@ -288,29 +292,12 @@ export class DataExtensionsEditorView extends AbstractWebview<
     // In application mode, we need the database of a specific library to generate
     // the modeled methods. In framework mode, we'll use the current database.
     if (this.mode === Mode.Application) {
-      const selectedDatabase = this.databaseManager.currentDatabaseItem;
-
-      // The external API methods are in the library source code, so we need to ask
-      // the user to import the library database. We need to have the database
-      // imported to the query server, so we need to register it to our workspace.
-      addedDatabase = await promptImportGithubDatabase(
-        this.app.commands,
-        this.databaseManager,
-        this.app.workspaceStoragePath ?? this.app.globalStoragePath,
-        this.app.credentials,
-        (update) => this.showProgress(update),
-        this.cliServer,
+      addedDatabase = await this.promptImportAndResetDatabase((update) =>
+        this.showProgress(update),
       );
       if (!addedDatabase) {
-        await this.clearProgress();
-        void this.app.logger.log("No database chosen");
-
         return;
       }
-
-      // The library database was set as the current database by importing it,
-      // but we need to set it back to the originally selected database.
-      await this.databaseManager.setCurrentDatabaseItem(selectedDatabase);
     }
 
     await this.showProgress({
@@ -435,30 +422,8 @@ export class DataExtensionsEditorView extends AbstractWebview<
 
   private async modelDependency(): Promise<void> {
     return withProgress(async (progress, token) => {
-      const selectedDatabase = this.databaseManager.currentDatabaseItem;
-
-      // The external API methods are in the library source code, so we need to ask
-      // the user to import the library database. We need to have the database
-      // imported to the query server, so we need to register it to our workspace.
-      const addedDatabase = await promptImportGithubDatabase(
-        this.app.commands,
-        this.databaseManager,
-        this.app.workspaceStoragePath ?? this.app.globalStoragePath,
-        this.app.credentials,
-        progress,
-        this.cliServer,
-        this.databaseItem.language,
-      );
-      if (!addedDatabase) {
-        void this.app.logger.log("No database chosen");
-        return;
-      }
-
-      // The library database was set as the current database by importing it,
-      // but we need to set it back to the originally selected database.
-      await this.databaseManager.setCurrentDatabaseItem(selectedDatabase);
-
-      if (token.isCancellationRequested) {
+      const addedDatabase = await this.promptImportAndResetDatabase(progress);
+      if (!addedDatabase || token.isCancellationRequested) {
         return;
       }
 
@@ -488,6 +453,35 @@ export class DataExtensionsEditorView extends AbstractWebview<
     });
   }
 
+  private async promptImportAndResetDatabase(
+    progress: ProgressCallback,
+  ): Promise<DatabaseItem | undefined> {
+    const selectedDatabase = this.databaseManager.currentDatabaseItem;
+
+    // The external API methods are in the library source code, so we need to ask
+    // the user to import the library database. We need to have the database
+    // imported to the query server, so we need to register it to our workspace.
+    const addedDatabase = await promptImportGithubDatabase(
+      this.app.commands,
+      this.databaseManager,
+      this.app.workspaceStoragePath ?? this.app.globalStoragePath,
+      this.app.credentials,
+      progress,
+      this.cliServer,
+      this.databaseItem.language,
+    );
+    if (!addedDatabase) {
+      void this.app.logger.log("No database chosen");
+      return undefined;
+    }
+
+    // The library database was set as the current database by importing it,
+    // but we need to set it back to the originally selected database.
+    await this.databaseManager.setCurrentDatabaseItem(selectedDatabase);
+
+    return addedDatabase;
+  }
+
   /*
    * Progress in this class is a bit weird. Most of the progress is based on running the query.
    * Query progress is always between 0 and 1000. However, we still have some steps that need