#130 AsyncIterator.return()

packages/aws-lambda-graphql/src/DynamoDBConnectionManager.ts

@@ -7,6 +7,7 @@ import {
   IConnectionManager,
   ISubscriptionManager,
   IConnectionData,
+  ISubscriber,
   HydrateConnectionOptions,
 } from './types';
 import { computeTTL } from './helpers';
@@ -211,8 +212,10 @@ export class DynamoDBConnectionManager implements IConnectionManager {
     }
   };
 
-  unregisterConnection = async ({ id }: DynamoDBConnection): Promise<void> => {
-    await Promise.all([
+  unregisterConnection = async ({
+    id,
+  }: DynamoDBConnection): Promise<ISubscriber[]> => {
+    const [, subscribers] = await Promise.all([
       this.db
         .delete({
           Key: {
@@ -223,6 +226,8 @@ export class DynamoDBConnectionManager implements IConnectionManager {
         .promise(),
       this.subscriptions.unsubscribeAllByConnectionId(id),
     ]);
+
+    return subscribers;
   };
 
   closeConnection = async ({ id, data }: DynamoDBConnection): Promise<void> => {

packages/aws-lambda-graphql/src/DynamoDBSubscriptionManager.ts

@@ -309,7 +309,10 @@ export class DynamoDBSubscriptionManager implements ISubscriptionManager {
     }
   };
 
-  unsubscribeAllByConnectionId = async (connectionId: string) => {
+  unsubscribeAllByConnectionId = async (
+    connectionId: string,
+  ): Promise<ISubscriber[]> => {
+    const subscribers: ISubscriber[] = [];
     let cursor: DynamoDB.DocumentClient.Key | undefined;
 
     do {
@@ -324,9 +327,10 @@ export class DynamoDBSubscriptionManager implements ISubscriptionManager {
           Limit: 12, // Maximum of 25 request items sent to DynamoDB a time
         })
         .promise();
+      subscribers.push(...(Items as ISubscriber[]));
 
       if (Items == null || (LastEvaluatedKey == null && Items.length === 0)) {
-        return;
+        return subscribers;
       }
 
       if (Items.length > 0) {
@@ -353,6 +357,8 @@ export class DynamoDBSubscriptionManager implements ISubscriptionManager {
 
       cursor = LastEvaluatedKey;
     } while (cursor);
+
+    return subscribers;
   };
 
   generateSubscriptionId = (

packages/aws-lambda-graphql/src/MemorySubscriptionManager.ts

@@ -136,6 +136,6 @@ export class MemorySubscriptionManager implements ISubscriptionManager {
       );
     }
 
-    return Promise.resolve();
+    return Promise.resolve([]);
   };
 }

packages/aws-lambda-graphql/src/RedisConnectionManager.ts

@@ -144,12 +144,13 @@ export class RedisConnectionManager implements IConnectionManager {
     }
   };
 
-  unregisterConnection = async ({ id }: IConnection): Promise<void> => {
+  unregisterConnection = async ({ id }: IConnection) => {
     const key = prefixRedisKey(`connection:${id}`);
-    await Promise.all([
+    const [, subscribers] = await Promise.all([
       this.redisClient.del(key),
       this.subscriptions.unsubscribeAllByConnectionId(id),
     ]);
+    return subscribers;
   };
 
   closeConnection = async ({ id, data }: IConnection): Promise<void> => {

packages/aws-lambda-graphql/src/RedisSubscriptionManager.ts

@@ -219,6 +219,7 @@ export class RedisSubscriptionManager implements ISubscriptionManager {
   };
 
   unsubscribeAllByConnectionId = async (connectionId: string) => {
+    const subscribers: ISubscriber[] = [];
     let done = false;
     const limit = 50;
     let offset = 0;
@@ -244,6 +245,7 @@ export class RedisSubscriptionManager implements ISubscriptionManager {
               const result = await this.redisClient.get(key);
               if (result) {
                 subscriber = JSON.parse(result);
+                subscribers.push(subscriber as ISubscriber);
                 const subscriptionId = this.generateSubscriptionId(
                   connectionId,
                   subscriber.operationId,
@@ -270,6 +272,7 @@ export class RedisSubscriptionManager implements ISubscriptionManager {
       }
     } while (!done);
     await this.redisClient.del(subscriptionListKey);
+    return subscribers;
   };
 
   generateSubscriptionId = (

packages/aws-lambda-graphql/src/Server.ts

@@ -11,7 +11,7 @@ import {
   Context as LambdaContext,
   Handler as LambdaHandler,
 } from 'aws-lambda';
-import { isAsyncIterable } from 'iterall';
+import { isAsyncIterable, getAsyncIterator } from 'iterall';
 import { ExecutionResult } from 'graphql';
 import { PubSub } from 'graphql-subscriptions';
 import {
@@ -357,7 +357,43 @@ export class Server<
               onDisconnect(connection);
             }
 
-            await this.connectionManager.unregisterConnection(connection);
+            const subscribers = await this.connectionManager.unregisterConnection(
+              connection,
+            );
+
+            const promises = subscribers.map(async (subscriber) => {
+              const pubSub = new PubSub();
+
+              const options = await this.createGraphQLServerOptions(
+                event,
+                lambdaContext,
+                {
+                  connection,
+                  operation: subscriber.operation,
+                  pubSub,
+                },
+              );
+
+              const iterable = await execute({
+                ...options,
+                connection,
+                connectionManager: this.connectionManager,
+                event,
+                lambdaContext,
+                operation: subscriber.operation,
+                pubSub,
+                registerSubscriptions: false,
+                subscriptionManager: this.subscriptionManager,
+              });
+
+              if (!isAsyncIterable(iterable)) {
+                return;
+              }
+
+              const iterator = getAsyncIterator(iterable);
+              await iterator.return?.();
+            });
+            await Promise.all(promises);
 
             return {
               body: '',

packages/aws-lambda-graphql/src/WebSocketConnectionManager.ts

@@ -77,8 +77,9 @@ export class WebSocketConnectionManager implements IConnectionManager {
     });
   };
 
-  unregisterConnection = async (connection: IConnection): Promise<void> => {
+  unregisterConnection = async (connection: IConnection) => {
     this.connections.delete(connection.id);
+    return [];
   };
 
   closeConnection = async (connection: WSConnection): Promise<void> => {

packages/aws-lambda-graphql/src/__tests__/Server.test.ts

@@ -437,7 +437,7 @@ describe('Server', () => {
           {},
         );
         (connectionManager.unregisterConnection as jest.Mock).mockResolvedValueOnce(
-          undefined,
+          [],
         );
 
         await expect(
@@ -478,6 +478,9 @@ describe('Server', () => {
         (connectionManager.hydrateConnection as jest.Mock).mockResolvedValueOnce(
           {},
         );
+        (connectionManager.unregisterConnection as jest.Mock).mockResolvedValueOnce(
+          [],
+        );
 
         await expect(
           handlerWithOnDisconnect(

packages/aws-lambda-graphql/src/types/connection.ts

@@ -0,0 +1,28 @@
+export interface IConnection {
+  /**
+   * Unique connection id
+   */
+  readonly id: string;
+
+  /**
+   * Extra connection data, this data is stored only upon registration
+   * All values should be JSON serializable
+   */
+  readonly data: IConnectionData;
+}
+
+export interface IConnectionData {
+  [key: string]: any;
+
+  /**
+   * Connection context data provided from GQL_CONNECTION_INIT message or from onConnect method
+   * This data is passed to graphql resolvers' context
+   * All values should be JSON serializable
+   */
+  context: Object;
+
+  /**
+   * Indicates whether connection sent GQL_CONNECTION_INIT message or
+   */
+  readonly isInitialized: boolean;
+}

packages/aws-lambda-graphql/src/types/connections.ts

@@ -1,31 +1,7 @@
-export interface IConnection {
-  /**
-   * Unique connection id
-   */
-  readonly id: string;
-
-  /**
-   * Extra connection data, this data is stored only upon registration
-   * All values should be JSON serializable
-   */
-  readonly data: IConnectionData;
-}
-
-export interface IConnectionData {
-  [key: string]: any;
+import { IConnection, IConnectionData } from './connection';
+import { ISubscriber } from './subscriptions';
 
-  /**
-   * Connection context data provided from GQL_CONNECTION_INIT message or from onConnect method
-   * This data is passed to graphql resolvers' context
-   * All values should be JSON serializable
-   */
-  context: Object;
-
-  /**
-   * Indicates whether connection sent GQL_CONNECTION_INIT message or
-   */
-  readonly isInitialized: boolean;
-}
+export { IConnection, IConnectionData };
 
 export interface HydrateConnectionOptions {
   /**
@@ -58,6 +34,6 @@ export interface IConnectionManager {
     connection: IConnection,
     payload: string | Buffer,
   ): Promise<void>;
-  unregisterConnection(connection: IConnection): Promise<void>;
+  unregisterConnection(connection: IConnection): Promise<ISubscriber[]>;
   closeConnection(connection: IConnection): Promise<void>;
 }

packages/aws-lambda-graphql/src/types/subscriptions.ts

@@ -1,4 +1,4 @@
-import { IConnection } from './connections';
+import { IConnection } from './connection';
 import { IdentifiedOperationRequest, OperationRequest } from './operations';
 
 export interface ISubscriptionEvent {
@@ -51,5 +51,5 @@ export interface ISubscriptionManager {
    *
    * @param connectionId
    */
-  unsubscribeAllByConnectionId(connectionId: string): Promise<any>;
+  unsubscribeAllByConnectionId(connectionId: string): Promise<ISubscriber[]>;
 }