chore(Queries): store [YTFRONT-4612]

ytsaurus · Feb 4, 2025 · 533f16c · 533f16c
1 parent ced7ff9
commit 533f16c
Show file tree

Hide file tree

Showing 6 changed files with 85 additions and 40 deletions.
diff --git a/packages/ui/src/ui/pages/query-tracker/module/index.ts b/packages/ui/src/ui/pages/query-tracker/module/index.ts
@@ -7,6 +7,7 @@ import {queryFilesFormReducer} from './queryFilesForm/queryFilesFormSlice';
 import {vcsReducer} from './vcs/vcsSlice';
 import {queryNavigationReducer} from './queryNavigation/queryNavigationSlice';
 import {queryChartReducer} from './queryChart/queryChartSlice';
+import {neuralNetworkReducer} from './neuralNetwork/neuralNetworkSlice';
 
 export const queryTracker = combineReducers({
     list: listReducer,
@@ -17,4 +18,5 @@ export const queryTracker = combineReducers({
     vcs: vcsReducer,
     queryNavigation: queryNavigationReducer,
     queryChart: queryChartReducer,
+    neuralNetwork: neuralNetworkReducer,
 });
diff --git a/packages/ui/src/ui/pages/query-tracker/module/neuralNetwork/actions.ts b/packages/ui/src/ui/pages/query-tracker/module/neuralNetwork/actions.ts
@@ -0,0 +1,28 @@
+import {ThunkAction} from 'redux-thunk';
+import {RootState} from '../../../../store/reducers';
+import {wrapApiPromiseByToaster} from '../../../../utils/utils';
+import {YTApiId, ytApiV4Id} from '../../../../rum/rum-wrap-api';
+import {setContextId} from './neuralNetworkSlice';
+
+export const createContextId =
+    (queryId: string): ThunkAction<void, RootState, unknown, any> =>
+    async (dispatch) => {
+        const newContextId = crypto.randomUUID();
+
+        await wrapApiPromiseByToaster(
+            ytApiV4Id.alterQuery(YTApiId.alterQuery, {
+                parameters: {
+                    query_id: queryId,
+                    annotations: {
+                        contextId: newContextId,
+                    },
+                },
+            }),
+            {
+                toasterName: `context_id`,
+                skipSuccessToast: true,
+                errorTitle: `Failed to set query contect id`,
+            },
+        );
+        dispatch(setContextId(newContextId));
+    };
diff --git a/packages/ui/src/ui/pages/query-tracker/module/neuralNetwork/neuralNetworkSlice.ts b/packages/ui/src/ui/pages/query-tracker/module/neuralNetwork/neuralNetworkSlice.ts
@@ -0,0 +1,33 @@
+import {PayloadAction, createSlice} from '@reduxjs/toolkit';
+
+type NeuralNetworkState = {
+    suggestions: string[];
+    requestId: string;
+    contextId: string;
+};
+
+const initialState: NeuralNetworkState = {
+    suggestions: [],
+    requestId: '',
+    contextId: '',
+};
+
+const neuralNetworkSlice = createSlice({
+    initialState,
+    name: 'neuralNetwork',
+    reducers: {
+        setSuggestions: (state, {payload}: PayloadAction<string[]>) => {
+            state.suggestions = payload;
+        },
+        setRequestId: (state, {payload}: PayloadAction<string>) => {
+            state.requestId = payload;
+        },
+        setContextId: (state, {payload}: PayloadAction<string>) => {
+            state.contextId = payload;
+        },
+    },
+});
+
+export const {setSuggestions, setRequestId, setContextId} = neuralNetworkSlice.actions;
+
+export const neuralNetworkReducer = neuralNetworkSlice.reducer;
diff --git a/packages/ui/src/ui/pages/query-tracker/module/neuralNetwork/selectors.ts b/packages/ui/src/ui/pages/query-tracker/module/neuralNetwork/selectors.ts
@@ -0,0 +1,10 @@
+import {RootState} from '../../../../store/reducers';
+
+export const selectNeuralNetworkContextId = (state: RootState) =>
+    state.queryTracker.neuralNetwork.contextId;
+
+export const selectNeuralNetworkRequestId = (state: RootState) =>
+    state.queryTracker.neuralNetwork.requestId;
+
+export const selectNeuralNetworkSuggestions = (state: RootState) =>
+    state.queryTracker.neuralNetwork.suggestions;
diff --git a/packages/ui/src/ui/pages/query-tracker/module/query/actions.ts b/packages/ui/src/ui/pages/query-tracker/module/query/actions.ts
@@ -38,6 +38,8 @@ import {chytApiAction, spytApiAction} from '../../../../utils/strawberryControll
 import guid from '../../../../common/hammer/guid';
 import {getSettingQueryTrackerStage} from '../../../../store/selectors/settings/settings-ts';
 import {getDefaultQueryACO, selectIsMultipleAco} from '../query_aco/selectors';
+import {createContextId} from '../neuralNetwork/actions';
+import {setContextId} from '../neuralNetwork/neuralNetworkSlice';
 
 import {
     REQUEST_QUERY,
@@ -50,7 +52,7 @@ import {
     UPDATE_QUERY,
 } from '../query-tracker-contants';
 import {loadVisualization} from '../queryChart/actions';
-import {YTApiId, ytApiV4Id} from '../../../../rum/rum-wrap-api';
+import {selectNeuralNetworkContextId} from '../neuralNetwork/selectors';
 
 export const setCurrentClusterToQuery =
     (): ThunkAction<void, RootState, unknown, any> => (dispatch, getState) => {
@@ -109,42 +111,10 @@ export const loadCliqueByCluster =
             });
     };
 
-export const setContextId =
-    (queryId: string): ThunkAction<void, RootState, unknown, any> =>
-    async (dispatch, getState) => {
-        const draft = getQueryDraft(getState());
-        const newContextId = crypto.randomUUID();
-
-        await wrapApiPromiseByToaster(
-            ytApiV4Id.alterQuery(YTApiId.alterQuery, {
-                parameters: {
-                    query_id: queryId,
-                    annotations: {
-                        contextId: newContextId,
-                    },
-                },
-            }),
-            {
-                toasterName: `context_id`,
-                skipSuccessToast: true,
-                errorTitle: `Failed to set query contect id`,
-            },
-        );
-        dispatch(
-            updateQueryDraft({
-                ...draft,
-                annotations: {
-                    ...draft.annotations,
-                    contextId: newContextId,
-                },
-            }),
-        );
-    };
-
 export function loadQuery(
     queryId: string,
     config?: {dontReplaceQueryText?: boolean},
-): ThunkAction<any, RootState, any, SetQueryAction | RequestQueryAction | SetQueryErrorLoadAction> {
+): ThunkAction<any, RootState, any, any> {
     return async (dispatch, getState) => {
         const state = getState();
         const stage = getSettingQueryTrackerStage(state);
@@ -169,7 +139,9 @@ export function loadQuery(
             }
 
             if (!queryItem.annotations?.contextId) {
-                dispatch(setContextId(queryId));
+                dispatch(createContextId(queryId));
+            } else {
+                dispatch(setContextId(queryItem.annotations?.contextId));
             }
 
             dispatch({
@@ -248,18 +220,16 @@ export function createEmptyQuery(
     engine = QueryEngine.YQL,
     query?: string,
     settings?: Record<string, string>,
-): ThunkAction<any, RootState, any, SetQueryAction> {
+): ThunkAction<any, RootState, any, any> {
     return (dispatch, getState) => {
         const state = getState();
         const defaultQueryACO = getDefaultQueryACO(state);
 
+        dispatch(setContextId(crypto.randomUUID()));
         dispatch({
             type: SET_QUERY,
             data: {
                 initialQuery: {
-                    annotations: {
-                        contextId: crypto.randomUUID(),
-                    },
                     access_control_object: defaultQueryACO,
                     access_control_objects: [defaultQueryACO],
                     query: query || '',
@@ -279,8 +249,9 @@ export function runQuery(
     return async (dispatch, getState) => {
         const state = getState();
         const query = getQueryDraft(state);
+        const contextId = selectNeuralNetworkContextId(state);
 
-        const newQuery = {...query};
+        const newQuery = {...query, annotations: {...query.annotations, contextId}};
         if ('access_control_objects' in newQuery) {
             newQuery.access_control_objects = newQuery.access_control_objects?.filter(
                 (i) => i !== SHARED_QUERY_ACO,

diff --git a/packages/ui/src/ui/pages/query-tracker/module/query/selectors.ts b/packages/ui/src/ui/pages/query-tracker/module/query/selectors.ts
@@ -20,6 +20,7 @@ export const DEFAULT_QUERY_ACO = 'nobody';
 export const SHARED_QUERY_ACO = 'everyone-share';
 
 export const getQuery = (state: RootState) => getState(state).queryItem;
+export const getQueryId = (state: RootState) => getState(state).queryItem?.id;
 export const getQueryGetParams = (state: RootState) => getState(state).params;
 
 export const getQueryDraft = (state: RootState) => getState(state).draft;