Repository: mcollina/fastify-gql
Pull request #233 (2020-07-31 09:19:33)

feat: add subscription context
This is #153 synced with master, I don't know if I can edit @bvanneerven PR, so this PR is just a fork from his repo. 

This PR could allow onConnect from #167 to extend the subscription context with onConnect return value (will be addressed in an other PR once this one and #167 are merged)

Closes: #153 #151 

mcollina: I actually made a mess.

Forget my reviews, this is ok. I got confused with https://github.com/mcollina/fastify-gql/pull/167, which is landed now.
mcollina: This needs a rebase.
PacoDu: Rebased, I've also added the context merge from the onConnect handler

Patch:
diff --git a/README.md b/README.md
index dd7af71..f5fef77 100644
--- a/README.md
+++ b/README.md
@@ -362,6 +362,41 @@ app.register(GQL, {
 })
 ```
 
+### Build a custom GraphQL context object for subscriptions
+
+```js
+...
+const resolvers = {
+  Mutation: {
+    sendMessage: async (_, { message, userId }, { pubsub }) => {
+      await pubsub.publish({
+        topic: userId,
+        payload: message
+      })
+
+      return "OK"
+    }
+  },
+  Subscription: {
+    receivedMessage: {
+      // If someone calls the sendMessage mutation with the Id of the user that was added
+      // to the subscription context, that user receives the message.
+      subscribe: (root, args, { pubsub, user }) => pubsub.subscribe(user.id)
+    }
+  }
+}
+
+app.register(GQL, {
+  schema,
+  resolvers,
+  subscription: {
+      // Add the decoded JWT from the Authorization header to the subscription context.
+      context: (_, req) => ({ user: jwt.verify(req.headers["Authorization"].slice(7))})
+  }
+})
+...
+```
+
 ### Subscription support (with redis)
 
 ```js
@@ -596,11 +631,13 @@ app.listen(3000)
 A GraphQL server can act as a Gateway that composes the schemas of the underlying services into one federated schema and executes queries across the services. Every underlying service must be a GraphQL server that supports the federation.
 
 In Gateway mode the following options are not allowed (the plugin will throw an error if any of them are defined):
+
 * `schema`
 * `resolvers`
 * `loaders`
 
 Also, using the following decorator methods will throw:
+
 * `app.graphql.defineResolvers`
 * `app.graphql.defineLoaders`
 * `app.graphql.replaceSchema`
@@ -719,7 +756,8 @@ __fastify-gql__ supports the following options:
 * `subscription`: Boolean | Object. Enable subscriptions. It is uses [mqemitter](https://github.com/mcollina/mqemitter) when it is true. To use a custom emitter set the value to an object containing the emitter.
   * `subscription.emitter`: Custom emitter
   * `subscription.verifyClient`: `Function` A function which can be used to validate incoming connections.
-  * `subscription.onConnect`: `Function` A function which can be used to validate the `connection_init` payload. If defined it should return a truthy value to authorize the connection. 
+  * `subscription.context`: `Function` Result of function is passed to subscription resolvers as a custom GraphQL context. The function receives the `connection` and `request` as parameters.
+  * `subscription.onConnect`: `Function` A function which can be used to validate the `connection_init` payload. If defined it should return a truthy value to authorize the connection. If it returns an object the subscription context will be extended with the returned object.
 * `federationMetadata`: Boolean. Enable federation metadata support so the service can be deployed behind an Apollo Gateway
 * `gateway`: Object. Run the GraphQL server in gateway mode.
   * `gateway.services`: Service[] An array of GraphQL services that are part of the gateway
@@ -741,6 +779,7 @@ __fastify-gql__ supports the following options:
 * `allowBatchedQueries`: Boolean. Flag to control whether to allow batched queries. When `true`, the server supports recieving an array of queries and returns an array of results.
 
 #### queryDepth example
+
 ```
 query {
   dogs {
@@ -760,6 +799,7 @@ query {
   }
 }
 ```
+
 A `queryDepth` of `6` would allow this query. `5` or less would throw with the error - `unnamedQuery query exceeds the query depth limit of 5`
 
 ### HTTP endpoints
@@ -798,6 +838,7 @@ payload must conform to the following JSON schema:
 ```
 
 For code from [example](#example) use:
+
 ```sh
 curl -H "Content-Type:application/json" -XPOST -d '{"query": "query { add(x: 2, y: 2) }"}' http://localhost:3000/graphql
 ```
@@ -808,6 +849,7 @@ Executes the GraphQL query or mutation described in the body. `operationName` an
 payload contains the GraphQL query.
 
 For code from [example](#example) use:
+
 ```sh
 curl -H "Content-Type:application/graphql" -XPOST -d "query { add(x: 2, y: 2) }" http://localhost:3000/graphql
 ```
@@ -822,7 +864,6 @@ the options.
 Serves [GraphQL IDE](https://www.npmjs.com/package/graphql-playground-react) if enabled by
 the options.
 
-
 ### decorators
 
 __fastify-gql__ adds the following decorators.
@@ -1008,7 +1049,6 @@ GraphQL context, and it includes a `reply` object.
 
 Example:
 
-
 ```js
 const loaders = {
   Dog: {
@@ -1051,7 +1091,6 @@ app.register(GQL, {
 Disabling caching has the advantage to avoid the serialization at
 the cost of more objects to fetch in the resolvers.
 
-
 Internally, it uses
 [single-user-cache](http://npm.im/single-user-cache).
 
diff --git a/index.d.ts b/index.d.ts
index 4b53dc3..9d1e3a8 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -16,6 +16,7 @@ import {
   GraphQLScalarType,
   ValidationRule,
 } from "graphql";
+import { SocketStream } from "fastify-websocket"
 
 declare interface FastifyGQLPlugin {
   /**
@@ -176,7 +177,8 @@ export interface FastifyGQLCommonOptions {
         verifyClient?: (
           info: object,
           next: (result: boolean) => void
-        ) => void;
+        ) => void,
+        context?: (connection: SocketStream, request: FastifyRequest) => object | Promise<object>
       };
   /**
    * Enable federation metadata support so the service can be deployed behind an Apollo Gateway
diff --git a/index.js b/index.js
index f46de85..83fb100 100644
--- a/index.js
+++ b/index.js
@@ -105,11 +105,13 @@ const plugin = fp(async function (app, opts) {
 
   let subscriber
   let verifyClient
+  let subscriptionContextFn
   let onConnect
 
   if (typeof subscriptionOpts === 'object') {
     emitter = subscriptionOpts.emitter || mq()
     verifyClient = subscriptionOpts.verifyClient
+    subscriptionContextFn = subscriptionOpts.context
     onConnect = subscriptionOpts.onConnect
   } else if (subscriptionOpts === true) {
     emitter = mq()
@@ -186,7 +188,8 @@ const plugin = fp(async function (app, opts) {
       verifyClient,
       onConnect,
       lruGatewayResolvers,
-      entityResolversFactory
+      entityResolversFactory,
+      subscriptionContextFn
     })
   }
 
diff --git a/lib/routes.js b/lib/routes.js
index 77a9cd0..1557746 100644
--- a/lib/routes.js
+++ b/lib/routes.js
@@ -116,6 +116,7 @@ module.exports = async function (app, opts) {
     })
   }
   const contextFn = opts.context
+  const { subscriptionContextFn } = opts
 
   const {
     path: graphqlPath = '/graphql',
@@ -245,7 +246,8 @@ module.exports = async function (app, opts) {
       verifyClient,
       onConnect,
       lruGatewayResolvers,
-      entityResolversFactory
+      entityResolversFactory,
+      subscriptionContextFn
     })
   } else {
     app.route(getOptions)
diff --git a/lib/subscription-connection.js b/lib/subscription-connection.js
index 57a4519..d1f784b 100644
--- a/lib/subscription-connection.js
+++ b/lib/subscription-connection.js
@@ -23,6 +23,7 @@ module.exports = class SubscriptionConnection {
     fastify,
     lruGatewayResolvers,
     entityResolvers,
+    context,
     onConnect
   }) {
     this.fastify = fastify
@@ -33,6 +34,7 @@ module.exports = class SubscriptionConnection {
     this.subscriber = subscriber
     this.onConnect = onConnect
     this.subscriptionContexts = new Map()
+    this.context = context
 
     this.socket.on('error', this.handleConnectionClose.bind(this))
     this.handleConnection()
@@ -92,6 +94,14 @@ module.exports = class SubscriptionConnection {
         this.sendMessage(GQL_CONNECTION_ERROR, undefined, { message: 'Forbidden' })
         return this.handleConnectionClose()
       }
+
+      // Merge context returned from onConnect
+      if (typeof authorize === 'object') {
+        this.context = {
+          ...this.context,
+          ...authorize
+        }
+      }
     }
 
     this.sendMessage(GQL_CONNECTION_ACK)
@@ -115,6 +125,7 @@ module.exports = class SubscriptionConnection {
       document,
       {}, // rootValue
       {
+        ...this.context,
         pubsub: sc,
         lruGatewayResolvers: this.lruGatewayResolvers,
         reply: {
diff --git a/lib/subscription.js b/lib/subscription.js
index c7ba198..0b22f98 100644
--- a/lib/subscription.js
+++ b/lib/subscription.js
@@ -8,8 +8,8 @@ function handle (conn) {
   conn.end(JSON.stringify({ error: 'unknown route' }))
 }
 
-function createConnectionHandler ({ schema, subscriber, fastify, onConnect, lruGatewayResolvers, entityResolversFactory }) {
-  return (connection, request) => {
+function createConnectionHandler ({ schema, subscriber, fastify, onConnect, lruGatewayResolvers, entityResolversFactory, subscriptionContextFn }) {
+  return async (connection, request) => {
     const { socket } = connection
     if (socket.protocol === undefined ||
       (socket.protocol.indexOf(GRAPHQL_WS) === -1)) {
@@ -21,13 +21,26 @@ function createConnectionHandler ({ schema, subscriber, fastify, onConnect, lruG
       return
     }
 
+    let context = {}
+    if (subscriptionContextFn) {
+      try {
+        context = await subscriptionContextFn(connection, request)
+      } catch (err) {
+        fastify.log.error(err)
+        // 1011: Internal Error
+        socket.close(1011)
+        return
+      }
+    }
+
     const subscriptionConnection = new SubscriptionConnection(socket, {
       schema,
       subscriber,
       fastify,
       onConnect,
       lruGatewayResolvers,
-      entityResolvers: entityResolversFactory && entityResolversFactory.create()
+      entityResolvers: entityResolversFactory && entityResolversFactory.create(),
+      context
     })
 
     /* istanbul ignore next */
@@ -41,7 +54,7 @@ function createConnectionHandler ({ schema, subscriber, fastify, onConnect, lruG
 }
 
 module.exports = function (fastify, opts, next) {
-  const { getOptions, schema, subscriber, verifyClient, onConnect, lruGatewayResolvers, entityResolversFactory } = opts
+  const { getOptions, schema, subscriber, verifyClient, onConnect, lruGatewayResolvers, entityResolversFactory, subscriptionContextFn } = opts
   fastify.register(fastifyWebsocket, {
     handle,
     options: {
@@ -58,7 +71,8 @@ module.exports = function (fastify, opts, next) {
       fastify,
       onConnect,
       lruGatewayResolvers,
-      entityResolversFactory
+      entityResolversFactory,
+      subscriptionContextFn
     })
   })
 
diff --git a/package.json b/package.json
index 4831338..2c24b03 100644
--- a/package.json
+++ b/package.json
@@ -25,6 +25,7 @@
   "homepage": "https://github.com/mcollina/fastify-gql#readme",
   "devDependencies": {
     "@types/node": "^14.0.23",
+    "@types/ws": "^7.2.6",
     "@typescript-eslint/eslint-plugin": "^3.6.1",
     "@typescript-eslint/parser": "^3.6.1",
     "autocannon": "^5.0.0",


Test patch:
diff --git a/test/subscription-connection.js b/test/subscription-connection.js
index ee1ff6b..420a7c9 100644
--- a/test/subscription-connection.js
+++ b/test/subscription-connection.js
@@ -201,3 +201,23 @@ test('subscription connection handles GQL_START message correctly, when payload.
     payload: { }
   }))
 })
+
+test('subscription connection extends context with onConnect return value', async (t) => {
+  const context = {}
+
+  const sc = new SubscriptionConnection({
+    on () {},
+    close () {},
+    send (message) {
+      t.pass()
+    }
+  }, {
+    context,
+    onConnect: function () {
+      return { data: true }
+    }
+  })
+
+  await sc.handleConnectionInit()
+  t.is(sc.context.data, true)
+})
diff --git a/test/subscription.js b/test/subscription.js
index d35d8c8..49f6c2c 100644
--- a/test/subscription.js
+++ b/test/subscription.js
@@ -223,6 +223,182 @@ test('subscription server sends update to subscriptions', t => {
   })
 })
 
+test('subscription server sends update to subscriptions with custom context', t => {
+  const app = Fastify()
+  t.tearDown(() => app.close())
+
+  const sendTestQuery = () => {
+    app.inject({
+      method: 'POST',
+      url: '/graphql',
+      body: {
+        query: `
+          query {
+            notifications {
+              id
+              message
+            }
+          }
+        `
+      }
+    }, () => {
+      sendTestMutation()
+    })
+  }
+
+  const sendTestMutation = () => {
+    app.inject({
+      method: 'POST',
+      url: '/graphql',
+      body: {
+        query: `
+          mutation {
+            addNotification(message: "Hello World") {
+              id
+            }
+          }
+        `
+      }
+    }, () => {})
+  }
+
+  const emitter = mq()
+  const schema = `
+    type Notification {
+      id: ID!
+      message: String
+    }
+
+    type Query {
+      notifications: [Notification]
+    }
+
+    type Mutation {
+      addNotification(message: String): Notification
+    }
+
+    type Subscription {
+      notificationAdded: Notification
+    }
+  `
+
+  let idCount = 1
+  const notifications = [{
+    id: idCount,
+    message: 'Notification message'
+  }]
+
+  const resolvers = {
+    Query: {
+      notifications: () => notifications
+    },
+    Mutation: {
+      addNotification: async (_, { message }) => {
+        const id = idCount++
+        const notification = {
+          id,
+          message
+        }
+        notifications.push(notification)
+        await emitter.emit({
+          topic: 'NOTIFICATION_ADDED',
+          payload: {
+            notificationAdded: notification
+          }
+        })
+
+        return notification
+      }
+    },
+    Subscription: {
+      notificationAdded: {
+        subscribe: (root, args, { pubsub, topic }) => pubsub.subscribe(topic)
+      }
+    }
+  }
+
+  app.register(GQL, {
+    schema,
+    resolvers,
+    subscription: {
+      emitter,
+      context: () => ({ topic: 'NOTIFICATION_ADDED' })
+    }
+  })
+
+  app.listen(0, err => {
+    t.error(err)
+
+    const ws = new WebSocket('ws://localhost:' + (app.server.address()).port + '/graphql', 'graphql-ws')
+    const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8', objectMode: true })
+    t.tearDown(client.destroy.bind(client))
+    client.setEncoding('utf8')
+
+    client.write(JSON.stringify({
+      type: 'connection_init'
+    }))
+
+    client.write(JSON.stringify({
+      id: 1,
+      type: 'start',
+      payload: {
+        query: `
+          subscription {
+            notificationAdded {
+              id
+              message
+            }
+          }
+        `
+      }
+    }))
+
+    client.write(JSON.stringify({
+      id: 2,
+      type: 'start',
+      payload: {
+        query: `
+          subscription {
+            notificationAdded {
+              id
+              message
+            }
+          }
+        `
+      }
+    }))
+
+    client.write(JSON.stringify({
+      id: 2,
+      type: 'stop'
+    }))
+
+    client.on('data', chunk => {
+      const data = JSON.parse(chunk)
+
+      if (data.id === 1 && data.type === 'data') {
+        t.equal(chunk, JSON.stringify({
+          type: 'data',
+          id: 1,
+          payload: {
+            data: {
+              notificationAdded: {
+                id: '1',
+                message: 'Hello World'
+              }
+            }
+          }
+        }))
+
+        client.end()
+        t.end()
+      } else if (data.id === 2 && data.type === 'complete') {
+        sendTestQuery()
+      }
+    })
+  })
+})
+
 test('subscription server register handle function arg is not empty', t => {
   const app = Fastify()
   const schema = `
@@ -290,3 +466,45 @@ test('subscription socket protocol different than graphql-ws, protocol = foobar'
     })
   })
 })
+
+test('subscription connection is closed if context function throws', t => {
+  const app = Fastify()
+  t.tearDown(() => app.close())
+
+  const schema = `
+    type Query {
+      add(x: Int, y: Int): Int
+    }
+  `
+
+  const resolvers = {
+    Query: {
+      add: (parent, { x, y }) => x + y
+    }
+  }
+
+  app.register(GQL, {
+    schema,
+    resolvers,
+    subscription: {
+      context: function () {
+        throw new Error('kaboom')
+      }
+    }
+  })
+
+  app.listen(0, err => {
+    t.error(err)
+
+    const url = 'ws://localhost:' + (app.server.address()).port + '/graphql'
+    const ws = new WebSocket(url, 'graphql-ws')
+    const client = WebSocket.createWebSocketStream(ws, { encoding: 'utf8', objectMode: true })
+    t.tearDown(client.destroy.bind(client))
+
+    client.setEncoding('utf8')
+    ws.on('close', () => {
+      client.end()
+      t.end()
+    })
+  })
+})