import { QueryClient, QueryKey } from 'react-query' import { serialize } from 'superjson' import { isClient, isServer } from '../stdlib/index' import { requestIdleCallback } from '../client/request-idle-callback' import { ResolverType, RpcClient } from './rpc' import { AsyncFunc } from '../types/utils' import { Resolver } from '../server/api-utils' type MutateOptions = { refetch?: boolean } export const initializeQueryClient = () => { let suspenseEnabled = true if (!process.env.CLI_COMMAND_CONSOLE && !process.env.CLI_COMMAND_DB) { suspenseEnabled = Boolean(process.env.__BLITZ_SUSPENSE_ENABLED) } return new QueryClient({ defaultOptions: { queries: { ...(isServer && { cacheTime: 0 }), suspense: suspenseEnabled, retry: (failureCount, error: any) => { if (process.env.NODE_ENV !== 'production') return false // Retry (max. 3 times) only if network error detected if (error.message === 'Network request failed' && failureCount <= 3) return true return false }, }, }, }) } // Create internal QueryClient instance export const queryClient = initializeQueryClient() function isRpcClient(f: any): f is RpcClient { return !!f._isRpcClient } export interface QueryCacheFunctions { setQueryData: ( newData: T | ((oldData: T | undefined) => T), opts?: MutateOptions ) => ReturnType } export const getQueryCacheFunctions = ( resolver: T | Resolver | RpcClient, params: TInput ): QueryCacheFunctions => ({ setQueryData: (newData, opts = { refetch: true }) => { return setQueryData(resolver, params, newData, opts) }, }) export const emptyQueryFn: RpcClient = (() => { const fn = ((() => new Promise(() => {})) as any) as RpcClient fn._isRpcClient = true return fn })() const isNotInUserTestEnvironment = () => { if (process.env.JEST_WORKER_ID === undefined) return true if (process.env.BLITZ_TEST_ENVIRONMENT !== undefined) return true return false } export const validateQueryFn = ( queryFn: Resolver | RpcClient ) => { if (isClient && !isRpcClient(queryFn) && isNotInUserTestEnvironment()) { throw new Error( `Either the file path to your resolver is incorrect (must be in a "queries" or "mutations" folder that isn't nested inside "pages" or "api") or you are trying to use Blitz's useQuery to fetch from third-party APIs (to do that, import useQuery directly from "react-query")` ) } } const sanitize = (type: ResolverType) => ( queryFn: Resolver | RpcClient ): RpcClient => { if (isServer) return queryFn as any validateQueryFn(queryFn) const rpcClient = queryFn as RpcClient const queryFnName = type === 'mutation' ? 'useMutation' : 'useQuery' if (rpcClient._resolverType !== type && isNotInUserTestEnvironment()) { throw new Error( `"${queryFnName}" was expected to be called with a ${type} but was called with a "${rpcClient._resolverType}"` ) } return rpcClient } export const sanitizeQuery = sanitize('query') export const sanitizeMutation = sanitize('mutation') export const getQueryKeyFromUrlAndParams = (url: string, params: unknown) => { const queryKey = [url] const args = typeof params === 'function' ? (params as Function)() : params queryKey.push(serialize(args) as any) return queryKey as [string, any] } export function getQueryKey( resolver: T | Resolver | RpcClient, params?: TInput ) { if (typeof resolver === 'undefined') { throw new Error( 'getQueryKey is missing the first argument - it must be a resolver function' ) } return getQueryKeyFromUrlAndParams(sanitizeQuery(resolver)._routePath, params) } export function getInfiniteQueryKey( resolver: T | Resolver | RpcClient, params?: TInput ) { if (typeof resolver === 'undefined') { throw new Error( 'getInfiniteQueryKey is missing the first argument - it must be a resolver function' ) } const queryKey = getQueryKeyFromUrlAndParams( sanitizeQuery(resolver)._routePath, params ) return [...queryKey, 'infinite'] } export function invalidateQuery( resolver: T | Resolver | RpcClient, params?: TInput ) { if (typeof resolver === 'undefined') { throw new Error( 'invalidateQuery is missing the first argument - it must be a resolver function' ) } const fullQueryKey = getQueryKey(resolver, params) let queryKey: QueryKey if (params) { queryKey = fullQueryKey } else { // Params not provided, only use first query key item (url) queryKey = fullQueryKey[0] } return queryClient.invalidateQueries(queryKey) } export function setQueryData( resolver: T | Resolver | RpcClient, params: TInput, newData: TResult | ((oldData: TResult | undefined) => TResult), opts: MutateOptions = { refetch: true } ): Promise> { if (typeof resolver === 'undefined') { throw new Error( 'setQueryData is missing the first argument - it must be a resolver function' ) } const queryKey = getQueryKey(resolver, params) return new Promise((res) => { queryClient.setQueryData(queryKey, newData) let result: void | ReturnType if (opts.refetch) { result = invalidateQuery(resolver, params) } if (isClient) { // Fix for https://github.com/blitz-js/blitz/issues/1174 requestIdleCallback(() => { res(result) }) } else { res(result) } }) }