Change data flow generator for framework mode

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

@@ -327,30 +327,36 @@ export class DataExtensionsEditorView extends AbstractWebview<
   protected async generateModeledMethods(): Promise<void> {
     const tokenSource = new CancellationTokenSource();
 
-    const selectedDatabase = this.databaseManager.currentDatabaseItem;
+    let addedDatabase: DatabaseItem | undefined;
 
-    // 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 database = await promptImportGithubDatabase(
-      this.app.commands,
-      this.databaseManager,
-      this.app.workspaceStoragePath ?? this.app.globalStoragePath,
-      this.app.credentials,
-      (update) => this.showProgress(update),
-      this.cliServer,
-    );
-    if (!database) {
-      await this.clearProgress();
-      void this.app.logger.log("No database chosen");
+    // 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;
 
-      return;
+      // 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,
+      );
+      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);
     }
 
-    // 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({
       step: 0,
       maxStep: 4000,
@@ -362,7 +368,7 @@ export class DataExtensionsEditorView extends AbstractWebview<
         cliServer: this.cliServer,
         queryRunner: this.queryRunner,
         queryStorageDir: this.queryStorageDir,
-        databaseItem: database,
+        databaseItem: addedDatabase ?? this.databaseItem,
         onResults: async (results) => {
           const modeledMethodsByName: Record<string, ModeledMethod> = {};
 
@@ -389,14 +395,16 @@ export class DataExtensionsEditorView extends AbstractWebview<
       );
     }
 
-    // After the flow model has been generated, we can remove the temporary database
-    // which we used for generating the flow model.
-    await this.showProgress({
-      step: 3900,
-      maxStep: 4000,
-      message: "Removing temporary database",
-    });
-    await this.databaseManager.removeDatabaseItem(database);
+    if (addedDatabase) {
+      // After the flow model has been generated, we can remove the temporary database
+      // which we used for generating the flow model.
+      await this.showProgress({
+        step: 3900,
+        maxStep: 4000,
+        message: "Removing temporary database",
+      });
+      await this.databaseManager.removeDatabaseItem(addedDatabase);
+    }
 
     await this.clearProgress();
   }

extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx

@@ -220,7 +220,9 @@ export function DataExtensionsEditor({
             <ButtonsContainer>
               <VSCodeButton onClick={onApplyClick}>Apply</VSCodeButton>
               <VSCodeButton onClick={onGenerateClick}>
-                Download and generate
+                {viewState?.mode === Mode.Framework
+                  ? "Generate"
+                  : "Download and generate"}
               </VSCodeButton>
               {viewState?.showLlmButton && (
                 <>