Add headers for dataset request (#1974)

datalens-tech · Dec 25, 2024 · 0121f46 · 0121f46
1 parent aae92da
commit 0121f46
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 28 deletions.
diff --git a/src/server/components/charts-engine/components/processor/data-fetcher.ts b/src/server/components/charts-engine/components/processor/data-fetcher.ts
@@ -158,7 +158,7 @@ export function addZitadelHeaders({
     headers,
     zitadelParams,
 }: {
-    headers: IncomingHttpHeaders;
+    headers: OutgoingHttpHeaders;
     zitadelParams: ZitadelParams;
 }) {
     if (zitadelParams?.accessToken) {
@@ -237,7 +237,6 @@ export class DataFetcher {
                 fetchPromisesList.push(() =>
                     source
                         ? DataFetcher.fetchSource({
-                              req,
                               ctx,
                               sourceName,
                               source: isString(source) ? {url: source} : source,
@@ -458,7 +457,6 @@ export class DataFetcher {
     private static async fetchSource({
         sourceName,
         source,
-        req,
         ctx,
         chartsEngine,
         fetchingStartTime,
@@ -476,7 +474,6 @@ export class DataFetcher {
     }: {
         sourceName: string;
         source: Source;
-        req: Request;
         ctx: AppContext;
         chartsEngine: ChartsEngine;
         fetchingStartTime: number;
@@ -786,6 +783,10 @@ export class DataFetcher {
             requestOptions.headers['x-forwarded-for'] = originalReqHeaders.xForwardedFor;
         }
 
+        if (!requestOptions.headers['x-real-ip']) {
+            requestOptions.headers['x-real-ip'] = originalReqHeaders.xRealIP;
+        }
+
         if (isSourceWithMiddlewareUrl(source)) {
             const middlewareSourceConfig = DataFetcher.getSourceConfig({
                 chartsEngine,
@@ -794,15 +795,16 @@ export class DataFetcher {
 
             if (middlewareSourceConfig?.middlewareAdapter) {
                 source = await middlewareSourceConfig.middlewareAdapter({
+                    ctx,
                     source,
                     sourceName,
-                    req,
                     iamToken: iamToken ?? undefined,
                     workbookId,
                     ChartsEngine: chartsEngine,
                     userId: userId === undefined ? null : userId,
                     rejectFetchingSource,
                     zitadelParams,
+                    requestHeaders: requestOptions.headers,
                 });
             }
         }
@@ -824,10 +826,6 @@ export class DataFetcher {
         const publicTargetUri = hideSensitiveData(targetUri);
         const publicSourceData = hideSensitiveData(sourceData);
 
-        if (!requestOptions.headers['x-real-ip']) {
-            requestOptions.headers['x-real-ip'] = originalReqHeaders.xRealIP;
-        }
-
         const traceId = ctx.getTraceId();
         const tenantId = ctx.get('tenantId');
 

diff --git a/...erver/modes/charts/plugins/request-with-dataset/middlewareAdapters/charts-with-dataset.ts b/...erver/modes/charts/plugins/request-with-dataset/middlewareAdapters/charts-with-dataset.ts
@@ -13,20 +13,19 @@ export default async (
     },
 ) => {
     const {
+        ctx,
         source,
         sourceName,
-        req,
         ChartsEngine,
         userId,
         iamToken,
         workbookId,
         rejectFetchingSource,
         pluginOptions,
         zitadelParams,
+        requestHeaders,
     } = args;
 
-    const ctx = req.ctx;
-
     const cacheClient = ChartsEngine.cacheClient as Cache;
 
     const [datasetId, layerId] = getDatasetIdAndLayerIdFromKey(sourceName);
@@ -49,14 +48,14 @@ export default async (
         const datasetFieldsResponse = await getDatasetFields({
             datasetId,
             workbookId: workbookId ?? null,
-            req,
             ctx,
             cacheClient,
             userId,
             iamToken,
             rejectFetchingSource,
             pluginOptions,
             zitadelParams,
+            requestHeaders,
         });
 
         revisionId = datasetFieldsResponse.revisionId;

diff --git a/...ver/modes/charts/plugins/request-with-dataset/middlewareAdapters/controls-with-dataset.ts b/...ver/modes/charts/plugins/request-with-dataset/middlewareAdapters/controls-with-dataset.ts
@@ -11,34 +11,33 @@ export default async (
     },
 ) => {
     const {
+        ctx,
         source,
-        req,
         ChartsEngine,
         userId,
         iamToken,
         workbookId,
         rejectFetchingSource,
         pluginOptions,
         zitadelParams,
+        requestHeaders,
     } = args;
 
-    const ctx = req.ctx;
-
     const cacheClient = ChartsEngine.cacheClient as Cache;
 
     const datasetId = source.datasetId || '';
 
     const datasetFieldsResponse = await getDatasetFields({
         datasetId,
         workbookId: workbookId ?? null,
-        req,
         ctx,
         cacheClient,
         userId,
         iamToken,
         rejectFetchingSource,
         pluginOptions,
         zitadelParams,
+        requestHeaders,
     });
 
     const datasetFields = datasetFieldsResponse.datasetFields;

diff --git a/src/server/modes/charts/plugins/request-with-dataset/request-dataset.ts b/src/server/modes/charts/plugins/request-with-dataset/request-dataset.ts
@@ -1,4 +1,5 @@
-import type {Request} from '@gravity-ui/expresskit';
+import type {OutgoingHttpHeaders} from 'http';
+
 import type {AppContext} from '@gravity-ui/nodekit';
 import {REQUEST_ID_PARAM_NAME} from '@gravity-ui/nodekit';
 import {isObject} from 'lodash';
@@ -25,21 +26,21 @@ const getStatusFromError = (error: unknown) =>
 const getDatasetFieldsById = async ({
     datasetId,
     workbookId,
-    req,
     ctx,
     rejectFetchingSource,
     iamToken,
     pluginOptions,
     zitadelParams,
+    headers,
 }: {
     datasetId: string;
     workbookId: string | null;
-    req: Request;
     ctx: AppContext;
     rejectFetchingSource: (reason?: any) => void;
     iamToken?: string;
     pluginOptions?: ConfigurableRequestWithDatasetPluginOptions;
     zitadelParams: ZitadelParams | undefined;
+    headers: OutgoingHttpHeaders;
 }): Promise<GetDataSetFieldsByIdResponse> => {
     const {gatewayApi} = registry.getGatewayApi<DatalensGatewaySchemas>();
 
@@ -48,13 +49,11 @@ const getDatasetFieldsById = async ({
 
     const requestDatasetFieldsByToken = gatewayApi.bi.embedsGetDataSetFieldsById;
     try {
-        const headers = {...req.headers};
-
         if (zitadelParams) {
             addZitadelHeaders({headers, zitadelParams});
         }
 
-        const response = req.headers[DL_EMBED_TOKEN_HEADER]
+        const response = headers[DL_EMBED_TOKEN_HEADER]
             ? await requestDatasetFieldsByToken({
                   ctx,
                   headers,
@@ -98,26 +97,26 @@ const getDatasetFieldsById = async ({
 export const getDatasetFields = async (args: {
     datasetId: string;
     workbookId: WorkbookId;
-    req: Request;
     ctx: AppContext;
     iamToken?: string;
     cacheClient: Cache;
     userId: string | null;
     rejectFetchingSource: (reason: any) => void;
     pluginOptions?: ConfigurableRequestWithDatasetPluginOptions;
     zitadelParams: ZitadelParams | undefined;
+    requestHeaders: OutgoingHttpHeaders;
 }): Promise<{datasetFields: PartialDatasetField[]; revisionId: string}> => {
     const {
         datasetId,
         workbookId,
         cacheClient,
-        req,
         ctx,
         userId,
         iamToken,
         rejectFetchingSource,
         pluginOptions,
         zitadelParams,
+        requestHeaders,
     } = args;
 
     const cacheKey = `${datasetId}__${userId}`;
@@ -140,12 +139,12 @@ export const getDatasetFields = async (args: {
             const response = await getDatasetFieldsById({
                 datasetId,
                 workbookId,
-                req,
                 ctx,
                 rejectFetchingSource,
                 iamToken,
                 pluginOptions,
                 zitadelParams,
+                headers: requestHeaders,
             });
             datasetFields = response.fields;
             revisionId = response.revision_id;
@@ -173,12 +172,12 @@ export const getDatasetFields = async (args: {
         const response = await getDatasetFieldsById({
             datasetId,
             workbookId,
-            req,
             ctx,
             rejectFetchingSource,
             iamToken,
             pluginOptions,
             zitadelParams,
+            headers: requestHeaders,
         });
         datasetFields = response.fields;
         revisionId = response.revision_id;

diff --git a/src/server/modes/charts/plugins/types.ts b/src/server/modes/charts/plugins/types.ts
@@ -1,4 +1,7 @@
+import type {OutgoingHttpHeaders} from 'http';
+
 import type {Request} from '@gravity-ui/expresskit';
+import type {AppContext} from '@gravity-ui/nodekit';
 
 import type {WorkbookId} from '../../../../shared';
 import type {ChartsEngine} from '../../../components/charts-engine';
@@ -42,19 +45,20 @@ export interface MiddlewareUrl {
 }
 
 export interface MiddlewareSourceAdapterArgs {
+    ctx: AppContext;
     sourceName: string;
     source: {
         sourceArgs: SourcesArgs;
         middlewareUrl: MiddlewareUrl;
         datasetId?: string;
     };
-    req: Request;
     iamToken?: string;
     workbookId?: WorkbookId;
     ChartsEngine: ChartsEngine;
     userId: string | null;
     rejectFetchingSource: (reason: any) => void;
     zitadelParams: ZitadelParams | undefined;
+    requestHeaders: OutgoingHttpHeaders;
 }
 
 export interface ProcessorHookInitArgs {