RTK Query and endpoints with continuation header

[LuisGAngel]

I'm working with an API that retrieves entities. The response contains some entities and a continuation header "x-ms-continuation": "someContinuationToken" if there are more entities to be retrieved. To get all entities one needs to call the same endpoint until continuation header is not present.

I'm not using pokeapi.co, but it should work as an example:
First request
url =>https://pokeapi.co/api/v2/pokemon/all
Response
data => ["Pikachu", "Charmander", ...]
response headers => "x-ms-continuation": "TokenAAA"

Second request
url => https://pokeapi.co/api/v2/pokemon/all
request headers => "x-ms-continuation": "TokenAAA"
Response
data => ["Snorlax", "Ditto", ...]
// no more pokemons available so no continuation header with token

Currently, I'm achieving this task using axios/slice/thunk:

// The continuation header for data requests/responses
const CONTINUATION_TOKEN_HEADER = "x-ms-continuation";
// Sanity check for data retrieval
const MAX_CONTINUATIONS = 100;
const url = 'https://pokeapi.co/api/v2/pokemon/all'
let response = await axios.get(url);
let { data } = response;
// real code includes auth header as well
let continuationHeader = response.headers[CONTINUATION_TOKEN_HEADER]; 
let iterationCount = 0;
while (continuationHeader && iterationCount < MAX_CONTINUATIONS) {
    // real code includes auth header as well
    response = await axios.get(url, { headers: { [CONTINUATION_TOKEN_HEADER]: continuationHeader } });
    data = data.concat(response.data);
    continuationHeader = response.headers[CONTINUATION_TOKEN_HEADER];
    iterationCount += 1;
}
return data; // This is the data returned from the thunk as payload

Now, I'm trying to migrate to RTK Query. Please let me know if you have any suggestions on how to achieve it. I don't mind using the fetch wrapper and ditch axios if I can pass the auth header:

export const pokemonApi = createApi({
    reducerPath: 'pokemonApi',
    baseQuery: fetchBaseQuery({
        baseUrl: 'https://pokeapi.co/api/v2/',
        prepareHeaders: (headers) => {
            headers.set("authorization", `Bearer ${AUTH_TOKEN}`);
            return headers;
        },
    }),
    endpoints: (builder) => ({
        getPokemon: builder.query<Pokemon>({
            query: () => `pokemon/all`,
        }),
    }),
});

Thanks for reading so far. I'm happy to answer any questions.

[markerikson]

I think your best option here is to cut and paste the fetching logic and put it into a queryFn, rather than trying to use the query option. You should be able to use the baseQuery to do the actual request part.

[LuisGAngel]

I agree, thanks for your swift response, Mark.
Based on your suggestion I came up with a solution. I'm leaving it below in case anyone finds it useful. Beware that I haven't tested it so take it with a grain of salt:

export const pokemonApi = createApi({
    reducerPath: 'pokemonApi',
    baseQuery: fetchBaseQuery({
        baseUrl: 'https://pokeapi.co/api/v2/',
        prepareHeaders: (headers) => {
            headers.set("authorization", `Bearer ${AUTH_TOKEN}`);
            return headers;
        },
    }),
    endpoints: (builder) => ({
        getAllPokemons: builder.query<Pokemon[]>({
            queryFn: async (_arg, _queryApi, _extraOptions, fetchWithBQ) => {
                const url = `pokemon/all`;
                return getContinuousData<Pokemon>(url, someParams, someHeaders, fetchWithBQ);
            },
        }),
        getAllPokemonsTrainers: builder.query<PokemonTrainer[]>({
            queryFn: async (_arg, _queryApi, _extraOptions, fetchWithBQ) => {
                const url = `trainers/all`;
                return getContinuousData<PokemonTrainer>(url, undefined, undefined, fetchWithBQ);
            },
        }),
    }),
});

async function getContinuousData<T>(
    url: string,
    params: Record<string, any> | undefined,
    headers: Record<string, string | undefined>,
    fetch: (arg: string | FetchArgs) => MaybePromise<QueryReturnValue<unknown, FetchBaseQueryError, FetchBaseQueryMeta>>,
) {
    // The continuation header for data requests/responses
    const CONTINUATION_TOKEN_HEADER = "x-ms-continuation";
    // Sanity check for data retrieval
    const MAX_CONTINUATIONS = 100;
    const data: T[] = [];
    let continuationHeader: string | undefined;
    let iterationCount = 0;
    do {
        const response = await fetch({
            url,
            params,
            headers: {
                ...headers,
                ...(continuationHeader && { [CONTINUATION_TOKEN_HEADER]: continuationHeader }),
            },
        });
        if (response.error) return { error: response.error as FetchBaseQueryError };
        data.concat(response.data as T[]);
        continuationHeader = response.meta?.response?.headers.get(CONTINUATION_TOKEN_HEADER);
        iterationCount += 1;
    } while (continuationHeader && iterationCount < MAX_CONTINUATIONS);
    return { data };
}