Merge remote-tracking branch 'origin/main' into koesie10/data-extension-editor-cli-tests

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

@@ -7,12 +7,14 @@ import { DatabaseManager } from "../local-databases";
 import { extLogger } from "../common";
 import { ensureDir } from "fs-extra";
 import { join } from "path";
+import { App } from "../common/app";
 
 export class DataExtensionsEditorModule {
   private readonly queryStorageDir: string;
 
   private constructor(
     private readonly ctx: ExtensionContext,
+    private readonly app: App,
     private readonly databaseManager: DatabaseManager,
     private readonly cliServer: CodeQLCliServer,
     private readonly queryRunner: QueryRunner,
@@ -26,6 +28,7 @@ export class DataExtensionsEditorModule {
 
   public static async initialize(
     ctx: ExtensionContext,
+    app: App,
     databaseManager: DatabaseManager,
     cliServer: CodeQLCliServer,
     queryRunner: QueryRunner,
@@ -33,6 +36,7 @@ export class DataExtensionsEditorModule {
   ): Promise<DataExtensionsEditorModule> {
     const dataExtensionsEditorModule = new DataExtensionsEditorModule(
       ctx,
+      app,
       databaseManager,
       cliServer,
       queryRunner,
@@ -54,6 +58,8 @@ export class DataExtensionsEditorModule {
 
         const view = new DataExtensionsEditorView(
           this.ctx,
+          this.app,
+          this.databaseManager,
           this.cliServer,
           this.queryRunner,
           this.queryStorageDir,

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

@@ -5,6 +5,7 @@ import {
   ViewColumn,
   window,
   workspace,
+  WorkspaceFolder,
 } from "vscode";
 import { AbstractWebview, WebviewPanelConfig } from "../abstract-webview";
 import {
@@ -20,9 +21,12 @@ import {
 import { extLogger } from "../common";
 import { readFile, writeFile } from "fs-extra";
 import { load as loadYaml } from "js-yaml";
-import { DatabaseItem } from "../local-databases";
+import { DatabaseItem, DatabaseManager } from "../local-databases";
 import { CodeQLCliServer } from "../cli";
 import { asError, assertNever, getErrorMessage } from "../pure/helpers-pure";
+import { generateFlowModel } from "./generate-flow-model";
+import { promptImportGithubDatabase } from "../databaseFetcher";
+import { App } from "../common/app";
 import { ResolvableLocationValue } from "../pure/bqrs-cli-types";
 import { showResolvableLocation } from "../interface-utils";
 import { decodeBqrsToExternalApiUsages } from "./bqrs";
@@ -32,12 +36,27 @@ import { createDataExtensionYaml, loadDataExtensionYaml } from "./yaml";
 import { ExternalApiUsage } from "./external-api-usage";
 import { ModeledMethod } from "./modeled-method";
 
+function getQlSubmoduleFolder(): WorkspaceFolder | undefined {
+  const workspaceFolder = workspace.workspaceFolders?.find(
+    (folder) => folder.name === "ql",
+  );
+  if (!workspaceFolder) {
+    void extLogger.log("No workspace folder 'ql' found");
+
+    return;
+  }
+
+  return workspaceFolder;
+}
+
 export class DataExtensionsEditorView extends AbstractWebview<
   ToDataExtensionsEditorMessage,
   FromDataExtensionsEditorMessage
 > {
   public constructor(
     ctx: ExtensionContext,
+    private readonly app: App,
+    private readonly databaseManager: DatabaseManager,
     private readonly cliServer: CodeQLCliServer,
     private readonly queryRunner: QueryRunner,
     private readonly queryStorageDir: string,
@@ -86,6 +105,10 @@ export class DataExtensionsEditorView extends AbstractWebview<
         );
         await this.loadExternalApiUsages();
 
+        break;
+      case "generateExternalApi":
+        await this.generateModeledMethods();
+
         break;
       default:
         assertNever(msg);
@@ -158,8 +181,8 @@ export class DataExtensionsEditorView extends AbstractWebview<
       }
 
       await this.postMessage({
-        t: "setExistingModeledMethods",
-        existingModeledMethods,
+        t: "addModeledMethods",
+        modeledMethods: existingModeledMethods,
       });
     } catch (e: unknown) {
       void extLogger.log(`Unable to read data extension YAML: ${e}`);
@@ -225,6 +248,92 @@ export class DataExtensionsEditorView extends AbstractWebview<
     }
   }
 
+  protected async generateModeledMethods(): Promise<void> {
+    const tokenSource = new CancellationTokenSource();
+
+    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 database = await promptImportGithubDatabase(
+      this.app.commands,
+      this.databaseManager,
+      this.app.workspaceStoragePath ?? this.app.globalStoragePath,
+      this.app.credentials,
+      (update) => this.showProgress(update),
+      tokenSource.token,
+      this.cliServer,
+    );
+    if (!database) {
+      await this.clearProgress();
+      void extLogger.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);
+
+    const workspaceFolder = getQlSubmoduleFolder();
+    if (!workspaceFolder) {
+      return;
+    }
+
+    await this.showProgress({
+      step: 0,
+      maxStep: 4000,
+      message: "Generating modeled methods for library",
+    });
+
+    try {
+      await generateFlowModel({
+        cliServer: this.cliServer,
+        queryRunner: this.queryRunner,
+        queryStorageDir: this.queryStorageDir,
+        qlDir: workspaceFolder.uri.fsPath,
+        databaseItem: database,
+        onResults: async (results) => {
+          const modeledMethodsByName: Record<string, ModeledMethod> = {};
+
+          for (const result of results) {
+            modeledMethodsByName[result.signature] = result.modeledMethod;
+          }
+
+          await this.postMessage({
+            t: "addModeledMethods",
+            modeledMethods: modeledMethodsByName,
+            overrideNone: true,
+          });
+        },
+        progress: (update) => this.showProgress(update),
+        token: tokenSource.token,
+      });
+    } catch (e: unknown) {
+      void showAndLogExceptionWithTelemetry(
+        redactableError(
+          asError(e),
+        )`Failed to generate flow model: ${getErrorMessage(e)}`,
+      );
+    }
+
+    // After the flow model has been generated, we can remove the temporary database
+    // which we used for generating the flow model.
+    await this.databaseManager.removeDatabaseItem(
+      () =>
+        this.showProgress({
+          step: 3900,
+          maxStep: 4000,
+          message: "Removing temporary database",
+        }),
+      tokenSource.token,
+      database,
+    );
+
+    await this.clearProgress();
+  }
+
   /*
    * 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
@@ -232,6 +341,13 @@ export class DataExtensionsEditorView extends AbstractWebview<
    * that there's 1000 steps of the query progress since that takes the most time, and then
    * an additional 500 steps for the rest of the work. The progress doesn't need to be 100%
    * accurate, so this is just a rough estimate.
+   *
+   * For generating the modeled methods for an external library, the max step is 4000. This is
+   * based on the following steps:
+   * - 1000 for the summary model
+   * - 1000 for the sink model
+   * - 1000 for the source model
+   * - 1000 for the neutral model
    */
   private async showProgress(update: ProgressUpdate, maxStep?: number) {
     await this.postMessage({
@@ -251,12 +367,8 @@ export class DataExtensionsEditorView extends AbstractWebview<
   }
 
   private calculateModelFilename(): string | undefined {
-    const workspaceFolder = workspace.workspaceFolders?.find(
-      (folder) => folder.name === "ql",
-    );
+    const workspaceFolder = getQlSubmoduleFolder();
     if (!workspaceFolder) {
-      void extLogger.log("No workspace folder 'ql' found");
-
       return;
     }
 

extensions/ql-vscode/src/data-extensions-editor/generate-flow-model.ts

@@ -0,0 +1,141 @@
+import { CancellationToken } from "vscode";
+import { DatabaseItem } from "../local-databases";
+import { join } from "path";
+import { QueryRunner } from "../queryRunner";
+import { CodeQLCliServer } from "../cli";
+import { TeeLogger } from "../common";
+import { extensiblePredicateDefinitions } from "./yaml";
+import { ProgressCallback } from "../progress";
+import { getOnDiskWorkspaceFolders } from "../helpers";
+import {
+  ModeledMethodType,
+  ModeledMethodWithSignature,
+} from "./modeled-method";
+
+type FlowModelOptions = {
+  cliServer: CodeQLCliServer;
+  queryRunner: QueryRunner;
+  queryStorageDir: string;
+  qlDir: string;
+  databaseItem: DatabaseItem;
+  progress: ProgressCallback;
+  token: CancellationToken;
+  onResults: (results: ModeledMethodWithSignature[]) => void | Promise<void>;
+};
+
+async function getModeledMethodsFromFlow(
+  type: Exclude<ModeledMethodType, "none">,
+  queryName: string,
+  queryStep: number,
+  {
+    cliServer,
+    queryRunner,
+    queryStorageDir,
+    qlDir,
+    databaseItem,
+    progress,
+    token,
+  }: Omit<FlowModelOptions, "onResults">,
+): Promise<ModeledMethodWithSignature[]> {
+  const definition = extensiblePredicateDefinitions[type];
+
+  const query = join(
+    qlDir,
+    databaseItem.language,
+    "ql/src/utils/modelgenerator",
+    queryName,
+  );
+
+  const queryRun = queryRunner.createQueryRun(
+    databaseItem.databaseUri.fsPath,
+    { queryPath: query, quickEvalPosition: undefined },
+    false,
+    getOnDiskWorkspaceFolders(),
+    undefined,
+    queryStorageDir,
+    undefined,
+    undefined,
+  );
+
+  const queryResult = await queryRun.evaluate(
+    ({ step, message }) =>
+      progress({
+        message: `Generating ${type} model: ${message}`,
+        step: queryStep * 1000 + step,
+        maxStep: 4000,
+      }),
+    token,
+    new TeeLogger(queryRunner.logger, queryRun.outputDir.logPath),
+  );
+
+  const bqrsPath = queryResult.outputDir.bqrsPath;
+
+  const bqrsInfo = await cliServer.bqrsInfo(bqrsPath);
+  if (bqrsInfo["result-sets"].length !== 1) {
+    throw new Error(
+      `Expected exactly one result set, got ${bqrsInfo["result-sets"].length}`,
+    );
+  }
+
+  const resultSet = bqrsInfo["result-sets"][0];
+
+  const decodedResults = await cliServer.bqrsDecode(bqrsPath, resultSet.name);
+
+  const results = decodedResults.tuples;
+
+  return (
+    results
+      // This is just a sanity check. The query should only return strings.
+      .filter((result) => typeof result[0] === "string")
+      .map((result) => {
+        const row = result[0] as string;
+
+        return definition.readModeledMethod(row.split(";"));
+      })
+  );
+}
+
+export async function generateFlowModel({
+  onResults,
+  ...options
+}: FlowModelOptions) {
+  const summaryResults = await getModeledMethodsFromFlow(
+    "summary",
+    "CaptureSummaryModels.ql",
+    0,
+    options,
+  );
+  if (summaryResults) {
+    await onResults(summaryResults);
+  }
+
+  const sinkResults = await getModeledMethodsFromFlow(
+    "sink",
+    "CaptureSinkModels.ql",
+    1,
+    options,
+  );
+  if (sinkResults) {
+    await onResults(sinkResults);
+  }
+
+  const sourceResults = await getModeledMethodsFromFlow(
+    "source",
+    "CaptureSourceModels.ql",
+    2,
+    options,
+  );
+  if (sourceResults) {
+    await onResults(sourceResults);
+  }
+
+  const neutralResults = await getModeledMethodsFromFlow(
+    "neutral",
+    "CaptureNeutralModels.ql",
+    3,
+    options,
+  );
+  if (neutralResults) {
+    await onResults(neutralResults);
+  }
+}

extensions/ql-vscode/src/data-extensions-editor/modeled-method.ts

@@ -11,3 +11,8 @@ export type ModeledMethod = {
   output: string;
   kind: string;
 };
+
+export type ModeledMethodWithSignature = {
+  signature: string;
+  modeledMethod: ModeledMethod;
+};

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

@@ -1,27 +1,31 @@
 import { ExternalApiUsage } from "./external-api-usage";
-import { ModeledMethod, ModeledMethodType } from "./modeled-method";
+import {
+  ModeledMethod,
+  ModeledMethodType,
+  ModeledMethodWithSignature,
+} from "./modeled-method";
 
 type ExternalApiUsageByType = {
   externalApiUsage: ExternalApiUsage;
   modeledMethod: ModeledMethod;
 };
 
-type DataExtensionDefinition = {
-  extensible: string;
+type ExtensiblePredicateDefinition = {
+  extensiblePredicate: string;
   generateMethodDefinition: (method: ExternalApiUsageByType) => any[];
-  readModeledMethod: (row: any[]) => [string, ModeledMethod] | undefined;
+  readModeledMethod: (row: any[]) => ModeledMethodWithSignature;
 };
 
 function readRowToMethod(row: any[]): string {
   return `${row[0]}.${row[1]}#${row[3]}${row[4]}`;
 }
 
-const definitions: Record<
+export const extensiblePredicateDefinitions: Record<
   Exclude<ModeledMethodType, "none">,
-  DataExtensionDefinition
+  ExtensiblePredicateDefinition
 > = {
   source: {
-    extensible: "sourceModel",
+    extensiblePredicate: "sourceModel",
     // extensible predicate sourceModel(
     //   string package, string type, boolean subtypes, string name, string signature, string ext,
     //   string output, string kind, string provenance
@@ -37,18 +41,18 @@ const definitions: Record<
       method.modeledMethod.kind,
       "manual",
     ],
-    readModeledMethod: (row) => [
-      readRowToMethod(row),
-      {
+    readModeledMethod: (row) => ({
+      signature: readRowToMethod(row),
+      modeledMethod: {
         type: "source",
         input: "",
         output: row[6],
         kind: row[7],
       },
-    ],
+    }),
   },
   sink: {
-    extensible: "sinkModel",
+    extensiblePredicate: "sinkModel",
     // extensible predicate sinkModel(
     //   string package, string type, boolean subtypes, string name, string signature, string ext,
     //   string input, string kind, string provenance
@@ -64,18 +68,18 @@ const definitions: Record<
       method.modeledMethod.kind,
       "manual",
     ],
-    readModeledMethod: (row) => [
-      readRowToMethod(row),
-      {
+    readModeledMethod: (row) => ({
+      signature: readRowToMethod(row),
+      modeledMethod: {
         type: "sink",
         input: row[6],
         output: "",
         kind: row[7],
       },
-    ],
+    }),
   },
   summary: {
-    extensible: "summaryModel",
+    extensiblePredicate: "summaryModel",
     // extensible predicate summaryModel(
     //   string package, string type, boolean subtypes, string name, string signature, string ext,
     //   string input, string output, string kind, string provenance
@@ -92,18 +96,18 @@ const definitions: Record<
       method.modeledMethod.kind,
       "manual",
     ],
-    readModeledMethod: (row) => [
-      readRowToMethod(row),
-      {
+    readModeledMethod: (row) => ({
+      signature: readRowToMethod(row),
+      modeledMethod: {
         type: "summary",
         input: row[6],
         output: row[7],
         kind: row[8],
       },
-    ],
+    }),
   },
   neutral: {
-    extensible: "neutralModel",
+    extensiblePredicate: "neutralModel",
     // extensible predicate neutralModel(
     //   string package, string type, string name, string signature, string provenance
     // );
@@ -114,21 +118,21 @@ const definitions: Record<
       method.externalApiUsage.methodParameters,
       "manual",
     ],
-    readModeledMethod: (row) => [
-      `${row[0]}.${row[1]}#${row[2]}${row[3]}`,
-      {
+    readModeledMethod: (row) => ({
+      signature: `${row[0]}.${row[1]}#${row[2]}${row[3]}`,
+      modeledMethod: {
         type: "neutral",
         input: "",
         output: "",
         kind: "",
       },
-    ],
+    }),
   },
 };
 
 function createDataProperty(
   methods: ExternalApiUsageByType[],
-  definition: DataExtensionDefinition,
+  definition: ExtensiblePredicateDefinition,
 ) {
   if (methods.length === 0) {
     return " []";
@@ -169,10 +173,10 @@ export function createDataExtensionYaml(
     }
   }
 
-  const extensions = Object.entries(definitions).map(
+  const extensions = Object.entries(extensiblePredicateDefinitions).map(
     ([type, definition]) => `  - addsTo:
       pack: codeql/java-all
-      extensible: ${definition.extensible}
+      extensible: ${definition.extensiblePredicate}
     data:${createDataProperty(
       methodsByType[type as Exclude<ModeledMethodType, "none">],
       definition,
@@ -214,8 +218,8 @@ export function loadDataExtensionYaml(
       continue;
     }
 
-    const definition = Object.values(definitions).find(
-      (definition) => definition.extensible === extensible,
+    const definition = Object.values(extensiblePredicateDefinitions).find(
+      (definition) => definition.extensiblePredicate === extensible,
     );
     if (!definition) {
       continue;
@@ -227,9 +231,9 @@ export function loadDataExtensionYaml(
         continue;
       }
 
-      const [apiInfo, modeledMethod] = result;
+      const { signature, modeledMethod } = result;
 
-      modeledMethods[apiInfo] = modeledMethod;
+      modeledMethods[signature] = modeledMethod;
     }
   }
 

extensions/ql-vscode/src/extension.ts

@@ -868,6 +868,7 @@ async function activateWithInstalledDistribution(
   const dataExtensionsEditorModule =
     await DataExtensionsEditorModule.initialize(
       ctx,
+      app,
       dbm,
       cliServer,
       qs,

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

@@ -493,6 +493,19 @@ export interface ShowProgressMessage {
   message: string;
 }
 
+export interface AddModeledMethodsMessage {
+  t: "addModeledMethods";
+  modeledMethods: Record<string, ModeledMethod>;
+
+  /**
+   * If true, then any existing modeled methods set to "none" will be
+   * overwritten by the new modeled methods. Otherwise, the "none" modeled
+   * methods will not be overwritten, even if the new modeled methods
+   * contain a better model.
+   */
+  overrideNone?: boolean;
+}
+
 export interface JumpToUsageMessage {
   t: "jumpToUsage";
   location: ResolvableLocationValue;
@@ -509,12 +522,17 @@ export interface SaveModeledMethods {
   modeledMethods: Record<string, ModeledMethod>;
 }
 
+export interface GenerateExternalApiMessage {
+  t: "generateExternalApi";
+}
+
 export type ToDataExtensionsEditorMessage =
   | SetExternalApiUsagesMessage
   | ShowProgressMessage
-  | SetExistingModeledMethods;
+  | AddModeledMethodsMessage;
 
 export type FromDataExtensionsEditorMessage =
   | ViewLoadedMsg
   | JumpToUsageMessage
-  | SaveModeledMethods;
+  | SaveModeledMethods
+  | GenerateExternalApiMessage;

extensions/ql-vscode/src/variant-analysis/store/repo-states-data-types.ts

@@ -0,0 +1,12 @@
+export interface VariantAnalysisScannedRepositoryStateData {
+  repositoryId: number;
+  downloadStatus: VariantAnalysisScannedRepositoryDownloadData;
+  downloadPercentage?: number;
+}
+
+export enum VariantAnalysisScannedRepositoryDownloadData {
+  Pending = "pending",
+  InProgress = "inProgress",
+  Succeeded = "succeeded",
+  Failed = "failed",
+}

extensions/ql-vscode/src/variant-analysis/store/repo-states-store.ts

@@ -1,17 +1,44 @@
 import { outputJson, readJson } from "fs-extra";
 import { VariantAnalysisScannedRepositoryState } from "../shared/variant-analysis";
+import { VariantAnalysisScannedRepositoryStateData } from "./repo-states-data-types";
+import { mapRepoStateToData } from "./repo-states-to-data-mapper";
+import { mapRepoStateToDomain } from "./repo-states-to-domain-mapper";
 
 export const REPO_STATES_FILENAME = "repo_states.json";
 
 export async function writeRepoStates(
   storagePath: string,
-  repoStates: Record<number, VariantAnalysisScannedRepositoryState> | undefined,
+  repoStates: Record<number, VariantAnalysisScannedRepositoryState>,
 ): Promise<void> {
-  return await outputJson(storagePath, repoStates);
+  // Map from repoStates Domain type to the repoStates Data type
+  const repoStatesData = Object.fromEntries(
+    Object.entries(repoStates).map(([key, value]) => {
+      return [key, mapRepoStateToData(value)];
+    }),
+  );
+
+  await outputJson(storagePath, repoStatesData);
 }
 
 export async function readRepoStates(
   storagePath: string,
-): Promise<Record<number, VariantAnalysisScannedRepositoryState>> {
-  return await readJson(storagePath);
+): Promise<Record<number, VariantAnalysisScannedRepositoryState> | undefined> {
+  try {
+    const repoStatesData: Record<
+      number,
+      VariantAnalysisScannedRepositoryStateData
+    > = await readJson(storagePath);
+
+    // Map from repoStates Data type to the repoStates Domain type
+    const repoStates = Object.fromEntries(
+      Object.entries(repoStatesData).map(([key, value]) => {
+        return [key, mapRepoStateToDomain(value)];
+      }),
+    );
+
+    return repoStates;
+  } catch (e) {
+    // Ignore this error, we simply might not have downloaded anything yet
+    return undefined;
+  }
 }

extensions/ql-vscode/src/variant-analysis/store/repo-states-to-data-mapper.ts

@@ -0,0 +1,36 @@
+import { assertNever } from "../../pure/helpers-pure";
+import {
+  VariantAnalysisScannedRepositoryDownloadStatus,
+  VariantAnalysisScannedRepositoryState,
+} from "../shared/variant-analysis";
+import {
+  VariantAnalysisScannedRepositoryDownloadData,
+  VariantAnalysisScannedRepositoryStateData,
+} from "./repo-states-data-types";
+
+export function mapRepoStateToData(
+  repoState: VariantAnalysisScannedRepositoryState,
+): VariantAnalysisScannedRepositoryStateData {
+  return {
+    repositoryId: repoState.repositoryId,
+    downloadStatus: processDownloadStatus(repoState.downloadStatus),
+    downloadPercentage: repoState.downloadPercentage,
+  };
+}
+
+function processDownloadStatus(
+  downloadedStatus: VariantAnalysisScannedRepositoryDownloadStatus,
+) {
+  switch (downloadedStatus) {
+    case VariantAnalysisScannedRepositoryDownloadStatus.Pending:
+      return VariantAnalysisScannedRepositoryDownloadData.Pending;
+    case VariantAnalysisScannedRepositoryDownloadStatus.InProgress:
+      return VariantAnalysisScannedRepositoryDownloadData.InProgress;
+    case VariantAnalysisScannedRepositoryDownloadStatus.Succeeded:
+      return VariantAnalysisScannedRepositoryDownloadData.Succeeded;
+    case VariantAnalysisScannedRepositoryDownloadStatus.Failed:
+      return VariantAnalysisScannedRepositoryDownloadData.Failed;
+    default:
+      assertNever(downloadedStatus);
+  }
+}

extensions/ql-vscode/src/variant-analysis/store/repo-states-to-domain-mapper.ts

@@ -0,0 +1,36 @@
+import { assertNever } from "../../pure/helpers-pure";
+import {
+  VariantAnalysisScannedRepositoryState,
+  VariantAnalysisScannedRepositoryDownloadStatus,
+} from "../shared/variant-analysis";
+import {
+  VariantAnalysisScannedRepositoryStateData,
+  VariantAnalysisScannedRepositoryDownloadData,
+} from "./repo-states-data-types";
+
+export function mapRepoStateToDomain(
+  repoState: VariantAnalysisScannedRepositoryStateData,
+): VariantAnalysisScannedRepositoryState {
+  return {
+    repositoryId: repoState.repositoryId,
+    downloadStatus: processDownloadStatus(repoState.downloadStatus),
+    downloadPercentage: repoState.downloadPercentage,
+  };
+}
+
+function processDownloadStatus(
+  downloadedStatus: VariantAnalysisScannedRepositoryDownloadData,
+) {
+  switch (downloadedStatus) {
+    case VariantAnalysisScannedRepositoryDownloadData.Pending:
+      return VariantAnalysisScannedRepositoryDownloadStatus.Pending;
+    case VariantAnalysisScannedRepositoryDownloadData.InProgress:
+      return VariantAnalysisScannedRepositoryDownloadStatus.InProgress;
+    case VariantAnalysisScannedRepositoryDownloadData.Succeeded:
+      return VariantAnalysisScannedRepositoryDownloadStatus.Succeeded;
+    case VariantAnalysisScannedRepositoryDownloadData.Failed:
+      return VariantAnalysisScannedRepositoryDownloadStatus.Failed;
+    default:
+      assertNever(downloadedStatus);
+  }
+}

extensions/ql-vscode/src/variant-analysis/store/repo-task-data-types.ts

@@ -0,0 +1,25 @@
+export interface VariantAnalysisRepositoryTaskData {
+  repository: RepositoryData;
+  analysisStatus: VariantAnalysisRepoStatusData;
+  resultCount?: number;
+  artifactSizeInBytes?: number;
+  failureMessage?: string;
+  databaseCommitSha?: string;
+  sourceLocationPrefix?: string;
+  artifactUrl?: string;
+}
+
+interface RepositoryData {
+  id: number;
+  fullName: string;
+  private: boolean;
+}
+
+export enum VariantAnalysisRepoStatusData {
+  Pending = "pending",
+  InProgress = "inProgress",
+  Succeeded = "succeeded",
+  Failed = "failed",
+  Canceled = "canceled",
+  TimedOut = "timedOut",
+}

extensions/ql-vscode/src/variant-analysis/store/repo-task-store.ts

@@ -0,0 +1,24 @@
+import { outputJson, readJson } from "fs-extra";
+import { join } from "path";
+import { VariantAnalysisRepositoryTask } from "../shared/variant-analysis";
+import { mapRepoTaskToData } from "./repo-task-to-data-mapper";
+import { mapRepoTaskToDomain } from "./repo-task-to-domain-mapper";
+
+export const REPO_TASK_FILENAME = "repo_task.json";
+
+export async function writeRepoTask(
+  storageDirectory: string,
+  repoTask: VariantAnalysisRepositoryTask,
+): Promise<void> {
+  const repoTaskData = mapRepoTaskToData(repoTask);
+  await outputJson(join(storageDirectory, REPO_TASK_FILENAME), repoTaskData);
+}
+
+export async function readRepoTask(
+  storageDirectory: string,
+): Promise<VariantAnalysisRepositoryTask> {
+  const repoTaskData = await readJson(
+    join(storageDirectory, REPO_TASK_FILENAME),
+  );
+  return mapRepoTaskToDomain(repoTaskData);
+}

extensions/ql-vscode/src/variant-analysis/store/repo-task-to-data-mapper.ts

@@ -0,0 +1,49 @@
+import { assertNever } from "../../pure/helpers-pure";
+import {
+  VariantAnalysisRepositoryTask,
+  VariantAnalysisRepoStatus,
+} from "../shared/variant-analysis";
+import {
+  VariantAnalysisRepositoryTaskData,
+  VariantAnalysisRepoStatusData,
+} from "./repo-task-data-types";
+
+export function mapRepoTaskToData(
+  repoTask: VariantAnalysisRepositoryTask,
+): VariantAnalysisRepositoryTaskData {
+  return {
+    repository: {
+      id: repoTask.repository.id,
+      fullName: repoTask.repository.fullName,
+      private: repoTask.repository.private,
+    },
+    analysisStatus: mapRepoTaskAnalysisStatusToData(repoTask.analysisStatus),
+    resultCount: repoTask.resultCount,
+    artifactSizeInBytes: repoTask.artifactSizeInBytes,
+    failureMessage: repoTask.failureMessage,
+    databaseCommitSha: repoTask.databaseCommitSha,
+    sourceLocationPrefix: repoTask.sourceLocationPrefix,
+    artifactUrl: repoTask.artifactUrl,
+  };
+}
+
+function mapRepoTaskAnalysisStatusToData(
+  analysisStatus: VariantAnalysisRepoStatus,
+): VariantAnalysisRepoStatusData {
+  switch (analysisStatus) {
+    case VariantAnalysisRepoStatus.Pending:
+      return VariantAnalysisRepoStatusData.Pending;
+    case VariantAnalysisRepoStatus.InProgress:
+      return VariantAnalysisRepoStatusData.InProgress;
+    case VariantAnalysisRepoStatus.Succeeded:
+      return VariantAnalysisRepoStatusData.Succeeded;
+    case VariantAnalysisRepoStatus.Failed:
+      return VariantAnalysisRepoStatusData.Failed;
+    case VariantAnalysisRepoStatus.Canceled:
+      return VariantAnalysisRepoStatusData.Canceled;
+    case VariantAnalysisRepoStatus.TimedOut:
+      return VariantAnalysisRepoStatusData.TimedOut;
+    default:
+      assertNever(analysisStatus);
+  }
+}

extensions/ql-vscode/src/variant-analysis/store/repo-task-to-domain-mapper.ts

@@ -0,0 +1,49 @@
+import { assertNever } from "../../pure/helpers-pure";
+import {
+  VariantAnalysisRepositoryTask,
+  VariantAnalysisRepoStatus,
+} from "../shared/variant-analysis";
+import {
+  VariantAnalysisRepositoryTaskData,
+  VariantAnalysisRepoStatusData,
+} from "./repo-task-data-types";
+
+export function mapRepoTaskToDomain(
+  repoTask: VariantAnalysisRepositoryTaskData,
+): VariantAnalysisRepositoryTask {
+  return {
+    repository: {
+      id: repoTask.repository.id,
+      fullName: repoTask.repository.fullName,
+      private: repoTask.repository.private,
+    },
+    analysisStatus: mapRepoTaskAnalysisStatusToDomain(repoTask.analysisStatus),
+    resultCount: repoTask.resultCount,
+    artifactSizeInBytes: repoTask.artifactSizeInBytes,
+    failureMessage: repoTask.failureMessage,
+    databaseCommitSha: repoTask.databaseCommitSha,
+    sourceLocationPrefix: repoTask.sourceLocationPrefix,
+    artifactUrl: repoTask.artifactUrl,
+  };
+}
+
+function mapRepoTaskAnalysisStatusToDomain(
+  analysisStatus: VariantAnalysisRepoStatusData,
+): VariantAnalysisRepoStatus {
+  switch (analysisStatus) {
+    case VariantAnalysisRepoStatusData.Pending:
+      return VariantAnalysisRepoStatus.Pending;
+    case VariantAnalysisRepoStatusData.InProgress:
+      return VariantAnalysisRepoStatus.InProgress;
+    case VariantAnalysisRepoStatusData.Succeeded:
+      return VariantAnalysisRepoStatus.Succeeded;
+    case VariantAnalysisRepoStatusData.Failed:
+      return VariantAnalysisRepoStatus.Failed;
+    case VariantAnalysisRepoStatusData.Canceled:
+      return VariantAnalysisRepoStatus.Canceled;
+    case VariantAnalysisRepoStatusData.TimedOut:
+      return VariantAnalysisRepoStatus.TimedOut;
+    default:
+      assertNever(analysisStatus);
+  }
+}

extensions/ql-vscode/src/variant-analysis/variant-analysis-manager.ts

@@ -264,15 +264,11 @@ export class VariantAnalysisManager
     } else {
       await this.setVariantAnalysis(variantAnalysis);
 
-      try {
-        const repoStates = await readRepoStates(
-          this.getRepoStatesStoragePath(variantAnalysis.id),
-        );
-        this.repoStates.set(variantAnalysis.id, repoStates);
-      } catch (e) {
-        // Ignore this error, we simply might not have downloaded anything yet
-        this.repoStates.set(variantAnalysis.id, {});
-      }
+      const repoStatesFromDisk = await readRepoStates(
+        this.getRepoStatesStoragePath(variantAnalysis.id),
+      );
+
+      this.repoStates.set(variantAnalysis.id, repoStatesFromDisk || {});
 
       if (
         !(await isVariantAnalysisComplete(
@@ -598,10 +594,13 @@ export class VariantAnalysisManager
       VariantAnalysisScannedRepositoryDownloadStatus.Succeeded;
     await this.onRepoStateUpdated(variantAnalysis.id, repoState);
 
-    await writeRepoStates(
-      this.getRepoStatesStoragePath(variantAnalysis.id),
-      this.repoStates.get(variantAnalysis.id),
-    );
+    const repoStates = this.repoStates.get(variantAnalysis.id);
+    if (repoStates) {
+      await writeRepoStates(
+        this.getRepoStatesStoragePath(variantAnalysis.id),
+        repoStates,
+      );
+    }
   }
 
   public async enqueueDownload(

extensions/ql-vscode/src/variant-analysis/variant-analysis-results-manager.ts

@@ -1,4 +1,4 @@
-import { appendFile, pathExists, mkdir, outputJson, readJson } from "fs-extra";
+import { appendFile, pathExists } from "fs-extra";
 import fetch from "node-fetch";
 import { EOL } from "os";
 import { join } from "path";
@@ -17,6 +17,7 @@ import {
 import { DisposableObject, DisposeHandler } from "../pure/disposable-object";
 import { EventEmitter } from "vscode";
 import { unzipFile } from "../pure/zip";
+import { readRepoTask, writeRepoTask } from "./store/repo-task-store";
 
 type CacheKey = `${number}/${string}`;
 
@@ -37,7 +38,6 @@ export type LoadResultsOptions = {
 };
 
 export class VariantAnalysisResultsManager extends DisposableObject {
-  private static readonly REPO_TASK_FILENAME = "repo_task.json";
   private static readonly RESULTS_DIRECTORY = "results";
 
   private readonly cachedResults: Map<
@@ -78,14 +78,7 @@ export class VariantAnalysisResultsManager extends DisposableObject {
       repoTask.repository.fullName,
     );
 
-    if (!(await pathExists(resultDirectory))) {
-      await mkdir(resultDirectory, { recursive: true });
-    }
-
-    await outputJson(
-      join(resultDirectory, VariantAnalysisResultsManager.REPO_TASK_FILENAME),
-      repoTask,
-    );
+    await writeRepoTask(resultDirectory, repoTask);
 
     const zipFilePath = join(resultDirectory, "results.zip");
 
@@ -184,8 +177,8 @@ export class VariantAnalysisResultsManager extends DisposableObject {
       repositoryFullName,
     );
 
-    const repoTask: VariantAnalysisRepositoryTask = await readJson(
-      join(storageDirectory, VariantAnalysisResultsManager.REPO_TASK_FILENAME),
+    const repoTask: VariantAnalysisRepositoryTask = await readRepoTask(
+      storageDirectory,
     );
 
     if (!repoTask.databaseCommitSha || !repoTask.sourceLocationPrefix) {

extensions/ql-vscode/src/variant-analysis/variant-analysis-view-serializer.ts

@@ -37,6 +37,14 @@ export class VariantAnalysisViewSerializer implements WebviewPanelSerializer {
       return;
     }
 
+    // Between the time the webview is deserialized and the time the extension
+    // is fully activated, the user may close the webview. In this case, we
+    // should not attempt to restore the view.
+    let disposed = false;
+    const unregisterOnDidDispose = webviewPanel.onDidDispose(() => {
+      disposed = true;
+    });
+
     const variantAnalysisState: VariantAnalysisState =
       state as VariantAnalysisState;
 
@@ -46,11 +54,16 @@ export class VariantAnalysisViewSerializer implements WebviewPanelSerializer {
       variantAnalysisState.variantAnalysisId,
     );
     if (existingView) {
+      unregisterOnDidDispose.dispose();
       await existingView.openView();
       webviewPanel.dispose();
       return;
     }
 
+    if (disposed) {
+      return;
+    }
+
     const view = new VariantAnalysisView(
       this.ctx,
       this.app,
@@ -58,6 +71,8 @@ export class VariantAnalysisViewSerializer implements WebviewPanelSerializer {
       manager,
     );
     await view.restoreView(webviewPanel);
+
+    unregisterOnDidDispose.dispose();
   }
 
   private waitForExtensionFullyLoaded(): Promise<

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

@@ -57,14 +57,21 @@ export function DataExtensionsEditor(): JSX.Element {
           case "showProgress":
             setProgress(msg);
             break;
-          case "setExistingModeledMethods":
+          case "addModeledMethods":
             setModeledMethods((oldModeledMethods) => {
+              const filteredOldModeledMethods = msg.overrideNone
+                ? Object.fromEntries(
+                    Object.entries(oldModeledMethods).filter(
+                      ([, value]) => value.type !== "none",
+                    ),
+                  )
+                : oldModeledMethods;
+
               return {
-                ...msg.existingModeledMethods,
-                ...oldModeledMethods,
+                ...msg.modeledMethods,
+                ...filteredOldModeledMethods,
               };
             });
-
             break;
           default:
             assertNever(msg);
@@ -107,6 +114,12 @@ export function DataExtensionsEditor(): JSX.Element {
     });
   }, [externalApiUsages, modeledMethods]);
 
+  const onGenerateClick = useCallback(() => {
+    vscode.postMessage({
+      t: "generateExternalApi",
+    });
+  }, []);
+
   return (
     <DataExtensionsEditorContainer>
       {progress.maxStep > 0 && (
@@ -128,6 +141,12 @@ export function DataExtensionsEditor(): JSX.Element {
           <div>
             <h3>External API modelling</h3>
             <VSCodeButton onClick={onApplyClick}>Apply</VSCodeButton>
+            &nbsp;
+            <VSCodeButton onClick={onGenerateClick}>
+              Download and generate
+            </VSCodeButton>
+            <br />
+            <br />
             <VSCodeDataGrid>
               <VSCodeDataGridRow rowType="header">
                 <VSCodeDataGridCell cellType="columnheader" gridColumn={1}>

extensions/ql-vscode/test/vscode-tests/activated-extension/variant-analysis/variant-analysis-manager.test.ts

@@ -45,6 +45,10 @@ import { App } from "../../../../src/common/app";
 import { ExtensionApp } from "../../../../src/common/vscode/vscode-app";
 import { DbConfigStore } from "../../../../src/databases/config/db-config-store";
 import { mockedObject } from "../../utils/mocking.helpers";
+import {
+  REPO_STATES_FILENAME,
+  writeRepoStates,
+} from "../../../../src/variant-analysis/store/repo-states-store";
 
 // up to 3 minutes per test
 jest.setTimeout(3 * 60 * 1000);
@@ -119,8 +123,12 @@ describe("Variant Analysis Manager", () => {
       });
 
       it("should read in the repo states if it exists", async () => {
-        await fs.writeJson(
-          join(storagePath, variantAnalysis.id.toString(), "repo_states.json"),
+        await writeRepoStates(
+          join(
+            storagePath,
+            variantAnalysis.id.toString(),
+            REPO_STATES_FILENAME,
+          ),
           {
             [scannedRepos[0].repository.id]: {
               repositoryId: scannedRepos[0].repository.id,
@@ -177,7 +185,7 @@ describe("Variant Analysis Manager", () => {
       repoStatesPath = join(
         storagePath,
         variantAnalysis.id.toString(),
-        "repo_states.json",
+        REPO_STATES_FILENAME,
       );
     });