Extend plugins to also apply their logic to subscriptions

Author: CarsonF

src/core/graphql/gql-context.host.ts

@@ -52,6 +52,17 @@ export class GqlContextHostImpl implements GqlContextHost, OnModuleDestroy {
     });
   };
 
+  onSubscribe: Plugin['onSubscribe'] = ({
+    subscribeFn,
+    setSubscribeFn,
+    args,
+  }) => {
+    const ctx = args.contextValue;
+    setSubscribeFn((...args) => {
+      return this.als.run(ctx, subscribeFn, ...args);
+    });
+  };
+
   onModuleDestroy() {
     this.als.disable();
   }

src/core/graphql/graphql-error-formatter.ts

@@ -51,6 +51,13 @@ export class GraphqlErrorFormatter {
       }),
   });
 
+  onSubscribe: Plugin['onSubscribe'] = () => ({
+    onSubscribeError: ({ error, setError }) => {
+      const formatted = this.formatError(error);
+      setError(formatted);
+    },
+  });
+
   formatError = (error: unknown) => {
     if (!(error instanceof GraphQLError)) {
       // I don't think this happens.

src/core/graphql/graphql-session.plugin.ts

@@ -9,4 +9,11 @@ export class GraphqlSessionPlugin {
       args.contextValue.session$.complete();
     },
   });
+  onSubscribe: Plugin['onSubscribe'] = () => ({
+    onSubscribeResult: ({ args }) => ({
+      onEnd: () => {
+        args.contextValue.session$.complete();
+      },
+    }),
+  });
 }