Pass query args to prepareHeaders function

Author: kyletsang

docs/rtk-query/api/fetchBaseQuery.mdx

@@ -65,6 +65,7 @@ type FetchBaseQueryArgs = {
       BaseQueryApi,
       'getState' | 'extra' | 'endpoint' | 'type' | 'forced'
     >,
+    args: string | FetchArgs
   ) => MaybePromise<Headers | void>
   fetchFn?: (
     input: RequestInfo,
@@ -105,7 +106,7 @@ Typically a string like `https://api.your-really-great-app.com/v1/`. If you don'
 
 _(optional)_
 
-Allows you to inject headers on every request. You can specify headers at the endpoint level, but you'll typically want to set common headers like `authorization` here. As a convenience mechanism, the second argument allows you to use `getState` to access your redux store in the event you store information you'll need there such as an auth token. Additionally, it provides access to `extra`, `endpoint`, `type`, and `forced` to unlock more granular conditional behaviors.
+Allows you to inject headers on every request. You can specify headers at the endpoint level, but you'll typically want to set common headers like `authorization` here. As a convenience mechanism, the second argument allows you to use `getState` to access your redux store in the event you store information you'll need there such as an auth token. Additionally, it provides access to `extra`, `endpoint`, `type`, and `forced` to unlock more granular conditional behaviors. Arguments passed to the query function is passed as the third argument in case you need those params in the context of the `prepareHeaders` function.
 
 You can mutate the `headers` argument directly, and returning it is optional.
 
@@ -119,6 +120,7 @@ type prepareHeaders = (
     type: 'query' | 'mutation'
     forced: boolean | undefined
   },
+  args: string | FetchArgs
 ) => Headers | void
 ```
 

packages/toolkit/src/query/fetchBaseQuery.ts

@@ -113,6 +113,7 @@ export type FetchBaseQueryArgs = {
       BaseQueryApi,
       'getState' | 'extra' | 'endpoint' | 'type' | 'forced'
     >,
+    args: string | FetchArgs
   ) => MaybePromise<Headers | void>
   fetchFn?: (
     input: RequestInfo,
@@ -164,9 +165,9 @@ export type FetchBaseQueryMeta = { request: Request; response?: Response }
  * The base URL for an API service.
  * Typically in the format of https://example.com/
  *
- * @param {(headers: Headers, api: { getState: () => unknown; extra: unknown; endpoint: string; type: 'query' | 'mutation'; forced: boolean; }) => Headers} prepareHeaders
+ * @param {(headers: Headers, api: { getState: () => unknown; extra: unknown; endpoint: string; type: 'query' | 'mutation'; forced: boolean; }, args: string | FetchArgs) => Headers} prepareHeaders
  * An optional function that can be used to inject headers on requests.
- * Provides a Headers object, as well as most of the `BaseQueryApi` (`dispatch` is not available).
+ * Provides a Headers object, most of the `BaseQueryApi` (`dispatch` is not available), and the args passed into the query function.
  * Useful for setting authentication or headers that need to be set conditionally.
  *
  * @link https://developer.mozilla.org/en-US/docs/Web/API/Headers
@@ -212,7 +213,7 @@ export function fetchBaseQuery({
       'Warning: `fetch` is not available. Please supply a custom `fetchFn` property to use `fetchBaseQuery` on SSR environments.',
     )
   }
-  return async (arg, api) => {
+  return async (args, api) => {
     const { signal, getState, extra, endpoint, forced, type } = api
     let meta: FetchBaseQueryMeta | undefined
     let {
@@ -223,7 +224,7 @@ export function fetchBaseQuery({
       validateStatus = globalValidateStatus ?? defaultValidateStatus,
       timeout = defaultTimeout,
       ...rest
-    } = typeof arg == 'string' ? { url: arg } : arg
+    } = typeof args == 'string' ? { url: args } : args
     let config: RequestInit = {
       ...baseFetchOptions,
       signal,
@@ -238,7 +239,7 @@ export function fetchBaseQuery({
         endpoint,
         forced,
         type,
-      })) || headers
+      }, args)) || headers
 
     // Only set the content-type to json if appropriate. Will not be true for FormData, ArrayBuffer, Blob, etc.
     const isJsonifiable = (body: any) =>

packages/toolkit/src/query/tests/fetchBaseQuery.test.tsx

@@ -850,6 +850,48 @@ describe('fetchBaseQuery', () => {
       expect(_forced).toBe(true)
       expect(_extra).toBe(fakeAuth0Client)
     })
+
+    test('prepareHeaders provides args', async () => {
+      let _args
+
+      const baseQuery = fetchBaseQuery({
+        baseUrl,
+        fetchFn: fetchFn as any,
+        prepareHeaders: (headers, _api, args) => {
+          _args = args
+
+          return headers
+        },
+      })
+
+      const fakeAuth0Client = {
+        getTokenSilently: async () => 'fakeToken',
+      }
+
+      const doRequest = async () => {
+        const abortController = new AbortController()
+        return baseQuery(
+          { url: '/echo' },
+          {
+            signal: abortController.signal,
+            abort: (reason) =>
+              // @ts-ignore
+              abortController.abort(reason),
+            dispatch: storeRef.store.dispatch,
+            getState: storeRef.store.getState,
+            extra: fakeAuth0Client,
+            type: 'query',
+            forced: true,
+            endpoint: 'someEndpointName',
+          },
+          {},
+        )
+      }
+
+      await doRequest()
+
+      expect(_args!.url).toBe('/echo')
+    })
   })
 
   test('can pass `headers` into `fetchBaseQuery`', async () => {