From f3796ac5bd2cbb78ba76292e1b02a982659595b4 Mon Sep 17 00:00:00 2001
From: shati-patel <42641846+shati-patel@users.noreply.github.com>
Date: Tue, 13 Feb 2024 13:43:29 +0000
Subject: [PATCH] Add feature flag for Python model editor

This makes Python available in the model editor if you have `"codeQL.model.enablePython": true` in your (user or workspace) settings.
---
 extensions/ql-vscode/src/config.ts                    |  6 ++++++
 .../ql-vscode/src/model-editor/model-editor-module.ts |  5 ++++-
 .../ql-vscode/src/model-editor/supported-languages.ts | 11 ++++++++++-
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/extensions/ql-vscode/src/config.ts b/extensions/ql-vscode/src/config.ts
index e081c09fcdd..50937cf15b7 100644
--- a/extensions/ql-vscode/src/config.ts
+++ b/extensions/ql-vscode/src/config.ts
@@ -716,6 +716,7 @@ const LLM_GENERATION_DEV_ENDPOINT = new Setting(
   MODEL_SETTING,
 );
 const EXTENSIONS_DIRECTORY = new Setting("extensionsDirectory", MODEL_SETTING);
+const ENABLE_PYTHON = new Setting("enablePython", MODEL_SETTING);
 const ENABLE_ACCESS_PATH_SUGGESTIONS = new Setting(
   "enableAccessPathSuggestions",
   MODEL_SETTING,
@@ -725,6 +726,7 @@ export interface ModelConfig {
   flowGeneration: boolean;
   llmGeneration: boolean;
   getExtensionsDirectory(languageId: string): string | undefined;
+  enablePython: boolean;
   enableAccessPathSuggestions: boolean;
 }
 
@@ -763,6 +765,10 @@ export class ModelConfigListener extends ConfigListener implements ModelConfig {
     });
   }
 
+  public get enablePython(): boolean {
+    return !!ENABLE_PYTHON.getValue<boolean>();
+  }
+
   public get enableAccessPathSuggestions(): boolean {
     return !!ENABLE_ACCESS_PATH_SUGGESTIONS.getValue<boolean>();
   }
diff --git a/extensions/ql-vscode/src/model-editor/model-editor-module.ts b/extensions/ql-vscode/src/model-editor/model-editor-module.ts
index 123a7680ab2..5a531c1b2b7 100644
--- a/extensions/ql-vscode/src/model-editor/model-editor-module.ts
+++ b/extensions/ql-vscode/src/model-editor/model-editor-module.ts
@@ -143,7 +143,10 @@ export class ModelEditorModule extends DisposableObject {
 
       const language = db.language;
 
-      if (!isQueryLanguage(language) || !isSupportedLanguage(language)) {
+      if (
+        !isQueryLanguage(language) ||
+        !isSupportedLanguage(language, this.modelConfig)
+      ) {
         void showAndLogErrorMessage(
           this.app.logger,
           `The CodeQL Model Editor is not supported for ${language} databases.`,
diff --git a/extensions/ql-vscode/src/model-editor/supported-languages.ts b/extensions/ql-vscode/src/model-editor/supported-languages.ts
index fa71c1beeb1..2bf3a0f81a6 100644
--- a/extensions/ql-vscode/src/model-editor/supported-languages.ts
+++ b/extensions/ql-vscode/src/model-editor/supported-languages.ts
@@ -1,4 +1,5 @@
 import { QueryLanguage } from "../common/query-language";
+import type { ModelConfig } from "../config";
 import { isCanary } from "../config";
 
 /**
@@ -10,7 +11,10 @@ export const SUPPORTED_LANGUAGES: QueryLanguage[] = [
   QueryLanguage.CSharp,
 ];
 
-export function isSupportedLanguage(language: QueryLanguage) {
+export function isSupportedLanguage(
+  language: QueryLanguage,
+  modelConfig: ModelConfig,
+) {
   if (SUPPORTED_LANGUAGES.includes(language)) {
     return true;
   }
@@ -20,5 +24,10 @@ export function isSupportedLanguage(language: QueryLanguage) {
     return isCanary();
   }
 
+  if (language === QueryLanguage.Python) {
+    // Python is only enabled when the config setting is set
+    return modelConfig.enablePython;
+  }
+
   return false;
 }