Convert ALS for gql context to yoga plugin

Author: CarsonF

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

@@ -1,20 +1,8 @@
-import {
-  ApolloServerPlugin as ApolloPlugin,
-  GraphQLRequestContext as RequestContext,
-  GraphQLRequestListener as RequestListener,
-} from '@apollo/server';
-import { Plugin } from '@nestjs/apollo';
-import {
-  CallHandler,
-  ExecutionContext,
-  Injectable,
-  NestInterceptor,
-  OnModuleDestroy,
-} from '@nestjs/common';
+import { Injectable, OnModuleDestroy } from '@nestjs/common';
 import { AsyncLocalStorage } from 'async_hooks';
 import { GqlContextType as ContextType } from '~/common';
-import { HttpMiddleware as NestMiddleware } from '~/core/http';
 import { AsyncLocalStorageNoContextException } from '../async-local-storage-no-context.exception';
+import { Plugin } from './plugin.decorator';
 
 export const isGqlContext = (object: unknown): object is ContextType =>
   object != null && typeof object === 'object' && isGqlContext.KEY in object;
@@ -30,85 +18,29 @@ export abstract class GqlContextHost {
   readonly context: ContextType;
 }
 
-/**
- * This is necessary to allow global pipes to have access to GraphQL request context.
- * At least until this is resolved: https://github.com/nestjs/graphql/issues/325
- */
 @Injectable()
 @Plugin()
-export class GqlContextHostImpl
-  implements
-    GqlContextHost,
-    NestMiddleware,
-    NestInterceptor,
-    OnModuleDestroy,
-    ApolloPlugin<ContextType>
-{
-  als = new AsyncLocalStorage<{
-    ctx?: ContextType;
-    execution?: ExecutionContext;
-  }>();
+export class GqlContextHostImpl implements GqlContextHost, OnModuleDestroy {
+  als = new AsyncLocalStorage<ContextType>();
 
-  /**
-   * Unwrap the ALS store or throw error if called incorrectly.
-   */
   get context() {
-    const store = this.als.getStore();
-    if (store?.ctx) {
-      return store.ctx;
-    }
-
-    const message = 'The GraphQL context is not available yet.';
-    if (!store) {
-      throw new AsyncLocalStorageNoContextException(message);
-    }
-    if (
-      !store.ctx &&
-      store.execution &&
-      store.execution.getType() !== 'graphql'
-    ) {
-      throw new NotGraphQLContext(message);
+    const context = this.als.getStore();
+    if (context) {
+      return context;
     }
-    throw new Error(message);
+    throw new AsyncLocalStorageNoContextException(
+      'The GraphQL context is not available yet.',
+    );
   }
 
-  use: NestMiddleware['use'] = (req, res, next) => {
-    // Connect middleware is the only place we get a function where we can
-    // completely wrap the request for the use of an ALS context.
-    this.attachScope(next);
+  onExecute: Plugin['onExecute'] = ({ executeFn, setExecuteFn, args }) => {
+    const ctx = args.contextValue;
+    setExecuteFn((...args) => {
+      return this.als.run(ctx, executeFn, ...args);
+    });
   };
 
-  attachScope<R>(fn: () => R): R {
-    // Just give it a placeholder object for now which we populate below.
-    return this.als.run({}, fn);
-  }
-
-  intercept(context: ExecutionContext, next: CallHandler) {
-    const store = this.als.getStore();
-    if (store) {
-      store.execution = context;
-    }
-
-    return next.handle();
-  }
-
-  /**
-   * Attach GQL context to the ALS store now that we have it.
-   */
-  async requestDidStart({
-    contextValue: context,
-  }: RequestContext<ContextType>): Promise<RequestListener<ContextType>> {
-    const store = this.als.getStore();
-    if (!store) {
-      return {};
-    }
-    store.ctx = context;
-    return {};
-  }
-
   onModuleDestroy() {
     this.als.disable();
   }
 }
-
-export class NotGraphQLContext extends Error {}

src/core/graphql/graphql.module.ts

@@ -1,7 +1,5 @@
-import { MiddlewareConsumer, Module, NestModule } from '@nestjs/common';
-import { APP_INTERCEPTOR } from '@nestjs/core';
+import { Module } from '@nestjs/common';
 import { GraphQLModule as NestGraphqlModule } from '@nestjs/graphql';
-import { HttpAdapterHost } from '~/core/http';
 import { TracingModule } from '../tracing';
 import { Driver } from './driver';
 import { GqlContextHost, GqlContextHostImpl } from './gql-context.host';
@@ -37,19 +35,7 @@ export class GraphqlOptionsModule {}
   providers: [
     GqlContextHostImpl,
     { provide: GqlContextHost, useExisting: GqlContextHostImpl },
-    { provide: APP_INTERCEPTOR, useExisting: GqlContextHostImpl },
   ],
   exports: [NestGraphqlModule, GqlContextHost],
 })
-export class GraphqlModule implements NestModule {
-  constructor(
-    private readonly middleware: GqlContextHostImpl,
-    private readonly app: HttpAdapterHost,
-  ) {}
-
-  configure(consumer: MiddlewareConsumer) {
-    // Always attach our GQL Context middleware.
-    // It has its own logic to handle non-gql requests.
-    consumer.apply(this.middleware.use).forRoutes('*');
-  }
-}
+export class GraphqlModule {}

src/core/graphql/index.ts

@@ -1,7 +1,3 @@
 export * from './graphql.module';
-export {
-  GqlContextHost,
-  isGqlContext,
-  NotGraphQLContext,
-} from './gql-context.host';
+export { GqlContextHost, isGqlContext } from './gql-context.host';
 export * from './plugin.decorator';