feat: specify resource pools when launching a query (#3198)

astandrik · web-flow · commit 6c682649d17b · 2025-12-11T20:22:11.000Z
diff --git a/src/containers/Heatmap/Heatmap.tsx b/src/containers/Heatmap/Heatmap.tsx
@@ -30,7 +30,7 @@ interface HeatmapProps {
 export const Heatmap = ({path, database, databaseFullPath}: HeatmapProps) => {
     const dispatch = useTypedDispatch();
 
-    const itemsContainer = React.createRef<HTMLDivElement>();
+    const itemsContainer = React.useRef<HTMLDivElement | null>(null);
 
     const [autoRefreshInterval] = useAutoRefreshInterval();
 
diff --git a/src/containers/Tenant/Query/QueryEditor/QueryEditor.tsx b/src/containers/Tenant/Query/QueryEditor/QueryEditor.tsx
@@ -32,6 +32,7 @@ import {
     useEventHandler,
     useQueryExecutionSettings,
     useQueryStreamingSetting,
+    useResourcePools,
     useSetting,
     useTypedDispatch,
     useTypedSelector,
@@ -88,7 +89,7 @@ export default function QueryEditor({theme, changeUserInput, queriesHistory}: Qu
 
     const isResultLoaded = Boolean(result);
 
-    const [querySettings] = useQueryExecutionSettings();
+    const [querySettings, setQuerySettings] = useQueryExecutionSettings();
     const enableTracingLevel = useTracingLevelOptionAvailable();
     const [lastQueryExecutionSettings, setLastQueryExecutionSettings] =
         useLastQueryExecutionSettings();
@@ -104,6 +105,12 @@ export default function QueryEditor({theme, changeUserInput, queriesHistory}: Qu
         SETTING_KEYS.BINARY_DATA_IN_PLAIN_TEXT_DISPLAY,
     );
 
+    const {
+        resourcePools,
+        normalizedResourcePool,
+        isLoading: isResourcePoolsLoading,
+    } = useResourcePools(database, querySettings.resourcePool);
+
     const encodeTextWithBase64 = !binaryDataInPlainTextDisplay;
 
     const isStreamingEnabled =
@@ -114,6 +121,28 @@ export default function QueryEditor({theme, changeUserInput, queriesHistory}: Qu
     const [sendQuery] = queryApi.useUseSendQueryMutation();
     const [streamQuery] = queryApi.useUseStreamQueryMutation();
 
+    // Normalize stored resourcePool if it's not available for current database
+    React.useEffect(() => {
+        if (isResourcePoolsLoading) {
+            return;
+        }
+
+        if (querySettings.resourcePool === normalizedResourcePool) {
+            return;
+        }
+
+        setQuerySettings({
+            ...querySettings,
+            resourcePool: normalizedResourcePool,
+        });
+    }, [
+        isResourcePoolsLoading,
+        normalizedResourcePool,
+        querySettings,
+        resourcePools.length,
+        setQuerySettings,
+    ]);
+
     const tableSettings = React.useMemo(() => {
         return isStreamingEnabled
             ? {
diff --git a/src/containers/Tenant/Query/QueryEditorControls/utils/getChangedQueryExecutionSettings.test.ts b/src/containers/Tenant/Query/QueryEditorControls/utils/getChangedQueryExecutionSettings.test.ts
@@ -36,6 +36,7 @@ describe('getChangedQueryExecutionSettings', () => {
             statisticsMode: STATISTICS_MODES.basic,
             tracingLevel: TRACING_LEVELS.basic,
             pragmas: 'PRAGMA TestPragma;',
+            resourcePool: DEFAULT_QUERY_SETTINGS.resourcePool,
         };
         const result = getChangedQueryExecutionSettings(currentSettings, DEFAULT_QUERY_SETTINGS);
         expect(result).toEqual([
diff --git a/src/containers/Tenant/Query/QueryEditorControls/utils/getChangedQueryExecutionSettingsDescription.test.ts b/src/containers/Tenant/Query/QueryEditorControls/utils/getChangedQueryExecutionSettingsDescription.test.ts
@@ -58,6 +58,7 @@ describe('getChangedQueryExecutionSettingsDescription', () => {
             statisticsMode: STATISTICS_MODES.profile,
             tracingLevel: TRACING_LEVELS.diagnostic,
             pragmas: 'PRAGMA TestPragma;',
+            resourcePool: DEFAULT_QUERY_SETTINGS.resourcePool,
         };
 
         const result = getChangedQueryExecutionSettingsDescription({
diff --git a/src/containers/Tenant/Query/QuerySettingsDialog/QuerySettingsDialog.tsx b/src/containers/Tenant/Query/QuerySettingsDialog/QuerySettingsDialog.tsx
@@ -15,11 +15,18 @@ import {cn} from '../../../../utils/cn';
 import {
     useQueryExecutionSettings,
     useQueryStreamingSetting,
+    useResourcePools,
     useSetting,
     useTypedDispatch,
     useTypedSelector,
 } from '../../../../utils/hooks';
-import {QUERY_MODES, querySettingsValidationSchema} from '../../../../utils/query';
+import type {ResourcePoolValue} from '../../../../utils/query';
+import {
+    QUERY_MODES,
+    RESOURCE_POOL_NO_OVERRIDE_VALUE,
+    querySettingsValidationSchema,
+} from '../../../../utils/query';
+import {useCurrentSchema} from '../../TenantContext';
 
 import {QuerySettingsSelect} from './QuerySettingsSelect';
 import {QuerySettingsTimeout} from './QuerySettingsTimeout';
@@ -86,10 +93,44 @@ function QuerySettingsForm({initialValues, onSubmit, onClose}: QuerySettingsForm
     const [useShowPlanToSvg] = useSetting<boolean>(SETTING_KEYS.USE_SHOW_PLAN_SVG);
     const enableTracingLevel = useTracingLevelOptionAvailable();
     const [isQueryStreamingEnabled] = useQueryStreamingSetting();
+    const {database} = useCurrentSchema();
+    const {resourcePools, isLoading: isResourcePoolsLoading} = useResourcePools(
+        database,
+        initialValues.resourcePool,
+    );
+
+    const resourcePoolOptions = React.useMemo(
+        () => [
+            {
+                value: RESOURCE_POOL_NO_OVERRIDE_VALUE,
+                content: i18n('form.resource-pool.no-override'),
+            },
+            ...resourcePools.map((name) => ({
+                value: name,
+                content: name,
+            })),
+        ],
+        [resourcePools],
+    );
 
     const timeout = watch('timeout');
+    const resourcePool = watch('resourcePool');
     const queryMode = watch('queryMode');
 
+    React.useEffect(() => {
+        if (isResourcePoolsLoading) {
+            return;
+        }
+
+        if (!resourcePool || resourcePool === RESOURCE_POOL_NO_OVERRIDE_VALUE) {
+            return;
+        }
+
+        if (!resourcePools.length || !resourcePools.includes(resourcePool)) {
+            setValue('resourcePool', RESOURCE_POOL_NO_OVERRIDE_VALUE);
+        }
+    }, [isResourcePoolsLoading, resourcePools, resourcePool, setValue]);
+
     return (
         <form onSubmit={handleSubmit(onSubmit)}>
             <Dialog.Body className={b('dialog-body')}>
@@ -255,6 +296,26 @@ function QuerySettingsForm({initialValues, onSubmit, onClose}: QuerySettingsForm
                         </div>
                     </Flex>
                 )}
+                <Flex direction="row" alignItems="flex-start" className={b('dialog-row')}>
+                    <label htmlFor="resourcePool" className={b('field-title')}>
+                        {QUERY_SETTINGS_FIELD_SETTINGS.resourcePool.title}
+                    </label>
+                    <div className={b('control-wrapper', {resourcePool: true})}>
+                        <Controller
+                            name="resourcePool"
+                            control={control}
+                            render={({field}) => (
+                                <QuerySettingsSelect<ResourcePoolValue>
+                                    id="resourcePool"
+                                    setting={field.value ?? RESOURCE_POOL_NO_OVERRIDE_VALUE}
+                                    disabled={isResourcePoolsLoading || !resourcePools.length}
+                                    onUpdateSetting={(value) => field.onChange(value)}
+                                    settingOptions={resourcePoolOptions}
+                                />
+                            )}
+                        />
+                    </div>
+                </Flex>
                 <Flex direction="row" alignItems="flex-start" className={b('dialog-row')}>
                     <label htmlFor="pragmas" className={b('field-title')}>
                         {QUERY_SETTINGS_FIELD_SETTINGS.pragmas.title}
diff --git a/src/containers/Tenant/Query/QuerySettingsDialog/QuerySettingsSelect.tsx b/src/containers/Tenant/Query/QuerySettingsDialog/QuerySettingsSelect.tsx
@@ -8,6 +8,7 @@ import type {
     TransactionMode,
 } from '../../../../types/store/query';
 import {cn} from '../../../../utils/cn';
+import type {ResourcePoolValue} from '../../../../utils/query';
 
 import i18n from './i18n';
 
@@ -19,7 +20,7 @@ export const getOptionHeight = () => -1;
 
 const b = cn('ydb-query-settings-select');
 
-type SelectType = QueryMode | TransactionMode | StatisticsMode | TracingLevel;
+type SelectType = QueryMode | TransactionMode | StatisticsMode | TracingLevel | ResourcePoolValue;
 type QuerySettingSelectOption<T> = SelectOption<T> & {isDefault?: boolean};
 
 interface QuerySettingsSelectProps<T extends SelectType> {
diff --git a/src/containers/Tenant/Query/QuerySettingsDialog/constants.ts b/src/containers/Tenant/Query/QuerySettingsDialog/constants.ts
@@ -157,4 +157,7 @@ export const QUERY_SETTINGS_FIELD_SETTINGS = {
     pragmas: {
         title: formI18n('form.pragmas'),
     },
+    resourcePool: {
+        title: formI18n('form.resource-pool'),
+    },
 } as const;
diff --git a/src/containers/Tenant/Query/QuerySettingsDialog/i18n/en.json b/src/containers/Tenant/Query/QuerySettingsDialog/i18n/en.json
@@ -9,6 +9,8 @@
   "form.output.chunk.max.size": "Output chunk max size",
   "form.output.chunk.max.size.bytes": "bytes",
   "form.pragmas": "Pragmas",
+  "form.resource-pool": "Resource pool",
+  "form.resource-pool.no-override": "No pool override",
   "button-done": "Save",
   "tooltip_plan-to-svg-statistics": "Statistics option is set to \"Full\" due to the enabled \"Execution plan\" experiment.\n To disable it, go to the \"Experiments\" section in the user settings.",
   "button-cancel": "Cancel",
diff --git a/src/containers/Tenant/Query/QuerySettingsDialog/i18n/ru.json b/src/containers/Tenant/Query/QuerySettingsDialog/i18n/ru.json
@@ -9,6 +9,8 @@
   "form.output.chunk.max.size": "Максимальный размер чанка",
   "form.output.chunk.max.size.bytes": "байт",
   "form.pragmas": "Прагмы",
+  "form.resource-pool": "Пул ресурсов",
+  "form.resource-pool.no-override": "Без переопределения пула",
   "tooltip_plan-to-svg-statistics": "Опция статистики установлена в значение \"Full\" из-за включенного эксперимента \"Execution plan\".\n Чтобы отключить его, перейдите в раздел \"Experiments\" в настройках пользователя.",
   "button-done": "Готово",
   "button-cancel": "Отменить",
diff --git a/src/store/reducers/query/query.ts b/src/store/reducers/query/query.ts
@@ -5,8 +5,16 @@ import {TracingLevelNumber} from '../../../types/api/query';
 import type {QueryAction, QueryRequestParams, QuerySettings} from '../../../types/store/query';
 import type {StreamDataChunk} from '../../../types/store/streaming';
 import {loadFromSessionStorage, saveToSessionStorage} from '../../../utils';
-import {QUERY_EDITOR_CURRENT_QUERY_KEY, QUERY_EDITOR_DIRTY_KEY} from '../../../utils/constants';
-import {isQueryErrorResponse} from '../../../utils/query';
+import {
+    QUERY_EDITOR_CURRENT_QUERY_KEY,
+    QUERY_EDITOR_DIRTY_KEY,
+    QUERY_TECHNICAL_MARK,
+} from '../../../utils/constants';
+import {
+    RESOURCE_POOL_NO_OVERRIDE_VALUE,
+    isQueryErrorResponse,
+    parseQueryAPIResponse,
+} from '../../../utils/query';
 import {isNumeric} from '../../../utils/utils';
 import type {RootState} from '../../defaultStore';
 import {api} from '../api';
@@ -119,6 +127,15 @@ export const {
     selectResultTab,
 } = slice.selectors;
 
+const getResourcePoolsQueryText = () => {
+    return `${QUERY_TECHNICAL_MARK}
+SELECT
+    Name
+FROM \`.sys/resource_pools\`
+ORDER BY Name
+`;
+};
+
 interface SendQueryParams extends QueryRequestParams {
     actionType?: QueryAction;
     queryId: string;
@@ -196,6 +213,11 @@ export const queryApi = api.injectEndpoints({
                                 : undefined,
                             concurrent_results: DEFAULT_CONCURRENT_RESULTS || undefined,
                             base64,
+                            resource_pool:
+                                querySettings.resourcePool === RESOURCE_POOL_NO_OVERRIDE_VALUE ||
+                                !querySettings.resourcePool
+                                    ? undefined
+                                    : querySettings.resourcePool,
                         },
                         {
                             signal,
@@ -300,6 +322,11 @@ export const queryApi = api.injectEndpoints({
                                 : undefined,
                             query_id: queryId,
                             base64,
+                            resource_pool:
+                                querySettings.resourcePool === RESOURCE_POOL_NO_OVERRIDE_VALUE ||
+                                !querySettings.resourcePool
+                                    ? undefined
+                                    : querySettings.resourcePool,
                         },
                         {signal},
                     );
@@ -366,6 +393,35 @@ export const queryApi = api.injectEndpoints({
                 }
             },
         }),
+        getResourcePools: build.query<string[], {database: string}>({
+            queryFn: async ({database}, {signal}) => {
+                try {
+                    const response = await window.api.viewer.sendQuery(
+                        {
+                            query: getResourcePoolsQueryText(),
+                            database,
+                            action: 'execute-query',
+                            internal_call: true,
+                        },
+                        {signal, withRetries: true},
+                    );
+
+                    if (isQueryErrorResponse(response)) {
+                        return {error: response};
+                    }
+
+                    const data = parseQueryAPIResponse(response);
+                    const rows = data.resultSets?.[0]?.result || [];
+                    const pools = rows
+                        .map((row) => row && row.Name)
+                        .filter((name): name is string => Boolean(name));
+
+                    return {data: pools};
+                } catch (error) {
+                    return {error};
+                }
+            },
+        }),
     }),
     overrideExisting: 'throw',
 });
diff --git a/src/types/api/query.ts b/src/types/api/query.ts
@@ -328,6 +328,7 @@ export interface SendQueryParams<Action extends Actions> {
     limit_rows?: number;
     internal_call?: boolean;
     base64?: boolean;
+    resource_pool?: string;
 }
 
 export interface StreamQueryParams<Action extends Actions> extends SendQueryParams<Action> {
diff --git a/src/utils/hooks/index.ts b/src/utils/hooks/index.ts
@@ -10,3 +10,4 @@ export * from './useEventHandler';
 export * from './useDelayed';
 export * from './useAclSyntax';
 export * from './useScrollPosition';
+export * from './useResourcePools';
diff --git a/src/utils/hooks/useResourcePools.ts b/src/utils/hooks/useResourcePools.ts
@@ -0,0 +1,47 @@
+import React from 'react';
+
+import {queryApi} from '../../store/reducers/query/query';
+import {RESOURCE_POOL_NO_OVERRIDE_VALUE} from '../query';
+import type {ResourcePoolValue} from '../query';
+
+type UseResourcePoolsResult = {
+    resourcePools: string[];
+    normalizedResourcePool: ResourcePoolValue;
+    isLoading: boolean;
+    isFetching: boolean;
+    error?: unknown;
+};
+
+export const useResourcePools = (
+    database: string | undefined,
+    resourcePool: ResourcePoolValue | undefined,
+): UseResourcePoolsResult => {
+    const {
+        data: resourcePools = [],
+        isLoading,
+        isFetching,
+        error,
+    } = queryApi.useGetResourcePoolsQuery(
+        {database: database || ''},
+        {
+            skip: !database,
+        },
+    );
+
+    const normalizedResourcePool = React.useMemo<ResourcePoolValue>(() => {
+        if (!resourcePool || resourcePool === RESOURCE_POOL_NO_OVERRIDE_VALUE) {
+            return RESOURCE_POOL_NO_OVERRIDE_VALUE;
+        }
+
+        if (!resourcePools.length) {
+            // If we couldn't get pools for any reason, fall back to no override
+            return RESOURCE_POOL_NO_OVERRIDE_VALUE;
+        }
+
+        return resourcePools.includes(resourcePool)
+            ? resourcePool
+            : RESOURCE_POOL_NO_OVERRIDE_VALUE;
+    }, [resourcePool, resourcePools]);
+
+    return {resourcePools, normalizedResourcePool, isLoading, isFetching, error};
+};
diff --git a/src/utils/query.ts b/src/utils/query.ts
diff --git a/tests/suites/tenant/queryEditor/models/SettingsDialog.ts b/tests/suites/tenant/queryEditor/models/SettingsDialog.ts
diff --git a/tests/suites/tenant/queryEditor/querySettings.test.ts b/tests/suites/tenant/queryEditor/querySettings.test.ts
