feat: add subscribe and complete to the context, improve their docs (#56)

Author: reconbot

README.md

@@ -47,7 +47,7 @@ const subscriptionServer = makeServer({
 ### Export the handler
 
 ```ts
-export const handler = subscriptionServer.webSocketHandler;
+export const handler = subscriptionServer.webSocketHandler
 ```
 
 ### Configure API Gateway
@@ -86,7 +86,7 @@ functions:
 
 </details>
 
-### Create DynanmoDB tables for state
+### Create DynamoDB tables for state
 
 In-flight connections and subscriptions need to be persisted.
 
@@ -305,7 +305,7 @@ resource "aws_dynamodb_table" "subscriptions-table" {
 Use the [`subscribe`](docs/README.md#subscribe) function to associate incoming subscriptions with a topic.
 
 ```ts
-import { subscribe } from 'graphql-lambda-subscriptions';
+import { subscribe } from 'graphql-lambda-subscriptions'
 
 export const resolver = {
   Subscribe: {
@@ -326,7 +326,7 @@ Use the [`subscribe`](docs/README.md#subscribe) with [`SubscribeOptions`](docs/i
 > Note: If a function is provided, it will be called **on subscription start** and must return a serializable object.
 
 ```ts
-import { subscribe } from 'graphql-lambda-subscriptions';
+import { subscribe } from 'graphql-lambda-subscriptions'
 
 // Subscription agnostic filter
 subscribe('MY_TOPIC', {
@@ -350,9 +350,9 @@ subscribe('MY_TOPIC',{
 
 #### Publishing events
 
-Use the `publish` on your graphql-lambda-subscriptions server to publish events to active subscriptions. Payloads must be of type `Record<string, any>` so they can be filtered and stored.
+Use the [`publish()`](docs/interfaces/SubscriptionServer.md#publish) function on your graphql-lambda-subscriptions server to publish events to active subscriptions. Payloads must be of type `Record<string, any>` so they can be filtered and stored.
 
-```tsx
+```ts
 subscriptionServer.publish({
   type: 'MY_TOPIC',
   payload: {
@@ -363,7 +363,7 @@ subscriptionServer.publish({
 
 Events can come from many sources
 
-```tsx
+```ts
 // SNS Event
 export const snsHandler = (event) =>
   Promise.all(
@@ -373,17 +373,17 @@ export const snsHandler = (event) =>
         payload: JSON.parse(r.Sns.Message),
       })
     )
-  );
+  )
 
 // Manual Invocation
-export const invocationHandler = (payload) => subscriptionServer.publish({ topic: 'MY_TOPIC', payload });
+export const invocationHandler = (payload) => subscriptionServer.publish({ topic: 'MY_TOPIC', payload })
 ```
 
 #### Completing Subscriptions
 
 Use the `complete` on your graphql-lambda-subscriptions server to complete active subscriptions. Payloads are optional and match against filters like events do.
 
-```tsx
+```ts
 subscriptionServer.complete({
   type: 'MY_TOPIC',
   // optional payload
@@ -395,33 +395,11 @@ subscriptionServer.complete({
 
 ### Context
 
-Context values are accessible in all callback and resolver functions (`resolve`, `filter`, `onAfterSubscribe`, `onSubscribe` and `onComplete`).
-
-<details>
-
-<summary>📖 Default value</summary>
-
-Assuming no `context` argument is provided, the default value is an object containing a `connectionInitPayload` attribute.
-
-This attribute contains the [(optionally parsed)](#events) payload from `connection_init`.
-
-```ts
-export const resolver = {
-  Subscribe: {
-    mySubscription: {
-      resolve: (event, args, context) => {
-        console.log(context.connectionInitPayload); // payload from connection_init
-      },
-    },
-  },
-};
-```
-
-</details>
+[Context](docs/interfaces/ServerArgs.md#context) is provided on the [`ServerArgs`](docs/interfaces/ServerArgs.md) object when creating a server. The values are accessible in all callback and resolver functions (eg. `resolve`, `filter`, `onAfterSubscribe`, `onSubscribe` and `onComplete`).
 
-<details>
+Assuming no `context` argument is provided when creating the server, the default value is an object with `connectionInitPayload`, `connectionId` properties and the [`publish()`](docs/interfaces/SubscriptionServer.md#publish) and [`complete()`](docs/interfaces/SubscriptionServer.md#complete) functions. These properties are merged into a provided object or passed into a provided function.
 
-<summary>📖 Setting static context value</summary>
+#### Setting static context value
 
 An object can be provided via the `context` attribute when calling `makeServer`.
 
@@ -431,16 +409,12 @@ const instance = makeServer({
   context: {
     myAttr: 'hello',
   },
-});
+})
 ```
 
 The default values (above) will be appended to this object prior to execution.
 
-</details>
-
-<details>
-
-<summary>📖 Setting dynamic context value</summary>
+#### Setting dynamic context value
 
 A function (optionally async) can be provided via the `context` attribute when calling `makeServer`.
 
@@ -453,10 +427,31 @@ const instance = makeServer({
     myAttr: 'hello',
     user: connectionInitPayload.user,
   }),
-});
+})
 ```
 
-</details>
+#### Using the context
+
+```ts
+export const resolver = {
+  Subscribe: {
+    mySubscription: {
+      subscribe: subscribe('GREETINGS', {
+        filter(_, _, context) {
+          console.log(context.connectionId) // the connectionId
+        },
+        async onAfterSubscribe(_, _, { connectionId, publish }) {
+          await publish('GREETINGS', { message: `HI from ${connectionId}!` })
+        }
+      })
+      resolve: (event, args, context) => {
+        console.log(context.connectionInitPayload) // payload from connection_init
+        return event.payload.message
+      },
+    },
+  },
+}
+```
 
 ### Side effects
 
@@ -481,7 +476,7 @@ export const resolver = {
       }),
     },
   },
-};
+}
 ```
 
 </details>
@@ -502,7 +497,7 @@ const instance = makeServer({
   onConnect: ({ event }) => {
     /* */
   },
-});
+})
 ```
 
 </details>
@@ -519,7 +514,7 @@ const instance = makeServer({
   onDisconnect: ({ event }) => {
     /* */
   },
-});
+})
 ```
 
 </details>
@@ -536,19 +531,19 @@ const instance = makeServer({
 const instance = makeServer({
   /* ... */
   onConnectionInit: ({ message }) => {
-    const token = message.payload.token;
+    const token = message.payload.token
 
     if (!myValidation(token)) {
-      throw Error('Token validation failed');
+      throw Error('Token validation failed')
     }
 
     // Prevent sensitive data from being written to DB
     return {
       ...message.payload,
       token: undefined,
-    };
+    }
   },
-});
+})
 ```
 
 By default, the (optionally parsed) payload will be accessible via [context](#context).
@@ -569,7 +564,7 @@ const instance = makeServer({
   onSubscribe: ({ event, message }) => {
     /* */
   },
-});
+})
 ```
 
 </details>
@@ -586,7 +581,7 @@ const instance = makeServer({
   onComplete: ({ event, message }) => {
     /* */
   },
-});
+})
 ```
 
 </details>
@@ -603,7 +598,7 @@ const instance = makeServer({
   onError: (error, context) => {
     /* */
   },
-});
+})
 ```
 
 </details>

docs/interfaces/ServerArgs.md

@@ -35,7 +35,7 @@ ___
 
 ### context
 
-• `Optional` **context**: `object` \| (`arg`: { `connectionId`: `string` ; `connectionInitPayload`: `any`  }) => [`MaybePromise`](../README.md#maybepromise)<`object`\>
+• `Optional` **context**: `object` \| (`arg`: { `complete`: (`event`: { `payload?`: `Record`<`string`, `any`\> ; `topic`: `string`  }) => `Promise`<`void`\> ; `connectionId`: `string` ; `connectionInitPayload`: `any` ; `publish`: (`event`: { `payload`: `Record`<`string`, `any`\> ; `topic`: `string`  }) => `Promise`<`void`\>  }) => [`MaybePromise`](../README.md#maybepromise)<`object`\>
 
 Makes the context object for all operations defaults to { connectionInitPayload, connectionId }
 
@@ -59,7 +59,7 @@ ___
 
 • `Optional` **pingpong**: [`MaybePromise`](../README.md#maybepromise)<`Object`\>
 
-If set
+If set you can use the `stepFunctionsHandler` and a step function to setup a per connection ping/pong cycle to detect disconnects sooner than the 10 minute idle timeout.
 
 ___
 

lib/messages/complete.ts

@@ -5,7 +5,7 @@ import { buildExecutionContext } from 'graphql/execution/execute'
 import { collect } from 'streaming-iterables'
 import { SubscribePseudoIterable, MessageHandler, PubSubEvent } from '../types'
 import { deleteConnection } from '../utils/deleteConnection'
-import { constructContext } from '../utils/constructContext'
+import { buildContext } from '../utils/buildContext'
 import { getResolverAndArgs } from '../utils/getResolverAndArgs'
 import { isArray } from '../utils/isArray'
 
@@ -26,7 +26,7 @@ export const complete: MessageHandler<CompleteMessage> =
         server.schema,
         parse(sub.subscription.query),
         undefined,
-        await constructContext({ server, connectionInitPayload: sub.connectionInitPayload, connectionId: sub.connectionId }),
+        await buildContext({ server, connectionInitPayload: sub.connectionInitPayload, connectionId: sub.connectionId }),
         sub.subscription.variables,
         sub.subscription.operationName,
         undefined,

lib/messages/disconnect.ts

@@ -2,7 +2,7 @@ import AggregateError from 'aggregate-error'
 import { parse } from 'graphql'
 import { equals } from '@aws/dynamodb-expressions'
 import { buildExecutionContext } from 'graphql/execution/execute'
-import { constructContext } from '../utils/constructContext'
+import { buildContext } from '../utils/buildContext'
 import { getResolverAndArgs } from '../utils/getResolverAndArgs'
 import { SubscribePseudoIterable, MessageHandler, PubSubEvent } from '../types'
 import { isArray } from '../utils/isArray'
@@ -36,7 +36,7 @@ export const disconnect: MessageHandler<null> = async ({ server, event }) => {
               server.schema,
               parse(sub.subscription.query),
               undefined,
-              await constructContext({ server, connectionInitPayload: sub.connectionInitPayload, connectionId: sub.connectionId }),
+              await buildContext({ server, connectionInitPayload: sub.connectionInitPayload, connectionId: sub.connectionId }),
               sub.subscription.variables,
               sub.subscription.operationName,
               undefined,

lib/messages/subscribe.ts

@@ -6,7 +6,7 @@ import {
   execute,
 } from 'graphql/execution/execute'
 import { APIGatewayWebSocketEvent, ServerClosure, MessageHandler, SubscribePseudoIterable, PubSubEvent } from '../types'
-import { constructContext } from '../utils/constructContext'
+import { buildContext } from '../utils/buildContext'
 import { getResolverAndArgs } from '../utils/getResolverAndArgs'
 import { postToConnection } from '../utils/postToConnection'
 import { deleteConnection } from '../utils/deleteConnection'
@@ -48,7 +48,7 @@ const setupSubscription: MessageHandler<SubscribeMessage> = async ({ server, eve
     })
   }
 
-  const contextValue = await constructContext({ server, connectionInitPayload: connection.payload, connectionId })
+  const contextValue = await buildContext({ server, connectionInitPayload: connection.payload, connectionId })
 
   const execContext = buildExecutionContext(
     server.schema,

lib/pubsub/complete.ts

@@ -4,12 +4,12 @@ import { CompleteMessage, MessageType } from 'graphql-ws'
 import { buildExecutionContext } from 'graphql/execution/execute'
 import { ServerClosure, PubSubEvent, SubscribePseudoIterable, SubscriptionServer } from '../types'
 import { postToConnection } from '../utils/postToConnection'
-import { constructContext } from '../utils/constructContext'
+import { buildContext } from '../utils/buildContext'
 import { getResolverAndArgs } from '../utils/getResolverAndArgs'
 import { isArray } from '../utils/isArray'
 import { getFilteredSubs } from './getFilteredSubs'
 
-export const complete = (serverPromise: Promise<ServerClosure>): SubscriptionServer['complete'] => async event => {
+export const complete = (serverPromise: Promise<ServerClosure> | ServerClosure): SubscriptionServer['complete'] => async event => {
   const server = await serverPromise
   const subscriptions = await getFilteredSubs({ server, event })
   server.log('pubsub:complete %j', { event, subscriptions })
@@ -29,7 +29,7 @@ export const complete = (serverPromise: Promise<ServerClosure>): SubscriptionSer
       server.schema,
       parse(sub.subscription.query),
       undefined,
-      await constructContext({ server, connectionInitPayload: sub.connectionInitPayload, connectionId: sub.connectionId }),
+      await buildContext({ server, connectionInitPayload: sub.connectionInitPayload, connectionId: sub.connectionId }),
       sub.subscription.variables,
       sub.subscription.operationName,
       undefined,

lib/pubsub/publish.ts

@@ -2,10 +2,10 @@ import { parse, execute } from 'graphql'
 import { MessageType, NextMessage } from 'graphql-ws'
 import { ServerClosure, SubscriptionServer } from '../types'
 import { postToConnection } from '../utils/postToConnection'
-import { constructContext } from '../utils/constructContext'
+import { buildContext } from '../utils/buildContext'
 import { getFilteredSubs } from './getFilteredSubs'
 
-export const publish = (serverPromise: Promise<ServerClosure>): SubscriptionServer['publish'] => async event => {
+export const publish = (serverPromise: Promise<ServerClosure> | ServerClosure): SubscriptionServer['publish'] => async event => {
   const server = await serverPromise
   server.log('pubsub:publish %j', { event })
   const subscriptions = await getFilteredSubs({ server, event })
@@ -16,7 +16,7 @@ export const publish = (serverPromise: Promise<ServerClosure>): SubscriptionServ
       server.schema,
       parse(sub.subscription.query),
       event,
-      await constructContext({ server, connectionInitPayload: sub.connectionInitPayload, connectionId: sub.connectionId }),
+      await buildContext({ server, connectionInitPayload: sub.connectionInitPayload, connectionId: sub.connectionId }),
       sub.subscription.variables,
       sub.subscription.operationName,
       undefined,

lib/types.ts

@@ -38,9 +38,9 @@ export interface ServerArgs {
   /**
    * Makes the context object for all operations defaults to { connectionInitPayload, connectionId }
    */
-  context?: ((arg: { connectionInitPayload: any, connectionId: string }) => MaybePromise<object>) | object
+  context?: ((arg: { connectionInitPayload: any, connectionId: string, publish: SubscriptionServer['publish'], complete: SubscriptionServer['complete'] }) => MaybePromise<object>) | object
   /**
-   * If set
+   * If set you can use the `stepFunctionsHandler` and a step function to setup a per connection ping/pong cycle to detect disconnects sooner than the 10 minute idle timeout.
    */
   pingpong?: MaybePromise<{
     /**

lib/utils/buildContext.ts

@@ -0,0 +1,14 @@
+/* eslint-disable @typescript-eslint/ban-types */
+import { complete as completeFactory } from '../pubsub/complete'
+import { publish as publishFactory } from '../pubsub/publish'
+import { ServerClosure } from '../types'
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export const buildContext = ({ server, connectionInitPayload, connectionId }: { connectionInitPayload: object, server: ServerClosure, connectionId: string }): any => {
+  const publish = publishFactory(server)
+  const complete = completeFactory(server)
+  if (typeof server.context === 'function') {
+    return server.context({ connectionInitPayload, connectionId, publish, complete })
+  }
+  return { ...server.context, connectionInitPayload, connectionId, publish, complete }
+}