sentry/static/app/views/settings/dynamicSampling/utils/useSamplingProjectRates.tsx

import type {Client} from 'sentry/api';
import type {Organization} from 'sentry/types/organization';
import parseLinkHeader from 'sentry/utils/parseLinkHeader';
import {
  type ApiQueryKey,
  setApiQueryData,
  useMutation,
  useQuery,
  useQueryClient,
} from 'sentry/utils/queryClient';
import type RequestError from 'sentry/utils/requestError/requestError';
import useApi from 'sentry/utils/useApi';
import useOrganization from 'sentry/utils/useOrganization';

interface SamplingProjectRate {
  id: number;
  sampleRate: number;
}

function getEndpoint(organization: Organization) {
  return `/organizations/${organization.slug}/sampling/project-rates/`;
}

function getQueryKey(organization: Organization): ApiQueryKey {
  return [getEndpoint(organization)];
}

/**
 * Fetches all sampling rates for the organization by looping through
 * the paginated results.
 */
const fetchAllSamplingRates = async (
  api: Client,
  organization: Organization
): Promise<SamplingProjectRate[]> => {
  const endpoint = getEndpoint(organization);
  let cursor: string | null = '';
  let result: SamplingProjectRate[] = [];

  while (cursor !== null) {
    const [data, _, response] = await api.requestPromise(endpoint, {
      method: 'GET',
      includeAllArgs: true,
      query: {cursor},
    });

    result = result.concat(data);

    cursor = null;
    const linkHeader = response?.getResponseHeader('Link');

    if (linkHeader) {
      const links = parseLinkHeader(linkHeader);
      cursor = (links.next.results && links.next.cursor) || null;
    }
  }

  return result;
};

export function useGetSamplingProjectRates() {
  const api = useApi();
  const organization = useOrganization();
  return useQuery<SamplingProjectRate[]>({
    queryKey: getQueryKey(organization),
    queryFn: () => fetchAllSamplingRates(api, organization),
    staleTime: 0,
  });
}

export function useUpdateSamplingProjectRates() {
  const api = useApi();
  const queryClient = useQueryClient();
  const organization = useOrganization();

  return useMutation<SamplingProjectRate[], RequestError, SamplingProjectRate[]>({
    mutationFn: variables => {
      return api.requestPromise(getEndpoint(organization), {
        method: 'PUT',
        data: variables,
      });
    },
    onSuccess: data => {
      setApiQueryData<SamplingProjectRate[]>(
        queryClient,
        getQueryKey(organization),
        previous => {
          if (!previous) {
            return data;
          }
          const newDataById = data.reduce(
            (acc, item) => {
              acc[item.id] = item;
              return acc;
            },
            {} as Record<number, SamplingProjectRate>
          );

          return previous.map(item => {
            const newItem = newDataById[item.id];
            if (newItem) {
              return newItem;
            }
            return item;
          });
        }
      );
    },
    onSettled: () => {
      queryClient.invalidateQueries({queryKey: getQueryKey(organization)});
    },
  });
}