feat: pass context to onParams hook (#3464)

* feat: pass context to inParams hook

* add changeset

* feat: pass server context to persisted document handler

* add changeset

* add test

* let it be

.changeset/chatty-lobsters-exercise.md

@@ -0,0 +1,5 @@
+---
+'@graphql-yoga/plugin-persisted-operations': minor
+---
+
+Forward server context into `extractPersistedOperationId` and `getPersistedOperation` handlers.

.changeset/quiet-cougars-camp.md

@@ -0,0 +1,5 @@
+---
+'graphql-yoga': minor
+---
+
+Inject initial context into `onParams` hook.

packages/graphql-yoga/__tests__/plugin-hooks.spec.ts

@@ -1,4 +1,4 @@
-import { createYoga, type Plugin } from '../src';
+import { createSchema, createYoga, type Plugin } from '../src';
 import { eventStream } from './utilities';
 
 test('onParams -> setResult to single execution result', async () => {
@@ -59,3 +59,49 @@ test('onParams -> setResult to event stream execution result', async () => {
   }
   expect(counter).toBe(2);
 });
+
+test('context value identity stays the same in all hooks', async () => {
+  const contextValues = [] as Array<unknown>;
+  const yoga = createYoga({
+    schema: createSchema({ typeDefs: `type Query {a:String}` }),
+    plugins: [
+      {
+        onParams(ctx) {
+          contextValues.push(ctx.context);
+        },
+        onParse(ctx) {
+          contextValues.push(ctx.context);
+        },
+        onValidate(ctx) {
+          contextValues.push(ctx.context);
+        },
+        onContextBuilding(ctx) {
+          contextValues.push(ctx.context);
+          // mutate context
+          ctx.extendContext({ a: 1 } as Record<string, unknown>);
+          contextValues.push(ctx.context);
+        },
+        onExecute(ctx) {
+          contextValues.push(ctx.args.contextValue);
+        },
+        onResponse(ctx) {
+          contextValues.push(ctx.serverContext);
+        },
+      } satisfies Plugin,
+    ],
+  });
+
+  const response = await yoga.fetch('http://localhost/graphql', {
+    method: 'POST',
+    headers: {
+      'content-type': 'application/json',
+    },
+    body: JSON.stringify({ query: '{__typename}' }),
+  });
+  expect(response.status).toEqual(200);
+  expect(await response.json()).toEqual({ data: { __typename: 'Query' } });
+  expect(contextValues).toHaveLength(7);
+  for (const value of contextValues) {
+    expect(value).toBe(contextValues[0]);
+  }
+});

packages/graphql-yoga/src/plugins/types.ts

@@ -55,7 +55,7 @@ export type Plugin<
      * Use this hook with your own risk. It is still experimental and may change in the future.
      * @internal
      */
-    onParams?: OnParamsHook;
+    onParams?: OnParamsHook<TServerContext>;
     /**
      * Use this hook with your own risk. It is still experimental and may change in the future.
      * @internal
@@ -105,14 +105,17 @@ export interface OnRequestParseDoneEventPayload {
   setRequestParserResult: (params: GraphQLParams | GraphQLParams[]) => void;
 }
 
-export type OnParamsHook = (payload: OnParamsEventPayload) => PromiseOrValue<void>;
+export type OnParamsHook<TServerContext> = (
+  payload: OnParamsEventPayload<TServerContext>,
+) => PromiseOrValue<void>;
 
-export interface OnParamsEventPayload {
+export interface OnParamsEventPayload<TServerContext = Record<string, unknown>> {
   params: GraphQLParams;
   request: Request;
   setParams: (params: GraphQLParams) => void;
   setResult: (result: ExecutionResult | AsyncIterable<ExecutionResult>) => void;
   fetchAPI: FetchAPI;
+  context: TServerContext;
 }
 
 export type OnResultProcess<TServerContext> = (

packages/graphql-yoga/src/server.ts

@@ -212,7 +212,7 @@ export class YogaServer<
     Plugin<TUserContext & TServerContext & YogaInitialContext, TServerContext, TUserContext>
   >;
   private onRequestParseHooks: OnRequestParseHook<TServerContext>[];
-  private onParamsHooks: OnParamsHook[];
+  private onParamsHooks: OnParamsHook<TServerContext>[];
   private onExecutionResultHooks: OnExecutionResultHook<TServerContext>[];
   private onResultProcessHooks: OnResultProcess<TServerContext>[];
   private maskedErrorsOpts: YogaMaskedErrorOpts | null;
@@ -471,7 +471,9 @@ export class YogaServer<
     serverContext: TServerContext,
   ) {
     let result: ExecutionResult | AsyncIterable<ExecutionResult> | undefined;
-    let context = serverContext as TServerContext & YogaInitialContext;
+    const context: TServerContext & YogaInitialContext =
+      batched === true ? Object.create(serverContext) : serverContext;
+
     try {
       for (const onParamsHook of this.onParamsHooks) {
         await onParamsHook({
@@ -484,6 +486,7 @@ export class YogaServer<
             result = newResult;
           },
           fetchAPI: this.fetchAPI,
+          context,
         });
       }
 
@@ -498,10 +501,7 @@ export class YogaServer<
                 params,
               };
 
-        context = Object.assign(
-          batched ? Object.create(serverContext) : serverContext,
-          additionalContext,
-        );
+        Object.assign(context, additionalContext);
 
         const enveloped = this.getEnveloped(context);
 

packages/plugins/persisted-operations/src/index.ts

@@ -18,9 +18,10 @@ export interface GraphQLErrorOptions {
   extensions?: Maybe<GraphQLErrorExtensions>;
 }
 
-export type ExtractPersistedOperationId = (
+export type ExtractPersistedOperationId<TPluginContext = Record<string, unknown>> = (
   params: GraphQLParams,
   request: Request,
+  context: TPluginContext,
 ) => null | string;
 
 export const defaultExtractPersistedOperationId: ExtractPersistedOperationId = (
@@ -41,13 +42,14 @@ export const defaultExtractPersistedOperationId: ExtractPersistedOperationId = (
 
 type AllowArbitraryOperationsHandler = (request: Request) => PromiseOrValue<boolean>;
 
-export type UsePersistedOperationsOptions = {
+export type UsePersistedOperationsOptions<TPluginContext = Record<string, unknown>> = {
   /**
    * A function that fetches the persisted operation
    */
   getPersistedOperation(
     key: string,
     request: Request,
+    context: TPluginContext,
   ): PromiseOrValue<DocumentNode | string | null>;
   /**
    * Whether to allow execution of arbitrary GraphQL operations aside from persisted operations.
@@ -100,7 +102,7 @@ export function usePersistedOperations<
   getPersistedOperation,
   skipDocumentValidation = false,
   customErrors,
-}: UsePersistedOperationsOptions): Plugin<TPluginContext> {
+}: UsePersistedOperationsOptions<TPluginContext>): Plugin<TPluginContext> {
   const operationASTByRequest = new WeakMap<Request, DocumentNode>();
   const persistedOperationRequest = new WeakSet<Request>();
 
@@ -129,13 +131,18 @@ export function usePersistedOperations<
         return;
       }
 
-      const persistedOperationKey = extractPersistedOperationId(params, request);
+      const persistedOperationKey = extractPersistedOperationId(params, request, payload.context);
 
       if (persistedOperationKey == null) {
         throw keyNotFoundErrorFactory(payload);
       }
 
-      const persistedQuery = await getPersistedOperation(persistedOperationKey, request);
+      const persistedQuery = await getPersistedOperation(
+        persistedOperationKey,
+        request,
+        payload.context as TPluginContext,
+      );
+
       if (persistedQuery == null) {
         throw notFoundErrorFactory(payload);
       }