Merge pull request #2623 from github/robertbrignull/data-model-dependency

Implement "model dependency" button in application mode

extensions/ql-vscode/src/common/interface-types.ts

@@ -553,6 +553,10 @@ export interface GenerateExternalApiFromLlmMessage {
   modeledMethods: Record<string, ModeledMethod>;
 }
 
+export interface ModelDependencyMessage {
+  t: "modelDependency";
+}
+
 export type ToDataExtensionsEditorMessage =
   | SetExtensionPackStateMessage
   | SetExternalApiUsagesMessage
@@ -568,4 +572,5 @@ export type FromDataExtensionsEditorMessage =
   | JumpToUsageMessage
   | SaveModeledMethods
   | GenerateExternalApiMessage
-  | GenerateExternalApiFromLlmMessage;
+  | GenerateExternalApiFromLlmMessage
+  | ModelDependencyMessage;

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 } from "../common/vscode/progress";
+import {
+  ProgressCallback,
+  ProgressUpdate,
+  withProgress,
+} from "../common/vscode/progress";
 import { QueryRunner } from "../query-server";
 import {
   showAndLogExceptionWithTelemetry,
@@ -44,6 +48,7 @@ import { getAutoModelUsages } from "./auto-model-usages-query";
 import { Mode } from "./shared/mode";
 import { loadModeledMethods, saveModeledMethods } from "./modeled-method-fs";
 import { join } from "path";
+import { pickExtensionPack } from "./extension-pack-picker";
 
 export class DataExtensionsEditorView extends AbstractWebview<
   ToDataExtensionsEditorMessage,
@@ -139,6 +144,10 @@ export class DataExtensionsEditorView extends AbstractWebview<
           msg.modeledMethods,
         );
 
+        break;
+      case "modelDependency":
+        await this.modelDependency();
+
         break;
       case "switchMode":
         this.mode = msg.mode;
@@ -284,29 +293,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({
@@ -429,6 +421,68 @@ export class DataExtensionsEditorView extends AbstractWebview<
     await this.clearProgress();
   }
 
+  private async modelDependency(): Promise<void> {
+    return withProgress(async (progress, token) => {
+      const addedDatabase = await this.promptImportAndResetDatabase(progress);
+      if (!addedDatabase || token.isCancellationRequested) {
+        return;
+      }
+
+      const modelFile = await pickExtensionPack(
+        this.cliServer,
+        addedDatabase,
+        this.app.logger,
+        progress,
+        token,
+      );
+      if (!modelFile) {
+        return;
+      }
+
+      const view = new DataExtensionsEditorView(
+        this.ctx,
+        this.app,
+        this.databaseManager,
+        this.cliServer,
+        this.queryRunner,
+        this.queryStorageDir,
+        addedDatabase,
+        modelFile,
+        Mode.Framework,
+      );
+      await view.openView();
+    });
+  }
+
+  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

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

@@ -85,6 +85,7 @@ export async function promptImportInternetDatabase(
  * @param credentials the credentials to use to authenticate with GitHub
  * @param progress the progress callback
  * @param cli the CodeQL CLI server
+ * @param language the language to download. If undefined, the user will be prompted to choose a language.
  */
 export async function promptImportGithubDatabase(
   commandManager: AppCommandManager,
@@ -93,6 +94,7 @@ export async function promptImportGithubDatabase(
   credentials: Credentials | undefined,
   progress: ProgressCallback,
   cli?: CodeQLCliServer,
+  language?: string,
 ): Promise<DatabaseItem | undefined> {
   const githubRepo = await askForGitHubRepo(progress);
   if (!githubRepo) {
@@ -106,6 +108,7 @@ export async function promptImportGithubDatabase(
     credentials,
     progress,
     cli,
+    language,
   );
 
   if (databaseItem) {

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

@@ -215,6 +215,12 @@ export function DataExtensionsEditor({
     });
   }, [externalApiUsages, modeledMethods]);
 
+  const onModelDependencyClick = useCallback(() => {
+    vscode.postMessage({
+      t: "modelDependency",
+    });
+  }, []);
+
   const onGenerateFromLlmClick = useCallback(
     (
       externalApiUsages: ExternalApiUsage[],
@@ -323,6 +329,7 @@ export function DataExtensionsEditor({
               onSaveModelClick={onSaveModelClick}
               onGenerateFromLlmClick={onGenerateFromLlmClick}
               onGenerateFromSourceClick={onGenerateFromSourceClick}
+              onModelDependencyClick={onModelDependencyClick}
             />
           </EditorContainer>
         </>

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

@@ -87,6 +87,7 @@ type Props = {
     modeledMethods: Record<string, ModeledMethod>,
   ) => void;
   onGenerateFromSourceClick: () => void;
+  onModelDependencyClick: () => void;
 };
 
 export const LibraryRow = ({
@@ -100,6 +101,7 @@ export const LibraryRow = ({
   onSaveModelClick,
   onGenerateFromLlmClick,
   onGenerateFromSourceClick,
+  onModelDependencyClick,
 }: Props) => {
   const modeledPercentage = useMemo(() => {
     return calculateModeledPercentage(externalApiUsages);
@@ -129,10 +131,14 @@ export const LibraryRow = ({
     [onGenerateFromSourceClick],
   );
 
-  const handleModelDependency = useCallback(async (e: React.MouseEvent) => {
-    e.stopPropagation();
-    e.preventDefault();
-  }, []);
+  const handleModelDependency = useCallback(
+    async (e: React.MouseEvent) => {
+      onModelDependencyClick();
+      e.stopPropagation();
+      e.preventDefault();
+    },
+    [onModelDependencyClick],
+  );
 
   const handleSave = useCallback(
     async (e: React.MouseEvent) => {

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

@@ -29,6 +29,7 @@ type Props = {
     modeledMethods: Record<string, ModeledMethod>,
   ) => void;
   onGenerateFromSourceClick: () => void;
+  onModelDependencyClick: () => void;
 };
 
 const libraryNameOverrides: Record<string, string> = {
@@ -44,6 +45,7 @@ export const ModeledMethodsList = ({
   onSaveModelClick,
   onGenerateFromLlmClick,
   onGenerateFromSourceClick,
+  onModelDependencyClick,
 }: Props) => {
   const grouped = useMemo(
     () => groupMethods(externalApiUsages, viewState.mode),
@@ -85,6 +87,7 @@ export const ModeledMethodsList = ({
           onSaveModelClick={onSaveModelClick}
           onGenerateFromLlmClick={onGenerateFromLlmClick}
           onGenerateFromSourceClick={onGenerateFromSourceClick}
+          onModelDependencyClick={onModelDependencyClick}
         />
       ))}
     </>