From 09eff80f83bbcd63eff2d1c820c66a94fdf09834 Mon Sep 17 00:00:00 2001
From: Viktor Rozaev <vrozaev@tracto.ai>
Date: Wed, 30 Oct 2024 14:37:43 +0100
Subject: [PATCH] feat(PoolTreeSuggestControl): now we could select multiply
 pool trees on clique creation/edit [#841]

---
 .../PoolSuggestControl/PoolSuggestControl.tsx | 24 +++++----
 .../PoolTreeSuggestControl.tsx                | 14 ++++--
 .../ui/components/Dialog/df-dialog-utils.ts   | 50 +++++++++++++++++--
 packages/ui/src/ui/hooks/global-pool-trees.ts |  2 +-
 .../ChytPageClique/ChytPageCliqueSpeclet.tsx  | 36 +++++++++----
 .../ui/src/ui/pages/chyt/ChytPageTopRow.tsx   |  9 ++--
 .../TableMergeSortModal/TableMergeModal.tsx   |  8 +--
 .../TableMergeSortModal/TableSortModal.tsx    |  8 +--
 8 files changed, 109 insertions(+), 42 deletions(-)

diff --git a/packages/ui/src/ui/components/Dialog/controls/PoolSuggestControl/PoolSuggestControl.tsx b/packages/ui/src/ui/components/Dialog/controls/PoolSuggestControl/PoolSuggestControl.tsx
index 1d4e84094..8db8701e8 100644
--- a/packages/ui/src/ui/components/Dialog/controls/PoolSuggestControl/PoolSuggestControl.tsx
+++ b/packages/ui/src/ui/components/Dialog/controls/PoolSuggestControl/PoolSuggestControl.tsx
@@ -57,11 +57,13 @@ export function PoolSuggestControl(props: Props) {
 
     const [poolNames, setPoolNames] = React.useState<Array<string>>([]);
 
-    const loadedPools = useLoadedPools(cluster, poolTrees);
+    const names = useLoadedPools(cluster, poolTrees);
 
     React.useEffect(
         function onPoolLoaded() {
-            const {names} = loadedPools;
+            if (!names) {
+                return;
+            }
             const noRoot = filter_(names, (item) => '<Root>' !== item);
             const valueIndex = indexOf_(noRoot, value);
             if (value && -1 === valueIndex) {
@@ -73,7 +75,7 @@ export function PoolSuggestControl(props: Props) {
             }
         },
         // value should not affect the useEffect
-        [loadedPools, setPoolNames, onChange, calculateValueOnPoolsLoaded /*, value */],
+        [names && names.join(), setPoolNames, onChange, calculateValueOnPoolsLoaded /*, value */],
     );
 
     const getItems = React.useCallback(
@@ -121,13 +123,19 @@ PoolSuggestControl.isEmpty = (value: Props['value']) => {
     return !value;
 };
 
-function useLoadedPools(cluster?: string, poolTrees: string[] = []): {names: Array<string>} {
-    const [poolNames, setPoolNames] = React.useState<Array<string>>([]);
+function useLoadedPools(cluster?: string, poolTrees?: string[]): Array<string> | null {
+    poolTrees = poolTrees || [];
+
+    const [poolNames, setPoolNames] = React.useState<Array<string> | null>(null);
 
     const defaultPoolTree = useDefaultPoolTree();
 
     React.useMemo(() => {
-        const localPoolTrees: string[] = poolTrees.length ? poolTrees : defaultPoolTree ? [defaultPoolTree] : [];
+        const localPoolTrees: string[] = poolTrees.length
+            ? poolTrees
+            : defaultPoolTree
+              ? [defaultPoolTree]
+              : [];
 
         if (!localPoolTrees.length) {
             return;
@@ -150,9 +158,7 @@ function useLoadedPools(cluster?: string, poolTrees: string[] = []): {names: Arr
         });
     }, [cluster, defaultPoolTree, poolTrees.join()]);
 
-    return {
-        names: poolNames,
-    };
+    return poolNames;
 }
 
 function fetchPoolNamesByCluster(cluster: string, poolTrees: string[]): Promise<string[]> {
diff --git a/packages/ui/src/ui/components/Dialog/controls/PoolTreeSuggestControl/PoolTreeSuggestControl.tsx b/packages/ui/src/ui/components/Dialog/controls/PoolTreeSuggestControl/PoolTreeSuggestControl.tsx
index 809234302..253e98bba 100644
--- a/packages/ui/src/ui/components/Dialog/controls/PoolTreeSuggestControl/PoolTreeSuggestControl.tsx
+++ b/packages/ui/src/ui/components/Dialog/controls/PoolTreeSuggestControl/PoolTreeSuggestControl.tsx
@@ -6,8 +6,9 @@ import {Select} from '@gravity-ui/uikit';
 
 import {DialogControlProps} from '../../../../components/Dialog/Dialog.types';
 import {getAllPoolTreeNames} from '../../../../store/selectors/global';
+import {usePoolTreesLoaded} from '../../../../hooks/global-pool-trees';
 
-type Props = DialogControlProps<string> & {
+type Props = DialogControlProps<string[]> & {
     disabled?: boolean;
 };
 
@@ -15,6 +16,8 @@ export function PoolTreeSuggestControl(props: Props) {
     const {value, onChange, disabled, placeholder} = props;
     const treeNames = useSelector(getAllPoolTreeNames);
 
+    usePoolTreesLoaded();
+
     const items = React.useMemo(() => {
         return map_(treeNames, (value) => {
             return {value, content: value};
@@ -24,20 +27,21 @@ export function PoolTreeSuggestControl(props: Props) {
     return (
         <Select
             disabled={disabled}
-            value={[value]}
+            value={value}
             options={items}
-            onUpdate={(values) => onChange(values[0])}
+            onUpdate={onChange}
             placeholder={placeholder}
             width="max"
             filterable={items?.length > 5}
+            multiple
         />
     );
 }
 
 PoolTreeSuggestControl.getDefaultValue = () => {
-    return '';
+    return [];
 };
 
 PoolTreeSuggestControl.isEmpty = (value: Props['value']) => {
-    return !value;
+    return !value.length;
 };
diff --git a/packages/ui/src/ui/components/Dialog/df-dialog-utils.ts b/packages/ui/src/ui/components/Dialog/df-dialog-utils.ts
index db41569f0..38584c740 100644
--- a/packages/ui/src/ui/components/Dialog/df-dialog-utils.ts
+++ b/packages/ui/src/ui/components/Dialog/df-dialog-utils.ts
@@ -17,7 +17,8 @@ export type OptionDescription =
           min_value?: number;
       })
     | Option<'yson', JsonAsString>
-    | Option<'path' | 'pool', string>;
+    | Option<'path' | 'pool', string>
+    | Option<'pool_trees', string[]>;
 
 export type JsonAsString = string;
 
@@ -31,7 +32,7 @@ export type Option<TypeName extends string, T> = {
 
 export function descriptionToDialogField<T = unknown>(
     item: OptionDescription,
-    {unipikaSettings, allowEdit}: MakeDialogFieldsOptions,
+    {unipikaSettings, allowEdit, defaultPoolTree}: MakeDialogFieldsOptions,
 ): DialogField<T> & {initialValue?: unknown; converter: Converter} {
     const common = {
         name: item.name,
@@ -115,7 +116,19 @@ export function descriptionToDialogField<T = unknown>(
         case 'path':
             return {...common, type: 'path', extras};
         case 'pool':
-            return {...common, type: 'pool', extras: {...extras, allowEmpty: true}};
+            return {
+                ...common,
+                type: 'pool',
+                extras: {...extras, allowEmpty: true},
+                initialValue: item.current_value ?? item.default_value,
+            };
+        case 'pool_trees':
+            return {
+                ...common,
+                type: 'pool-tree',
+                extras,
+                initialValue: item.current_value ?? item.default_value ?? [defaultPoolTree],
+            };
         default:
             return {...common, type: 'plain'};
     }
@@ -195,6 +208,7 @@ function makeDialogField<FormValues = any>(
 type MakeDialogFieldsOptions = {
     allowEdit: boolean;
     unipikaSettings: UnipikaSettings;
+    defaultPoolTree: string;
 };
 
 export function makeDialogFieldsFromDescription<
@@ -253,3 +267,33 @@ export function makeTabbedDialogFieldsFromDescription<
         }),
     };
 }
+
+export function linkPoolWithPoolTree<
+    FormValues extends Record<string, Record<string, unknown>> = Record<
+        string,
+        Record<string, unknown>
+    >,
+>(data: ReturnType<typeof makeTabbedDialogFieldsFromDescription>) {
+    for (const group of Object.values(data.fields)) {
+        let pool: any;
+        let poolTrees: any;
+
+        (group.fields as Array<any>).some((field) => {
+            if (field.name === 'pool_trees') {
+                poolTrees = field;
+            }
+            if (field.name === 'pool') {
+                pool = field;
+            }
+            return Boolean(pool && poolTrees);
+        });
+
+        if (pool && poolTrees) {
+            const extras = pool.extras;
+            pool.extras = (values: FormValues) => ({
+                ...extras,
+                poolTrees: values[group.name].pool_trees,
+            });
+        }
+    }
+}
diff --git a/packages/ui/src/ui/hooks/global-pool-trees.ts b/packages/ui/src/ui/hooks/global-pool-trees.ts
index 781cbe7c6..fa53ce2a4 100644
--- a/packages/ui/src/ui/hooks/global-pool-trees.ts
+++ b/packages/ui/src/ui/hooks/global-pool-trees.ts
@@ -4,7 +4,7 @@ import {useDispatch} from 'react-redux';
 import {loadPoolTreesIfNotLoaded} from '../store/actions/global';
 import {loadDefaultPoolTree} from '../utils/poolTrees';
 
-function usePoolTreesLoaded() {
+export function usePoolTreesLoaded() {
     const dispatch = useDispatch();
 
     React.useEffect(() => {
diff --git a/packages/ui/src/ui/pages/chyt/ChytPageClique/ChytPageCliqueSpeclet.tsx b/packages/ui/src/ui/pages/chyt/ChytPageClique/ChytPageCliqueSpeclet.tsx
index e92e90981..695673328 100644
--- a/packages/ui/src/ui/pages/chyt/ChytPageClique/ChytPageCliqueSpeclet.tsx
+++ b/packages/ui/src/ui/pages/chyt/ChytPageClique/ChytPageCliqueSpeclet.tsx
@@ -9,7 +9,10 @@ import Icon from '../../../components/Icon/Icon';
 import {YTDFDialog} from '../../../components/Dialog';
 import Yson from '../../../components/Yson/Yson';
 import {UnipikaSettings} from '../../../components/Yson/StructuredYson/StructuredYsonTypes';
-import {makeTabbedDialogFieldsFromDescription} from '../../../components/Dialog/df-dialog-utils';
+import {
+    linkPoolWithPoolTree,
+    makeTabbedDialogFieldsFromDescription,
+} from '../../../components/Dialog/df-dialog-utils';
 
 import {useUpdater} from '../../../hooks/use-updater';
 
@@ -31,6 +34,7 @@ import {
 } from '../../../store/selectors/chyt/speclet';
 import {useThunkDispatch} from '../../../store/thunkDispatch';
 import {YTError} from '../../../../@types/types';
+import {WaitForDefaultPoolTree} from '../../../hooks/global-pool-trees';
 
 import './ChytPageCliqueSpeclet.scss';
 
@@ -124,14 +128,19 @@ export function ChytSpecletEditButton({
     return (
         <React.Fragment>
             {!visible || !specletData ? null : (
-                <ChytSpecletEditDialog
-                    key={dataAlias}
-                    data={specletData}
-                    alias={alias}
-                    unipikaSettings={unipikaSettings}
-                    allowEdit={dataAlias === alias}
-                    onClose={() => setVisible(false)}
-                />
+                <WaitForDefaultPoolTree>
+                    {({defaultPoolTree}) => (
+                        <ChytSpecletEditDialog
+                            key={dataAlias}
+                            data={specletData}
+                            alias={alias}
+                            unipikaSettings={unipikaSettings}
+                            allowEdit={dataAlias === alias}
+                            onClose={() => setVisible(false)}
+                            defaultPoolTree={defaultPoolTree}
+                        />
+                    )}
+                </WaitForDefaultPoolTree>
             )}
             <Button
                 view={compact ? 'outlined' : undefined}
@@ -153,21 +162,28 @@ function ChytSpecletEditDialog({
     allowEdit,
     unipikaSettings,
     onClose,
+    defaultPoolTree,
 }: {
     allowEdit: boolean;
     alias: string;
     data: ChytCliqueOptionsState['data'];
     unipikaSettings: UnipikaSettings;
     onClose: () => void;
+    defaultPoolTree: string;
 }) {
     const dispatch = useThunkDispatch();
     const [error, setError] = React.useState<YTError | undefined>();
 
     const {fields, initialValues, fieldTypeByName} = React.useMemo(() => {
-        return makeTabbedDialogFieldsFromDescription(data ?? [], {
+        const groups = makeTabbedDialogFieldsFromDescription(data ?? [], {
             allowEdit,
             unipikaSettings,
+            defaultPoolTree,
         });
+
+        linkPoolWithPoolTree(groups);
+
+        return groups;
     }, [data, allowEdit, unipikaSettings]);
 
     return (
diff --git a/packages/ui/src/ui/pages/chyt/ChytPageTopRow.tsx b/packages/ui/src/ui/pages/chyt/ChytPageTopRow.tsx
index 6e32ac4a5..8a358869f 100644
--- a/packages/ui/src/ui/pages/chyt/ChytPageTopRow.tsx
+++ b/packages/ui/src/ui/pages/chyt/ChytPageTopRow.tsx
@@ -178,7 +178,7 @@ function ChytAliasSuggest({
 type FormValues = {
     alias: string;
     instance_count: number;
-    tree: string;
+    tree: string[];
     pool: string;
     runAfterCreation: boolean;
 };
@@ -246,9 +246,6 @@ function CreateChytButton() {
                                     name: 'tree',
                                     type: 'pool-tree',
                                     caption: 'Pool tree',
-                                    extras: {
-                                        disabled: true,
-                                    },
                                 },
                                 {
                                     name: 'pool',
@@ -256,7 +253,7 @@ function CreateChytButton() {
                                     caption: 'Pool',
                                     extras: ({tree}: FormValues) => {
                                         return {
-                                            poolTrees: [tree],
+                                            poolTrees: tree,
                                             placeholder: 'Pool name...',
                                             allowEmpty: true,
                                         };
@@ -294,7 +291,7 @@ function CreateChytButton() {
                             ]}
                             initialValues={{
                                 instance_count: 1,
-                                tree: defaultPoolTree,
+                                tree: [defaultPoolTree],
                                 runAfterCreation: true,
                             }}
                         />
diff --git a/packages/ui/src/ui/pages/navigation/modals/TableMergeSortModal/TableMergeModal.tsx b/packages/ui/src/ui/pages/navigation/modals/TableMergeSortModal/TableMergeModal.tsx
index d194d2eb7..5d4c663e8 100644
--- a/packages/ui/src/ui/pages/navigation/modals/TableMergeSortModal/TableMergeModal.tsx
+++ b/packages/ui/src/ui/pages/navigation/modals/TableMergeSortModal/TableMergeModal.tsx
@@ -54,7 +54,7 @@ export default function TableMergeModal() {
                 } = values;
                 const chunkSizeBytes = parseBytes(chunkSize);
                 const data_size_per_job = isNaN(chunkSizeBytes) ? undefined : chunkSizeBytes;
-                const pool_trees = poolTree ? [poolTree] : undefined;
+                const pool_trees = poolTree.length ? poolTree : undefined;
                 await dispatch(
                     runTableMerge(
                         pickBy_(
@@ -114,7 +114,7 @@ export default function TableMergeModal() {
                         outputPath,
                         columns: [],
                         force_transform: true,
-                        poolTree: defaultPoolTree,
+                        poolTree: [defaultPoolTree],
                         combine_chunks: true,
                     }}
                     fields={[
@@ -194,7 +194,7 @@ export default function TableMergeModal() {
                             ),
                             extras: ({poolTree}: FormValues) => {
                                 return {
-                                    poolTrees: [poolTree],
+                                    poolTrees: poolTree,
                                     placeholder: login,
                                     allowEphemeral: true,
                                 };
@@ -237,7 +237,7 @@ export default function TableMergeModal() {
 interface FormValues {
     mode: string;
     pool: string;
-    poolTree: string;
+    poolTree: string[];
 
     paths: Array<string>;
     outputPath: string;
diff --git a/packages/ui/src/ui/pages/navigation/modals/TableMergeSortModal/TableSortModal.tsx b/packages/ui/src/ui/pages/navigation/modals/TableMergeSortModal/TableSortModal.tsx
index 575cf8fec..28c630f8d 100644
--- a/packages/ui/src/ui/pages/navigation/modals/TableMergeSortModal/TableSortModal.tsx
+++ b/packages/ui/src/ui/pages/navigation/modals/TableMergeSortModal/TableSortModal.tsx
@@ -63,7 +63,7 @@ export default function TableSortModal() {
                         }),
                     },
                     pool ? {pool} : {},
-                    poolTree ? {pool_trees: [poolTree]} : {},
+                    poolTree.length ? {pool_trees: poolTree} : {},
                 );
 
                 await dispatch(runTableSort(spec));
@@ -125,7 +125,7 @@ export default function TableSortModal() {
                         paths,
                         outputPath,
                         columns: [],
-                        poolTree: defaultPoolTree,
+                        poolTree: [defaultPoolTree],
                     }}
                     fields={[
                         {
@@ -183,7 +183,7 @@ export default function TableSortModal() {
                                 const {poolTree} = values;
                                 return {
                                     placeholder: login,
-                                    poolTrees: [poolTree],
+                                    poolTrees: poolTree,
                                     allowEphemeral: true,
                                 };
                             },
@@ -200,6 +200,6 @@ interface FormValues {
     paths: Array<string>;
     outputPath: string;
     columns: Array<string>;
-    poolTree: string;
+    poolTree: string[];
     pool: string;
 }