Automatic rate limit restoration w/ admin API calls

Author: shauns

packages/app/src/cli/api/admin-as-app.ts

@@ -13,6 +13,7 @@ interface AdminAsAppRequestOptions<TResult, TVariables extends Variables> {
   query: TypedDocumentNode<TResult, TVariables>
   session: AdminSession
   variables?: TVariables
+  autoRateLimitRestore?: boolean
 }
 
 /**
@@ -33,6 +34,10 @@ async function setupAdminAsAppRequest(session: AdminSession) {
 /**
  * Executes a GraphQL query against the Shopify Admin API, on behalf of the app. Uses typed documents.
  *
+ * If `autoRateLimitRestore` is true, the function will wait for a period of time such that the rate limit consumed by
+ * the query is restored back to its original value. This means this function is suitable for use in loops with
+ * multiple queries performed.
+ *
  * @param options - The options for the request.
  * @returns The response of the query of generic type <T>.
  */
@@ -43,5 +48,6 @@ export async function adminAsAppRequestDoc<TResult, TVariables extends Variables
     query: options.query,
     ...(await setupAdminAsAppRequest(options.session)),
     variables: options.variables,
+    autoRateLimitRestore: options.autoRateLimitRestore,
   })
 }

packages/cli-kit/src/public/node/api/graphql.test.ts

@@ -9,6 +9,7 @@ import {LocalStorage} from '../local-storage.js'
 import {ConfSchema, GraphQLRequestKey} from '../../../private/node/conf-store.js'
 import {nonRandomUUID} from '../crypto.js'
 import {CLI_KIT_VERSION} from '../../common/version.js'
+import * as system from '../system.js'
 import {test, vi, describe, expect, beforeEach, beforeAll, afterAll, afterEach} from 'vitest'
 import {TypedDocumentNode} from '@graphql-typed-document-node/core'
 import {setupServer} from 'msw/node'
@@ -18,6 +19,8 @@ let mockedRequestId = 'request-id-123'
 
 vi.spyOn(debugRequest, 'debugLogRequestInfo').mockResolvedValue(undefined)
 
+vi.spyOn(system, 'sleep').mockImplementation(async () => {})
+
 const mockedAddress = 'https://shopify.example/graphql'
 const mockVariables = {some: 'variables'}
 const mockToken = 'token'
@@ -36,6 +39,14 @@ const handlers = [
         data: {
           QueryName: {example: 'hello'},
         },
+        extensions: {
+          cost: {
+            actualQueryCost: 10,
+            throttleStatus: {
+              restoreRate: 10000,
+            },
+          },
+        },
       },
       {
         headers: {
@@ -365,6 +376,47 @@ describe('graphqlRequestDoc', () => {
       expect.anything(),
     )
   })
+
+  test('applies rate limit restoration', async () => {
+    const document = {
+      kind: 'Document',
+      definitions: [
+        {
+          kind: 'OperationDefinition',
+          operation: 'query',
+          name: {kind: 'Name', value: 'QueryName'},
+          selectionSet: {
+            kind: 'SelectionSet',
+            selections: [
+              {
+                kind: 'Field',
+                name: {kind: 'Name', value: 'example'},
+              },
+            ],
+          },
+        },
+      ],
+    } as unknown as TypedDocumentNode<unknown, unknown>
+
+    // When
+    const res = await graphqlRequestDoc({
+      query: document,
+      api: 'mockApi',
+      url: mockedAddress,
+      token: mockToken,
+      addedHeaders: mockedAddedHeaders,
+      variables: mockVariables,
+      autoRateLimitRestore: true,
+    })
+    expect(res).toMatchInlineSnapshot(`
+      {
+        "QueryName": {
+          "example": "hello",
+        },
+      }
+    `)
+    expect(system.sleep).toHaveBeenCalledWith(0.001)
+  })
 })
 
 describe('sanitizeVariables', () => {

packages/cli-kit/src/public/node/api/graphql.ts

@@ -12,8 +12,10 @@ import {
   timeIntervalToMilliseconds,
 } from '../../../private/node/conf-store.js'
 import {LocalStorage} from '../local-storage.js'
-import {abortSignalFromRequestBehaviour, requestMode, RequestModeInput} from '../http.js'
+import {abortSignalFromRequestBehaviour, RequestBehaviour, requestMode, RequestModeInput} from '../http.js'
 import {CLI_KIT_VERSION} from '../../common/version.js'
+import {sleep} from '../system.js'
+import {outputContent, outputDebug} from '../output.js'
 import {
   GraphQLClient,
   rawRequest,
@@ -65,6 +67,7 @@ type PerformGraphQLRequestOptions<TResult> = GraphQLRequestBaseOptions<TResult>
   queryAsString: string
   variables?: Variables
   unauthorizedHandler?: UnauthorizedHandler
+  autoRateLimitRestore?: boolean
 }
 
 export type GraphQLRequestOptions<T> = GraphQLRequestBaseOptions<T> & {
@@ -77,13 +80,28 @@ export type GraphQLRequestDocOptions<TResult, TVariables> = GraphQLRequestBaseOp
   query: TypedDocumentNode<TResult, TVariables> | TypedDocumentNode<TResult, Exact<{[key: string]: never}>>
   variables?: TVariables
   unauthorizedHandler?: UnauthorizedHandler
+  autoRateLimitRestore?: boolean
+}
+
+interface RunRawGraphQLRequestOptions<TResult> {
+  client: {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    setAbortSignal: (signal: any) => void
+    rawRequest: (query: string, variables?: Variables) => Promise<GraphQLResponse<TResult>>
+  }
+  behaviour: RequestBehaviour
+  queryAsString: string
+  variables?: Variables
+  autoRateLimitRestore: boolean
 }
 
 export interface GraphQLResponseOptions<T> {
   handleErrors?: boolean
   onResponse?: (response: GraphQLResponse<T>) => void
 }
 
+const MAX_RATE_LIMIT_RESTORE_DELAY_SECONDS = 0.3
+
 async function createGraphQLClient({
   url,
   addedHeaders,
@@ -105,38 +123,74 @@ async function createGraphQLClient({
   }
 }
 
-/**
- * Handles execution of a GraphQL query.
- *
- * @param options - GraphQL request options.
- */
+async function runSingleRawGraphQLRequest<TResult>(
+  options: RunRawGraphQLRequestOptions<TResult>,
+): Promise<GraphQLResponse<TResult>> {
+  const {client, behaviour, queryAsString, variables, autoRateLimitRestore} = options
+  let fullResponse: GraphQLResponse<TResult>
+  // there is a errorPolicy option which returns rather than throwing on errors, but we _do_ ultimately want to
+  // throw.
+  try {
+    client.setAbortSignal(abortSignalFromRequestBehaviour(behaviour))
+    fullResponse = await client.rawRequest(queryAsString, variables)
+    await logLastRequestIdFromResponse(fullResponse)
+
+    if (autoRateLimitRestore) {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const cost: any = (fullResponse.extensions as any)?.cost
+      const actualQueryCost: unknown = cost?.actualQueryCost
+      const restoreRate: unknown = cost?.throttleStatus?.restoreRate
+
+      if (actualQueryCost && typeof actualQueryCost === 'number' && restoreRate && typeof restoreRate === 'number') {
+        const secondsToRestoreRate = actualQueryCost / restoreRate
+        outputDebug(outputContent`Sleeping for ${secondsToRestoreRate.toString()} seconds to restore the rate limit.`)
+        await sleep(Math.min(secondsToRestoreRate, MAX_RATE_LIMIT_RESTORE_DELAY_SECONDS))
+      }
+    }
+
+    return fullResponse
+  } catch (error) {
+    if (error instanceof ClientError) {
+      // error.response does have a headers property like a normal response, but it's not typed as such.
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      await logLastRequestIdFromResponse(error.response as any)
+    }
+    throw error
+  }
+}
+
 async function performGraphQLRequest<TResult>(options: PerformGraphQLRequestOptions<TResult>) {
-  const {token, addedHeaders, queryAsString, variables, api, url, responseOptions, unauthorizedHandler, cacheOptions} =
-    options
+  const {
+    token,
+    addedHeaders,
+    queryAsString,
+    variables,
+    api,
+    url,
+    responseOptions,
+    unauthorizedHandler,
+    cacheOptions,
+    autoRateLimitRestore,
+  } = options
   const behaviour = requestMode(options.preferredBehaviour ?? 'default')
 
   let {headers, client} = await createGraphQLClient({url, addedHeaders, token})
   debugLogRequestInfo(api, queryAsString, url, variables, headers)
 
   const rawGraphQLRequest = async () => {
-    let fullResponse: GraphQLResponse<TResult>
-    // there is a errorPolicy option which returns rather than throwing on errors, but we _do_ ultimately want to
-    // throw.
-    try {
-      // mapping signal to any due to polyfill meaning types don't exactly match (but are functionally equivalent)
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      client.requestConfig.signal = abortSignalFromRequestBehaviour(behaviour) as any
-      fullResponse = await client.rawRequest<TResult>(queryAsString, variables)
-      await logLastRequestIdFromResponse(fullResponse)
-      return fullResponse
-    } catch (error) {
-      if (error instanceof ClientError) {
-        // error.response does have a headers property like a normal response, but it's not typed as such.
+    return runSingleRawGraphQLRequest({
+      client: {
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        await logLastRequestIdFromResponse(error.response as any)
-      }
-      throw error
-    }
+        setAbortSignal: (signal: any) => {
+          client.requestConfig.signal = signal
+        },
+        rawRequest: (query: string, variables?: Variables) => client.rawRequest<TResult>(query, variables),
+      },
+      behaviour,
+      queryAsString,
+      variables,
+      autoRateLimitRestore: autoRateLimitRestore ?? false,
+    })
   }
 
   const tokenRefreshHandler = unauthorizedHandler?.handler

packages/cli-kit/src/public/node/http.ts

@@ -38,7 +38,7 @@ type AutomaticCancellationBehaviour =
       useAbortSignal: AbortSignal | (() => AbortSignal)
     }
 
-type RequestBehaviour = NetworkRetryBehaviour & AutomaticCancellationBehaviour
+export type RequestBehaviour = NetworkRetryBehaviour & AutomaticCancellationBehaviour
 
 export type RequestModeInput = PresetFetchBehaviour | RequestBehaviour