Context per emit/next

[intellix]

Story

As a server

I want the ability to put something on the context object, for use PER emission

So that I can have a dataloader for fixing N+1 queries for each emission

Acceptance criteria

• server type is able to create a dataloader within context and use it for all of the socket emission

Description

I feel like a robot writing out the stories/acceptance etc so I'll write it like a human. Basically a typical problem in any graphql server is the N+1 problem, you want to fetch a list of Chat Messages and each of them have a user:

{
  messages {
    id
    body
    user {
      id
      name
      avatar
    }
  }
}

If you have 3x messages, your resolver will fire off 3x SELECT * FROM user WHERE id = ? so when using GET/POST we get around this by adding a context.dataloader for the whole request, so that we can batch those operations and fire off 1x query: SELECT * FROM user WHERE id IN (?, ?, ?). This works fine today.

Our problem with websockets, is that the context object is shared and re-used for the whole socket (or subscription, I can't remember) and that would mean we're batching over a long period of time as opposed to resolving a single operation. Consider this subscription where a dataloader would fix N+1:

subscription GameCreated {
  gameCreated {
    id
    users {
      id
      name
      avatar
    }
  }
}

That single emission may contain an array of users, and without a dataloader, we fire off 3x queries. I believe that the correct way to use a dataloader in this scenario, is to only batch those queries for that 1x emission.

Looking at the hooks, I'm seeing these available to me:

• context
• onClose
• onComplete
• onConnect
• onDisconnect
• onNext
• onOperation
• onSubscribe
• onError

I was sort of able to hack it by rewriting the dataloader context within onNext, which works for 2x tabs creating a subscription over 2x websockets... but if I have 1x socket open and create 2x subscriptions, I'm noticing that the 2nd emission only ever uses the base context (created using the context hook).

So if I do this (as an example):

useServer<any, ServerExtra>({
  context({ extra }) {
    return new Promise(resolve => {
      const baseContext: any = {
        random: Math.random(),
      };

      extra.context = baseContext;
      return resolve(baseContext);
    });
  },
  async onNext(ctx, message, args, result) {
    console.log('onNext');
    (ctx.extra.context as any).random = Math.random();
  },

Inside my resolver I'm logging what's inside the context:

  user: ({ userId }, _, context: any) => {
    console.log('Random: ', context.random);
    return { ... };
  }

I see this behaviour:

2x sockets (1 subscription per)

Random: 0.351 // Socket 1
Random: 0.351 // Socket 1
Random: 0.579 // Socket 2
Random: 0.579 // Socket 2
onNext
onNext

Random: 0.983 // Socket 1
Random: 0.983 // Socket 1
Random: 0.746 // Socket 2
Random: 0.746 // Socket 2
onNext
onNext

1x socket (2 subscriptions):

Random: 0.351 // Socket 1, Sub 1
Random: 0.351 // Socket 1, Sub 1
Random: 0.579 // Socket 1, Sub 2
Random: 0.579 // Socket 1, Sub 2
onNext
onNext

Random: 0.983 // Socket 1, Sub 1
Random: 0.983 // Socket 1, Sub 1
Random: 0.579 // Socket 1, Sub 2 <-- Notice this didn't change since context() hook
Random: 0.579 // Socket 1, Sub 2 <-- same
onNext
onNext

So, I'm not sure if my replacing of the random number (or dataloader context) within onNext is correct and that there's a bug, or if there should be another hook available to me to do this properly.... or if you disagree with the idea completely. LMK

[enisdenjo]

This feels more like a question, so I've moved it to a discussion.

onNext is invoked after the subscription event, so after the resolvers; meaning, you cant see what onNext changed in the context in a resolver.

[enisdenjo]

Maybe I am not understanding your question right. A reproduction or a runnable example could help.

[intellix]

we've managed to get what we wanted with a patch to add another hook:

patches/graphql-ws+5.12.1.patch

diff --git a/node_modules/graphql-ws/lib/server.d.ts b/node_modules/graphql-ws/lib/server.d.ts
index adfa2b5..80319cd 100644
--- a/node_modules/graphql-ws/lib/server.d.ts
+++ b/node_modules/graphql-ws/lib/server.d.ts
@@ -266,6 +266,10 @@ export interface ServerOptions<P extends ConnectionInitMessage['payload'] = Conn
      * will still be called.
      */
     onComplete?: (ctx: Context<P, E>, message: CompleteMessage) => Promise<void> | void;
+
+    /** Custom Hook to execute before every subscription emit */
+   beforeEmitExecute?: (args: ExecutionArgs) => void | Promise<void>;
+
     /**
      * An optional override for the JSON.parse function used to hydrate
      * incoming messages to this server. Useful for parsing custom datatypes
diff --git a/node_modules/graphql-ws/lib/server.js b/node_modules/graphql-ws/lib/server.js
index 4b5cfd3..0a5baf8 100644
--- a/node_modules/graphql-ws/lib/server.js
+++ b/node_modules/graphql-ws/lib/server.js
@@ -26,7 +26,7 @@ const utils_1 = require("./utils");
  * @category Server
  */
 function makeServer(options) {
-    const { schema, context, roots, validate, execute, subscribe, connectionInitWaitTimeout = 3000, // 3 seconds
+    const { schema, beforeEmitExecute, context, roots, validate, execute, subscribe, connectionInitWaitTimeout = 3000, // 3 seconds
     onConnect, onDisconnect, onClose, onSubscribe, onOperation, onNext, onError, onComplete, jsonMessageReviver: reviver, jsonMessageReplacer: replacer, } = options;
     return {
         opened(socket, extra) {
@@ -185,10 +185,30 @@ function makeServer(options) {
                                         ? await context(ctx, message, execArgs)
                                         : context;
                             // the execution arguments have been prepared
+
+
+
                             // perform the operation and act accordingly
                             let operationResult;
-                            if (operationAST.operation === 'subscription')
-                                operationResult = await (subscribe !== null && subscribe !== void 0 ? subscribe : graphql_1.subscribe)(execArgs);
+                            if (operationAST.operation === 'subscription') {
+                              const rawAsyncIterator = await (subscribe !== null && subscribe !== void 0 ? subscribe : graphql_1.subscribe)(execArgs);
+
+                              if ((0, utils_1.isAsyncIterable)(rawAsyncIterator)) {
+                                const originalNext = rawAsyncIterator.next.bind(rawAsyncIterator);
+
+                                operationResult = {
+                                  ...rawAsyncIterator,
+                                  async next(...args) {
+                                    if (typeof beforeEmitExecute === 'function') {
+                                      await beforeEmitExecute(execArgs);
+                                    }
+                                    return originalNext(...args);
+                                  },
+                                };
+                              } else {
+                                operationResult = rawAsyncIterator;
+                              }
+                            }
                             // operation === 'query' || 'mutation'
                             else
                                 operationResult = await (execute !== null && execute !== void 0 ? execute : graphql_1.execute)(execArgs);

And then we can have a new dataloader with a hook like so:

// Patched to add fresh dataloader context on every emit
beforeEmitExecute: (args: ExecutionArgs) => {
  args.contextValue = {
    ...((args.contextValue || {}) as Context),
    dataloaderContext: createContext(sequelize),
  };
},