feat: support union type for launch plan

packages/zapp/console/src/components/Launch/LaunchForm/LaunchFormInputs.tsx

@@ -8,16 +8,20 @@ import { MapInput } from './MapInput';
 import { NoInputsNeeded } from './NoInputsNeeded';
 import { SimpleInput } from './SimpleInput';
 import { StructInput } from './StructInput';
+import { UnionInput } from './UnionInput';
 import { useStyles } from './styles';
 import { BaseInterpretedLaunchState, InputProps, InputType, LaunchFormInputsRef } from './types';
 import { UnsupportedInput } from './UnsupportedInput';
 import { UnsupportedRequiredInputsError } from './UnsupportedRequiredInputsError';
 import { useFormInputsState } from './useFormInputsState';
 import { isEnterInputsState } from './utils';
 
-function getComponentForInput(input: InputProps, showErrors: boolean) {
+export function getComponentForInput(input: InputProps, showErrors: boolean) {
   const props = { ...input, error: showErrors ? input.error : undefined };
+
   switch (input.typeDefinition.type) {
+    case InputType.Union:
+      return <UnionInput {...props} />;
     case InputType.Blob:
       return <BlobInput {...props} />;
     case InputType.Collection:

packages/zapp/console/src/components/Launch/LaunchForm/SimpleInput.tsx

@@ -68,6 +68,7 @@ export const SimpleInput: React.FC<InputProps> = (props) => {
         <TextField
           error={hasError}
           id={getLaunchInputId(name)}
+          key={getLaunchInputId(name)}
           helperText={helperText}
           fullWidth={true}
           label={label}

packages/zapp/console/src/components/Launch/LaunchForm/UnionInput.tsx

@@ -0,0 +1,156 @@
+import { Typography } from '@material-ui/core';
+import { makeStyles, Theme } from '@material-ui/core/styles';
+import * as React from 'react';
+import Card from '@material-ui/core/Card';
+import CardContent from '@material-ui/core/CardContent';
+import { InputProps, InputType, InputTypeDefinition, UnionValue, InputValue } from './types';
+import { formatType } from './utils';
+import { getComponentForInput } from './LaunchFormInputs';
+import { getHelperForInput } from './inputHelpers/getHelperForInput';
+import { SearchableSelector, SearchableSelectorOption } from './SearchableSelector';
+
+const useStyles = makeStyles((theme: Theme) => ({
+  inlineTitle: {
+    display: 'flex',
+    gap: theme.spacing(1),
+    alignItems: 'center',
+    paddingBottom: theme.spacing(3),
+  },
+}));
+
+const generateInputTypeToValueMap = (
+  listOfSubTypes: InputTypeDefinition[] | undefined,
+  initialInputValue: UnionValue | undefined,
+  initialType: InputTypeDefinition,
+): Record<InputType, InputValue> | {} => {
+  if (!listOfSubTypes?.length) {
+    return {};
+  }
+
+  return listOfSubTypes.reduce(function (map, subType) {
+    if (initialInputValue && subType === initialType) {
+      map[subType.type] = initialInputValue;
+    } else {
+      map[subType.type] = { value: '', typeDefinition: subType };
+    }
+    return map;
+  }, {});
+};
+
+const generateSearchableSelectorOption = (
+  inputTypeDefinition: InputTypeDefinition,
+): SearchableSelectorOption<InputType> => {
+  return {
+    id: inputTypeDefinition.type,
+    data: inputTypeDefinition.type,
+    name: formatType(inputTypeDefinition),
+  } as SearchableSelectorOption<InputType>;
+};
+
+const generateListOfSearchableSelectorOptions = (
+  listOfInputTypeDefinition: InputTypeDefinition[],
+): SearchableSelectorOption<InputType>[] => {
+  return listOfInputTypeDefinition.map((inputTypeDefinition) =>
+    generateSearchableSelectorOption(inputTypeDefinition),
+  );
+};
+
+export const UnionInput = (props: InputProps) => {
+  const { initialValue, required, label, onChange, typeDefinition, error, description } = props;
+
+  const classes = useStyles();
+
+  const listOfSubTypes = typeDefinition?.listOfSubTypes;
+
+  if (!listOfSubTypes) {
+    return <></>;
+  }
+
+  const inputTypeToInputTypeDefinition = listOfSubTypes.reduce(
+    (previous, current) => ({ ...previous, [current.type]: current }),
+    {},
+  );
+
+  const initialInputValue =
+    initialValue &&
+    (getHelperForInput(typeDefinition.type).fromLiteral(
+      initialValue,
+      typeDefinition,
+    ) as UnionValue);
+
+  const initialInputTypeDefinition = initialInputValue?.typeDefinition ?? listOfSubTypes[0];
+
+  if (!initialInputTypeDefinition) {
+    return <></>;
+  }
+
+  const [inputTypeToValueMap, setInputTypeToValueMap] = React.useState<
+    Record<InputType, InputValue> | {}
+  >(generateInputTypeToValueMap(listOfSubTypes, initialInputValue, initialInputTypeDefinition));
+
+  const [selectedInputType, setSelectedInputType] = React.useState<InputType>(
+    initialInputTypeDefinition.type,
+  );
+
+  const selectedInputTypeDefintion = inputTypeToInputTypeDefinition[
+    selectedInputType
+  ] as InputTypeDefinition;
+
+  const handleTypeOnSelectionChanged = (value: SearchableSelectorOption<InputType>) => {
+    setSelectedInputType(value.data);
+  };
+
+  const handleSubTypeOnChange = (input: InputValue) => {
+    onChange({
+      value: input,
+      typeDefinition: selectedInputTypeDefintion,
+    } as UnionValue);
+    setInputTypeToValueMap({
+      ...inputTypeToValueMap,
+      [selectedInputType]: {
+        value: input,
+        typeDefinition: selectedInputTypeDefintion,
+      } as UnionValue,
+    });
+  };
+
+  return (
+    <Card
+      variant="outlined"
+      style={{
+        overflow: 'visible',
+      }}
+    >
+      <CardContent>
+        <div className={classes.inlineTitle}>
+          <Typography variant="body1" component="label">
+            {label}
+          </Typography>
+
+          <SearchableSelector
+            label="Type"
+            options={generateListOfSearchableSelectorOptions(listOfSubTypes)}
+            selectedItem={generateSearchableSelectorOption(selectedInputTypeDefintion)}
+            onSelectionChanged={handleTypeOnSelectionChanged}
+          />
+        </div>
+
+        <div>
+          {getComponentForInput(
+            {
+              description: description,
+              name: `${formatType(selectedInputTypeDefintion)}`,
+              label: '',
+              required: required,
+              typeDefinition: selectedInputTypeDefintion,
+              onChange: handleSubTypeOnChange,
+              value: inputTypeToValueMap[selectedInputType]?.value,
+              error: error,
+            } as InputProps,
+            true,
+          )}
+        </div>
+      </CardContent>
+    </Card>
+  );
+};

packages/zapp/console/src/components/Launch/LaunchForm/constants.ts

@@ -45,6 +45,7 @@ export const typeLabels: { [k in InputType]: string } = {
   [InputType.Schema]: 'schema - uri',
   [InputType.String]: 'string',
   [InputType.Struct]: 'struct',
+  [InputType.Union]: '',
   [InputType.Unknown]: 'unknown',
 };
 

packages/zapp/console/src/components/Launch/LaunchForm/inputHelpers/getHelperForInput.ts

@@ -11,6 +11,7 @@ import { noneHelper } from './none';
 import { schemaHelper } from './schema';
 import { stringHelper } from './string';
 import { structHelper } from './struct';
+import { unionHelper } from './union';
 import { InputHelper } from './types';
 
 const unsupportedHelper = noneHelper;
@@ -32,6 +33,7 @@ const inputHelpers: Record<InputType, InputHelper> = {
   [InputType.Schema]: schemaHelper,
   [InputType.String]: stringHelper,
   [InputType.Struct]: structHelper,
+  [InputType.Union]: unionHelper,
   [InputType.Unknown]: unsupportedHelper,
 };
 

packages/zapp/console/src/components/Launch/LaunchForm/inputHelpers/union.ts

@@ -0,0 +1,62 @@
+import { Core } from 'flyteidl';
+import { isObject } from 'lodash';
+import { InputTypeDefinition, InputValue, UnionValue } from '../types';
+import { getHelperForInput } from './getHelperForInput';
+import { ConverterInput, InputHelper, InputValidatorParams } from './types';
+import t from '../../../common/strings';
+
+function fromLiteral(literal: Core.ILiteral, inputTypeDefinition: InputTypeDefinition): InputValue {
+  const { listOfSubTypes } = inputTypeDefinition;
+  if (!listOfSubTypes?.length) {
+    throw new Error(t('missingUnionListOfSubType'));
+  }
+
+  // loop though the subtypes to find the correct match literal type
+  for (let i = 0; i < listOfSubTypes.length; i++) {
+    try {
+      const value = getHelperForInput(listOfSubTypes[i].type).fromLiteral(
+        literal,
+        listOfSubTypes[i],
+      );
+      return { value, typeDefinition: listOfSubTypes[i] } as UnionValue;
+    } catch (error) {
+      // do nothing here it's expected to have error
+    }
+  }
+  throw new Error(t('noMatchingResults'));
+}
+
+function toLiteral({ value, typeDefinition: { listOfSubTypes } }: ConverterInput): Core.ILiteral {
+  if (!listOfSubTypes) {
+    throw new Error(t('missingUnionListOfSubType'));
+  }
+
+  if (!isObject(value)) {
+    throw new Error(t('valueMustBeObject'));
+  }
+
+  const { value: unionValue, typeDefinition } = value as UnionValue;
+
+  return getHelperForInput(typeDefinition.type).toLiteral({
+    value: unionValue,
+    typeDefinition: typeDefinition,
+  } as ConverterInput);
+}
+
+function validate({ value, typeDefinition: { listOfSubTypes } }: InputValidatorParams) {
+  if (!value) {
+    throw new Error(t('valueRequired'));
+  }
+  if (!isObject(value)) {
+    throw new Error(t('valueMustBeObject'));
+  }
+
+  const { typeDefinition } = value as UnionValue;
+  getHelperForInput(typeDefinition.type).validate(value as InputValidatorParams);
+}
+
+export const unionHelper: InputHelper = {
+  fromLiteral,
+  toLiteral,
+  validate,
+};

packages/zapp/console/src/components/Launch/LaunchForm/inputHelpers/utils.ts

@@ -1,18 +1,17 @@
 import { assertNever, stringifyValue } from 'common/utils';
 import { Core } from 'flyteidl';
-import { get } from 'lodash';
+import { get, has } from 'lodash';
 import { BlobDimensionality } from 'models/Common/types';
 import { InputType, InputTypeDefinition } from '../types';
 
 /** Performs a deep get of `path` on the given `Core.ILiteral`. Will throw
  * if the given property doesn't exist.
  */
 export function extractLiteralWithCheck<T>(literal: Core.ILiteral, path: string): T {
-  const value = get(literal, path);
-  if (value === undefined) {
+  if (!has(literal, path)) {
     throw new Error(`Failed to extract literal value with path ${path}`);
   }
-  return value as T;
+  return get(literal, path) as T;
 }
 
 /** Converts a value within a collection to the appropriate string
@@ -29,7 +28,7 @@ export function collectionChildToString(type: InputType, value: any) {
  * supported for use in the Launch form.
  */
 export function typeIsSupported(typeDefinition: InputTypeDefinition): boolean {
-  const { type, subtype } = typeDefinition;
+  const { type, subtype, listOfSubTypes } = typeDefinition;
   switch (type) {
     case InputType.Binary:
     case InputType.Error:
@@ -47,6 +46,17 @@ export function typeIsSupported(typeDefinition: InputTypeDefinition): boolean {
     case InputType.String:
     case InputType.Struct:
       return true;
+    case InputType.Union:
+      if (listOfSubTypes?.length) {
+        var isSupported = true;
+        listOfSubTypes.forEach((subtype) => {
+          if (!typeIsSupported(subtype)) {
+            isSupported = false;
+          }
+        });
+        return isSupported;
+      }
+      return false;
     case InputType.Map:
       if (!subtype) {
         return false;

packages/zapp/console/src/components/Launch/LaunchForm/types.ts

@@ -159,13 +159,15 @@ export enum InputType {
   Schema = 'SCHEMA',
   String = 'STRING',
   Struct = 'STRUCT',
+  Union = 'Union',
   Unknown = 'UNKNOWN',
 }
 
 export interface InputTypeDefinition {
   literalType: LiteralType;
   type: InputType;
   subtype?: InputTypeDefinition;
+  listOfSubTypes?: InputTypeDefinition[];
 }
 
 export interface BlobValue {
@@ -174,7 +176,12 @@ export interface BlobValue {
   uri: string;
 }
 
-export type InputValue = string | number | boolean | Date | BlobValue;
+export interface UnionValue {
+  value: InputValue;
+  typeDefinition: InputTypeDefinition;
+}
+
+export type InputValue = string | number | boolean | Date | BlobValue | UnionValue;
 export type InputChangeHandler = (newValue: InputValue) => void;
 
 export interface InputProps {

packages/zapp/console/src/components/Launch/LaunchForm/utils.ts

@@ -75,13 +75,25 @@ export function getTaskInputs(task: Task): Record<string, Variable> {
 /** Returns a formatted string based on an InputTypeDefinition.
  * ex. `string`, `string[]`, `map<string, number>`
  */
-export function formatType({ type, subtype }: InputTypeDefinition): string {
+export function formatType({ type, subtype, listOfSubTypes }: InputTypeDefinition): string {
   if (type === InputType.Collection) {
     return subtype ? `${formatType(subtype)}[]` : 'collection';
   }
   if (type === InputType.Map) {
     return subtype ? `map<string, ${formatType(subtype)}>` : 'map';
   }
+  if (type === InputType.Union) {
+    if (!listOfSubTypes) return 'union';
+
+    const concatListOfSubTypes = listOfSubTypes.reduce((previous, currentValue, currentIndex) => {
+      if (currentIndex === listOfSubTypes.length - 1) {
+        return previous + `${formatType(currentValue)}`;
+      }
+      return previous + `${formatType(currentValue)} |`;
+    }, '');
+
+    return `union [${concatListOfSubTypes}]`;
+  }
   return typeLabels[type];
 }
 
@@ -157,6 +169,11 @@ export function getInputDefintionForLiteralType(literalType: LiteralType): Input
     result.type = simpleTypeToInputType[literalType.simple];
   } else if (literalType.enumType) {
     result.type = InputType.Enum;
+  } else if (literalType.unionType) {
+    result.type = InputType.Union;
+    result.listOfSubTypes = literalType.unionType.variants?.map((variant) =>
+      getInputDefintionForLiteralType(variant as LiteralType),
+    );
   }
   return result;
 }