Add add/remove buttons for models

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

@@ -6,7 +6,10 @@ import { MethodRow as MethodRowComponent } from "../../view/model-editor/MethodR
 import { CallClassification, Method } from "../../model-editor/method";
 import { ModeledMethod } from "../../model-editor/modeled-method";
 import { VSCodeDataGrid } from "@vscode/webview-ui-toolkit/react";
-import { GRID_TEMPLATE_COLUMNS } from "../../view/model-editor/ModeledMethodDataGrid";
+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";
@@ -16,11 +19,16 @@ export default {
   component: MethodRowComponent,
 } as Meta<typeof MethodRowComponent>;
 
-const Template: StoryFn<typeof MethodRowComponent> = (args) => (
+const Template: StoryFn<typeof MethodRowComponent> = (args) => {
-  <VSCodeDataGrid gridTemplateColumns={GRID_TEMPLATE_COLUMNS}>
+  const gridTemplateColumns = args.viewState?.showMultipleModels
-    <MethodRowComponent {...args} />
+    ? MULTIPLE_MODELS_GRID_TEMPLATE_COLUMNS
-  </VSCodeDataGrid>
+    : SINGLE_MODEL_GRID_TEMPLATE_COLUMNS;
-);
+  return (
+    <VSCodeDataGrid gridTemplateColumns={gridTemplateColumns}>
+      <MethodRowComponent {...args} />
+    </VSCodeDataGrid>
+  );
+};
 
 const method: Method = {
   library: "sql2o-1.6.0.jar",

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

@@ -1,4 +1,5 @@
 import {
+  VSCodeButton,
   VSCodeDataGridCell,
   VSCodeDataGridRow,
   VSCodeLink,
@@ -22,6 +23,7 @@ import { ModelTypeDropdown } from "./ModelTypeDropdown";
 import { ModelInputDropdown } from "./ModelInputDropdown";
 import { ModelOutputDropdown } from "./ModelOutputDropdown";
 import { ModelEditorViewState } from "../../model-editor/shared/view-state";
+import { Codicon } from "../common";
 
 const MultiModelColumn = styled(VSCodeDataGridCell)`
   display: flex;
@@ -55,6 +57,11 @@ const ProgressRing = styled(VSCodeProgressRing)`
   margin-left: auto;
 `;
 
+const CodiconRow = styled(VSCodeButton)`
+  min-height: calc(var(--input-height) * 1px);
+  align-items: center;
+`;
+
 const DataGridRow = styled(VSCodeDataGridRow)<{ focused?: boolean }>`
   outline: ${(props) =>
     props.focused ? "1px solid var(--vscode-focusBorder)" : "none"};
@@ -159,6 +166,13 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
             <VSCodeDataGridCell gridColumn={5}>
               <InProgressDropdown />
             </VSCodeDataGridCell>
+            {viewState.showMultipleModels && (
+              <VSCodeDataGridCell gridColumn={6}>
+                <CodiconRow appearance="icon" disabled={true}>
+                  <Codicon name="add" label="Add new model" />
+                </CodiconRow>
+              </VSCodeDataGridCell>
+            )}
           </>
         )}
         {!props.modelingInProgress && (
@@ -203,6 +217,19 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
                 />
               ))}
             </MultiModelColumn>
+            {viewState.showMultipleModels && (
+              <MultiModelColumn gridColumn={6}>
+                {modeledMethods.map((_, index) => (
+                  <CodiconRow key={index} appearance="icon" disabled={false}>
+                    {index === modeledMethods.length - 1 ? (
+                      <Codicon name="add" label="Add new model" />
+                    ) : (
+                      <Codicon name="trash" label="Remove model" />
+                    )}
+                  </CodiconRow>
+                ))}
+              </MultiModelColumn>
+            )}
           </>
         )}
       </DataGridRow>

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

@@ -13,7 +13,10 @@ import { InProgressMethods } from "../../model-editor/shared/in-progress-methods
 import { HiddenMethodsRow } from "./HiddenMethodsRow";
 import { ModelEditorViewState } from "../../model-editor/shared/view-state";
 
-export const GRID_TEMPLATE_COLUMNS = "0.5fr 0.125fr 0.125fr 0.125fr 0.125fr";
+export const SINGLE_MODEL_GRID_TEMPLATE_COLUMNS =
+  "0.5fr 0.125fr 0.125fr 0.125fr 0.125fr";
+export const MULTIPLE_MODELS_GRID_TEMPLATE_COLUMNS =
+  "0.5fr 0.125fr 0.125fr 0.125fr 0.125fr max-content";
 
 export type ModeledMethodDataGridProps = {
   packageName: string;
@@ -64,8 +67,12 @@ export const ModeledMethodDataGrid = ({
 
   const someMethodsAreVisible = methodsWithModelability.length > 0;
 
+  const gridTemplateColumns = viewState.showMultipleModels
+    ? MULTIPLE_MODELS_GRID_TEMPLATE_COLUMNS
+    : SINGLE_MODEL_GRID_TEMPLATE_COLUMNS;
+
   return (
-    <VSCodeDataGrid gridTemplateColumns={GRID_TEMPLATE_COLUMNS}>
+    <VSCodeDataGrid gridTemplateColumns={gridTemplateColumns}>
       {someMethodsAreVisible && (
         <>
           <VSCodeDataGridRow rowType="header">
@@ -84,6 +91,9 @@ export const ModeledMethodDataGrid = ({
             <VSCodeDataGridCell cellType="columnheader" gridColumn={5}>
               Kind
             </VSCodeDataGridCell>
+            {viewState.showMultipleModels && (
+              <VSCodeDataGridCell cellType="columnheader" gridColumn={6} />
+            )}
           </VSCodeDataGridRow>
           {methodsWithModelability.map(({ method, methodCanBeModeled }) => {
             const modeledMethods = modeledMethodsMap[method.signature] ?? [];