Repository: mcollina/fastify-gql
Pull request #265 (2020-08-26 14:45:59)

fix: handle gateway.service.wsConnectionParams
This PR fixes the `gateway.service.wsConnectionParams` that is currently defined in the documentation but not handled.

It also adds the `connection_error` message handling for the `SubscriptionClient` and the `connectionInitPayload` option.

PacoDu: A little change so that `gateway.service.wsConnectionParams.connectionInitPayload` can be an async function.
PacoDu: is there anything that prevent this PR from being merged ?

Patch:
diff --git a/README.md b/README.md
index 240a116..2702238 100644
--- a/README.md
+++ b/README.md
@@ -900,6 +900,12 @@ __fastify-gql__ supports the following options:
     * `service.rewriteHeaders`: `Function` A function that gets the original headers as a parameter and returns an object containing values that should be added to the headers
     * `service.wsUrl`: The url of the websocket endpoint
     * `service.wsConnectionParams`: `Function` or `Object`
+      * `wsConnectionParams.connectionInitPayload`: `Function` or `Object` An object or a function that returns the `connection_init` payload sent to the service.
+      * `wsConnectionParams.reconnect`: `Boolean` Enable reconnect on connection close (Default: `false`)
+      * `wsConnectionParams.maxReconnectAttempts`: `Number` Defines the maximum reconnect attempts if reconnect is enabled (Default: `Infinity`)
+      * `wsConnectionParams.connectionCallback`: `Function` A function called after a `connection_ack` message is received.
+      * `wsConnectionParams.failedReconnectCallback`: `Function` A function called if reconnect is enabled and maxReconnectAttempts is reached.
+
 * `persistedQueries`: A hash/query map to resolve the full query text using it's unique hash. Overrides `persistedQueryProvider`.
 * `onlyPersisted`: Boolean. Flag to control whether to allow graphql queries other than persisted. When `true`, it'll make the server reject any queries that are not present in the `persistedQueries` option above. It will also disable any ide available (playground/graphiql). Requires `persistedQueries` to be set, and overrides `persistedQueryProvider`.
 * `persistedQueryProvider`
diff --git a/lib/gateway/service-map.js b/lib/gateway/service-map.js
index beaf725..103bbad 100644
--- a/lib/gateway/service-map.js
+++ b/lib/gateway/service-map.js
@@ -50,6 +50,19 @@ function createTypeMap (schemaDefinition) {
   return { typeMap, types, extensionTypeMap }
 }
 
+async function getWsOpts (service) {
+  let opts = {
+    serviceName: service.name
+  }
+  if (typeof service.wsConnectionParams === 'object') {
+    opts = { ...opts, ...service.wsConnectionParams }
+  } else if (typeof service.wsConnectionParams === 'function') {
+    opts = { ...opts, ...(await service.wsConnectionParams()) }
+  }
+
+  return opts
+}
+
 async function buildServiceMap (services) {
   const serviceMap = {}
 
@@ -76,9 +89,8 @@ async function buildServiceMap (services) {
         if (serviceConfig.client) {
           serviceConfig.client.close()
 
-          const client = new SubscriptionClient(service.wsUrl, {
-            serviceName: service.name
-          })
+          const wsOpts = await getWsOpts(service)
+          const client = new SubscriptionClient(service.wsUrl, wsOpts)
 
           serviceConfig.client = client
           serviceConfig.createSubscription = client.createSubscription.bind(
@@ -121,9 +133,8 @@ async function buildServiceMap (services) {
     }
 
     if (service.wsUrl) {
-      const client = new SubscriptionClient(service.wsUrl, {
-        serviceName: service.name
-      })
+      const wsOpts = await getWsOpts(service)
+      const client = new SubscriptionClient(service.wsUrl, wsOpts)
 
       serviceConfig.client = client
       serviceConfig.createSubscription = client.createSubscription.bind(client)
diff --git a/lib/subscription-client.js b/lib/subscription-client.js
index 63b91ce..5f44ac8 100644
--- a/lib/subscription-client.js
+++ b/lib/subscription-client.js
@@ -4,6 +4,7 @@ const WebSocket = require('ws')
 const {
   GQL_CONNECTION_INIT,
   GQL_CONNECTION_ACK,
+  GQL_CONNECTION_ERROR,
   GQL_START,
   GQL_DATA,
   GQL_COMPLETE,
@@ -26,7 +27,8 @@ class SubscriptionClient {
       maxReconnectAttempts = Infinity,
       serviceName,
       connectionCallback,
-      failedReconnectCallback
+      failedReconnectCallback,
+      connectionInitPayload
     } = config
 
     this.protocols = [GRAPHQL_WS, ...protocols]
@@ -36,6 +38,7 @@ class SubscriptionClient {
     this.reconnectAttempts = 0
     this.connectionCallback = connectionCallback
     this.failedReconnectCallback = failedReconnectCallback
+    this.connectionInitPayload = connectionInitPayload
 
     this.connect()
   }
@@ -43,10 +46,16 @@ class SubscriptionClient {
   connect () {
     this.socket = new WebSocket(this.uri, this.protocols)
 
-    this.socket.onopen = () => {
+    this.socket.onopen = async () => {
       /* istanbul ignore else */
       if (this.socket && this.socket.readyState === WebSocket.OPEN) {
-        this.sendMessage(null, GQL_CONNECTION_INIT)
+        try {
+          const payload = typeof this.connectionInitPayload === 'function'
+            ? await this.connectionInitPayload() : this.connectionInitPayload
+          this.sendMessage(null, GQL_CONNECTION_INIT, payload)
+        } catch (err) {
+          this.close(this.tryReconnect, false)
+        }
       }
     }
 
@@ -191,6 +200,9 @@ class SubscriptionClient {
           this.operations.delete(operationId)
         }
 
+        break
+      case GQL_CONNECTION_ERROR:
+        this.close(this.tryReconnect, false)
         break
       /* istanbul ignore next */
       default:


Test patch:
diff --git a/test/gateway/subscription.js b/test/gateway/subscription.js
index 9043827..cc4da55 100644
--- a/test/gateway/subscription.js
+++ b/test/gateway/subscription.js
@@ -181,7 +181,6 @@ test('gateway subscription handling works correctly', t => {
       messageService.close()
       userService.close()
       gateway.close()
-      process.exit(0)
     })
     client.setEncoding('utf8')
 
@@ -285,3 +284,103 @@ test('gateway subscription handling works correctly', t => {
 
   createUserService(startMessageService)
 })
+
+test('gateway wsConnectionParams object is passed to SubscriptionClient', t => {
+  function onConnect (data) {
+    t.deepEqual(data.payload, connectionInitPayload)
+    t.end()
+  }
+
+  const connectionInitPayload = {
+    hello: 'world'
+  }
+  const testService = Fastify()
+  t.tearDown(() => {
+    testService.close()
+  })
+
+  testService.register(GQL, {
+    schema: `
+      type Query {
+        test: String
+      }
+    `,
+    federationMetadata: true,
+    subscription: { onConnect }
+  })
+
+  testService.listen(0, async err => {
+    t.error(err)
+    const testServicePort = testService.server.address().port
+
+    const gateway = Fastify()
+    t.tearDown(() => {
+      gateway.close()
+    })
+    gateway.register(GQL, {
+      subscription: true,
+      gateway: {
+        services: [{
+          name: 'test',
+          url: `http://localhost:${testServicePort}/graphql`,
+          wsUrl: `ws://localhost:${testServicePort}/graphql`,
+          wsConnectionParams: {
+            connectionInitPayload
+          }
+        }]
+      }
+    })
+    await gateway.ready()
+  })
+})
+
+test('gateway wsConnectionParams function is passed to SubscriptionClient', t => {
+  function onConnect (data) {
+    t.deepEqual(data.payload, connectionInitPayload)
+    t.end()
+  }
+
+  const connectionInitPayload = {
+    hello: 'world'
+  }
+  const testService = Fastify()
+  t.tearDown(() => {
+    testService.close()
+  })
+
+  testService.register(GQL, {
+    schema: `
+      type Query {
+        test: String
+      }
+    `,
+    federationMetadata: true,
+    subscription: { onConnect }
+  })
+
+  testService.listen(0, async err => {
+    t.error(err)
+    const testServicePort = testService.server.address().port
+
+    const gateway = Fastify()
+    t.tearDown(() => {
+      gateway.close()
+    })
+    gateway.register(GQL, {
+      subscription: true,
+      gateway: {
+        services: [{
+          name: 'test',
+          url: `http://localhost:${testServicePort}/graphql`,
+          wsUrl: `ws://localhost:${testServicePort}/graphql`,
+          wsConnectionParams: async function () {
+            return {
+              connectionInitPayload
+            }
+          }
+        }]
+      }
+    })
+    await gateway.ready()
+  })
+})
diff --git a/test/subscription-client.js b/test/subscription-client.js
index 92a21c4..216c96e 100644
--- a/test/subscription-client.js
+++ b/test/subscription-client.js
@@ -215,3 +215,77 @@ test('subscription client multiple subscriptions unsubscribe removes only one su
     }
   })
 })
+
+test('subscription client closes the connection after GQL_CONNECTION_ERROR type payload received', (t) => {
+  const server = new WS.Server({ port: 0 })
+  const port = server.address().port
+
+  server.on('connection', function connection (ws) {
+    ws.on('message', function incoming (message) {
+      const data = JSON.parse(message)
+      if (data.type === 'connection_init') {
+        ws.send(JSON.stringify({ id: '1', type: 'connection_error' }))
+      }
+    })
+
+    ws.on('close', function () {
+      client.close()
+      server.close()
+      t.end()
+    })
+  })
+
+  const client = new SubscriptionClient(`ws://localhost:${port}`, {
+    reconnect: false,
+    serviceName: 'test-service'
+  })
+})
+
+test('subscription client connectionInitPayload is correctly passed', (t) => {
+  const connectionInitPayload = {
+    hello: 'world'
+  }
+  const server = new WS.Server({ port: 0 })
+  const port = server.address().port
+
+  server.on('connection', function connection (ws) {
+    ws.on('message', function incoming (message) {
+      const data = JSON.parse(message)
+      if (data.type === 'connection_init') {
+        t.deepEqual(data.payload, connectionInitPayload)
+        client.close()
+        server.close()
+        t.end()
+      }
+    })
+  })
+
+  const client = new SubscriptionClient(`ws://localhost:${port}`, {
+    reconnect: false,
+    serviceName: 'test-service',
+    connectionInitPayload: async function () {
+      return connectionInitPayload
+    }
+  })
+})
+
+test('subscription client closes the connection if connectionInitPayload throws', (t) => {
+  const server = new WS.Server({ port: 0 })
+  const port = server.address().port
+
+  server.on('connection', function connection (ws) {
+    ws.on('close', function () {
+      client.close()
+      server.close()
+      t.end()
+    })
+  })
+
+  const client = new SubscriptionClient(`ws://localhost:${port}`, {
+    reconnect: false,
+    serviceName: 'test-service',
+    connectionInitPayload: async function () {
+      throw new Error('kaboom')
+    }
+  })
+})