@apollo/federation-internals 2.1.0-alpha.1 → 2.1.0-alpha.4

package/src/definitions.ts

@@ -15,9 +15,11 @@ import {
   TypeNode,
   VariableDefinitionNode,
   VariableNode,
-  TypeSystemDefinitionNode,
   SchemaDefinitionNode,
-  TypeDefinitionNode
+  TypeDefinitionNode,
+  DefinitionNode,
+  DirectiveDefinitionNode,
+  DirectiveNode,
 } from "graphql";
 import {
   CoreImport,
@@ -32,7 +34,7 @@ import {
 import { assert, mapValues, MapWithCachedArrays, setValues } from "./utils";
 import { withDefaultValues, valueEquals, valueToString, valueToAST, variablesInValue, valueFromAST, valueNodeToConstValueNode, argumentsEquals } from "./values";
 import { removeInaccessibleElements } from "./inaccessibleSpec";
-import { printSchema } from './print';
+import { printDirectiveDefinition, printSchema } from './print';
 import { sameType } from './types';
 import { addIntrospectionFields, introspectionFieldNames, isIntrospectionName } from "./introspection";
 import { err } from '@apollo/core-schema';
@@ -291,6 +293,32 @@ export const executableDirectiveLocations: DirectiveLocation[] = [
   DirectiveLocation.VARIABLE_DEFINITION,
 ];
 
+const executableDirectiveLocationsSet = new Set(executableDirectiveLocations);
+
+export function isExecutableDirectiveLocation(loc: DirectiveLocation): boolean {
+  return executableDirectiveLocationsSet.has(loc);
+}
+
+export const typeSystemDirectiveLocations: DirectiveLocation[] = [
+  DirectiveLocation.SCHEMA,
+  DirectiveLocation.SCALAR,
+  DirectiveLocation.OBJECT,
+  DirectiveLocation.FIELD_DEFINITION,
+  DirectiveLocation.ARGUMENT_DEFINITION,
+  DirectiveLocation.INTERFACE,
+  DirectiveLocation.UNION,
+  DirectiveLocation.ENUM,
+  DirectiveLocation.ENUM_VALUE,
+  DirectiveLocation.INPUT_OBJECT,
+  DirectiveLocation.INPUT_FIELD_DEFINITION,
+];
+
+const typeSystemDirectiveLocationsSet = new Set(typeSystemDirectiveLocations);
+
+export function isTypeSystemDirectiveLocation(loc: DirectiveLocation): boolean {
+  return typeSystemDirectiveLocationsSet.has(loc);
+}
+
 /**
  * Converts a type to an AST of a "reference" to that type, one corresponding to the type `toString()` (and thus never a type definition).
  *
@@ -495,11 +523,37 @@ export class Extension<TElement extends ExtendableElement> {
   }
 }
 
+type UnappliedDirective = {
+  nameOrDef: DirectiveDefinition<Record<string, any>> | string,
+  args: Record<string, any>,
+  extension?: Extension<any>,
+  directive: DirectiveNode,
+};
+
 // TODO: ideally, we should hide the ctor of this class as we rely in places on the fact the no-one external defines new implementations.
 export abstract class SchemaElement<TOwnType extends SchemaElement<any, TParent>, TParent extends SchemaElement<any, any> | Schema> extends Element<TParent> {
   protected readonly _appliedDirectives: Directive<TOwnType>[] = [];
+  protected _unappliedDirectives: UnappliedDirective[] = [];
   description?: string;
 
+  addUnappliedDirective({ nameOrDef, args, extension, directive }: UnappliedDirective) {
+    this._unappliedDirectives.push({
+      nameOrDef,
+      args: args ?? {},
+      extension,
+      directive,
+    });
+  }
+
+  processUnappliedDirectives() {
+    for (const { nameOrDef, args, extension, directive } of this._unappliedDirectives) {
+      const d = this.applyDirective(nameOrDef, args);
+      d.setOfExtension(extension);
+      d.sourceAST = directive;
+    }
+    this._unappliedDirectives = [];
+  }
+
   get appliedDirectives(): readonly Directive<TOwnType>[] {
     return this._appliedDirectives;
   }
@@ -697,7 +751,7 @@ abstract class BaseNamedType<TReferencer, TOwnType extends NamedType & NamedSche
   }
 
   hasNonExtensionElements(): boolean {
-    return this.preserveEmptyDefinition 
+    return this.preserveEmptyDefinition
       || this._appliedDirectives.some(d => d.ofExtension() === undefined)
       || this.hasNonExtensionInnerElements();
   }
@@ -916,6 +970,10 @@ export class SchemaBlueprint {
   onUnknownDirectiveValidationError(_schema: Schema, _unknownDirectiveName: string, error: GraphQLError): GraphQLError {
     return error;
   }
+
+  applyDirectivesAfterParsing() {
+    return false;
+  }
 }
 
 export const defaultSchemaBlueprint = new SchemaBlueprint();
@@ -1009,24 +1067,37 @@ export class CoreFeatures {
     this.byIdentity.set(feature.url.identity, feature);
   }
 
-  sourceFeature(element: DirectiveDefinition | Directive | NamedType): CoreFeature | undefined {
+  sourceFeature(element: DirectiveDefinition | Directive | NamedType): { feature: CoreFeature, nameInFeature: string, isImported: boolean } | undefined {
     const isDirective = element instanceof DirectiveDefinition || element instanceof Directive;
     const splitted = element.name.split('__');
     if (splitted.length > 1) {
-      return this.byAlias.get(splitted[0]);
+      const feature = this.byAlias.get(splitted[0]);
+      return feature ? {
+        feature,
+        nameInFeature: splitted[1],
+        isImported: false,
+      } : undefined;
     } else {
       const directFeature = this.byAlias.get(element.name);
       if (directFeature && isDirective) {
-        return directFeature;
+        return {
+          feature: directFeature,
+          nameInFeature: directFeature.imports.find(imp => imp.as === `@${element.name}`)?.name.slice(1) ?? element.name,
+          isImported: true,
+        };
       }
 
       // Let's see if it's an import. If not, it's not associated to a declared feature.
       const importName = isDirective ? '@' + element.name : element.name;
       const allFeatures = [this.coreItself, ...this.byIdentity.values()];
       for (const feature of allFeatures) {
-        for (const { as } of feature.imports) {
-          if (as === importName) {
-            return feature;
+        for (const { as, name } of feature.imports) {
+          if ((as ?? name) === importName) {
+            return {
+              feature,
+              nameInFeature: name.slice(1),
+              isImported: true,
+            };
           }
         }
       }
@@ -1059,8 +1130,50 @@ const graphQLBuiltInDirectivesSpecifications: readonly DirectiveSpecification[]
     locations: [DirectiveLocation.SCALAR],
     argumentFct: (schema) => ({ args: [{ name: 'url', type: new NonNullType(schema.stringType()) }], errors: [] })
   }),
+  // TODO: currently inconditionally adding @defer as the list of built-in. It's probably fine, but double check if we want to not do so when @defer-support is
+  // not enabled or something (it would probably be hard to handle it at that point anyway but well...).
+  createDirectiveSpecification({
+    name: 'defer',
+    locations: [DirectiveLocation.FRAGMENT_SPREAD, DirectiveLocation.INLINE_FRAGMENT],
+    argumentFct: (schema) => ({
+      args: [
+        { name: 'label', type: schema.stringType() },
+        { name: 'if', type: schema.booleanType() },
+      ],
+      errors: [],
+    })
+  }),
+  // Adding @stream too so that it's know and we don't error out if it is queries. It feels like it would be weird to do so for @stream but not
+  // @defer when both are defined in the same spec. That said, that does *not* mean we currently _implement_ @stream, we don't, and so putting
+  // it in a query will be a no-op at the moment (which technically is valid according to the spec so ...).
+  createDirectiveSpecification({
+    name: 'stream',
+    locations: [DirectiveLocation.FIELD],
+    argumentFct: (schema) => ({
+      args: [
+        { name: 'label', type: schema.stringType() },
+        { name: 'initialCount', type: schema.intType(), defaultValue: 0 },
+        { name: 'if', type: schema.booleanType() },
+      ],
+      errors: [],
+    })
+  }),
 ];
 
+export type DeferDirectiveArgs = {
+  // TODO: we currently do not support variables for the defer label. Passing a label in a variable
+  // feels like a weird use case in the first place, but we should probably fix this nonetheless (or
+  // if we decide to have it be a known limitations, we should at least reject it cleanly).
+  label?: string,
+  if?: boolean | Variable,
+}
+
+export type StreamDirectiveArgs = {
+  label?: string,
+  initialCount: number,
+  if?: boolean,
+}
+
 
 // A coordinate is up to 3 "graphQL name" ([_A-Za-z][_0-9A-Za-z]*).
 const coordinateRegexp = /^@?[_A-Za-z][_0-9A-Za-z]*(\.[_A-Za-z][_0-9A-Za-z]*)?(\([_A-Za-z][_0-9A-Za-z]*:\))?$/;
@@ -1162,7 +1275,7 @@ export class Schema {
     return this.apiSchema;
   }
 
-  private emptyASTDefinitionsForExtensionsWithoutDefinition(): TypeSystemDefinitionNode[] {
+  private emptyASTDefinitionsForExtensionsWithoutDefinition(): DefinitionNode[] {
     const nodes = [];
     if (this.schemaDefinition.hasExtensionElements() && !this.schemaDefinition.hasNonExtensionElements()) {
       const node: SchemaDefinitionNode = { kind: Kind.SCHEMA_DEFINITION, operationTypes: [] };
@@ -1180,7 +1293,9 @@ export class Schema {
     return nodes;
   }
 
-  toGraphQLJSSchema(): GraphQLSchema {
+  toGraphQLJSSchema(config?: { includeDefer?: boolean }): GraphQLSchema {
+    const includeDefer = config?.includeDefer ?? false;
+
     let ast = this.toAST();
 
     // Note that AST generated by `this.toAST()` may not be fully graphQL valid because, in federation subgraphs, we accept
@@ -1188,6 +1303,9 @@ export class Schema {
     // we need to "fix" that problem. For that, we add empty definitions for every element that has extensions without
     // definitions (which is also what `fed1` was effectively doing).
     const additionalNodes = this.emptyASTDefinitionsForExtensionsWithoutDefinition();
+    if (includeDefer) {
+      additionalNodes.push(this.deferDirective().toAST());
+    }
     if (additionalNodes.length > 0) {
       ast = {
         kind: Kind.DOCUMENT,
@@ -1464,28 +1582,35 @@ export class Schema {
   }
 
   private getBuiltInDirective<TApplicationArgs extends {[key: string]: any}>(
-    schema: Schema,
     name: string
   ): DirectiveDefinition<TApplicationArgs> {
-    const directive = schema.directive(name);
+    const directive = this.directive(name);
     assert(directive, `The provided schema has not be built with the ${name} directive built-in`);
     return directive as DirectiveDefinition<TApplicationArgs>;
   }
 
-  includeDirective(schema: Schema): DirectiveDefinition<{if: boolean}> {
-    return this.getBuiltInDirective(schema, 'include');
+  includeDirective(): DirectiveDefinition<{if: boolean}> {
+    return this.getBuiltInDirective('include');
+  }
+
+  skipDirective(): DirectiveDefinition<{if: boolean}> {
+    return this.getBuiltInDirective('skip');
   }
 
-  skipDirective(schema: Schema): DirectiveDefinition<{if: boolean}> {
-    return this.getBuiltInDirective(schema, 'skip');
+  deprecatedDirective(): DirectiveDefinition<{reason?: string}> {
+    return this.getBuiltInDirective('deprecated');
   }
 
-  deprecatedDirective(schema: Schema): DirectiveDefinition<{reason?: string}> {
-    return this.getBuiltInDirective(schema, 'deprecated');
+  specifiedByDirective(): DirectiveDefinition<{url: string}> {
+    return this.getBuiltInDirective('specifiedBy');
   }
 
-  specifiedByDirective(schema: Schema): DirectiveDefinition<{url: string}> {
-    return this.getBuiltInDirective(schema, 'specifiedBy');
+  deferDirective(): DirectiveDefinition<DeferDirectiveArgs> {
+    return this.getBuiltInDirective('defer');
+  }
+
+  streamDirective(): DirectiveDefinition<StreamDirectiveArgs> {
+    return this.getBuiltInDirective('stream');
   }
 
   /**
@@ -1653,7 +1778,7 @@ export class SchemaDefinition extends SchemaElement<SchemaDefinition, Schema>  {
   }
 
   hasNonExtensionElements(): boolean {
-    return this.preserveEmptyDefinition 
+    return this.preserveEmptyDefinition
       || this._appliedDirectives.some((d) => d.ofExtension() === undefined)
       || this.roots().some((r) => r.ofExtension() === undefined);
   }
@@ -2776,6 +2901,9 @@ export class DirectiveDefinition<TApplicationArgs extends {[key: string]: any} =
     return this.addLocations(...Object.values(DirectiveLocation));
   }
 
+  /**
+  * Adds the subset of type system locations that correspond to type definitions.
+  */
   addAllTypeLocations(): DirectiveDefinition {
     return this.addLocations(
       DirectiveLocation.SCALAR,
@@ -2802,6 +2930,14 @@ export class DirectiveDefinition<TApplicationArgs extends {[key: string]: any} =
     return this;
   }
 
+  hasExecutableLocations(): boolean {
+    return this.locations.some((loc) => isExecutableDirectiveLocation(loc));
+  }
+
+  hasTypeSystemLocations(): boolean {
+    return this.locations.some((loc) => isTypeSystemDirectiveLocation(loc));
+  }
+
   applications(): readonly Directive<SchemaElement<any, any>, TApplicationArgs>[] {
     return setValues(this._referencers);
   }
@@ -2858,6 +2994,11 @@ export class DirectiveDefinition<TApplicationArgs extends {[key: string]: any} =
     this.remove().forEach(ref => ref.remove());
   }
 
+  toAST(): DirectiveDefinitionNode {
+    const doc = parse(printDirectiveDefinition(this));
+    return doc.definitions[0] as DirectiveDefinitionNode;
+  }
+
   toString(): string {
     return `@${this.name}`;
   }
@@ -3054,7 +3195,7 @@ export function sameDirectiveApplications(applications1: Directive<any, any>[],
 /**
  * Checks whether a given array of directive applications (`maybeSubset`) is a sub-set of another array of directive applications (`applications`).
  *
- * Sub-set here means that all of the applications in `maybeSubset` appears in `applications`. 
+ * Sub-set here means that all of the applications in `maybeSubset` appears in `applications`.
  */
 export function isDirectiveApplicationsSubset(applications: Directive<any, any>[], maybeSubset: Directive<any, any>[]): boolean {
   if (maybeSubset.length > applications.length) {
@@ -3306,6 +3447,34 @@ function *directivesToCopy(source: Schema, dest: Schema): Generator<DirectiveDef
   yield* source.directives();
 }
 
+/**
+ * Creates, in the provided schema, a directive definition equivalent to the provided one.
+ *
+ * Note that this method assumes that:
+ *  - the provided schema does not already have a directive with the name of the definition to copy.
+ *  - if the copied definition has arguments, then the provided schema has existing types with
+ *    names matching any type used in copied definition.
+ */
+export function copyDirectiveDefinitionToSchema({
+  definition,
+  schema,
+  copyDirectiveApplicationsInArguments = true,
+  locationFilter,
+}: {
+  definition: DirectiveDefinition,
+  schema: Schema,
+  copyDirectiveApplicationsInArguments: boolean,
+  locationFilter?: (loc: DirectiveLocation) => boolean,
+}
+) {
+  copyDirectiveDefinitionInner(
+    definition,
+    schema.addDirectiveDefinition(definition.name),
+    copyDirectiveApplicationsInArguments,
+    locationFilter,
+  );
+}
+
 function copy(source: Schema, dest: Schema) {
   // We shallow copy types first so any future reference to any of them can be dereferenced.
   for (const type of typesToCopy(source, dest)) {
@@ -3458,22 +3627,41 @@ function copyWrapperTypeOrTypeRef(source: Type | undefined, destParent: Schema):
   }
 }
 
-function copyArgumentDefinitionInner<P extends FieldDefinition<any> | DirectiveDefinition>(source: ArgumentDefinition<P>, dest: ArgumentDefinition<P>) {
+function copyArgumentDefinitionInner<P extends FieldDefinition<any> | DirectiveDefinition>(
+  source: ArgumentDefinition<P>,
+  dest: ArgumentDefinition<P>,
+  copyDirectiveApplications: boolean = true,
+) {
   const type = copyWrapperTypeOrTypeRef(source.type, dest.schema()) as InputType;
   dest.type = type;
   dest.defaultValue = source.defaultValue;
-  copyAppliedDirectives(source, dest);
+  if (copyDirectiveApplications) {
+    copyAppliedDirectives(source, dest);
+  }
   dest.description = source.description;
   dest.sourceAST = source.sourceAST;
 }
 
-function copyDirectiveDefinitionInner(source: DirectiveDefinition, dest: DirectiveDefinition) {
+function copyDirectiveDefinitionInner(
+  source: DirectiveDefinition,
+  dest: DirectiveDefinition,
+  copyDirectiveApplicationsInArguments: boolean = true,
+  locationFilter?: (loc: DirectiveLocation) => boolean,
+) {
+  let locations = source.locations;
+  if (locationFilter) {
+    locations = locations.filter((loc) => locationFilter(loc));
+  }
+  if (locations.length === 0) {
+    return;
+  }
+
   for (const arg of source.arguments()) {
     const type = copyWrapperTypeOrTypeRef(arg.type, dest.schema());
-    copyArgumentDefinitionInner(arg, dest.addArgument(arg.name, type as InputType));
+    copyArgumentDefinitionInner(arg, dest.addArgument(arg.name, type as InputType), copyDirectiveApplicationsInArguments);
   }
   dest.repeatable = source.repeatable;
-  dest.addLocations(...source.locations);
+  dest.addLocations(...locations);
   dest.sourceAST = source.sourceAST;
   dest.description = source.description;
 }

package/src/error.ts

@@ -459,6 +459,12 @@ const DOWNSTREAM_SERVICE_ERROR = makeCodeDefinition(
   { addedIn: FED1_CODE },
 );
 
+const DIRECTIVE_COMPOSITION_ERROR = makeCodeDefinition(
+  'DIRECTIVE_COMPOSITION_ERROR',
+  'Error when composing custom directives.',
+  { addedIn: '2.1.0' },
+);
+
 export const ERROR_CATEGORIES = {
   DIRECTIVE_FIELDS_MISSING_EXTERNAL,
   DIRECTIVE_UNSUPPORTED_ON_INTERFACE,
@@ -540,6 +546,7 @@ export const ERRORS = {
   KEY_HAS_DIRECTIVE_IN_FIELDS_ARGS,
   PROVIDES_HAS_DIRECTIVE_IN_FIELDS_ARGS,
   REQUIRES_HAS_DIRECTIVE_IN_FIELDS_ARGS,
+  DIRECTIVE_COMPOSITION_ERROR,
 };
 
 const codeDefByCode = Object.values(ERRORS).reduce((obj: {[code: string]: ErrorCodeDefinition}, codeDef: ErrorCodeDefinition) => { obj[codeDef.code] = codeDef; return obj; }, {});

package/src/extractSubgraphsFromSupergraph.ts

@@ -1,11 +1,13 @@
 import {
   baseType,
   CompositeType,
+  copyDirectiveDefinitionToSchema,
   Directive,
   FieldDefinition,
   InputFieldDefinition,
   InputObjectType,
   InterfaceType,
+  isExecutableDirectiveLocation,
   isEnumType,
   isInterfaceType,
   isObjectType,
@@ -234,6 +236,7 @@ export function extractSubgraphsFromSupergraph(supergraph: Schema): Subgraphs {
       }
     }
 
+    const allExecutableDirectives = supergraph.directives().filter((def) => def.hasExecutableLocations());
     for (const subgraph of subgraphs) {
       if (isFed1) {
         // The join spec in fed1 was not including external fields. Let's make sure we had them or we'll get validation
@@ -266,6 +269,23 @@ export function extractSubgraphsFromSupergraph(supergraph: Schema): Subgraphs {
             break;
         }
       }
+
+      // Lastly, we add all the "executable" directives from the supergraph to each subgraphs, as those may be part
+      // of a query and end up in any subgraph fetches. We do this "last" to make sure that if one of the directive
+      // use a type for an argument, that argument exists.
+      // Note that we don't bother with non-executable directives at the moment since we've don't extract their
+      // applications. It might become something we need later, but we don't so far.
+      for (const definition of allExecutableDirectives) {
+        // Note that we skip any potentially applied directives in the argument of the copied definition, because as said
+        // in the comment above, we haven't copied type-system directives. And so far, we really don't care about those
+        // applications.
+        copyDirectiveDefinitionToSchema({
+          definition,
+          schema: subgraph.schema,
+          copyDirectiveApplicationsInArguments: false,
+          locationFilter: (loc) => isExecutableDirectiveLocation(loc),
+        });
+      }
     }
 
     // TODO: Not sure that code is needed anymore (any field necessary to validate an interface will have been marked

package/src/federation.ts

@@ -76,6 +76,7 @@ import {
   extendsDirectiveSpec,
   shareableDirectiveSpec,
   overrideDirectiveSpec,
+  composeDirectiveSpec,
   FEDERATION2_SPEC_DIRECTIVES,
   ALL_FEDERATION_DIRECTIVES_DEFAULT_NAMES,
   FEDERATION2_ONLY_SPEC_DIRECTIVES,
@@ -599,6 +600,10 @@ export class FederationMetadata {
     return this.getFederationDirective(tagSpec.tagDirectiveSpec.name);
   }
 
+  composeDirective(): DirectiveDefinition<{name: string}> {
+    return this.getFederationDirective(composeDirectiveSpec.name);
+  }
+
   inaccessibleDirective(): DirectiveDefinition<{}> {
     return this.getFederationDirective(
       inaccessibleSpec.inaccessibleDirectiveSpec.name
@@ -615,7 +620,7 @@ export class FederationMetadata {
       this.extendsDirective(),
     ];
     return this.isFed2Schema()
-      ? baseDirectives.concat(this.shareableDirective(), this.inaccessibleDirective(), this.overrideDirective())
+      ? baseDirectives.concat(this.shareableDirective(), this.inaccessibleDirective(), this.overrideDirective(), this.composeDirective())
       : baseDirectives;
   }
 
@@ -698,7 +703,9 @@ export class FederationBlueprint extends SchemaBlueprint {
   }
 
   onDirectiveDefinitionAndSchemaParsed(schema: Schema): GraphQLError[] {
-    return completeSubgraphSchema(schema);
+    const errors = completeSubgraphSchema(schema);
+    schema.schemaDefinition.processUnappliedDirectives();
+    return errors;
   }
 
   onInvalidation(schema: Schema) {
@@ -873,6 +880,10 @@ export class FederationBlueprint extends SchemaBlueprint {
     }
     return error;
   }
+
+  applyDirectivesAfterParsing() {
+    return true;
+  }
 }
 
 function findUnusedNamedForLinkDirective(schema: Schema): string | undefined {
@@ -927,7 +938,7 @@ export function setSchemaAsFed2Subgraph(schema: Schema) {
 
 // This is the full @link declaration as added by `asFed2SubgraphDocument`. It's here primarily for uses by tests that print and match
 // subgraph schema to avoid having to update 20+ tests every time we use a new directive or the order of import changes ...
-export const FEDERATION2_LINK_WTH_FULL_IMPORTS = '@link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@key", "@requires", "@provides", "@external", "@tag", "@extends", "@shareable", "@inaccessible", "@override"])';
+export const FEDERATION2_LINK_WTH_FULL_IMPORTS = '@link(url: "https://specs.apollo.dev/federation/v2.1", import: ["@key", "@requires", "@provides", "@external", "@tag", "@extends", "@shareable", "@inaccessible", "@override", "@composeDirective"])';
 
 export function asFed2SubgraphDocument(document: DocumentNode): DocumentNode {
   const fed2LinkExtension: SchemaExtensionNode = {
@@ -1422,7 +1433,7 @@ export class Subgraph {
     }
 
     const core = this.schema.coreFeatures;
-    return !core || core.sourceFeature(d)?.url.identity !== linkIdentity;
+    return !core || core.sourceFeature(d)?.feature.url.identity !== linkIdentity;
   }
 
   private isPrintedType(t: NamedType): boolean {
@@ -1437,7 +1448,7 @@ export class Subgraph {
     }
 
     const core = this.schema.coreFeatures;
-    return !core || core.sourceFeature(t)?.url.identity !== linkIdentity;
+    return !core || core.sourceFeature(t)?.feature.url.identity !== linkIdentity;
   }
 
   private isPrintedDirectiveApplication(d: Directive): boolean {

package/src/federationSpec.ts

@@ -8,6 +8,7 @@ import {
   ArgumentSpecification,
   createDirectiveSpecification,
   createScalarTypeSpecification,
+  DirectiveSpecification,
 } from "./directiveAndTypeSpecification";
 import { DirectiveLocation, GraphQLError } from "graphql";
 import { assert } from "./utils";
@@ -79,6 +80,16 @@ export const overrideDirectiveSpec = createDirectiveSpecification({
   }),
 });
 
+export const composeDirectiveSpec = createDirectiveSpecification({
+  name: 'composeDirective',
+  locations: [DirectiveLocation.SCHEMA],
+  repeatable: true,
+  argumentFct: (schema) => ({
+    args: [{ name: 'name', type: schema.stringType() }],
+    errors: [],
+  }),
+})
+
 function fieldsArgument(schema: Schema): ArgumentSpecification {
   return { name: 'fields', type: fieldSetType(schema) };
 }
@@ -95,15 +106,24 @@ export const FEDERATION2_ONLY_SPEC_DIRECTIVES = [
   overrideDirectiveSpec,
 ];
 
-// Note that this is only used for federation 2+ (federation 1 adds the same directive, but not through a core spec).
-export const FEDERATION2_SPEC_DIRECTIVES = [
+export const FEDERATION2_1_ONLY_SPEC_DIRECTIVES = [
+  composeDirectiveSpec,
+];
+
+const PRE_FEDERATION2_SPEC_DIRECTIVES = [
   keyDirectiveSpec,
   requiresDirectiveSpec,
   providesDirectiveSpec,
   externalDirectiveSpec,
   TAG_VERSIONS.latest().tagDirectiveSpec,
   extendsDirectiveSpec, // TODO: should we stop supporting that?
+];
+
+// Note that this is only used for federation 2+ (federation 1 adds the same directive, but not through a core spec).
+export const FEDERATION2_SPEC_DIRECTIVES = [
+  ...PRE_FEDERATION2_SPEC_DIRECTIVES,
   ...FEDERATION2_ONLY_SPEC_DIRECTIVES,
+  ...FEDERATION2_1_ONLY_SPEC_DIRECTIVES,
 ];
 
 // Note that this is meant to contain _all_ federation directive names ever supported, regardless of which version.
@@ -120,6 +140,12 @@ export class FederationSpecDefinition extends FeatureDefinition {
     super(new FeatureUrl(federationIdentity, 'federation', version));
   }
 
+  private allFedDirectives(): DirectiveSpecification[] {
+    return PRE_FEDERATION2_SPEC_DIRECTIVES
+      .concat(FEDERATION2_ONLY_SPEC_DIRECTIVES)
+      .concat(this.url.version >= (new FeatureVersion(2, 1)) ? FEDERATION2_1_ONLY_SPEC_DIRECTIVES : []);
+  }
+
   addElementsToSchema(schema: Schema): GraphQLError[] {
     const feature = this.featureInSchema(schema);
     assert(feature, 'The federation specification should have been added to the schema before this is called');
@@ -127,20 +153,21 @@ export class FederationSpecDefinition extends FeatureDefinition {
     let errors: GraphQLError[] = [];
     errors = errors.concat(this.addTypeSpec(schema, fieldSetTypeSpec));
 
-    for (const directive of FEDERATION2_SPEC_DIRECTIVES) {
+    for (const directive of this.allFedDirectives()) {
       errors = errors.concat(this.addDirectiveSpec(schema, directive));
     }
     return errors;
   }
 
   allElementNames(): string[] {
-    return FEDERATION2_SPEC_DIRECTIVES.map((spec) => `@${spec.name}`).concat([
+    return this.allFedDirectives().map((spec) => `@${spec.name}`).concat([
       fieldSetTypeSpec.name,
     ])
   }
 }
 
 export const FEDERATION_VERSIONS = new FeatureDefinitions<FederationSpecDefinition>(federationIdentity)
-  .add(new FederationSpecDefinition(new FeatureVersion(2, 0)));
+  .add(new FederationSpecDefinition(new FeatureVersion(2, 0)))
+  .add(new FederationSpecDefinition(new FeatureVersion(2, 1)));
 
 registerKnownFeature(FEDERATION_VERSIONS);

package/src/inaccessibleSpec.ts

@@ -6,7 +6,6 @@ import {
   EnumType,
   EnumValue,
   ErrGraphQLAPISchemaValidationFailed,
-  executableDirectiveLocations,
   FieldDefinition,
   InputFieldDefinition,
   InputObjectType,
@@ -17,6 +16,7 @@ import {
   isListType,
   isNonNullType,
   isScalarType,
+  isTypeSystemDirectiveLocation,
   isVariable,
   NamedType,
   ObjectType,
@@ -682,11 +682,8 @@ function validateInaccessibleElements(
     }
   }
 
-  const executableDirectiveLocationSet = new Set(executableDirectiveLocations);
   for (const directive of schema.allDirectives()) {
-    const typeSystemLocations = directive.locations.filter((loc) =>
-      !executableDirectiveLocationSet.has(loc)
-    );
+    const typeSystemLocations = directive.locations.filter((loc) => isTypeSystemDirectiveLocation(loc));
     if (hasBuiltInName(directive)) {
       // Built-in directives (and their descendants) aren't allowed to be
       // @inaccessible, regardless of shadowing.