Merge pull request #2535 from github/koesie10/framework-mode

Add initial implementation of framework mode

Author: koesie10

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

@@ -20,6 +20,7 @@ import { DataFlowPaths } from "../variant-analysis/shared/data-flow-paths";
 import { ExternalApiUsage } from "../data-extensions-editor/external-api-usage";
 import { ModeledMethod } from "../data-extensions-editor/modeled-method";
 import { DataExtensionEditorViewState } from "../data-extensions-editor/shared/view-state";
+import { Mode } from "../data-extensions-editor/shared/mode";
 
 /**
  * This module contains types and code that are shared between
@@ -521,6 +522,11 @@ export interface AddModeledMethodsMessage {
   overrideNone?: boolean;
 }
 
+export interface SwitchModeMessage {
+  t: "switchMode";
+  mode: Mode;
+}
+
 export interface JumpToUsageMessage {
   t: "jumpToUsage";
   location: ResolvableLocationValue;
@@ -559,6 +565,7 @@ export type ToDataExtensionsEditorMessage =
 
 export type FromDataExtensionsEditorMessage =
   | ViewLoadedMsg
+  | SwitchModeMessage
   | OpenModelFileMessage
   | OpenExtensionPackMessage
   | JumpToUsageMessage

extensions/ql-vscode/src/config.ts

@@ -714,6 +714,7 @@ export function showQueriesPanel(): boolean {
 
 const DATA_EXTENSIONS = new Setting("dataExtensions", ROOT_SETTING);
 const LLM_GENERATION = new Setting("llmGeneration", DATA_EXTENSIONS);
+const FRAMEWORK_MODE = new Setting("frameworkMode", DATA_EXTENSIONS);
 const DISABLE_AUTO_NAME_EXTENSION_PACK = new Setting(
   "disableAutoNameExtensionPack",
   DATA_EXTENSIONS,
@@ -723,6 +724,10 @@ export function showLlmGeneration(): boolean {
   return !!LLM_GENERATION.getValue<boolean>();
 }
 
+export function enableFrameworkMode(): boolean {
+  return !!FRAMEWORK_MODE.getValue<boolean>();
+}
+
 export function disableAutoNameExtensionPack(): boolean {
   return !!DISABLE_AUTO_NAME_EXTENSION_PACK.getValue<boolean>();
 }

extensions/ql-vscode/src/data-extensions-editor/auto-model-usages-query.ts

@@ -32,7 +32,7 @@ export async function getAutoModelUsages({
   // This will re-run the query that was already run when opening the data extensions editor. This
   // might be unnecessary, but this makes it really easy to get the path to the BQRS file which we
   // need to interpret the results.
-  const queryResult = await runQuery({
+  const queryResult = await runQuery("applicationModeQuery", {
     cliServer,
     queryRunner,
     queryStorageDir,

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

@@ -36,7 +36,8 @@ import { decodeBqrsToExternalApiUsages } from "./bqrs";
 import { redactableError } from "../common/errors";
 import { readQueryResults, runQuery } from "./external-api-usage-query";
 import {
-  createDataExtensionYamlsPerLibrary,
+  createDataExtensionYamlsForApplicationMode,
+  createDataExtensionYamlsForFrameworkMode,
   createFilenameForLibrary,
   loadDataExtensionYaml,
 } from "./yaml";
@@ -48,9 +49,10 @@ import {
   createAutoModelRequest,
   parsePredictedClassifications,
 } from "./auto-model";
-import { showLlmGeneration } from "../config";
+import { enableFrameworkMode, showLlmGeneration } from "../config";
 import { getAutoModelUsages } from "./auto-model-usages-query";
 import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
+import { Mode } from "./shared/mode";
 
 export class DataExtensionsEditorView extends AbstractWebview<
   ToDataExtensionsEditorMessage,
@@ -65,6 +67,7 @@ export class DataExtensionsEditorView extends AbstractWebview<
     private readonly queryStorageDir: string,
     private readonly databaseItem: DatabaseItem,
     private readonly extensionPack: ExtensionPack,
+    private mode: Mode = Mode.Application,
   ) {
     super(ctx);
   }
@@ -138,6 +141,12 @@ export class DataExtensionsEditorView extends AbstractWebview<
           msg.modeledMethods,
         );
 
+        break;
+      case "switchMode":
+        this.mode = msg.mode;
+
+        await Promise.all([this.setViewState(), this.loadExternalApiUsages()]);
+
         break;
       default:
         assertNever(msg);
@@ -159,7 +168,9 @@ export class DataExtensionsEditorView extends AbstractWebview<
       t: "setDataExtensionEditorViewState",
       viewState: {
         extensionPack: this.extensionPack,
+        enableFrameworkMode: enableFrameworkMode(),
         showLlmButton: showLlmGeneration(),
+        mode: this.mode,
       },
     });
   }
@@ -188,11 +199,26 @@ export class DataExtensionsEditorView extends AbstractWebview<
     externalApiUsages: ExternalApiUsage[],
     modeledMethods: Record<string, ModeledMethod>,
   ): Promise<void> {
-    const yamls = createDataExtensionYamlsPerLibrary(
-      this.databaseItem.language,
-      externalApiUsages,
-      modeledMethods,
-    );
+    let yamls: Record<string, string>;
+    switch (this.mode) {
+      case Mode.Application:
+        yamls = createDataExtensionYamlsForApplicationMode(
+          this.databaseItem.language,
+          externalApiUsages,
+          modeledMethods,
+        );
+        break;
+      case Mode.Framework:
+        yamls = createDataExtensionYamlsForFrameworkMode(
+          this.databaseItem.name,
+          this.databaseItem.language,
+          externalApiUsages,
+          modeledMethods,
+        );
+        break;
+      default:
+        assertNever(this.mode);
+    }
 
     for (const [filename, yaml] of Object.entries(yamls)) {
       await outputFile(join(this.extensionPack.path, filename), yaml);
@@ -255,16 +281,21 @@ export class DataExtensionsEditorView extends AbstractWebview<
     const cancellationTokenSource = new CancellationTokenSource();
 
     try {
-      const queryResult = await runQuery({
-        cliServer: this.cliServer,
-        queryRunner: this.queryRunner,
-        databaseItem: this.databaseItem,
-        queryStorageDir: this.queryStorageDir,
-        progress: (progressUpdate: ProgressUpdate) => {
-          void this.showProgress(progressUpdate, 1500);
+      const queryResult = await runQuery(
+        this.mode === Mode.Framework
+          ? "frameworkModeQuery"
+          : "applicationModeQuery",
+        {
+          cliServer: this.cliServer,
+          queryRunner: this.queryRunner,
+          databaseItem: this.databaseItem,
+          queryStorageDir: this.queryStorageDir,
+          progress: (progressUpdate: ProgressUpdate) => {
+            void this.showProgress(progressUpdate, 1500);
+          },
+          token: cancellationTokenSource.token,
         },
-        token: cancellationTokenSource.token,
-      });
+      );
       if (!queryResult) {
         await this.clearProgress();
         return;
@@ -313,29 +344,35 @@ 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),
-      this.cliServer,
-    );
-    if (!database) {
-      await this.clearProgress();
-      void this.app.logger.log("No database chosen");
+    let addedDatabase: DatabaseItem | undefined;
+
+    // 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,
+      );
+      if (!addedDatabase) {
+        await this.clearProgress();
+        void this.app.logger.log("No database chosen");
 
-      return;
-    }
+        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,
@@ -348,7 +385,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> = {};
 
@@ -375,14 +412,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/data-extensions-editor/extension-pack-name.ts

@@ -37,7 +37,7 @@ export function autoNameExtensionPack(
   };
 }
 
-function sanitizeExtensionPackName(name: string) {
+export function sanitizeExtensionPackName(name: string) {
   // Lowercase everything
   name = name.toLowerCase();
 

extensions/ql-vscode/src/data-extensions-editor/external-api-usage-query.ts

@@ -15,6 +15,7 @@ import { QueryResultType } from "../query-server/new-messages";
 import { join } from "path";
 import { redactableError } from "../common/errors";
 import { telemetryListener } from "../common/vscode/telemetry";
+import { Query } from "./queries/query";
 
 export type RunQueryOptions = {
   cliServer: Pick<CodeQLCliServer, "resolveQlpacks">;
@@ -26,14 +27,17 @@ export type RunQueryOptions = {
   token: CancellationToken;
 };
 
-export async function runQuery({
-  cliServer,
-  queryRunner,
-  databaseItem,
-  queryStorageDir,
-  progress,
-  token,
-}: RunQueryOptions): Promise<CoreCompletedQuery | undefined> {
+export async function runQuery(
+  queryName: keyof Omit<Query, "dependencies">,
+  {
+    cliServer,
+    queryRunner,
+    databaseItem,
+    queryStorageDir,
+    progress,
+    token,
+  }: RunQueryOptions,
+): Promise<CoreCompletedQuery | undefined> {
   // The below code is temporary to allow for rapid prototyping of the queries. Once the queries are stabilized, we will
   // move these queries into the `github/codeql` repository and use them like any other contextual (e.g. AST) queries.
   // This is intentionally not pretty code, as it will be removed soon.
@@ -61,7 +65,7 @@ export async function runQuery({
 
   const queryDir = (await dir({ unsafeCleanup: true })).path;
   const queryFile = join(queryDir, "FetchExternalApis.ql");
-  await writeFile(queryFile, query.mainQuery, "utf8");
+  await writeFile(queryFile, query[queryName], "utf8");
 
   if (query.dependencies) {
     for (const [filename, contents] of Object.entries(query.dependencies)) {