fix(schema): use one cache entree for tenant path schemas (#1016)

ValeraS · web-flow · commit 7f283a0f595d · 2024-07-16T11:44:47.000+02:00
diff --git a/src/containers/Tenant/Diagnostics/HotKeys/HotKeys.tsx b/src/containers/Tenant/Diagnostics/HotKeys/HotKeys.tsx
@@ -8,7 +8,7 @@ import {Button, Card, Icon} from '@gravity-ui/uikit';
 import {ResponseError} from '../../../../components/Errors/ResponseError';
 import {ResizeableDataTable} from '../../../../components/ResizeableDataTable/ResizeableDataTable';
 import {hotKeysApi} from '../../../../store/reducers/hotKeys/hotKeys';
-import {schemaApi} from '../../../../store/reducers/schema/schema';
+import {useGetSchemaQuery} from '../../../../store/reducers/schema/schema';
 import type {HotKey} from '../../../../types/api/hotkeys';
 import {cn} from '../../../../utils/cn';
 import {DEFAULT_TABLE_SETTINGS, IS_HOTKEYS_HELP_HIDDEN_KEY} from '../../../../utils/constants';
@@ -61,9 +61,7 @@ export function HotKeys({path}: HotKeysProps) {
     const {currentData: data, isFetching, error} = hotKeysApi.useGetHotKeysQuery({path});
     const loading = isFetching && data === undefined;
 
-    const {currentData: schemaData, isFetching: schemaIsFetching} =
-        schemaApi.endpoints.getSchema.useQueryState({path});
-    const schemaLoading = schemaIsFetching && schemaData === undefined;
+    const {data: schemaData, isLoading: schemaLoading} = useGetSchemaQuery({path});
 
     const keyColumnsIds = schemaData?.PathDescription?.Table?.KeyColumnNames;
 
diff --git a/src/containers/Tenant/Diagnostics/Overview/Overview.tsx b/src/containers/Tenant/Diagnostics/Overview/Overview.tsx
@@ -8,7 +8,10 @@ import {TableIndexInfo} from '../../../../components/InfoViewer/schemaInfo';
 import {Loader} from '../../../../components/Loader';
 import {olapApi} from '../../../../store/reducers/olapStats';
 import {overviewApi} from '../../../../store/reducers/overview/overview';
-import {schemaApi, selectSchemaMergedChildrenPaths} from '../../../../store/reducers/schema/schema';
+import {
+    selectSchemaMergedChildrenPaths,
+    useGetSchemaQuery,
+} from '../../../../store/reducers/schema/schema';
 import {EPathType} from '../../../../types/api/schema';
 import {useAutoRefreshInterval, useTypedSelector} from '../../../../utils/hooks';
 import {ExternalDataSourceInfo} from '../../Info/ExternalDataSource/ExternalDataSource';
@@ -66,7 +69,7 @@ function Overview({type, path}: OverviewProps) {
     const overviewLoading = isFetching && currentData === undefined;
     const {data: rawData, additionalData} = currentData || {};
 
-    const {error: schemaError} = schemaApi.endpoints.getSchema.useQueryState({path});
+    const {error: schemaError} = useGetSchemaQuery({path});
 
     const entityLoading = overviewLoading || olapStatsLoading;
     const entityNotReady = isEntityWithMergedImpl && !mergedChildrenPaths;
diff --git a/src/containers/Tenant/ObjectSummary/ObjectSummary.tsx b/src/containers/Tenant/ObjectSummary/ObjectSummary.tsx
@@ -15,7 +15,7 @@ import {LinkWithIcon} from '../../../components/LinkWithIcon/LinkWithIcon';
 import {Loader} from '../../../components/Loader';
 import SplitPane from '../../../components/SplitPane';
 import routes, {createExternalUILink, createHref} from '../../../routes';
-import {schemaApi, setShowPreview} from '../../../store/reducers/schema/schema';
+import {setShowPreview, useGetSchemaQuery} from '../../../store/reducers/schema/schema';
 import {
     TENANT_PAGES_IDS,
     TENANT_QUERY_TABS_ID,
@@ -92,7 +92,7 @@ export function ObjectSummary({
         ignoreQueryPrefix: true,
     });
 
-    const {currentData: currentObjectData} = schemaApi.endpoints.getSchema.useQueryState({path});
+    const {data: currentObjectData} = useGetSchemaQuery({path});
     const currentSchemaData = currentObjectData?.PathDescription?.Self;
 
     React.useEffect(() => {
@@ -399,14 +399,14 @@ export function ObjectSummary({
 }
 
 function ObjectTree({tenantName, path}: {tenantName: string; path?: string}) {
-    const {currentData: tenantData = {}, isFetching} = schemaApi.useGetSchemaQuery({
+    const {data: tenantData = {}, isLoading} = useGetSchemaQuery({
         path: tenantName,
     });
     const pathData = tenantData?.PathDescription?.Self;
 
     const [, setCurrentPath] = useQueryParam('schema', StringParam);
 
-    if (!pathData && isFetching) {
+    if (!pathData && isLoading) {
         // If Loader isn't wrapped with div, SplitPane doesn't calculate panes height correctly
         return (
             <div>
diff --git a/src/containers/Tenant/Schema/SchemaTree/SchemaTree.tsx b/src/containers/Tenant/Schema/SchemaTree/SchemaTree.tsx
@@ -7,7 +7,7 @@ import {NavigationTree} from 'ydb-ui-components';
 
 import {USE_DIRECTORY_OPERATIONS} from '../../../../lib';
 import {schemaApi} from '../../../../store/reducers/schema/schema';
-import type {EPathType} from '../../../../types/api/schema';
+import type {EPathType, TEvDescribeSchemeResult} from '../../../../types/api/schema';
 import {useQueryModes, useSetting, useTypedDispatch} from '../../../../utils/hooks';
 import {isChildlessPathType, mapPathTypeToNavigationTreeType} from '../../utils/schema';
 import {getActions} from '../../utils/schemaActions';
@@ -33,15 +33,26 @@ export function SchemaTree(props: SchemaTreeProps) {
     const [schemaTreeKey, setSchemaTreeKey] = React.useState('');
 
     const fetchPath = async (path: string) => {
-        const promise = dispatch(
-            schemaApi.endpoints.getSchema.initiate({path}, {forceRefetch: true}),
-        );
-        const {data} = await promise;
-        promise.unsubscribe();
-        if (!data) {
+        let schemaData: TEvDescribeSchemeResult | undefined;
+        do {
+            const promise = dispatch(
+                schemaApi.endpoints.getSchema.initiate({path}, {forceRefetch: true}),
+            );
+            const {data, originalArgs} = await promise;
+            promise.unsubscribe();
+            // Check if the result from the current request is received. rtk-query may skip the current request and
+            // return data from a parallel request, due to the same cache key.
+            if (originalArgs?.path === path) {
+                schemaData = data?.[path];
+                break;
+            }
+            // eslint-disable-next-line no-constant-condition
+        } while (true);
+
+        if (!schemaData) {
             throw new Error(`no describe data about path ${path}`);
         }
-        const {PathDescription: {Children = []} = {}} = data;
+        const {PathDescription: {Children = []} = {}} = schemaData;
 
         const childItems = Children.map((childData) => {
             const {Name = '', PathType, PathSubType} = childData;
diff --git a/src/containers/Tenant/Schema/SchemaViewer/SchemaViewer.tsx b/src/containers/Tenant/Schema/SchemaViewer/SchemaViewer.tsx
@@ -5,7 +5,7 @@ import {skipToken} from '@reduxjs/toolkit/query';
 
 import {ResizeableDataTable} from '../../../../components/ResizeableDataTable/ResizeableDataTable';
 import {TableSkeleton} from '../../../../components/TableSkeleton/TableSkeleton';
-import {schemaApi} from '../../../../store/reducers/schema/schema';
+import {useGetSchemaQuery} from '../../../../store/reducers/schema/schema';
 import {viewSchemaApi} from '../../../../store/reducers/viewSchema/viewSchema';
 import type {EPathType} from '../../../../types/api/schema';
 import {DEFAULT_TABLE_SETTINGS} from '../../../../utils/constants';
@@ -37,10 +37,7 @@ interface SchemaViewerProps {
 }
 
 export const SchemaViewer = ({type, path, tenantName, extended = false}: SchemaViewerProps) => {
-    const {currentData: schemaData, isFetching} = schemaApi.endpoints.getSchema.useQueryState({
-        path,
-    });
-    const loading = isFetching && schemaData === undefined;
+    const {data: schemaData, isLoading: loading} = useGetSchemaQuery({path});
 
     const viewSchemaRequestParams = isViewType(type) ? {path, database: tenantName} : skipToken;
 
diff --git a/src/containers/Tenant/Tenant.tsx b/src/containers/Tenant/Tenant.tsx
@@ -7,7 +7,7 @@ import {AccessDenied} from '../../components/Errors/403';
 import {Loader} from '../../components/Loader';
 import SplitPane from '../../components/SplitPane';
 import {setHeaderBreadcrumbs} from '../../store/reducers/header/header';
-import {schemaApi} from '../../store/reducers/schema/schema';
+import {useGetSchemaQuery} from '../../store/reducers/schema/schema';
 import type {AdditionalNodesProps, AdditionalTenantsProps} from '../../types/additionalProps';
 import {cn} from '../../utils/cn';
 import {DEFAULT_IS_TENANT_SUMMARY_COLLAPSED, DEFAULT_SIZE_TENANT_KEY} from '../../utils/constants';
@@ -74,11 +74,7 @@ export function Tenant(props: TenantProps) {
 
     const path = schema ?? tenantName;
 
-    const {
-        currentData: currentItem,
-        error,
-        isLoading,
-    } = schemaApi.useGetSchemaQuery({path}, {refetchOnMountOrArgChange: true});
+    const {data: currentItem, error, isLoading} = useGetSchemaQuery({path});
     const {PathType: currentPathType, PathSubType: currentPathSubType} =
         currentItem?.PathDescription?.Self || {};
 
diff --git a/src/store/reducers/schema/schema.ts b/src/store/reducers/schema/schema.ts
@@ -1,3 +1,5 @@
+import React from 'react';
+
 import type {Reducer, Selector} from '@reduxjs/toolkit';
 import {createSelector} from '@reduxjs/toolkit';
 
@@ -51,61 +53,60 @@ export const schemaApi = api.injectEndpoints({
                 }
             },
         }),
-        getSchema: builder.query<TEvDescribeSchemeResult & {partial?: boolean}, {path: string}>({
+        getSchema: builder.query<
+            {[path: string]: TEvDescribeSchemeResult & {partial?: boolean}},
+            {path: string}
+        >({
             queryFn: async ({path}, {signal}) => {
                 try {
                     const data = await window.api.getSchema({path}, {signal});
-                    return {data: data ?? {}};
+                    return {data: data ? {[path]: data, ...getSchemaChildren(data)} : {}};
                 } catch (error) {
                     return {error};
                 }
             },
             keepUnusedDataFor: Infinity,
-            forceRefetch: ({endpointState}) => {
-                const data = endpointState?.data;
-                if (data && typeof data === 'object' && 'partial' in data && data.partial) {
-                    return true;
-                }
-                return false;
+            serializeQueryArgs: ({queryArgs: {path}}) => {
+                const parts = path.split('/');
+                return {path: parts[0] || parts[1]};
             },
-            onQueryStarted: async ({path}, {dispatch, getState, queryFulfilled}) => {
-                const {data} = await queryFulfilled;
+            merge: (existing, incoming, {arg: {path}}) => {
+                const {[path]: data, ...children} = incoming;
                 if (data) {
-                    const state = getState();
-                    const {PathDescription: {Children = []} = {}} = data;
-                    for (const child of Children) {
-                        const {Name = ''} = child;
-                        const childPath = `${path}/${Name}`;
-                        const cachedData = schemaApi.endpoints.getSchema.select({path: childPath})(
-                            state,
-                        ).data;
-                        if (!cachedData) {
-                            // not full data, but it contains PathType, which ensures seamless switch between nodes
-                            dispatch(
-                                schemaApi.util.upsertQueryData(
-                                    'getSchema',
-                                    {path: childPath},
-                                    {PathDescription: {Self: child}, partial: true},
-                                ),
-                            );
-                        }
-                    }
+                    return {
+                        ...children,
+                        ...existing,
+                        [path]: data,
+                    };
                 }
+                return existing;
             },
         }),
     }),
     overrideExisting: 'throw',
 });
 
+function getSchemaChildren(data: TEvDescribeSchemeResult) {
+    const children: {[path: string]: TEvDescribeSchemeResult & {partial?: boolean}} = {};
+    const {PathDescription: {Children = []} = {}, Path: path} = data;
+    for (const child of Children) {
+        const {Name = ''} = child;
+        const childPath = `${path}/${Name}`;
+        children[childPath] = {PathDescription: {Self: child}, Path: childPath, partial: true};
+    }
+    return children;
+}
+
 const getSchemaSelector = createSelector(
     (path: string) => path,
     (path) => schemaApi.endpoints.getSchema.select({path}),
 );
 
 const selectGetSchema = createSelector(
     (state: RootState) => state,
+    (_state: RootState, path: string) => path,
     (_state: RootState, path: string) => getSchemaSelector(path),
-    (state, selectTabletsInfo) => selectTabletsInfo(state).data,
+    (state, path, selectTabletsInfo) => selectTabletsInfo(state).data?.[path],
 );
 
 const selectSchemaChildren = (state: RootState, path: string) =>
@@ -127,3 +128,19 @@ export const selectSchemaMergedChildrenPaths: Selector<
             : undefined;
     },
 );
+
+export function useGetSchemaQuery({path}: {path: string}) {
+    const {currentData, isFetching, error, refetch} = schemaApi.useGetSchemaQuery({path});
+
+    const data = currentData?.[path];
+    const isLoading = isFetching && data === undefined;
+
+    const shouldLoad = !isLoading && ((!data && !error) || data?.partial);
+    React.useEffect(() => {
+        if (shouldLoad) {
+            refetch();
+        }
+    }, [refetch, path, shouldLoad]);
+
+    return {data, isLoading, error};
+}
