Exposes hooks to customize the workbench-view (#16749)

* Exposes hooks to customize the workbench-view

* addressed PR feedback

* naming

* auto -> formatInteger(maxNum)

web-console/src/views/workbench-view/max-tasks-button/max-tasks-button.tsx

@@ -39,14 +39,28 @@ const TASK_ASSIGNMENT_DESCRIPTION: Record<string, string> = {
   auto: `Use as few tasks as possible without exceeding 512 MiB or 10,000 files per task, unless exceeding these limits is necessary to stay within 'maxNumTasks'. When calculating the size of files, the weighted size is used, which considers the file format and compression format used if any. When file sizes cannot be determined through directory listing (for example: http), behaves the same as 'max'.`,
 };
 
+const DEFAULT_MAX_NUM_LABEL_FN = (maxNum: number) => {
+  if (maxNum === 2) return { text: formatInteger(maxNum), label: '(1 controller + 1 worker)' };
+  return { text: formatInteger(maxNum), label: `(1 controller + max ${maxNum - 1} workers)` };
+};
+
 export interface MaxTasksButtonProps extends Omit<ButtonProps, 'text' | 'rightIcon'> {
   clusterCapacity: number | undefined;
   queryContext: QueryContext;
   changeQueryContext(queryContext: QueryContext): void;
+  menuHeader?: JSX.Element;
+  maxNumLabelFn?: (maxNum: number) => { text: string; label?: string };
 }
 
 export const MaxTasksButton = function MaxTasksButton(props: MaxTasksButtonProps) {
-  const { clusterCapacity, queryContext, changeQueryContext, ...rest } = props;
+  const {
+    clusterCapacity,
+    queryContext,
+    changeQueryContext,
+    menuHeader,
+    maxNumLabelFn = DEFAULT_MAX_NUM_LABEL_FN,
+    ...rest
+  } = props;
   const [customMaxNumTasksDialogOpen, setCustomMaxNumTasksDialogOpen] = useState(false);
 
   const maxNumTasks = getMaxNumTasks(queryContext);
@@ -68,6 +82,7 @@ export const MaxTasksButton = function MaxTasksButton(props: MaxTasksButtonProps
         position={Position.BOTTOM_LEFT}
         content={
           <Menu>
+            {menuHeader}
             <MenuDivider title="Maximum number of tasks to launch" />
             {Boolean(fullClusterCapacity) && (
               <MenuItem
@@ -76,15 +91,19 @@ export const MaxTasksButton = function MaxTasksButton(props: MaxTasksButtonProps
                 onClick={() => changeQueryContext(changeMaxNumTasks(queryContext, undefined))}
               />
             )}
-            {shownMaxNumTaskOptions.map(m => (
+            {shownMaxNumTaskOptions.map(m => {
-              <MenuItem
+              const { text, label } = maxNumLabelFn(m);
-                key={String(m)}
+
-                icon={tickIcon(m === maxNumTasks)}
+              return (
-                text={formatInteger(m)}
+                <MenuItem
-                label={`(1 controller + ${m === 2 ? '1 worker' : `max ${m - 1} workers`})`}
+                  key={String(m)}
-                onClick={() => changeQueryContext(changeMaxNumTasks(queryContext, m))}
+                  icon={tickIcon(m === maxNumTasks)}
-              />
+                  text={text}
-            ))}
+                  label={label}
+                  onClick={() => changeQueryContext(changeMaxNumTasks(queryContext, m))}
+                />
+              );
+            })}
             <MenuItem
               icon={tickIcon(
                 typeof maxNumTasks === 'number' && !shownMaxNumTaskOptions.includes(maxNumTasks),

web-console/src/views/workbench-view/query-tab/query-tab.tsx

@@ -21,7 +21,7 @@ import { IconNames } from '@blueprintjs/icons';
 import type { QueryResult } from '@druid-toolkit/query';
 import { QueryRunner, SqlQuery } from '@druid-toolkit/query';
 import axios from 'axios';
-import type { JSX } from 'react';
+import type { ComponentProps, JSX } from 'react';
 import React, { useCallback, useEffect, useRef, useState } from 'react';
 import SplitterLayout from 'react-splitter-layout';
 import { useStore } from 'zustand';
@@ -82,6 +82,9 @@ export interface QueryTabProps {
   clusterCapacity: number | undefined;
   goToTask(taskId: string): void;
   getClusterCapacity: (() => Promise<CapacityInfo | undefined>) | undefined;
+  maxTaskMenuHeader?: JSX.Element;
+  enginesLabelFn?: ComponentProps<typeof RunPanel>['enginesLabelFn'];
+  maxTaskLabelFn?: ComponentProps<typeof RunPanel>['maxTaskLabelFn'];
 }
 
 export const QueryTab = React.memo(function QueryTab(props: QueryTabProps) {
@@ -98,6 +101,9 @@ export const QueryTab = React.memo(function QueryTab(props: QueryTabProps) {
     clusterCapacity,
     goToTask,
     getClusterCapacity,
+    maxTaskMenuHeader,
+    enginesLabelFn,
+    maxTaskLabelFn,
   } = props;
   const [alertElement, setAlertElement] = useState<JSX.Element | undefined>();
 
@@ -399,6 +405,9 @@ export const QueryTab = React.memo(function QueryTab(props: QueryTabProps) {
               queryEngines={queryEngines}
               clusterCapacity={clusterCapacity}
               moreMenu={runMoreMenu}
+              maxTaskMenuHeader={maxTaskMenuHeader}
+              enginesLabelFn={enginesLabelFn}
+              maxTaskLabelFn={maxTaskLabelFn}
             />
             {executionState.isLoading() && (
               <ExecutionTimerPanel

web-console/src/views/workbench-view/run-panel/run-panel.tsx

@@ -30,7 +30,7 @@ import {
   useHotkeys,
 } from '@blueprintjs/core';
 import { IconNames } from '@blueprintjs/icons';
-import type { JSX } from 'react';
+import type { ComponentProps, JSX } from 'react';
 import React, { useCallback, useMemo, useState } from 'react';
 
 import { MenuCheckbox, MenuTristate } from '../../../components';
@@ -111,6 +111,14 @@ const ARRAY_INGEST_MODE_DESCRIPTION: Record<ArrayIngestMode, JSX.Element> = {
   ),
 };
 
+const DEFAULT_ENGINES_LABEL_FN = (engine: DruidEngine | undefined) => {
+  if (!engine) return { text: 'auto' };
+  return {
+    text: engine,
+    label: engine === 'sql-msq-task' ? 'multi-stage-query' : undefined,
+  };
+};
+
 export interface RunPanelProps {
   query: WorkbenchQuery;
   onQueryChange(query: WorkbenchQuery): void;
@@ -120,11 +128,25 @@ export interface RunPanelProps {
   queryEngines: DruidEngine[];
   clusterCapacity: number | undefined;
   moreMenu?: JSX.Element;
+  maxTaskMenuHeader?: JSX.Element;
+  enginesLabelFn?: (engine: DruidEngine | undefined) => { text: string; label?: string };
+  maxTaskLabelFn?: ComponentProps<typeof MaxTasksButton>['maxNumLabelFn'];
 }
 
 export const RunPanel = React.memo(function RunPanel(props: RunPanelProps) {
-  const { query, onQueryChange, onRun, moreMenu, running, small, queryEngines, clusterCapacity } =
+  const {
-    props;
+    query,
+    onQueryChange,
+    onRun,
+    moreMenu,
+    running,
+    small,
+    queryEngines,
+    clusterCapacity,
+    maxTaskMenuHeader,
+    maxTaskLabelFn,
+    enginesLabelFn = DEFAULT_ENGINES_LABEL_FN,
+  } = props;
   const [editContextDialogOpen, setEditContextDialogOpen] = useState(false);
   const [editParametersDialogOpen, setEditParametersDialogOpen] = useState(false);
   const [customTimezoneDialogOpen, setCustomTimezoneDialogOpen] = useState(false);
@@ -186,25 +208,11 @@ export const RunPanel = React.memo(function RunPanel(props: RunPanelProps) {
   useHotkeys(hotkeys);
 
   const queryEngine = query.engine;
-  function renderQueryEngineMenuItem(e: DruidEngine | undefined) {
-    return (
-      <MenuItem
-        key={String(e)}
-        icon={tickIcon(e === queryEngine)}
-        text={typeof e === 'undefined' ? 'auto' : e}
-        label={e === 'sql-msq-task' ? 'multi-stage-query' : undefined}
-        onClick={() => onQueryChange(query.changeEngine(e))}
-        shouldDismissPopover={false}
-      />
-    );
-  }
 
   function changeQueryContext(queryContext: QueryContext) {
     onQueryChange(query.changeQueryContext(queryContext));
   }
 
-  const availableEngines = ([undefined] as (DruidEngine | undefined)[]).concat(queryEngines);
-
   function offsetOptions(): JSX.Element[] {
     const items: JSX.Element[] = [];
 
@@ -231,6 +239,9 @@ export const RunPanel = React.memo(function RunPanel(props: RunPanelProps) {
   const intent = overloadWarning ? Intent.WARNING : undefined;
 
   const effectiveEngine = query.getEffectiveEngine();
+
+  const autoEngineLabel = enginesLabelFn(undefined);
+
   return (
     <div className="run-panel">
       <Button
@@ -262,7 +273,29 @@ export const RunPanel = React.memo(function RunPanel(props: RunPanelProps) {
                 {queryEngines.length > 1 && (
                   <>
                     <MenuDivider title="Select engine" />
-                    {availableEngines.map(renderQueryEngineMenuItem)}
+                    <MenuItem
+                      key="auto"
+                      icon={tickIcon(queryEngine === undefined)}
+                      text={autoEngineLabel.text}
+                      label={autoEngineLabel.label}
+                      onClick={() => onQueryChange(query.changeEngine(undefined))}
+                      shouldDismissPopover={false}
+                    />
+                    {queryEngines.map(engine => {
+                      const { text, label } = enginesLabelFn(engine);
+
+                      return (
+                        <MenuItem
+                          key={String(engine)}
+                          icon={tickIcon(engine === queryEngine)}
+                          text={text}
+                          label={label}
+                          onClick={() => onQueryChange(query.changeEngine(engine))}
+                          shouldDismissPopover={false}
+                        />
+                      );
+                    })}
+
                     <MenuDivider />
                   </>
                 )}
@@ -485,7 +518,10 @@ export const RunPanel = React.memo(function RunPanel(props: RunPanelProps) {
             }
           >
             <Button
-              text={`Engine: ${queryEngine || `auto (${effectiveEngine})`}`}
+              text={`Engine: ${
+                (enginesLabelFn ? enginesLabelFn(queryEngine).text : queryEngine) ||
+                `auto (${enginesLabelFn ? enginesLabelFn(effectiveEngine) : effectiveEngine})`
+              }`}
               rightIcon={IconNames.CARET_DOWN}
               intent={intent}
             />
@@ -495,6 +531,8 @@ export const RunPanel = React.memo(function RunPanel(props: RunPanelProps) {
               clusterCapacity={clusterCapacity}
               queryContext={queryContext}
               changeQueryContext={changeQueryContext}
+              menuHeader={maxTaskMenuHeader}
+              maxNumLabelFn={maxTaskLabelFn}
             />
           )}
           {ingestMode && (

web-console/src/views/workbench-view/workbench-view.tsx

@@ -30,7 +30,7 @@ import type { SqlQuery } from '@druid-toolkit/query';
 import { SqlExpression } from '@druid-toolkit/query';
 import classNames from 'classnames';
 import copy from 'copy-to-clipboard';
-import React from 'react';
+import React, { ComponentProps } from 'react';
 
 import { SpecDialog, StringInputDialog } from '../../dialogs';
 import type {
@@ -101,6 +101,9 @@ export interface WorkbenchViewProps {
   allowExplain: boolean;
   goToTask(taskId: string): void;
   getClusterCapacity: (() => Promise<CapacityInfo | undefined>) | undefined;
+  maxTaskMenuHeader?: JSX.Element;
+  enginesLabelFn?: ComponentProps<typeof QueryTab>['enginesLabelFn'];
+  maxTaskLabelFn?: ComponentProps<typeof QueryTab>['maxTaskLabelFn'];
 }
 
 export interface WorkbenchViewState {
@@ -649,6 +652,9 @@ export class WorkbenchView extends React.PureComponent<WorkbenchViewProps, Workb
       allowExplain,
       goToTask,
       getClusterCapacity,
+      maxTaskMenuHeader,
+      enginesLabelFn,
+      maxTaskLabelFn,
     } = this.props;
     const { columnMetadataState } = this.state;
     const currentTabEntry = this.getCurrentTabEntry();
@@ -673,6 +679,9 @@ export class WorkbenchView extends React.PureComponent<WorkbenchViewProps, Workb
           clusterCapacity={capabilities.getMaxTaskSlots()}
           goToTask={goToTask}
           getClusterCapacity={getClusterCapacity}
+          maxTaskMenuHeader={maxTaskMenuHeader}
+          enginesLabelFn={enginesLabelFn}
+          maxTaskLabelFn={maxTaskLabelFn}
           runMoreMenu={
             <Menu>
               {allowExplain &&