Fix headers for controls requests (#1859)

datalens-tech · Nov 29, 2024 · e72c649 · e72c649
1 parent c7e2d18
commit e72c649
Show file tree

Hide file tree

Showing 9 changed files with 82 additions and 16 deletions.
diff --git a/src/shared/schema/bi/actions/datasets.ts b/src/shared/schema/bi/actions/datasets.ts
@@ -183,7 +183,10 @@ export const actions = {
         path: ({datasetId}) => `${API_DATA_V2}/datasets/${datasetId}/values/distinct`,
         params: ({datasetId: _datasetId, workbookId, ...body}, headers) => ({
             body,
-            headers: {...(workbookId ? {[WORKBOOK_ID_HEADER]: workbookId} : {}), ...headers},
+            headers: {
+                ...(workbookId ? {[WORKBOOK_ID_HEADER]: workbookId} : {}),
+                ...headers,
+            },
         }),
         transformResponseData: transformApiV2DistinctsResponse,
         timeout: TIMEOUT_95_SEC,

diff --git a/src/shared/schema/bi/types/datasets.ts b/src/shared/schema/bi/types/datasets.ts
@@ -189,6 +189,8 @@ export type GetDistinctsApiV2TransformedResponse = {
     };
 };
 
+export type GetDistinctsApiV2InfoHeadersArg = Record<string, string>;
+
 export type GetDistinctsApiV2Args = Omit<
     ApiV2RequestBody,
     'pivot' | 'order_by' | 'disable_group_by' | 'with_totals' | 'autofill_legend'

diff --git a/src/ui/components/DashKit/DashKit.ts b/src/ui/components/DashKit/DashKit.ts
@@ -34,17 +34,19 @@ const wrapPlugins = (plugins: Plugin[], pluginDefaultsGetter?: typeof currentDef
 };
 
 export const getConfiguredDashKit = (pluginDefaultsGetter: typeof currentDefaultsGetter = null) => {
-    const controlSettings = {
-        getDistincts: getDistinctsAction(),
-    };
-
     if (currentDefaultsGetter !== pluginDefaultsGetter || !isConfigured) {
+        const textSettings = {
+            apiHandler: MarkdownProvider.getMarkdown,
+        };
+
+        const controlSettings = {
+            getDistincts: getDistinctsAction(),
+        };
+
         const plugins = wrapPlugins(
             [
                 pluginTitle,
-                textPlugin.setSettings({
-                    apiHandler: MarkdownProvider.getMarkdown,
-                }),
+                textPlugin.setSettings(textSettings),
                 pluginControl.setSettings(controlSettings),
                 pluginGroupControl.setSettings(controlSettings),
                 widgetPlugin,

diff --git a/src/ui/components/DashKit/plugins/Control/Control.tsx b/src/ui/components/DashKit/plugins/Control/Control.tsx
@@ -103,6 +103,7 @@ class Control extends React.PureComponent<PluginControlProps, PluginControlState
     _cancelSource: any = null;
 
     adjustWidgetLayout = debounce(this.setAdjustWidgetLayout, CONTROL_LAYOUT_DEBOUNCE_TIME);
+    _getDistinctsMemo: ControlSettings['getDistincts'];
 
     resolve: ((value: unknown) => void) | null = null;
 
@@ -426,6 +427,7 @@ class Control extends React.PureComponent<PluginControlProps, PluginControlState
             const {workbookId} = this.props;
 
             const payloadCancellation = chartsDataProvider.getRequestCancellation();
+            const dataProviderContextGetter = this.context?.dataProviderContextGetter?.();
 
             const payload = {
                 data: {
@@ -445,6 +447,7 @@ class Control extends React.PureComponent<PluginControlProps, PluginControlState
                     ...(workbookId ? {workbookId} : {}),
                 },
                 cancelToken: payloadCancellation.token,
+                headers: dataProviderContextGetter,
             };
 
             if (data.sourceType !== DashTabItemControlSourceType.External) {
@@ -552,9 +555,29 @@ class Control extends React.PureComponent<PluginControlProps, PluginControlState
         return null;
     }
 
+    getDistinctsWithHeaders() {
+        if (this.props.getDistincts) {
+            this._getDistinctsMemo =
+                this._getDistinctsMemo ||
+                ((params) => {
+                    const {getDistincts} = this.props;
+                    const headers = this?.context?.dataProviderContextGetter?.();
+
+                    return (getDistincts as Exclude<ControlSettings['getDistincts'], void>)?.(
+                        params,
+                        headers,
+                    );
+                });
+        } else {
+            this._getDistinctsMemo = undefined;
+        }
+
+        return this._getDistinctsMemo;
+    }
+
     renderSelectControl() {
         const data = this.props.data as unknown as DashTabItemControlSingle;
-        const {id, defaults, getDistincts} = this.props;
+        const {id, defaults} = this.props;
         const {loadedData, status, loadingItems, errorData, validationError} = this.state;
 
         const {label, innerLabel} = getLabels(data);
@@ -575,7 +598,7 @@ class Control extends React.PureComponent<PluginControlProps, PluginControlState
                 validationError={validationError}
                 errorData={errorData}
                 validateValue={this.validateValue}
-                getDistincts={getDistincts}
+                getDistincts={this.getDistinctsWithHeaders()}
                 classMixin={b('item')}
                 selectProps={{label, innerLabel}}
             />

diff --git a/src/ui/components/DashKit/plugins/Control/types.ts b/src/ui/components/DashKit/plugins/Control/types.ts
@@ -4,6 +4,7 @@ import type {StringParams} from '@gravity-ui/dashkit';
 import type {TitlePlacementOption} from 'shared';
 import type {
     GetDistinctsApiV2Args,
+    GetDistinctsApiV2InfoHeadersArg,
     GetDistinctsApiV2TransformedResponse,
 } from 'shared/schema/types';
 import type {ServerFilter} from 'shared/types/config/wizard';
@@ -55,7 +56,10 @@ export type SelectorError = {
 };
 
 export interface ControlSettings {
-    getDistincts?: (params: GetDistinctsApiV2Args) => Promise<GetDistinctsApiV2TransformedResponse>;
+    getDistincts?: (
+        params: GetDistinctsApiV2Args,
+        headers?: GetDistinctsApiV2InfoHeadersArg,
+    ) => Promise<GetDistinctsApiV2TransformedResponse>;
 }
 
 export interface PluginControlState {

diff --git a/src/ui/components/DashKit/plugins/GroupControl/Control/Control.tsx b/src/ui/components/DashKit/plugins/GroupControl/Control/Control.tsx
@@ -99,6 +99,7 @@ type ControlProps = {
     }) => void;
     silentLoading: boolean;
     getDistincts?: ControlSettings['getDistincts'];
+    requestHeaders?: Record<string, string>;
     onChange: ({
         params,
         callChangeByClick,
@@ -121,6 +122,7 @@ export const Control = ({
     silentLoading,
     onStatusChanged,
     getDistincts,
+    requestHeaders,
     onChange,
     needReload,
     workbookId,
@@ -224,6 +226,7 @@ export const Control = ({
                     ...(workbookId ? {workbookId} : {}),
                 },
                 cancelToken: payloadCancellation.token,
+                headers: requestHeaders,
             };
 
             cancelCurrentRunRequest();

diff --git a/src/ui/components/DashKit/plugins/GroupControl/GroupControl.tsx b/src/ui/components/DashKit/plugins/GroupControl/GroupControl.tsx
@@ -86,6 +86,7 @@ class GroupControl extends React.PureComponent<PluginGroupControlProps, PluginGr
 
     // a quick loader for imitating action by clicking on apply button
     _quickActionTimer: ReturnType<typeof setTimeout> | null = null;
+    _getDistinctsMemo: ControlSettings['getDistincts'];
 
     // params of current dash state
     initialParams: Record<string, StringParams> = {};
@@ -679,9 +680,30 @@ class GroupControl extends React.PureComponent<PluginGroupControlProps, PluginGr
         }
     };
 
+    private getDistinctsWithHeaders() {
+        if (this.props.getDistincts) {
+            this._getDistinctsMemo =
+                this._getDistinctsMemo ||
+                ((params) => {
+                    const {getDistincts} = this.props;
+                    const headers = this?.context?.dataProviderContextGetter?.();
+
+                    return (getDistincts as Exclude<ControlSettings['getDistincts'], void>)?.(
+                        params,
+                        headers,
+                    );
+                });
+        } else {
+            this._getDistinctsMemo = undefined;
+        }
+
+        return this._getDistinctsMemo;
+    }
+
     private renderControl(item: DashTabItemControlSingle) {
-        const {getDistincts, workbookId} = this.props;
+        const {workbookId} = this.props;
         const {silentLoading} = this.state;
+        const dataProviderContextGetter = this?.context?.dataProviderContextGetter?.();
 
         return (
             <Control
@@ -691,12 +713,13 @@ class GroupControl extends React.PureComponent<PluginGroupControlProps, PluginGr
                 params={this.state.stateParams[item.id] || {}}
                 onStatusChanged={this.handleStatusChanged}
                 silentLoading={silentLoading}
-                getDistincts={getDistincts}
+                getDistincts={this.getDistinctsWithHeaders()}
                 onChange={this.onChange}
                 needReload={this.state.needReload}
                 workbookId={workbookId}
                 dependentSelectors={this.dependentSelectors}
                 groupId={this.props.id}
+                requestHeaders={dataProviderContextGetter}
             />
         );
     }

diff --git a/src/ui/registry/units/common/functions-map.ts b/src/ui/registry/units/common/functions-map.ts
@@ -10,6 +10,7 @@ import type {
     CollectDashStatsArgs,
     CollectDashStatsResponse,
     GetDistinctsApiV2Args,
+    GetDistinctsApiV2InfoHeadersArg,
     GetDistinctsApiV2TransformedResponse,
     GetEntryResponse,
 } from 'shared/schema';
@@ -120,7 +121,10 @@ export const commonFunctionsMap = {
         >(),
     fetchDistinctsByApi:
         makeFunctionTemplate<
-            (params: GetDistinctsApiV2Args) => Promise<GetDistinctsApiV2TransformedResponse>
+            (
+                params: GetDistinctsApiV2Args,
+                headers?: GetDistinctsApiV2InfoHeadersArg,
+            ) => Promise<GetDistinctsApiV2TransformedResponse>
         >(),
     requestCollectDashStats:
         makeFunctionTemplate<

diff --git a/src/ui/utils/sdkRequests.ts b/src/ui/utils/sdkRequests.ts
@@ -2,6 +2,7 @@ import type {
     CollectChartkitStatsArgs,
     CollectDashStatsArgs,
     GetDistinctsApiV2Args,
+    GetDistinctsApiV2InfoHeadersArg,
     GetDistinctsApiV2TransformedResponse,
 } from 'shared/schema/types';
 import {Feature} from 'shared/types';
@@ -12,10 +13,11 @@ import Utils from './utils';
 
 export const fetchDistinctsByApi = (
     params: GetDistinctsApiV2Args,
+    headers?: GetDistinctsApiV2InfoHeadersArg,
 ): Promise<GetDistinctsApiV2TransformedResponse> => {
     return Utils.isEnabledFeature(Feature.UsePublicDistincts)
-        ? getSdk().bi.getPublicDistinctsApiV2(params)
-        : getSdk().bi.getDistinctsApiV2(params);
+        ? getSdk().bi.getPublicDistinctsApiV2(params, {headers})
+        : getSdk().bi.getDistinctsApiV2(params, {headers});
 };
 
 export const fetchBatchRenderedMarkdown = (texts: Record<string, string>) => {