FR Support Subscription in Gateway / Federated Schema

Is there an existing issue that is already proposing this?

[X] I have searched the existing issues

Is your feature request related to a problem? Please describe it

Enable subscription in Apollo Gateway Driver (and maybe other) following the solution proposed by Apollo here https://github.com/apollosolutions/federation-subscription-tools.

The problem is to use the Gateway Schema Federation feature in order to use types defined in subgraphs in the subscription schema. Wonder if you are interested in a work on a solution in a PR based on the following POC.

Describe the solution you'd like

I've made a proof of concept for this the solution but there are various problems and limitation:

Subscription Schema actually cannot use code first approach
Subscriptions must be defined at the gateway level in dedicated resolvers
We need use internal queries to forward some fields resolutions to underlying schema
subscriptions are not integrated into the generated schema
works only with graphql-ws ATM

The ApolloGatewayDriver looks like the following:

import { ApolloBaseDriver } from '@nestjs/apollo/dist/drivers/apollo-base.driver';
import { Injectable } from '@nestjs/common';
import { loadPackage } from '@nestjs/common/utils/load-package.util';
import {
  extend,
  GqlSubscriptionService,
  SubscriptionConfig,
} from '@nestjs/graphql';
import {
  DocumentNode,
  getOperationAST,
  GraphQLError,
  parse,
  printSchema,
  validate,
} from 'graphql';
import { ApolloGatewayDriverConfig } from './apollo-gateway-driver.config';
import { typeDefs } from '../resolver/type-defs';
import {
  GraphQLWsSubscriptionsConfig,
  ResolversExplorerService,
} from '@nestjs/graphql/dist/services';
import { makeExecutableSchema } from 'graphql-tools';
import { gql } from 'apollo-server-express';

export function makeSubscriptionSchema({
  gatewaySchema,
  typeDefs,
  resolvers,
}: any) {
  if (!typeDefs || !resolvers) {
    throw new Error(
      'Both `typeDefs` and `resolvers` are required to make the executable subscriptions schema.',
    );
  }

  const gatewayTypeDefs = gatewaySchema
    ? gql(printSchema(gatewaySchema))
    : undefined;

  return makeExecutableSchema({
    typeDefs: [
      ...((gatewayTypeDefs && [gatewayTypeDefs]) as DocumentNode[]),
      typeDefs,
    ],
    resolvers,
  });
}

@Injectable()
export class ApolloGatewayDriver extends ApolloBaseDriver<ApolloGatewayDriverConfig> {
  private _subscriptionService?: GqlSubscriptionService;

  constructor(
    private readonly resolversExplorerService: ResolversExplorerService,
  ) {
    super();
  }

  public async start(options: ApolloGatewayDriverConfig): Promise<void> {
    options.server.plugins = extend(options.server.plugins || [], []);

    const { ApolloGateway } = loadPackage(
      '@apollo/gateway',
      'ApolloGateway',
      () => require('@apollo/gateway'),
    );

    const { server: serverOpts = {}, gateway: gatewayOpts = {} } = options;
    const gateway = new ApolloGateway(gatewayOpts);
    const apolloServerOptions = serverOpts;
    let schema;
    const resolvers = this.resolversExplorerService.explore();

    gateway.onSchemaLoadOrUpdate((schemaContext) => {
      schema = makeSubscriptionSchema({
        gatewaySchema: schemaContext.apiSchema,
        typeDefs,
        resolvers,
      });
    });

    await super.start({
      ...apolloServerOptions,
      gateway,
    });

    if (serverOpts.installSubscriptionHandlers || serverOpts.subscriptions) {
      const subscriptionsOptions: SubscriptionConfig =
        serverOpts.subscriptions || { 'graphql-ws': {} };

      this._subscriptionService = new GqlSubscriptionService(
        {
          schema: serverOpts.schema,
          path: serverOpts.path,
          context: serverOpts.context,
          ...subscriptionsOptions,
          'graphql-ws': {
            ...(subscriptionsOptions[
              'graphql-ws'
            ] as GraphQLWsSubscriptionsConfig),
            onSubscribe: (_ctx, msg) => {
              // Construct the execution arguments
              const args = {
                schema,
                operationName: msg.payload.operationName,
                document: parse(msg.payload.query),
                variableValues: msg.payload.variables,
              };

              const operationAST = getOperationAST(
                args.document,
                args.operationName,
              );

              // Stops the subscription and sends an error message
              if (!operationAST) {
                return [new GraphQLError('Unable to identify operation')];
              }

              // Handle mutation and query requests
              if (operationAST.operation !== 'subscription') {
                return [
                  new GraphQLError(
                    'Only subscription operations are supported',
                  ),
                ];
              }

              // Validate the operation document
              const errors = validate(args.schema, args.document);

              if (errors.length > 0) {
                return errors;
              }

              // Ready execution arguments
              return args;
            },
          },
        },
        this.httpAdapterHost.httpAdapter?.getHttpServer(),
      );
    }
  }

  public async mergeDefaultOptions(
    options: Record<string, any>,
  ): Promise<Record<string, any>> {
    return {
      ...options,
      server: await super.mergeDefaultOptions(options?.server ?? {}),
    };
  }

  public async stop() {
    await this._subscriptionService?.stop();
    await super.stop();
  }
}

I have removed some feature from the original driver due to some missing imports and the ApolloGatewayDriverConfig needs some evolutions / improvements.

Teachability, documentation, adoption, migration strategy

We may also see if the solution can be adapted for the mercurius driver. I may also have missed some underlying options, like updating the schema live, additional configuration options...

What is the motivation / use case for changing the behavior?

We need to be able to create an independent subscription schema without the need of redefining the whole GraphQL types. Anyways with this solution the subscriptions handling and management should be written in the Gateway and fowarded to the underlying schema.

Building an autonomous Subscription service may also be an alternative but it will reuqire to redefine all types or to wrap them into a dedicated module.

Let me know if you are interested

nestjs / graphql