feat: provide the connectionId in the context at all times

.eslintrc.js

@@ -16,22 +16,24 @@ module.exports = {
     '@typescript-eslint',
   ],
   rules: {
-    indent: [ 'error', 2 ],
-    'linebreak-style': [ 'error', 'unix' ],
-    quotes: [ 'error', 'single' ],
-    semi: 'off',
-    '@typescript-eslint/semi': ['error', 'never'],
-    'quote-props': ['error', 'as-needed'],
-    'no-param-reassign': 'error',
-    'comma-dangle': ['error', 'always-multiline'],
-    'space-infix-ops': ['error'],
-    'no-multi-spaces': ['error'],
-    'no-unused-vars': 'off',
-    '@typescript-eslint/no-unused-vars': ['error', { argsIgnorePattern: '^_' }],
     '@typescript-eslint/member-delimiter-style': ['error', {
       multiline: { delimiter: 'none' },
       singleline: { delimiter: 'comma', requireLast: false },
       multilineDetection: 'last-member',
     }],
+    '@typescript-eslint/no-unused-vars': ['error', { argsIgnorePattern: '^_' }],
+    '@typescript-eslint/semi': ['error', 'never'],
+    'array-bracket-spacing': ['error', 'never', { singleValue: false }],
+    'comma-dangle': ['error', 'always-multiline'],
+    'linebreak-style': ['error', 'unix'],
+    'no-multi-spaces': ['error'],
+    'no-param-reassign': 'error',
+    'no-unused-vars': 'off',
+    'object-curly-spacing': ['error', 'always'],
+    'quote-props': ['error', 'as-needed'],
+    'space-infix-ops': ['error'],
+    indent: ['error', 2],
+    quotes: ['error', 'single'],
+    semi: 'off',
   },
 }

lib/gateway.ts

@@ -12,81 +12,79 @@ import { subscribe } from './messages/subscribe'
 import { connection_init } from './messages/connection_init'
 import { pong } from './messages/pong'
 
-export const handleGatewayEvent =
-  (server: ServerClosure): ApiGatewayHandler<APIGatewayWebSocketEvent, WebsocketResponse> =>
-    async (event) => {
-      if (!event.requestContext) {
-        return {
-          statusCode: 200,
-          body: '',
-        }
-      }
-
-      if (event.requestContext.eventType === 'CONNECT') {
-        await server.onConnect?.({ event })
-        return {
-          statusCode: 200,
-          headers: {
-            'Sec-WebSocket-Protocol': GRAPHQL_TRANSPORT_WS_PROTOCOL,
-          },
-          body: '',
-        }
-      }
-
-      if (event.requestContext.eventType === 'MESSAGE') {
-        const message = JSON.parse(event.body!)
+export const handleGatewayEvent = (server: ServerClosure): ApiGatewayHandler<APIGatewayWebSocketEvent, WebsocketResponse> => async (event) => {
+  if (!event.requestContext) {
+    return {
+      statusCode: 200,
+      body: '',
+    }
+  }
 
-        if (message.type === MessageType.ConnectionInit) {
-          await connection_init({ server, event, message })
-          return {
-            statusCode: 200,
-            body: '',
-          }
-        }
+  if (event.requestContext.eventType === 'CONNECT') {
+    await server.onConnect?.({ event })
+    return {
+      statusCode: 200,
+      headers: {
+        'Sec-WebSocket-Protocol': GRAPHQL_TRANSPORT_WS_PROTOCOL,
+      },
+      body: '',
+    }
+  }
 
-        if (message.type === MessageType.Subscribe) {
-          await subscribe({ server, event, message })
-          return {
-            statusCode: 200,
-            body: '',
-          }
-        }
+  if (event.requestContext.eventType === 'MESSAGE') {
+    const message = event.body === null ? null : JSON.parse(event.body)
 
-        if (message.type === MessageType.Complete) {
-          await complete({ server, event, message })
-          return {
-            statusCode: 200,
-            body: '',
-          }
-        }
+    if (message.type === MessageType.ConnectionInit) {
+      await connection_init({ server, event, message })
+      return {
+        statusCode: 200,
+        body: '',
+      }
+    }
 
-        if (message.type === MessageType.Ping) {
-          await ping({ server, event, message })
-          return {
-            statusCode: 200,
-            body: '',
-          }
-        }
+    if (message.type === MessageType.Subscribe) {
+      await subscribe({ server, event, message })
+      return {
+        statusCode: 200,
+        body: '',
+      }
+    }
 
-        if (message.type === MessageType.Pong) {
-          await pong({ server, event, message })
-          return {
-            statusCode: 200,
-            body: '',
-          }
-        }
+    if (message.type === MessageType.Complete) {
+      await complete({ server, event, message })
+      return {
+        statusCode: 200,
+        body: '',
       }
+    }
 
-      if (event.requestContext.eventType === 'DISCONNECT') {
-        await disconnect({ server, event, message: null })
-        return {
-          statusCode: 200,
-          body: '',
-        }
+    if (message.type === MessageType.Ping) {
+      await ping({ server, event, message })
+      return {
+        statusCode: 200,
+        body: '',
       }
+    }
 
+    if (message.type === MessageType.Pong) {
+      await pong({ server, event, message })
       return {
         statusCode: 200,
         body: '',
       }
     }
+  }
+
+  if (event.requestContext.eventType === 'DISCONNECT') {
+    await disconnect({ server, event, message: null })
+    return {
+      statusCode: 200,
+      body: '',
+    }
+  }
+
+  return {
+    statusCode: 200,
+    body: '',
+  }
+}

lib/messages/complete.ts

@@ -2,53 +2,47 @@ import AggregateError from 'aggregate-error'
 import { parse } from 'graphql'
 import { CompleteMessage } from 'graphql-ws'
 import { buildExecutionContext } from 'graphql/execution/execute'
-import { SubscribePsuedoIterable, MessageHandler } from '../types'
-import { deleteConnection } from '../utils/aws'
-import { constructContext, getResolverAndArgs } from '../utils/graphql'
+import { collect } from 'streaming-iterables'
+import { SubscribePseudoIterable, MessageHandler } from '../types'
+import { deleteConnection } from '../utils/deleteConnection'
+import { constructContext } from '../utils/constructContext'
+import { getResolverAndArgs } from '../utils/getResolverAndArgs'
 import { isArray } from '../utils/isArray'
 
 /** Handler function for 'complete' message. */
 export const complete: MessageHandler<CompleteMessage> =
   async ({ server, event, message }) => {
     try {
-      const topicSubscriptions = server.mapper.query(server.model.Subscription, {
-        id: `${event.requestContext.connectionId!}|${message.id}`,
-      })
-      let deletions = [] as Promise<any>[]
-      for await (const entity of topicSubscriptions) {
-        deletions = [
-          ...deletions,
-          (async () => {
-            // only call onComplete per subscription
-            if (deletions.length === 0) {
-              const execContext = buildExecutionContext(
-                server.schema,
-                parse(entity.subscription.query),
-                undefined,
-                await constructContext(server)(entity),
-                entity.subscription.variables,
-                entity.subscription.operationName,
-                undefined,
-              )
-
-              if (isArray(execContext)) {
-                throw new AggregateError(execContext)
-              }
+      const topicSubscriptions = await collect(server.mapper.query(server.model.Subscription, {
+        id: `${event.requestContext.connectionId}|${message.id}`,
+      }))
+      if (topicSubscriptions.length === 0) {
+        return
+      }
+      // only call onComplete on the first one as any others are duplicates
+      const sub = topicSubscriptions[0]
+      const execContext = buildExecutionContext(
+        server.schema,
+        parse(sub.subscription.query),
+        undefined,
+        await constructContext({ server, connectionParams: sub.connectionParams, connectionId: sub.connectionId }),
+        sub.subscription.variables,
+        sub.subscription.operationName,
+        undefined,
+      )
 
-              const [field, root, args, context, info] = getResolverAndArgs(server)(execContext)
+      if (isArray(execContext)) {
+        throw new AggregateError(execContext)
+      }
 
-              const onComplete = (field?.subscribe as SubscribePsuedoIterable)?.onComplete
-              if (onComplete) {
-                await onComplete(root, args, context, info)
-              }
-            }
+      const [field, root, args, context, info] = getResolverAndArgs(server)(execContext)
 
-            await server.mapper.delete(entity)
-          })(),
-        ]
+      const onComplete = (field?.subscribe as SubscribePseudoIterable)?.onComplete
+      if (onComplete) {
+        await onComplete(root, args, context, info)
       }
 
-      await Promise.all(deletions)
+      await Promise.all(topicSubscriptions.map(sub => server.mapper.delete(sub)))
     } catch (err) {
       await server.onError?.(err, { event, message })
       await deleteConnection(server)(event.requestContext)

lib/messages/connection_init.ts

@@ -1,7 +1,8 @@
 import { StepFunctions } from 'aws-sdk'
 import { ConnectionInitMessage, MessageType } from 'graphql-ws'
 import { StateFunctionInput, MessageHandler } from '../types'
-import { deleteConnection, sendMessage } from '../utils/aws'
+import { sendMessage } from '../utils/sendMessage'
+import { deleteConnection } from '../utils/deleteConnection'
 
 /** Handler function for 'connection_init' message. */
 export const connection_init: MessageHandler<ConnectionInitMessage> =
@@ -15,10 +16,10 @@ export const connection_init: MessageHandler<ConnectionInitMessage> =
         await new StepFunctions()
           .startExecution({
             stateMachineArn: server.pingpong.machine,
-            name: event.requestContext.connectionId!,
+            name: event.requestContext.connectionId,
             input: JSON.stringify({
-              connectionId: event.requestContext.connectionId!,
-              domainName: event.requestContext.domainName!,
+              connectionId: event.requestContext.connectionId,
+              domainName: event.requestContext.domainName,
               stage: event.requestContext.stage,
               state: 'PING',
               choice: 'WAIT',
@@ -30,7 +31,7 @@ export const connection_init: MessageHandler<ConnectionInitMessage> =
 
       // Write to persistence
       const connection = Object.assign(new server.model.Connection(), {
-        id: event.requestContext.connectionId!,
+        id: event.requestContext.connectionId,
         requestContext: event.requestContext,
         payload: res,
       })

lib/messages/disconnect.ts

@@ -2,40 +2,43 @@ import AggregateError from 'aggregate-error'
 import { parse } from 'graphql'
 import { equals } from '@aws/dynamodb-expressions'
 import { buildExecutionContext } from 'graphql/execution/execute'
-import { constructContext, getResolverAndArgs } from '../utils/graphql'
-import { SubscribePsuedoIterable, MessageHandler } from '../types'
+import { constructContext } from '../utils/constructContext'
+import { getResolverAndArgs } from '../utils/getResolverAndArgs'
+import { SubscribePseudoIterable, MessageHandler } from '../types'
 import { isArray } from '../utils/isArray'
+import { collect } from 'streaming-iterables'
+import { Connection } from '../model/Connection'
 
 /** Handler function for 'disconnect' message. */
 export const disconnect: MessageHandler<null> =
   async ({ server, event }) => {
     try {
       await server.onDisconnect?.({ event })
 
-      const entities = server.mapper.query(
+      const topicSubscriptions = await collect(server.mapper.query(
         server.model.Subscription,
         {
           connectionId: equals(event.requestContext.connectionId),
         },
         { indexName: 'ConnectionIndex' },
-      )
+      ))
 
       const completed = {} as Record<string, boolean>
-      const deletions = [] as Promise<any>[]
-      for await (const entity of entities) {
+      const deletions = [] as Promise<void|Connection>[]
+      for (const sub of topicSubscriptions) {
         deletions.push(
           (async () => {
             // only call onComplete per subscription
-            if (!completed[entity.subscriptionId]) {
-              completed[entity.subscriptionId] = true
+            if (!completed[sub.subscriptionId]) {
+              completed[sub.subscriptionId] = true
 
               const execContext = buildExecutionContext(
                 server.schema,
-                parse(entity.subscription.query),
+                parse(sub.subscription.query),
                 undefined,
-                await constructContext(server)(entity),
-                entity.subscription.variables,
-                entity.subscription.operationName,
+                await constructContext({ server, connectionParams: sub.connectionParams, connectionId: sub.connectionId }),
+                sub.subscription.variables,
+                sub.subscription.operationName,
                 undefined,
               )
 
@@ -46,13 +49,13 @@ export const disconnect: MessageHandler<null> =
 
               const [field, root, args, context, info] = getResolverAndArgs(server)(execContext)
 
-              const onComplete = (field?.subscribe as SubscribePsuedoIterable)?.onComplete
+              const onComplete = (field?.subscribe as SubscribePseudoIterable)?.onComplete
               if (onComplete) {
                 await onComplete(root, args, context, info)
               }
             }
 
-            await server.mapper.delete(entity)
+            await server.mapper.delete(sub)
           })(),
         )
       }
@@ -63,7 +66,7 @@ export const disconnect: MessageHandler<null> =
         // Delete connection
         server.mapper.delete(
           Object.assign(new server.model.Connection(), {
-            id: event.requestContext.connectionId!,
+            id: event.requestContext.connectionId,
           }),
         ),
       ])