Merge pull request #3020 from github/koesie10/refactor-predicates

Refactor model editor predicates

extensions/ql-vscode/src/model-editor/auto-modeler.ts

@@ -18,6 +18,7 @@ import { Mode } from "./shared/mode";
 import { CancellationTokenSource } from "vscode";
 import { ModelingStore } from "./modeling-store";
 import { ModelConfigListener } from "../config";
+import { QueryLanguage } from "../common/query-language";
 
 /**
  * The auto-modeler holds state around auto-modeling jobs and allows
@@ -36,6 +37,7 @@ export class AutoModeler {
     private readonly modelingStore: ModelingStore,
     private readonly queryStorageDir: string,
     private readonly databaseItem: DatabaseItem,
+    private readonly language: QueryLanguage,
     private readonly addModeledMethods: (
       modeledMethods: Record<string, ModeledMethod[]>,
     ) => Promise<void>,
@@ -202,7 +204,7 @@ export class AutoModeler {
       filename: "auto-model.yml",
     });
 
-    const loadedMethods = loadDataExtensionYaml(models);
+    const loadedMethods = loadDataExtensionYaml(models, this.language);
     if (!loadedMethods) {
       return;
     }

extensions/ql-vscode/src/model-editor/flow-model-queries.ts

@@ -5,7 +5,7 @@ import { QueryRunner } from "../query-server";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { showAndLogExceptionWithTelemetry } from "../common/logging";
 import { extLogger } from "../common/logging/vscode";
-import { extensiblePredicateDefinitions } from "./predicates";
+import { getModelsAsDataLanguage } from "./languages";
 import { ProgressCallback } from "../common/vscode/progress";
 import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
 import { ModeledMethod, ModeledMethodType } from "./modeled-method";
@@ -13,12 +13,14 @@ import { redactableError } from "../common/errors";
 import { telemetryListener } from "../common/vscode/telemetry";
 import { runQuery } from "../local-queries/run-query";
 import { resolveQueries } from "../local-queries";
+import { QueryLanguage } from "../common/query-language";
 
 type FlowModelOptions = {
   cliServer: CodeQLCliServer;
   queryRunner: QueryRunner;
   queryStorageDir: string;
   databaseItem: DatabaseItem;
+  language: QueryLanguage;
   progress: ProgressCallback;
   token: CancellationToken;
   onResults: (results: ModeledMethod[]) => void | Promise<void>;
@@ -104,6 +106,7 @@ async function runSingleFlowQuery(
     queryRunner,
     queryStorageDir,
     databaseItem,
+    language,
     progress,
     token,
   }: Omit<FlowModelOptions, "onResults">,
@@ -140,7 +143,9 @@ async function runSingleFlowQuery(
   }
 
   // Interpret the results
-  const definition = extensiblePredicateDefinitions[type];
+  const modelsAsDataLanguage = getModelsAsDataLanguage(language);
+
+  const definition = modelsAsDataLanguage[type];
 
   const bqrsPath = completedQuery.outputDir.bqrsPath;
 

extensions/ql-vscode/src/model-editor/languages/index.ts

@@ -0,0 +1,2 @@
+export * from "./languages";
+export * from "./models-as-data";

extensions/ql-vscode/src/model-editor/languages/languages.ts

@@ -0,0 +1,18 @@
+import { QueryLanguage } from "../../common/query-language";
+import { ModelsAsDataLanguage } from "./models-as-data";
+import { staticLanguage } from "./static";
+
+const languages: Partial<Record<QueryLanguage, ModelsAsDataLanguage>> = {
+  [QueryLanguage.CSharp]: staticLanguage,
+  [QueryLanguage.Java]: staticLanguage,
+};
+
+export function getModelsAsDataLanguage(
+  language: QueryLanguage,
+): ModelsAsDataLanguage {
+  const definition = languages[language];
+  if (!definition) {
+    throw new Error(`No models-as-data definition for ${language}`);
+  }
+  return definition;
+}

extensions/ql-vscode/src/model-editor/languages/models-as-data.ts

@@ -0,0 +1,19 @@
+import { ModeledMethod, ModeledMethodType } from "../modeled-method";
+import { DataTuple } from "../model-extension-file";
+
+type GenerateMethodDefinition = (method: ModeledMethod) => DataTuple[];
+type ReadModeledMethod = (row: DataTuple[]) => ModeledMethod;
+
+export type ModelsAsDataLanguageModelType = Exclude<ModeledMethodType, "none">;
+
+export type ModelsAsDataLanguageModel = {
+  extensiblePredicate: string;
+  supportedKinds: string[];
+  generateMethodDefinition: GenerateMethodDefinition;
+  readModeledMethod: ReadModeledMethod;
+};
+
+export type ModelsAsDataLanguage = Record<
+  ModelsAsDataLanguageModelType,
+  ModelsAsDataLanguageModel
+>;

extensions/ql-vscode/src/model-editor/languages/shared.ts

@@ -0,0 +1,25 @@
+export const sharedExtensiblePredicates = {
+  source: "sourceModel",
+  sink: "sinkModel",
+  summary: "summaryModel",
+  neutral: "neutralModel",
+};
+
+export const sharedKinds = {
+  source: ["local", "remote"],
+  sink: [
+    "code-injection",
+    "command-injection",
+    "file-content-store",
+    "html-injection",
+    "js-injection",
+    "ldap-injection",
+    "log-injection",
+    "path-injection",
+    "request-forgery",
+    "sql-injection",
+    "url-redirection",
+  ],
+  summary: ["taint", "value"],
+  neutral: ["summary", "source", "sink"],
+};

extensions/ql-vscode/src/model-editor/languages/static.ts

@@ -1,24 +1,16 @@
-import { ModeledMethod, ModeledMethodType, Provenance } from "./modeled-method";
-import { DataTuple } from "./model-extension-file";
-
-export type ExtensiblePredicateDefinition = {
-  extensiblePredicate: string;
-  generateMethodDefinition: (method: ModeledMethod) => DataTuple[];
-  readModeledMethod: (row: DataTuple[]) => ModeledMethod;
-
-  supportedKinds?: string[];
-};
+import { ModelsAsDataLanguage } from "./models-as-data";
+import { ModeledMethodType, Provenance } from "../modeled-method";
+import { DataTuple } from "../model-extension-file";
+import { sharedExtensiblePredicates, sharedKinds } from "./shared";
 
 function readRowToMethod(row: DataTuple[]): string {
   return `${row[0]}.${row[1]}#${row[3]}${row[4]}`;
 }
 
-export const extensiblePredicateDefinitions: Record<
-  Exclude<ModeledMethodType, "none">,
-  ExtensiblePredicateDefinition
-> = {
+export const staticLanguage: ModelsAsDataLanguage = {
   source: {
-    extensiblePredicate: "sourceModel",
+    extensiblePredicate: sharedExtensiblePredicates.source,
+    supportedKinds: sharedKinds.source,
     // extensible predicate sourceModel(
     //   string package, string type, boolean subtypes, string name, string signature, string ext,
     //   string output, string kind, string provenance
@@ -35,7 +27,7 @@ export const extensiblePredicateDefinitions: Record<
       method.provenance,
     ],
     readModeledMethod: (row) => ({
-      type: "source",
+      type: "source" as ModeledMethodType,
       input: "",
       output: row[6] as string,
       kind: row[7] as string,
@@ -46,10 +38,10 @@ export const extensiblePredicateDefinitions: Record<
       methodName: row[3] as string,
       methodParameters: row[4] as string,
     }),
-    supportedKinds: ["local", "remote"],
   },
   sink: {
-    extensiblePredicate: "sinkModel",
+    extensiblePredicate: sharedExtensiblePredicates.sink,
+    supportedKinds: sharedKinds.sink,
     // extensible predicate sinkModel(
     //   string package, string type, boolean subtypes, string name, string signature, string ext,
     //   string input, string kind, string provenance
@@ -77,22 +69,10 @@ export const extensiblePredicateDefinitions: Record<
       methodName: row[3] as string,
       methodParameters: row[4] as string,
     }),
-    supportedKinds: [
-      "code-injection",
-      "command-injection",
-      "file-content-store",
-      "html-injection",
-      "js-injection",
-      "ldap-injection",
-      "log-injection",
-      "path-injection",
-      "request-forgery",
-      "sql-injection",
-      "url-redirection",
-    ],
   },
   summary: {
-    extensiblePredicate: "summaryModel",
+    extensiblePredicate: sharedExtensiblePredicates.summary,
+    supportedKinds: sharedKinds.summary,
     // extensible predicate summaryModel(
     //   string package, string type, boolean subtypes, string name, string signature, string ext,
     //   string input, string output, string kind, string provenance
@@ -121,10 +101,10 @@ export const extensiblePredicateDefinitions: Record<
       methodName: row[3] as string,
       methodParameters: row[4] as string,
     }),
-    supportedKinds: ["taint", "value"],
   },
   neutral: {
-    extensiblePredicate: "neutralModel",
+    extensiblePredicate: sharedExtensiblePredicates.neutral,
+    supportedKinds: sharedKinds.neutral,
     // extensible predicate neutralModel(
     //   string package, string type, string name, string signature, string kind, string provenance
     // );
@@ -148,6 +128,5 @@ export const extensiblePredicateDefinitions: Record<
       methodName: row[2] as string,
       methodParameters: row[3] as string,
     }),
-    supportedKinds: ["summary", "source", "sink"],
   },
 };

extensions/ql-vscode/src/model-editor/method-modeling/method-modeling-view-provider.ts

@@ -15,6 +15,10 @@ import { ModelEditorViewTracker } from "../model-editor-view-tracker";
 import { ModelConfigListener } from "../../config";
 import { DatabaseItem } from "../../databases/local-databases";
 import { ModelingEvents } from "../modeling-events";
+import {
+  QueryLanguage,
+  tryGetQueryLanguage,
+} from "../../common/query-language";
 
 export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
   ToMethodModelingMessage,
@@ -24,6 +28,7 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
 
   private method: Method | undefined = undefined;
   private databaseItem: DatabaseItem | undefined = undefined;
+  private language: QueryLanguage | undefined = undefined;
 
   constructor(
     app: App,
@@ -45,6 +50,7 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
     await this.postMessage({
       t: "setMethodModelingPanelViewState",
       viewState: {
+        language: this.language,
         showMultipleModels: this.modelConfig.showMultipleModels,
       },
     });
@@ -56,6 +62,7 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
   ): Promise<void> {
     this.method = method;
     this.databaseItem = databaseItem;
+    this.language = databaseItem && tryGetQueryLanguage(databaseItem.language);
 
     if (this.isShowingView) {
       await this.postMessage({
@@ -70,6 +77,9 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
       const selectedMethod = this.modelingStore.getSelectedMethodDetails();
       if (selectedMethod) {
         this.databaseItem = selectedMethod.databaseItem;
+        this.language = tryGetQueryLanguage(
+          selectedMethod.databaseItem.language,
+        );
         this.method = selectedMethod.method;
 
         await this.postMessage({
@@ -185,6 +195,7 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
         if (this.webviewView) {
           this.method = e.method;
           this.databaseItem = e.databaseItem;
+          this.language = tryGetQueryLanguage(e.databaseItem.language);
 
           await this.postMessage({
             t: "setSelectedMethod",

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

@@ -233,6 +233,7 @@ export class ModelEditorModule extends DisposableObject {
             queryDir,
             db,
             modelFile,
+            language,
           );
 
           this.modelingEvents.onDbClosed(async (dbUri) => {

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

@@ -38,7 +38,10 @@ import { ModelConfigListener } from "../config";
 import { INITIAL_MODE, Mode } from "./shared/mode";
 import { loadModeledMethods, saveModeledMethods } from "./modeled-method-fs";
 import { pickExtensionPack } from "./extension-pack-picker";
-import { getLanguageDisplayName } from "../common/query-language";
+import {
+  getLanguageDisplayName,
+  QueryLanguage,
+} from "../common/query-language";
 import { AutoModeler } from "./auto-modeler";
 import { telemetryListener } from "../common/vscode/telemetry";
 import { ModelingStore } from "./modeling-store";
@@ -64,6 +67,8 @@ export class ModelEditorView extends AbstractWebview<
     private readonly queryDir: string,
     private readonly databaseItem: DatabaseItem,
     private readonly extensionPack: ExtensionPack,
+    // The language is equal to databaseItem.language but is properly typed as QueryLanguage
+    private readonly language: QueryLanguage,
     initialMode: Mode = INITIAL_MODE,
   ) {
     super(app);
@@ -82,6 +87,7 @@ export class ModelEditorView extends AbstractWebview<
       modelingStore,
       queryStorageDir,
       databaseItem,
+      language,
       async (modeledMethods) => {
         this.addModeledMethods(modeledMethods);
       },
@@ -218,7 +224,7 @@ export class ModelEditorView extends AbstractWebview<
               });
               await saveModeledMethods(
                 this.extensionPack,
-                this.databaseItem.language,
+                this.language,
                 methods,
                 modeledMethods,
                 mode,
@@ -367,6 +373,7 @@ export class ModelEditorView extends AbstractWebview<
       t: "setModelEditorViewState",
       viewState: {
         extensionPack: this.extensionPack,
+        language: this.language,
         showFlowGeneration: this.modelConfig.flowGeneration,
         showLlmButton,
         showMultipleModels: this.modelConfig.showMultipleModels,
@@ -394,6 +401,7 @@ export class ModelEditorView extends AbstractWebview<
     try {
       const modeledMethods = await loadModeledMethods(
         this.extensionPack,
+        this.language,
         this.cliServer,
         this.app.logger,
       );
@@ -458,6 +466,14 @@ export class ModelEditorView extends AbstractWebview<
           if (!addedDatabase) {
             return;
           }
+
+          if (addedDatabase.language !== this.language) {
+            void showAndLogErrorMessage(
+              this.app.logger,
+              `The selected database is for ${addedDatabase.language}, but the current database is for ${this.language}.`,
+            );
+            return;
+          }
         }
 
         progress({
@@ -472,6 +488,7 @@ export class ModelEditorView extends AbstractWebview<
             queryRunner: this.queryRunner,
             queryStorageDir: this.queryStorageDir,
             databaseItem: addedDatabase ?? this.databaseItem,
+            language: this.language,
             onResults: async (modeledMethods) => {
               const modeledMethodsByName: Record<string, ModeledMethod[]> = {};
 
@@ -579,6 +596,7 @@ export class ModelEditorView extends AbstractWebview<
         this.queryDir,
         addedDatabase,
         modelFile,
+        this.language,
         Mode.Framework,
       );
       await view.openView();

extensions/ql-vscode/src/model-editor/modeled-method-fs.ts

@@ -10,10 +10,11 @@ import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
 import { load as loadYaml } from "js-yaml";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { pathsEqual } from "../common/files";
+import { QueryLanguage } from "../common/query-language";
 
 export async function saveModeledMethods(
   extensionPack: ExtensionPack,
-  language: string,
+  language: QueryLanguage,
   methods: readonly Method[],
   modeledMethods: Readonly<Record<string, readonly ModeledMethod[]>>,
   mode: Mode,
@@ -22,6 +23,7 @@ export async function saveModeledMethods(
 ): Promise<void> {
   const existingModeledMethods = await loadModeledMethodFiles(
     extensionPack,
+    language,
     cliServer,
     logger,
   );
@@ -43,6 +45,7 @@ export async function saveModeledMethods(
 
 async function loadModeledMethodFiles(
   extensionPack: ExtensionPack,
+  language: QueryLanguage,
   cliServer: CodeQLCliServer,
   logger: NotificationLogger,
 ): Promise<Record<string, Record<string, ModeledMethod[]>>> {
@@ -60,7 +63,7 @@ async function loadModeledMethodFiles(
       filename: modelFile,
     });
 
-    const modeledMethods = loadDataExtensionYaml(data);
+    const modeledMethods = loadDataExtensionYaml(data, language);
     if (!modeledMethods) {
       void showAndLogErrorMessage(
         logger,
@@ -76,6 +79,7 @@ async function loadModeledMethodFiles(
 
 export async function loadModeledMethods(
   extensionPack: ExtensionPack,
+  language: QueryLanguage,
   cliServer: CodeQLCliServer,
   logger: NotificationLogger,
 ): Promise<Record<string, ModeledMethod[]>> {
@@ -83,6 +87,7 @@ export async function loadModeledMethods(
 
   const modeledMethodsByFile = await loadModeledMethodFiles(
     extensionPack,
+    language,
     cliServer,
     logger,
   );

extensions/ql-vscode/src/model-editor/shared/view-state.ts

@@ -1,8 +1,10 @@
 import { ExtensionPack } from "./extension-pack";
 import { Mode } from "./mode";
+import { QueryLanguage } from "../../common/query-language";
 
 export interface ModelEditorViewState {
   extensionPack: ExtensionPack;
+  language: QueryLanguage;
   showFlowGeneration: boolean;
   showLlmButton: boolean;
   showMultipleModels: boolean;
@@ -11,5 +13,6 @@ export interface ModelEditorViewState {
 }
 
 export interface MethodModelingPanelViewState {
+  language: QueryLanguage | undefined;
   showMultipleModels: boolean;
 }

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

@@ -3,21 +3,22 @@ import Ajv from "ajv";
 import { Method } from "./method";
 import { ModeledMethod, ModeledMethodType } from "./modeled-method";
 import {
-  ExtensiblePredicateDefinition,
-  extensiblePredicateDefinitions,
-} from "./predicates";
+  getModelsAsDataLanguage,
+  ModelsAsDataLanguageModel,
+} from "./languages";
 
 import * as modelExtensionFileSchema from "./model-extension-file.schema.json";
 import { Mode } from "./shared/mode";
 import { assertNever } from "../common/helpers-pure";
 import { ModelExtensionFile } from "./model-extension-file";
+import { QueryLanguage } from "../common/query-language";
 
 const ajv = new Ajv({ allErrors: true, allowUnionTypes: true });
 const modelExtensionFileSchemaValidate = ajv.compile(modelExtensionFileSchema);
 
 function createDataProperty(
   methods: readonly ModeledMethod[],
-  definition: ExtensiblePredicateDefinition,
+  definition: ModelsAsDataLanguageModel,
 ) {
   if (methods.length === 0) {
     return " []";
@@ -34,9 +35,11 @@ function createDataProperty(
 }
 
 export function createDataExtensionYaml(
-  language: string,
+  language: QueryLanguage,
   modeledMethods: readonly ModeledMethod[],
 ) {
+  const modelsAsDataLanguage = getModelsAsDataLanguage(language);
+
   const methodsByType: Record<
     Exclude<ModeledMethodType, "none">,
     ModeledMethod[]
@@ -53,7 +56,7 @@ export function createDataExtensionYaml(
     }
   }
 
-  const extensions = Object.entries(extensiblePredicateDefinitions).map(
+  const extensions = Object.entries(modelsAsDataLanguage).map(
     ([type, definition]) => `  - addsTo:
       pack: codeql/${language}-all
       extensible: ${definition.extensiblePredicate}
@@ -69,7 +72,7 @@ ${extensions.join("\n")}`;
 }
 
 export function createDataExtensionYamls(
-  language: string,
+  language: QueryLanguage,
   methods: readonly Method[],
   newModeledMethods: Readonly<Record<string, readonly ModeledMethod[]>>,
   existingModeledMethods: Readonly<
@@ -98,7 +101,7 @@ export function createDataExtensionYamls(
 }
 
 function createDataExtensionYamlsByGrouping(
-  language: string,
+  language: QueryLanguage,
   methods: readonly Method[],
   newModeledMethods: Readonly<Record<string, readonly ModeledMethod[]>>,
   existingModeledMethods: Readonly<
@@ -153,7 +156,7 @@ function createDataExtensionYamlsByGrouping(
 }
 
 export function createDataExtensionYamlsForApplicationMode(
-  language: string,
+  language: QueryLanguage,
   methods: readonly Method[],
   newModeledMethods: Readonly<Record<string, readonly ModeledMethod[]>>,
   existingModeledMethods: Readonly<
@@ -170,7 +173,7 @@ export function createDataExtensionYamlsForApplicationMode(
 }
 
 export function createDataExtensionYamlsForFrameworkMode(
-  language: string,
+  language: QueryLanguage,
   methods: readonly Method[],
   newModeledMethods: Readonly<Record<string, readonly ModeledMethod[]>>,
   existingModeledMethods: Readonly<
@@ -240,11 +243,14 @@ function validateModelExtensionFile(data: unknown): data is ModelExtensionFile {
 
 export function loadDataExtensionYaml(
   data: unknown,
+  language: QueryLanguage,
 ): Record<string, ModeledMethod[]> | undefined {
   if (!validateModelExtensionFile(data)) {
     return undefined;
   }
 
+  const modelsAsDataLanguage = getModelsAsDataLanguage(language);
+
   const extensions = data.extensions;
 
   const modeledMethods: Record<string, ModeledMethod[]> = {};
@@ -254,7 +260,7 @@ export function loadDataExtensionYaml(
     const extensible = addsTo.extensible;
     const data = extension.data;
 
-    const definition = Object.values(extensiblePredicateDefinitions).find(
+    const definition = Object.values(modelsAsDataLanguage).find(
       (definition) => definition.extensiblePredicate === extensible,
     );
     if (!definition) {

extensions/ql-vscode/src/stories/model-editor/LibraryRow.stories.tsx

@@ -2,10 +2,9 @@ import * as React from "react";
 
 import { Meta, StoryFn } from "@storybook/react";
 
-import { Mode } from "../../model-editor/shared/mode";
 import { LibraryRow as LibraryRowComponent } from "../../view/model-editor/LibraryRow";
 import { CallClassification } from "../../model-editor/method";
-import { createMockExtensionPack } from "../../../test/factories/model-editor/extension-pack";
+import { createMockModelEditorViewState } from "../../../test/factories/model-editor/view-state";
 
 export default {
   title: "CodeQL Model Editor/Library Row",
@@ -219,13 +218,10 @@ LibraryRow.args = {
   },
   modifiedSignatures: new Set(["org.sql2o.Sql2o#Sql2o(String)"]),
   inProgressMethods: new Set(),
-  viewState: {
-    extensionPack: createMockExtensionPack(),
+  viewState: createMockModelEditorViewState({
     showFlowGeneration: true,
     showLlmButton: true,
     showMultipleModels: true,
-    mode: Mode.Application,
-    sourceArchiveAvailable: true,
-  },
+  }),
   hideModeledMethods: false,
 };

extensions/ql-vscode/src/stories/model-editor/MethodRow.stories.tsx

@@ -10,10 +10,8 @@ import {
   MULTIPLE_MODELS_GRID_TEMPLATE_COLUMNS,
   SINGLE_MODEL_GRID_TEMPLATE_COLUMNS,
 } from "../../view/model-editor/ModeledMethodDataGrid";
-import { ModelEditorViewState } from "../../model-editor/shared/view-state";
-import { createMockExtensionPack } from "../../../test/factories/model-editor/extension-pack";
-import { Mode } from "../../model-editor/shared/mode";
 import { DataGrid } from "../../view/common/DataGrid";
+import { createMockModelEditorViewState } from "../../../test/factories/model-editor/view-state";
 
 export default {
   title: "CodeQL Model Editor/Method Row",
@@ -99,14 +97,11 @@ const modeledMethod: ModeledMethod = {
   methodParameters: "()",
 };
 
-const viewState: ModelEditorViewState = {
-  extensionPack: createMockExtensionPack(),
+const viewState = createMockModelEditorViewState({
   showFlowGeneration: true,
   showLlmButton: true,
   showMultipleModels: true,
-  mode: Mode.Application,
-  sourceArchiveAvailable: true,
-};
+});
 
 export const Unmodeled = Template.bind({});
 Unmodeled.args = {

extensions/ql-vscode/src/stories/model-editor/ModelEditor.stories.tsx

@@ -2,9 +2,9 @@ import * as React from "react";
 
 import { Meta, StoryFn } from "@storybook/react";
 
-import { Mode } from "../../model-editor/shared/mode";
 import { ModelEditor as ModelEditorComponent } from "../../view/model-editor/ModelEditor";
 import { CallClassification } from "../../model-editor/method";
+import { createMockModelEditorViewState } from "../../../test/factories/model-editor/view-state";
 
 export default {
   title: "CodeQL Model Editor/CodeQL Model Editor",
@@ -17,7 +17,7 @@ const Template: StoryFn<typeof ModelEditorComponent> = (args) => (
 
 export const ModelEditor = Template.bind({});
 ModelEditor.args = {
-  initialViewState: {
+  initialViewState: createMockModelEditorViewState({
     extensionPack: {
       path: "/home/user/vscode-codeql-starter/codeql-custom-queries-java/sql2o",
       yamlPath:
@@ -31,9 +31,7 @@ ModelEditor.args = {
     showFlowGeneration: true,
     showLlmButton: true,
     showMultipleModels: true,
-    mode: Mode.Application,
-    sourceArchiveAvailable: true,
-  },
+  }),
   initialMethods: [
     {
       library: "sql2o",

extensions/ql-vscode/src/view/method-modeling/MethodModeling.tsx

@@ -8,6 +8,7 @@ import { ModeledMethod } from "../../model-editor/modeled-method";
 import { VSCodeTag } from "@vscode/webview-ui-toolkit/react";
 import { ReviewInEditorButton } from "./ReviewInEditorButton";
 import { ModeledMethodsPanel } from "./ModeledMethodsPanel";
+import { QueryLanguage } from "../../common/query-language";
 
 const Container = styled.div`
   padding-top: 0.5rem;
@@ -49,6 +50,7 @@ const UnsavedTag = ({ modelingStatus }: { modelingStatus: ModelingStatus }) => (
 );
 
 export type MethodModelingProps = {
+  language: QueryLanguage;
   modelingStatus: ModelingStatus;
   method: Method;
   modeledMethods: ModeledMethod[];
@@ -58,6 +60,7 @@ export type MethodModelingProps = {
 };
 
 export const MethodModeling = ({
+  language,
   modelingStatus,
   modeledMethods,
   method,
@@ -77,6 +80,7 @@ export const MethodModeling = ({
         <MethodName {...method} />
       </DependencyContainer>
       <ModeledMethodsPanel
+        language={language}
         method={method}
         modeledMethods={modeledMethods}
         showMultipleModels={showMultipleModels}

extensions/ql-vscode/src/view/method-modeling/MethodModelingInputs.tsx

@@ -7,6 +7,7 @@ import { ModelInputDropdown } from "../model-editor/ModelInputDropdown";
 import { ModelOutputDropdown } from "../model-editor/ModelOutputDropdown";
 import { ModelKindDropdown } from "../model-editor/ModelKindDropdown";
 import { InProgressDropdown } from "../model-editor/InProgressDropdown";
+import { QueryLanguage } from "../../common/query-language";
 
 const Container = styled.div`
   padding-top: 0.5rem;
@@ -23,6 +24,7 @@ const Name = styled.span`
 `;
 
 export type MethodModelingInputsProps = {
+  language: QueryLanguage;
   method: Method;
   modeledMethod: ModeledMethod | undefined;
   isModelingInProgress: boolean;
@@ -30,6 +32,7 @@ export type MethodModelingInputsProps = {
 };
 
 export const MethodModelingInputs = ({
+  language,
   method,
   modeledMethod,
   isModelingInProgress,
@@ -79,7 +82,7 @@ export const MethodModelingInputs = ({
           {isModelingInProgress ? (
             <InProgressDropdown />
           ) : (
-            <ModelKindDropdown {...inputProps} />
+            <ModelKindDropdown language={language} {...inputProps} />
           )}
         </Input>
       </Container>

extensions/ql-vscode/src/view/method-modeling/MethodModelingView.tsx

@@ -80,7 +80,7 @@ export function MethodModelingView({ initialViewState }: Props): JSX.Element {
     };
   }, []);
 
-  if (!inModelingMode) {
+  if (!inModelingMode || !viewState?.language) {
     return <NotInModelingMode />;
   }
 
@@ -105,6 +105,7 @@ export function MethodModelingView({ initialViewState }: Props): JSX.Element {
 
   return (
     <MethodModeling
+      language={viewState?.language}
       modelingStatus={modelingStatus}
       method={method}
       modeledMethods={modeledMethods}

extensions/ql-vscode/src/view/method-modeling/ModeledMethodsPanel.tsx

@@ -6,8 +6,10 @@ import { Method } from "../../model-editor/method";
 import { styled } from "styled-components";
 import { MultipleModeledMethodsPanel } from "./MultipleModeledMethodsPanel";
 import { convertToLegacyModeledMethod } from "../../model-editor/shared/modeled-methods-legacy";
+import { QueryLanguage } from "../../common/query-language";
 
 export type ModeledMethodsPanelProps = {
+  language: QueryLanguage;
   method: Method;
   modeledMethods: ModeledMethod[];
   isModelingInProgress: boolean;
@@ -20,6 +22,7 @@ const SingleMethodModelingInputs = styled(MethodModelingInputs)`
 `;
 
 export const ModeledMethodsPanel = ({
+  language,
   method,
   modeledMethods,
   isModelingInProgress,
@@ -36,6 +39,7 @@ export const ModeledMethodsPanel = ({
   if (!showMultipleModels) {
     return (
       <SingleMethodModelingInputs
+        language={language}
         method={method}
         modeledMethod={convertToLegacyModeledMethod(modeledMethods)}
         isModelingInProgress={isModelingInProgress}
@@ -46,6 +50,7 @@ export const ModeledMethodsPanel = ({
 
   return (
     <MultipleModeledMethodsPanel
+      language={language}
       method={method}
       modeledMethods={modeledMethods}
       isModelingInProgress={isModelingInProgress}

extensions/ql-vscode/src/view/method-modeling/MultipleModeledMethodsPanel.tsx

@@ -12,8 +12,10 @@ import { VSCodeButton } from "@vscode/webview-ui-toolkit/react";
 import { Codicon } from "../common";
 import { validateModeledMethods } from "../../model-editor/shared/validation";
 import { ModeledMethodAlert } from "./ModeledMethodAlert";
+import { QueryLanguage } from "../../common/query-language";
 
 export type MultipleModeledMethodsPanelProps = {
+  language: QueryLanguage;
   method: Method;
   modeledMethods: ModeledMethod[];
   isModelingInProgress: boolean;
@@ -53,6 +55,7 @@ const ModificationActions = styled.div`
 `;
 
 export const MultipleModeledMethodsPanel = ({
+  language,
   method,
   modeledMethods,
   isModelingInProgress,
@@ -150,6 +153,7 @@ export const MultipleModeledMethodsPanel = ({
       )}
       {modeledMethods.length > 0 ? (
         <MethodModelingInputs
+          language={language}
           method={method}
           modeledMethod={modeledMethods[selectedIndex]}
           isModelingInProgress={isModelingInProgress}
@@ -157,6 +161,7 @@ export const MultipleModeledMethodsPanel = ({
         />
       ) : (
         <MethodModelingInputs
+          language={language}
           method={method}
           modeledMethod={undefined}
           isModelingInProgress={isModelingInProgress}

extensions/ql-vscode/src/view/method-modeling/__tests__/MethodModeling.spec.tsx

@@ -3,6 +3,7 @@ import { render as reactRender, screen } from "@testing-library/react";
 import { MethodModeling, MethodModelingProps } from "../MethodModeling";
 import { createMethod } from "../../../../test/factories/model-editor/method-factories";
 import { createModeledMethod } from "../../../../test/factories/model-editor/modeled-method-factories";
+import { QueryLanguage } from "../../../common/query-language";
 
 describe(MethodModeling.name, () => {
   const render = (props: MethodModelingProps) =>
@@ -15,6 +16,7 @@ describe(MethodModeling.name, () => {
     const onChange = jest.fn();
 
     render({
+      language: QueryLanguage.Java,
       modelingStatus: "saved",
       method,
       modeledMethods: [modeledMethod],

extensions/ql-vscode/src/view/method-modeling/__tests__/MethodModelingInputs.spec.tsx

@@ -7,11 +7,13 @@ import {
 } from "../MethodModelingInputs";
 import { createMethod } from "../../../../test/factories/model-editor/method-factories";
 import { createModeledMethod } from "../../../../test/factories/model-editor/modeled-method-factories";
+import { QueryLanguage } from "../../../common/query-language";
 
 describe(MethodModelingInputs.name, () => {
   const render = (props: MethodModelingInputsProps) =>
     reactRender(<MethodModelingInputs {...props} />);
 
+  const language = QueryLanguage.Java;
   const method = createMethod();
   const modeledMethod = createModeledMethod();
   const isModelingInProgress = false;
@@ -19,6 +21,7 @@ describe(MethodModelingInputs.name, () => {
 
   it("renders the method modeling inputs", () => {
     render({
+      language,
       method,
       modeledMethod,
       isModelingInProgress,
@@ -44,6 +47,7 @@ describe(MethodModelingInputs.name, () => {
 
   it("allows changing the type", async () => {
     render({
+      language,
       method,
       modeledMethod,
       isModelingInProgress,
@@ -65,6 +69,7 @@ describe(MethodModelingInputs.name, () => {
 
   it("sets other dropdowns when model type is changed", () => {
     const { rerender } = render({
+      language,
       method,
       modeledMethod,
       isModelingInProgress,
@@ -77,6 +82,7 @@ describe(MethodModelingInputs.name, () => {
 
     rerender(
       <MethodModelingInputs
+        language={language}
         method={method}
         modeledMethod={updatedModeledMethod}
         isModelingInProgress={isModelingInProgress}
@@ -105,6 +111,7 @@ describe(MethodModelingInputs.name, () => {
 
   it("sets in progress dropdowns when modeling is in progress", () => {
     render({
+      language,
       method,
       modeledMethod,
       isModelingInProgress: true,

extensions/ql-vscode/src/view/method-modeling/__tests__/ModeledMethodsPanel.spec.tsx

@@ -6,11 +6,13 @@ import {
   ModeledMethodsPanel,
   ModeledMethodsPanelProps,
 } from "../ModeledMethodsPanel";
+import { QueryLanguage } from "../../../common/query-language";
 
 describe(ModeledMethodsPanel.name, () => {
   const render = (props: ModeledMethodsPanelProps) =>
     reactRender(<ModeledMethodsPanel {...props} />);
 
+  const language = QueryLanguage.Java;
   const method = createMethod();
   const modeledMethods = [createModeledMethod(), createModeledMethod()];
   const isModelingInProgress = false;
@@ -21,6 +23,7 @@ describe(ModeledMethodsPanel.name, () => {
 
     it("renders the method modeling inputs", () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -33,6 +36,7 @@ describe(ModeledMethodsPanel.name, () => {
 
     it("does not render the pagination", () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -52,6 +56,7 @@ describe(ModeledMethodsPanel.name, () => {
 
     it("renders the method modeling inputs once", () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -64,6 +69,7 @@ describe(ModeledMethodsPanel.name, () => {
 
     it("renders the pagination", () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,

extensions/ql-vscode/src/view/method-modeling/__tests__/MultipleModeledMethodsPanel.spec.tsx

@@ -8,11 +8,13 @@ import {
 } from "../MultipleModeledMethodsPanel";
 import userEvent from "@testing-library/user-event";
 import { ModeledMethod } from "../../../model-editor/modeled-method";
+import { QueryLanguage } from "../../../common/query-language";
 
 describe(MultipleModeledMethodsPanel.name, () => {
   const render = (props: MultipleModeledMethodsPanelProps) =>
     reactRender(<MultipleModeledMethodsPanel {...props} />);
 
+  const language = QueryLanguage.Java;
   const method = createMethod();
   const isModelingInProgress = false;
   const onChange = jest.fn<void, [string, ModeledMethod[]]>();
@@ -22,6 +24,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("renders the method modeling inputs once", () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -38,6 +41,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("disables all pagination", () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -58,6 +62,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("cannot add or delete modeling", () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -88,6 +93,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("renders the method modeling inputs once", () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -104,6 +110,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("disables all pagination", () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -123,6 +130,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("cannot delete modeling", () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -138,6 +146,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("can add modeling", async () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -165,6 +174,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("changes selection to the newly added modeling", async () => {
       const { rerender } = render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -175,6 +185,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
       rerender(
         <MultipleModeledMethodsPanel
+          language={language}
           method={method}
           modeledMethods={
             onChange.mock.calls[onChange.mock.calls.length - 1][1]
@@ -208,6 +219,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("renders the method modeling inputs once", () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -224,6 +236,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("renders the pagination", () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -237,6 +250,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("disables the correct pagination", async () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -255,6 +269,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("can use the pagination", async () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -290,6 +305,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("correctly updates selected pagination index when the number of models decreases", async () => {
       const { rerender } = render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -300,6 +316,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
       rerender(
         <MultipleModeledMethodsPanel
+          language={language}
           method={method}
           modeledMethods={[modeledMethods[1]]}
           isModelingInProgress={isModelingInProgress}
@@ -317,6 +334,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("does not show errors", () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -328,6 +346,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("can update the first modeling", async () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -359,6 +378,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("can update the second modeling", async () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -392,6 +412,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("can delete modeling", async () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -408,6 +429,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("can add modeling", async () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -435,6 +457,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("shows an error when adding a neutral modeling", async () => {
       const { rerender } = render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -445,6 +468,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
       rerender(
         <MultipleModeledMethodsPanel
+          language={language}
           method={method}
           modeledMethods={
             onChange.mock.calls[onChange.mock.calls.length - 1][1]
@@ -464,6 +488,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
       rerender(
         <MultipleModeledMethodsPanel
+          language={language}
           method={method}
           modeledMethods={
             onChange.mock.calls[onChange.mock.calls.length - 1][1]
@@ -481,6 +506,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
       rerender(
         <MultipleModeledMethodsPanel
+          language={language}
           method={method}
           modeledMethods={
             onChange.mock.calls[onChange.mock.calls.length - 1][1]
@@ -498,6 +524,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("changes selection to the newly added modeling", async () => {
       const { rerender } = render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -510,6 +537,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
       rerender(
         <MultipleModeledMethodsPanel
+          language={language}
           method={method}
           modeledMethods={
             onChange.mock.calls[onChange.mock.calls.length - 1][1]
@@ -550,6 +578,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("can use the pagination", async () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -639,6 +668,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("preserves selection when a modeling other than the selected modeling is removed", async () => {
       const { rerender } = render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -649,6 +679,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
       rerender(
         <MultipleModeledMethodsPanel
+          language={language}
           method={method}
           modeledMethods={modeledMethods.slice(0, 2)}
           isModelingInProgress={isModelingInProgress}
@@ -661,6 +692,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("reduces selection when the selected modeling is removed", async () => {
       const { rerender } = render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -673,6 +705,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
       rerender(
         <MultipleModeledMethodsPanel
+          language={language}
           method={method}
           modeledMethods={modeledMethods.slice(0, 2)}
           isModelingInProgress={isModelingInProgress}
@@ -704,6 +737,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("can add modeling", () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -717,6 +751,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("can delete first modeling", async () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -733,6 +768,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("can delete second modeling", async () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -750,6 +786,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("can add modeling after deleting second modeling", async () => {
       const { rerender } = render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -766,6 +803,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
       rerender(
         <MultipleModeledMethodsPanel
+          language={language}
           method={method}
           modeledMethods={modeledMethods.slice(0, 1)}
           isModelingInProgress={isModelingInProgress}
@@ -806,6 +844,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("shows errors", () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,
@@ -817,6 +856,7 @@ describe(MultipleModeledMethodsPanel.name, () => {
 
     it("shows the correct error message", async () => {
       render({
+        language,
         method,
         modeledMethods,
         isModelingInProgress,

extensions/ql-vscode/src/view/model-editor/MethodRow.tsx

@@ -262,6 +262,7 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
                 </DataGridCell>
                 <DataGridCell>
                   <ModelKindDropdown
+                    language={viewState.language}
                     method={method}
                     modeledMethod={modeledMethod}
                     onChange={modeledMethodChangedHandlers[index]}

extensions/ql-vscode/src/view/model-editor/ModelKindDropdown.tsx

@@ -6,24 +6,29 @@ import type {
 } from "../../model-editor/modeled-method";
 import { Dropdown } from "../common/Dropdown";
 import { Method } from "../../model-editor/method";
-import { extensiblePredicateDefinitions } from "../../model-editor/predicates";
+import { getModelsAsDataLanguage } from "../../model-editor/languages";
+import { QueryLanguage } from "../../common/query-language";
 
 type Props = {
+  language: QueryLanguage;
   method: Method;
   modeledMethod: ModeledMethod | undefined;
   onChange: (modeledMethod: ModeledMethod) => void;
 };
 
 export const ModelKindDropdown = ({
+  language,
   method,
   modeledMethod,
   onChange,
 }: Props) => {
   const predicate = useMemo(() => {
+    const modelsAsDataLanguage = getModelsAsDataLanguage(language);
+
     return modeledMethod?.type && modeledMethod.type !== "none"
-      ? extensiblePredicateDefinitions[modeledMethod.type]
+      ? modelsAsDataLanguage[modeledMethod.type]
       : undefined;
-  }, [modeledMethod?.type]);
+  }, [language, modeledMethod?.type]);
 
   const kinds = useMemo(() => predicate?.supportedKinds || [], [predicate]);
 

extensions/ql-vscode/src/view/model-editor/__tests__/HiddenMethodsRow.spec.tsx

@@ -1,19 +1,10 @@
 import * as React from "react";
 import { render, screen } from "@testing-library/react";
 import { HiddenMethodsRow } from "../HiddenMethodsRow";
-import { createMockExtensionPack } from "../../../../test/factories/model-editor/extension-pack";
-import { ModelEditorViewState } from "../../../model-editor/shared/view-state";
-import { Mode } from "../../../model-editor/shared/mode";
+import { createMockModelEditorViewState } from "../../../../test/factories/model-editor/view-state";
 
 describe(HiddenMethodsRow.name, () => {
-  const viewState: ModelEditorViewState = {
-    mode: Mode.Application,
-    showFlowGeneration: false,
-    showLlmButton: false,
-    showMultipleModels: false,
-    extensionPack: createMockExtensionPack(),
-    sourceArchiveAvailable: true,
-  };
+  const viewState = createMockModelEditorViewState();
 
   it("does not render with 0 hidden methods", () => {
     const { container } = render(

extensions/ql-vscode/src/view/model-editor/__tests__/LibraryRow.spec.tsx

@@ -2,10 +2,8 @@ import * as React from "react";
 import { render as reactRender, screen } from "@testing-library/react";
 import { createMethod } from "../../../../test/factories/model-editor/method-factories";
 import { LibraryRow, LibraryRowProps } from "../LibraryRow";
-import { createMockExtensionPack } from "../../../../test/factories/model-editor/extension-pack";
-import { Mode } from "../../../model-editor/shared/mode";
-import { ModelEditorViewState } from "../../../model-editor/shared/view-state";
 import userEvent from "@testing-library/user-event";
+import { createMockModelEditorViewState } from "../../../../test/factories/model-editor/view-state";
 
 describe(LibraryRow.name, () => {
   const method = createMethod();
@@ -15,14 +13,7 @@ describe(LibraryRow.name, () => {
   const onStopGenerateFromLlmClick = jest.fn();
   const onModelDependencyClick = jest.fn();
 
-  const viewState: ModelEditorViewState = {
-    mode: Mode.Application,
-    showFlowGeneration: false,
-    showLlmButton: false,
-    showMultipleModels: false,
-    extensionPack: createMockExtensionPack(),
-    sourceArchiveAvailable: true,
-  };
+  const viewState = createMockModelEditorViewState();
 
   const render = (props: Partial<LibraryRowProps> = {}) =>
     reactRender(

extensions/ql-vscode/src/view/model-editor/__tests__/MethodRow.spec.tsx

@@ -5,12 +5,10 @@ import {
   screen,
 } from "@testing-library/react";
 import { createMethod } from "../../../../test/factories/model-editor/method-factories";
-import { Mode } from "../../../model-editor/shared/mode";
 import { MethodRow, MethodRowProps } from "../MethodRow";
 import { ModeledMethod } from "../../../model-editor/modeled-method";
 import userEvent from "@testing-library/user-event";
-import { ModelEditorViewState } from "../../../model-editor/shared/view-state";
-import { createMockExtensionPack } from "../../../../test/factories/model-editor/extension-pack";
+import { createMockModelEditorViewState } from "../../../../test/factories/model-editor/view-state";
 
 describe(MethodRow.name, () => {
   const method = createMethod({
@@ -33,14 +31,7 @@ describe(MethodRow.name, () => {
   };
   const onChange = jest.fn();
 
-  const viewState: ModelEditorViewState = {
-    mode: Mode.Application,
-    showFlowGeneration: false,
-    showLlmButton: false,
-    showMultipleModels: false,
-    extensionPack: createMockExtensionPack(),
-    sourceArchiveAvailable: true,
-  };
+  const viewState = createMockModelEditorViewState();
 
   const render = (props: Partial<MethodRowProps> = {}) =>
     reactRender(

extensions/ql-vscode/src/view/model-editor/__tests__/ModelKindDropdown.spec.tsx

@@ -4,6 +4,7 @@ import { ModelKindDropdown } from "../ModelKindDropdown";
 import userEvent from "@testing-library/user-event";
 import { createMethod } from "../../../../test/factories/model-editor/method-factories";
 import { createModeledMethod } from "../../../../test/factories/model-editor/modeled-method-factories";
+import { QueryLanguage } from "../../../common/query-language";
 
 describe(ModelKindDropdown.name, () => {
   const onChange = jest.fn();
@@ -21,6 +22,7 @@ describe(ModelKindDropdown.name, () => {
 
     render(
       <ModelKindDropdown
+        language={QueryLanguage.Java}
         method={method}
         modeledMethod={modeledMethod}
         onChange={onChange}
@@ -45,6 +47,7 @@ describe(ModelKindDropdown.name, () => {
 
     const { rerender } = render(
       <ModelKindDropdown
+        language={QueryLanguage.Java}
         method={method}
         modeledMethod={modeledMethod}
         onChange={onChange}
@@ -62,6 +65,7 @@ describe(ModelKindDropdown.name, () => {
 
     rerender(
       <ModelKindDropdown
+        language={QueryLanguage.Java}
         method={method}
         modeledMethod={updatedModeledMethod}
         onChange={onChange}
@@ -79,6 +83,7 @@ describe(ModelKindDropdown.name, () => {
 
     render(
       <ModelKindDropdown
+        language={QueryLanguage.Java}
         method={method}
         modeledMethod={modeledMethod}
         onChange={onChange}
@@ -102,6 +107,7 @@ describe(ModelKindDropdown.name, () => {
 
     render(
       <ModelKindDropdown
+        language={QueryLanguage.Java}
         method={method}
         modeledMethod={modeledMethod}
         onChange={onChange}
@@ -120,6 +126,7 @@ describe(ModelKindDropdown.name, () => {
 
     render(
       <ModelKindDropdown
+        language={QueryLanguage.Java}
         method={method}
         modeledMethod={modeledMethod}
         onChange={onChange}

extensions/ql-vscode/src/view/model-editor/__tests__/ModeledMethodDataGrid.spec.tsx

@@ -1,13 +1,11 @@
 import * as React from "react";
 import { render as reactRender, screen } from "@testing-library/react";
 import { createMethod } from "../../../../test/factories/model-editor/method-factories";
-import { Mode } from "../../../model-editor/shared/mode";
 import {
   ModeledMethodDataGrid,
   ModeledMethodDataGridProps,
 } from "../ModeledMethodDataGrid";
-import { ModelEditorViewState } from "../../../model-editor/shared/view-state";
-import { createMockExtensionPack } from "../../../../test/factories/model-editor/extension-pack";
+import { createMockModelEditorViewState } from "../../../../test/factories/model-editor/view-state";
 
 describe(ModeledMethodDataGrid.name, () => {
   const method1 = createMethod({
@@ -42,14 +40,7 @@ describe(ModeledMethodDataGrid.name, () => {
   });
   const onChange = jest.fn();
 
-  const viewState: ModelEditorViewState = {
-    mode: Mode.Application,
-    showFlowGeneration: false,
-    showLlmButton: false,
-    showMultipleModels: false,
-    extensionPack: createMockExtensionPack(),
-    sourceArchiveAvailable: true,
-  };
+  const viewState = createMockModelEditorViewState();
 
   const render = (props: Partial<ModeledMethodDataGridProps> = {}) =>
     reactRender(

extensions/ql-vscode/src/view/model-editor/__tests__/ModeledMethodsList.spec.tsx

@@ -1,13 +1,11 @@
 import * as React from "react";
 import { render as reactRender, screen } from "@testing-library/react";
 import { createMethod } from "../../../../test/factories/model-editor/method-factories";
-import { createMockExtensionPack } from "../../../../test/factories/model-editor/extension-pack";
-import { Mode } from "../../../model-editor/shared/mode";
-import { ModelEditorViewState } from "../../../model-editor/shared/view-state";
 import {
   ModeledMethodsList,
   ModeledMethodsListProps,
 } from "../ModeledMethodsList";
+import { createMockModelEditorViewState } from "../../../../test/factories/model-editor/view-state";
 
 describe(ModeledMethodsList.name, () => {
   const method1 = createMethod({
@@ -43,14 +41,7 @@ describe(ModeledMethodsList.name, () => {
   const onStopGenerateFromLlmClick = jest.fn();
   const onModelDependencyClick = jest.fn();
 
-  const viewState: ModelEditorViewState = {
-    mode: Mode.Application,
-    showFlowGeneration: false,
-    showLlmButton: false,
-    showMultipleModels: false,
-    extensionPack: createMockExtensionPack(),
-    sourceArchiveAvailable: true,
-  };
+  const viewState = createMockModelEditorViewState();
 
   const render = (props: Partial<ModeledMethodsListProps> = {}) =>
     reactRender(

extensions/ql-vscode/test/factories/model-editor/view-state.ts

@@ -0,0 +1,19 @@
+import { ModelEditorViewState } from "../../../src/model-editor/shared/view-state";
+import { Mode } from "../../../src/model-editor/shared/mode";
+import { createMockExtensionPack } from "./extension-pack";
+import { QueryLanguage } from "../../../src/common/query-language";
+
+export function createMockModelEditorViewState(
+  data: Partial<ModelEditorViewState> = {},
+): ModelEditorViewState {
+  return {
+    language: QueryLanguage.Java,
+    mode: Mode.Application,
+    showFlowGeneration: false,
+    showLlmButton: false,
+    showMultipleModels: false,
+    extensionPack: createMockExtensionPack(),
+    sourceArchiveAvailable: true,
+    ...data,
+  };
+}

extensions/ql-vscode/test/unit-tests/model-editor/yaml.test.ts

@@ -7,10 +7,11 @@ import {
   loadDataExtensionYaml,
 } from "../../../src/model-editor/yaml";
 import { CallClassification } from "../../../src/model-editor/method";
+import { QueryLanguage } from "../../../src/common/query-language";
 
 describe("createDataExtensionYaml", () => {
   it("creates the correct YAML file", () => {
-    const yaml = createDataExtensionYaml("java", [
+    const yaml = createDataExtensionYaml(QueryLanguage.Java, [
       {
         type: "sink",
         input: "Argument[0]",
@@ -50,7 +51,7 @@ describe("createDataExtensionYaml", () => {
   });
 
   it("includes the correct language", () => {
-    const yaml = createDataExtensionYaml("csharp", []);
+    const yaml = createDataExtensionYaml(QueryLanguage.CSharp, []);
 
     expect(yaml).toEqual(`extensions:
   - addsTo:
@@ -79,7 +80,7 @@ describe("createDataExtensionYaml", () => {
 describe("createDataExtensionYamlsForApplicationMode", () => {
   it("creates the correct YAML files when there are no existing modeled methods", () => {
     const yaml = createDataExtensionYamlsForApplicationMode(
-      "java",
+      QueryLanguage.Java,
       [
         {
           library: "sql2o",
@@ -323,7 +324,7 @@ describe("createDataExtensionYamlsForApplicationMode", () => {
 
   it("creates the correct YAML files when there are existing modeled methods", () => {
     const yaml = createDataExtensionYamlsForApplicationMode(
-      "java",
+      QueryLanguage.Java,
       [
         {
           library: "sql2o",
@@ -618,7 +619,7 @@ describe("createDataExtensionYamlsForApplicationMode", () => {
 describe("createDataExtensionYamlsForFrameworkMode", () => {
   it("creates the correct YAML files when there are no existing modeled methods", () => {
     const yaml = createDataExtensionYamlsForFrameworkMode(
-      "java",
+      QueryLanguage.Java,
       [
         {
           library: "sql2o",
@@ -774,7 +775,7 @@ describe("createDataExtensionYamlsForFrameworkMode", () => {
 
   it("creates the correct YAML files when there are existing modeled methods", () => {
     const yaml = createDataExtensionYamlsForFrameworkMode(
-      "java",
+      QueryLanguage.Java,
       [
         {
           library: "sql2o",
@@ -980,38 +981,41 @@ describe("createDataExtensionYamlsForFrameworkMode", () => {
 
 describe("loadDataExtensionYaml", () => {
   it("loads the YAML file", () => {
-    const data = loadDataExtensionYaml({
-      extensions: [
-        {
-          addsTo: { pack: "codeql/java-all", extensible: "sourceModel" },
-          data: [],
-        },
-        {
-          addsTo: { pack: "codeql/java-all", extensible: "sinkModel" },
-          data: [
-            [
-              "org.sql2o",
-              "Connection",
-              true,
-              "createQuery",
-              "(String)",
-              "",
-              "Argument[0]",
-              "sql",
-              "manual",
+    const data = loadDataExtensionYaml(
+      {
+        extensions: [
+          {
+            addsTo: { pack: "codeql/java-all", extensible: "sourceModel" },
+            data: [],
+          },
+          {
+            addsTo: { pack: "codeql/java-all", extensible: "sinkModel" },
+            data: [
+              [
+                "org.sql2o",
+                "Connection",
+                true,
+                "createQuery",
+                "(String)",
+                "",
+                "Argument[0]",
+                "sql",
+                "manual",
+              ],
             ],
-          ],
-        },
-        {
-          addsTo: { pack: "codeql/java-all", extensible: "summaryModel" },
-          data: [],
-        },
-        {
-          addsTo: { pack: "codeql/java-all", extensible: "neutralModel" },
-          data: [],
-        },
-      ],
-    });
+          },
+          {
+            addsTo: { pack: "codeql/java-all", extensible: "summaryModel" },
+            data: [],
+          },
+          {
+            addsTo: { pack: "codeql/java-all", extensible: "neutralModel" },
+            data: [],
+          },
+        ],
+      },
+      QueryLanguage.Java,
+    );
 
     expect(data).toEqual({
       "org.sql2o.Connection#createQuery(String)": [
@@ -1033,13 +1037,16 @@ describe("loadDataExtensionYaml", () => {
 
   it("returns undefined if given a string", () => {
     expect(() =>
-      loadDataExtensionYaml(`extensions:
+      loadDataExtensionYaml(
+        `extensions:
   - addsTo:
       pack: codeql/java-all
       extensible: sinkModel
     data:
       - ["org.sql2o","Connection",true,"createQuery","(String)","","Argument[0]","sql","manual"]
-`),
+`,
+        QueryLanguage.Java,
+      ),
     ).toThrow("Invalid data extension YAML:  must be object");
   });
 });

extensions/ql-vscode/test/vscode-tests/cli-integration/model-editor/modeled-method-fs.test.ts

@@ -11,6 +11,7 @@ import { ExtensionPack } from "../../../../src/model-editor/shared/extension-pac
 import { join } from "path";
 import { extLogger } from "../../../../src/common/logging/vscode";
 import { homedir } from "os";
+import { QueryLanguage } from "../../../../src/common/query-language";
 
 const dummyExtensionPackContents = `
 name: dummy/pack
@@ -192,6 +193,7 @@ describe("modeled-method-fs", () => {
 
       const modeledMethods = await loadModeledMethods(
         makeExtensionPack(extensionPackPath),
+        QueryLanguage.Java,
         cli,
         extLogger,
       );

extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/model-editor-view.test.ts

@@ -12,6 +12,7 @@ import { createMockModelingStore } from "../../../__mocks__/model-editor/modelin
 import { createMockModelEditorViewTracker } from "../../../__mocks__/model-editor/modelEditorViewTrackerMock";
 import { ModelConfigListener } from "../../../../src/config";
 import { createMockModelingEvents } from "../../../__mocks__/model-editor/modelingEventsMock";
+import { QueryLanguage } from "../../../../src/common/query-language";
 
 describe("ModelEditorView", () => {
   const app = createMockApp({});
@@ -56,6 +57,7 @@ describe("ModelEditorView", () => {
       queryDir,
       databaseItem,
       extensionPack,
+      QueryLanguage.Java,
       mode,
     );
   });