Add support for enabling/disabling custom visualizations

frontend/server/server.ts

@@ -54,7 +54,9 @@ const {
   /** API service will listen to this port */
   ML_PIPELINE_SERVICE_PORT = '3001',
   /** path to viewer:tensorboard pod template spec */
-  VIEWER_TENSORBOARD_POD_TEMPLATE_SPEC_PATH
+  VIEWER_TENSORBOARD_POD_TEMPLATE_SPEC_PATH,
+  /** Whether custom visualizations are allowed to be generated by the frontend */
+  ALLOW_CUSTOM_VISUALIZATIONS = false
 } = process.env;
 
 /** construct minio endpoint from host and namespace (optional) */
@@ -328,6 +330,10 @@ const projectIdHandler = async (req, res) => {
   res.send(await response.text());
 };
 
+const allowCustomVisualizationsHandler = (req, res) => {
+  res.send(ALLOW_CUSTOM_VISUALIZATIONS === 'true');
+};
+
 app.get('/' + v1beta1Prefix + '/healthz', healthzHandler);
 app.get(BASEPATH + '/' + v1beta1Prefix + '/healthz', healthzHandler);
 
@@ -349,6 +355,9 @@ app.get(BASEPATH + '/system/cluster-name', clusterNameHandler);
 app.get('/system/project-id', projectIdHandler);
 app.get(BASEPATH + '/system/project-id', projectIdHandler);
 
+app.get('/visualizations/allowed', allowCustomVisualizationsHandler);
+app.get(BASEPATH + '/visualizations/allowed', allowCustomVisualizationsHandler);
+
 // Order matters here, since both handlers can match any proxied request with a referer,
 // and we prioritize the basepath-friendly handler
 proxyMiddleware(app, BASEPATH + '/' + v1beta1Prefix);

frontend/src/components/viewers/VisualizationCreator.tsx

@@ -29,7 +29,7 @@ import 'brace/ext/language_tools';
 import 'brace/mode/json';
 import 'brace/mode/python';
 import 'brace/theme/github';
-
+import { Apis } from 'src/lib/Apis';
 
 export interface VisualizationCreatorConfig extends ViewerConfig {
   // Whether there is currently a visualization being generated or not.
@@ -44,6 +44,7 @@ interface VisualizationCreatorProps {
 }
 
 interface VisualizationCreatorState {
+  allowCustomVisualizations: boolean;
   // arguments is expected to be a JSON object in string form.
   arguments: string;
   code: string;
@@ -52,34 +53,16 @@ interface VisualizationCreatorState {
 }
 
 class VisualizationCreator extends Viewer<VisualizationCreatorProps, VisualizationCreatorState> {
-  /*
-    Due to the swagger API definition generation, enum value that include
-    an _ (underscore) remove all _ from the enum key. Additionally, due to the
-    manner in which TypeScript is compiled to Javascript, enums are duplicated
-    iff they included an _ in the proto file. This filters out those duplicate
-    keys that are generated by the complication from TypeScript to JavaScript.
-  
-    For example:
-    export enum ApiVisualizationType {
-      ROCCURVE = <any> 'ROC_CURVE'
-    }
-  
-    Object.keys(ApiVisualizationType) = ['CURVE', 'ROC_CURVE'];
-  
-    Additional details can be found here:
-    https://www.typescriptlang.org/play/#code/KYOwrgtgBAggDgSwGoIM5gIYBsEC8MAuCA9iACoCecwUA3gLABQUUASgPIDCnAqq0gFEoAXigAeDCAoA+AOQdOAfV78BsgDRMAvkyYBjUqmJZgAOizEA5gAp4yNJhz4ipStQCUAbiA
-  */
-  private _types = Object.keys(ApiVisualizationType)
-    .map((key: string) => key.replace('_', ''))
-    .filter((key: string, i: number, arr: string[]) => arr.indexOf(key) === i);
-
   constructor(props: VisualizationCreatorProps) {
     super(props);
     this.state = {
+      allowCustomVisualizations: false,
       arguments: '',
       code: '',
       source: '',
     };
+    Apis.areCustomVisualizationsAllowed()
+      .then((allowed) => this.setState({ allowCustomVisualizations: allowed }));
   }
 
   public getDisplayName(): string {
@@ -132,12 +115,11 @@ class VisualizationCreator extends Viewer<VisualizationCreatorProps, Visualizati
           }}
           disabled={isBusy}
         >
-          {this._types
-            .map((key: string) => (
-              <MenuItem key={key} value={ApiVisualizationType[key]}>
-                {ApiVisualizationType[key]}
-              </MenuItem>
-            ))}
+          {this.getAvailableTypes().map((key: string) => (
+            <MenuItem key={key} value={ApiVisualizationType[key]}>
+              {ApiVisualizationType[key]}
+            </MenuItem>
+          ))}
         </Select>
       </FormControl>
 
@@ -182,6 +164,32 @@ class VisualizationCreator extends Viewer<VisualizationCreatorProps, Visualizati
     </div>;
   }
 
+  /*
+    Due to the swagger API definition generation, enum value that include
+    an _ (underscore) remove all _ from the enum key. Additionally, due to the
+    manner in which TypeScript is compiled to Javascript, enums are duplicated
+    iff they included an _ in the proto file. This filters out those duplicate
+    keys that are generated by the complication from TypeScript to JavaScript.
+  
+    For example:
+    export enum ApiVisualizationType {
+      ROCCURVE = <any> 'ROC_CURVE'
+    }
+  
+    Object.keys(ApiVisualizationType) = ['CURVE', 'ROC_CURVE'];
+  
+    Additional details can be found here:
+    https://www.typescriptlang.org/play/#code/KYOwrgtgBAggDgSwGoIM5gIYBsEC8MAuCA9iACoCecwUA3gLABQUUASgPIDCnAqq0gFEoAXigAeDCAoA+AOQdOAfV78BsgDRMAvkyYBjUqmJZgAOizEA5gAp4yNJhz4ipStQCUAbiA
+  */
+  private getAvailableTypes(): string[] {
+    return Object.keys(ApiVisualizationType)
+      .map((key: string) => key.replace('_', ''))
+      .filter((key: string, i: number, arr: string[]) => {
+        return (this.state.allowCustomVisualizations || key !== 'CUSTOM')
+          && arr.indexOf(key) === i;
+      });
+  }
+
   private getArgumentPlaceholderForType(type: ApiVisualizationType | undefined): string {
     let placeholder= 'Arguments, provided as JSON, to be used during visualization generation.';
     switch(type) {

frontend/src/lib/Apis.ts

@@ -37,8 +37,21 @@ export interface BuildInfo {
   frontendCommitHash?: string;
 }
 
+let customVisualizationsAllowed: boolean;
+
 export class Apis {
 
+  public static async areCustomVisualizationsAllowed(): Promise<boolean> {
+    // Result is cached to prevent excessive network calls for simple request.
+    // The value of customVisualizationsAllowed will only change if the
+    // deployment is updated and then the entire pod is restarted.
+    if (customVisualizationsAllowed === undefined) {
+      const result = await this._fetch('visualizations/allowed');
+      customVisualizationsAllowed = result === 'true';
+    }
+    return customVisualizationsAllowed;
+  }
+
   /**
    * Get pod logs
    */