Add the ability to reference elements further up the call chain as named parameters

composition-js/src/__tests__/compose.test.ts

@@ -70,7 +70,7 @@ describe('composition', () => {
     expect(result.supergraphSdl).toMatchString(`
       schema
         @link(url: "https://specs.apollo.dev/link/v1.0")
-        @link(url: "https://specs.apollo.dev/join/v0.4", for: EXECUTION)
+        @link(url: "https://specs.apollo.dev/join/v0.5", for: EXECUTION)
       {
         query: Query
       }
@@ -79,7 +79,7 @@ describe('composition', () => {
 
       directive @join__enumValue(graph: join__Graph!) repeatable on ENUM_VALUE
 
-      directive @join__field(graph: join__Graph, requires: join__FieldSet, provides: join__FieldSet, type: String, external: Boolean, override: String, usedOverridden: Boolean, overrideLabel: String) repeatable on FIELD_DEFINITION | INPUT_FIELD_DEFINITION
+      directive @join__field(graph: join__Graph, requires: join__FieldSet, provides: join__FieldSet, type: String, external: Boolean, override: String, usedOverridden: Boolean, overrideLabel: String, contextArguments: [join__ContextArgument!]) repeatable on FIELD_DEFINITION | INPUT_FIELD_DEFINITION
 
       directive @join__graph(name: String!, url: String!) on ENUM_VALUE
 
@@ -98,10 +98,19 @@ describe('composition', () => {
         V2 @join__enumValue(graph: SUBGRAPH2)
       }
 
+      input join__ContextArgument {
+        name: String!
+        type: String!
+        context: String!
+        selection: join__FieldValue!
+      }
+
       scalar join__DirectiveArguments
 
       scalar join__FieldSet
 
+      scalar join__FieldValue
+
       enum join__Graph {
         SUBGRAPH1 @join__graph(name: "Subgraph1", url: "https://Subgraph1")
         SUBGRAPH2 @join__graph(name: "Subgraph2", url: "https://Subgraph2")
@@ -222,7 +231,7 @@ describe('composition', () => {
     expect(result.supergraphSdl).toMatchString(`
       schema
         @link(url: "https://specs.apollo.dev/link/v1.0")
-        @link(url: "https://specs.apollo.dev/join/v0.4", for: EXECUTION)
+        @link(url: "https://specs.apollo.dev/join/v0.5", for: EXECUTION)
       {
         query: Query
       }
@@ -231,7 +240,7 @@ describe('composition', () => {
 
       directive @join__enumValue(graph: join__Graph!) repeatable on ENUM_VALUE
 
-      directive @join__field(graph: join__Graph, requires: join__FieldSet, provides: join__FieldSet, type: String, external: Boolean, override: String, usedOverridden: Boolean, overrideLabel: String) repeatable on FIELD_DEFINITION | INPUT_FIELD_DEFINITION
+      directive @join__field(graph: join__Graph, requires: join__FieldSet, provides: join__FieldSet, type: String, external: Boolean, override: String, usedOverridden: Boolean, overrideLabel: String, contextArguments: [join__ContextArgument!]) repeatable on FIELD_DEFINITION | INPUT_FIELD_DEFINITION
 
       directive @join__graph(name: String!, url: String!) on ENUM_VALUE
 
@@ -250,10 +259,19 @@ describe('composition', () => {
         V2 @join__enumValue(graph: SUBGRAPH2)
       }
 
+      input join__ContextArgument {
+        context: String!
+        name: String!
+        selection: join__FieldValue!
+        type: String!
+      }
+
       scalar join__DirectiveArguments
 
       scalar join__FieldSet
 
+      scalar join__FieldValue
+
       enum join__Graph {
         SUBGRAPH1 @join__graph(name: "Subgraph1", url: "https://Subgraph1")
         SUBGRAPH2 @join__graph(name: "Subgraph2", url: "https://Subgraph2")

composition-js/src/__tests__/override.compose.test.ts

@@ -974,7 +974,7 @@ describe("composition involving @override directive", () => {
       expect(result.supergraphSdl).toMatchInlineSnapshot(`
         "schema
           @link(url: \\"https://specs.apollo.dev/link/v1.0\\")
-          @link(url: \\"https://specs.apollo.dev/join/v0.4\\", for: EXECUTION)
+          @link(url: \\"https://specs.apollo.dev/join/v0.5\\", for: EXECUTION)
         {
           query: Query
         }
@@ -983,7 +983,7 @@ describe("composition involving @override directive", () => {
 
         directive @join__enumValue(graph: join__Graph!) repeatable on ENUM_VALUE
 
-        directive @join__field(graph: join__Graph, requires: join__FieldSet, provides: join__FieldSet, type: String, external: Boolean, override: String, usedOverridden: Boolean, overrideLabel: String) repeatable on FIELD_DEFINITION | INPUT_FIELD_DEFINITION
+        directive @join__field(graph: join__Graph, requires: join__FieldSet, provides: join__FieldSet, type: String, external: Boolean, override: String, usedOverridden: Boolean, overrideLabel: String, contextArguments: [join__ContextArgument!]) repeatable on FIELD_DEFINITION | INPUT_FIELD_DEFINITION
 
         directive @join__graph(name: String!, url: String!) on ENUM_VALUE
 
@@ -995,10 +995,19 @@ describe("composition involving @override directive", () => {
 
         directive @link(url: String, as: String, for: link__Purpose, import: [link__Import]) repeatable on SCHEMA
 
+        input join__ContextArgument {
+          name: String!
+          type: String!
+          context: String!
+          selection: join__FieldValue!
+        }
+
         scalar join__DirectiveArguments
 
         scalar join__FieldSet
 
+        scalar join__FieldValue
+
         enum join__Graph {
           SUBGRAPH1 @join__graph(name: \\"Subgraph1\\", url: \\"https://Subgraph1\\")
           SUBGRAPH2 @join__graph(name: \\"Subgraph2\\", url: \\"https://Subgraph2\\")

composition-js/src/composeDirectiveManager.ts

@@ -176,6 +176,7 @@ export class ComposeDirectiveManager {
       sg.metadata().authenticatedDirective(),
       sg.metadata().requiresScopesDirective(),
       sg.metadata().policyDirective(),
+      sg.metadata().contextDirective(),
     ].map(d => d.name);
     if (directivesComposedByDefault.includes(directive.name)) {
       this.pushHint(new CompositionHint(

composition-js/src/merging/merge.ts

@@ -74,8 +74,12 @@ import {
   LinkDirectiveArgs,
   sourceIdentity,
   FeatureUrl,
+  isFederationDirectiveDefinedInSchema,
+  parseSelectionSet,
+  parseContext,
   CoreFeature,
   Subgraph,
+  StaticArgumentsTransform,
 } from "@apollo/federation-internals";
 import { ASTNode, GraphQLError, DirectiveLocation } from "graphql";
 import {
@@ -291,7 +295,7 @@ class Merger {
   readonly merged: Schema = new Schema();
   readonly subgraphNamesToJoinSpecName: Map<string, string>;
   readonly mergedFederationDirectiveNames = new Set<string>();
-  readonly mergedFederationDirectiveInSupergraph = new Map<string, { definition: DirectiveDefinition, argumentsMerger?: ArgumentMerger }>();
+  readonly mergedFederationDirectiveInSupergraph = new Map<string, { definition: DirectiveDefinition, argumentsMerger?: ArgumentMerger, staticArgumentTransform?: StaticArgumentsTransform }>();
   readonly enumUsages = new Map<string, EnumTypeUsage>();
   private composeDirectiveManager: ComposeDirectiveManager;
   private mismatchReporter: MismatchReporter;
@@ -306,6 +310,7 @@ class Merger {
   private latestFedVersionUsed: FeatureVersion;
   private joinDirectiveIdentityURLs = new Set<string>();
   private schemaToImportNameToFeatureUrl = new Map<Schema, Map<string, FeatureUrl>>();
+  private contextToTypeMap = new Map<string, { types: Set<CompositeType>, usages: { usage: string, argumentDefinition: ArgumentDefinition<FieldDefinition<ObjectType | InterfaceType>> }[] }>();
 
   constructor(readonly subgraphs: Subgraphs, readonly options: CompositionOptions) {
     this.latestFedVersionUsed = this.getLatestFederationVersionUsed();
@@ -455,6 +460,7 @@ class Merger {
       this.mergedFederationDirectiveInSupergraph.set(specInSupergraph.url.name, {
         definition: this.merged.directive(nameInSupergraph)!,
         argumentsMerger,
+        staticArgumentTransform: compositionSpec.staticArgumentTransform,
       });
     }
   }
@@ -872,6 +878,7 @@ class Merger {
       const isInterfaceObject = sourceMetadata.isInterfaceObjectType(source) ? true : undefined;
       const keys = source.appliedDirectivesOf(sourceMetadata.keyDirective());
       const name = this.joinSpecName(idx);
+
       if (!keys.length) {
         dest.applyDirective(joinTypeDirective, { graph: name, isInterfaceObject });
       } else {
@@ -1709,6 +1716,37 @@ class Merger {
         continue;
       }
 
+      const fromContextDirective = this.subgraphs.values()[idx].metadata().fromContextDirective();
+
+      const contextArguments = (source.kind === 'FieldDefinition' ? source.arguments() : [])
+        .map((arg): {
+          context: string,
+          name: string,
+          type: string,
+          selection: string,
+        } | undefined => {
+          if (!isFederationDirectiveDefinedInSchema(fromContextDirective)) {
+            return undefined;
+          }
+          const appliedDirectives = arg.appliedDirectivesOf(fromContextDirective);
+          if (appliedDirectives.length === 0) {
+            return undefined;
+          }
+          assert(appliedDirectives.length === 1, 'There should be at most one @fromContext directive applied to an argument');
+          const directive = appliedDirectives[0];
+          const { context, selection } = parseContext(directive.arguments().field);
+          // these are validated in the subgraph validation phase
+          assert(context, 'Context should be defined');
+          assert(selection, 'Selection should be defined');
+          return {
+            context: `${this.subgraphs.values()[idx].name}__${context}`,
+            name: arg.name,
+            type: arg.type!.toString(),
+            selection,
+          };
+        })
+        .filter(isDefined);
+
       const external = this.isExternal(idx, source);
       const sourceMeta = this.subgraphs.values()[idx].metadata();
       const name = this.joinSpecName(idx);
@@ -1721,10 +1759,10 @@ class Merger {
         external: external ? true : undefined,
         usedOverridden: usedOverridden ? true : undefined,
         overrideLabel: mergeContext.overrideLabel(idx),
+        contextArguments: contextArguments.length > 0 ? contextArguments : undefined,
       });
     }
   }
-
   private getFieldSet(element: SchemaElement<any, any>, directive: DirectiveDefinition<{fields: string}>): string | undefined {
     const applications = element.appliedDirectivesOf(directive);
     assert(applications.length <= 1, () => `Found more than one application of ${directive} on ${element}`);
@@ -1847,6 +1885,29 @@ class Merger {
       // some path. Done because this helps reusing our "reportMismatchHint" method
       // in those cases.
       const arg = dest.addArgument(argName);
+
+      const isContextualArg = (index: number, arg: ArgumentDefinition<DirectiveDefinition<any>> | ArgumentDefinition<FieldDefinition<any>>) => {
+        const fromContextDirective = this.metadata(index).fromContextDirective();
+        return fromContextDirective && isFederationDirectiveDefinedInSchema(fromContextDirective) && arg.appliedDirectivesOf(fromContextDirective).length >= 1;
+      }
+      const hasContextual = sources.map((s, idx) => {
+        const arg = s?.argument(argName);
+        return arg && isContextualArg(idx, arg);
+      });
+
+      if (hasContextual.some((c) => c === true)) {
+        // If any of the sources has a contextual argument, then we need to remove it from the supergraph
+        // and ensure that all the sources have it.
+        if (hasContextual.some((c) => c === false)) {
+          this.errors.push(ERRORS.CONTEXTUAL_ARGUMENT_NOT_CONTEXTUAL_IN_ALL_SUBGRAPHS.err(
+            `Argument "${arg.coordinate}" is contextual in some subgraphs but not in all subgraphs: it is contextual in ${printSubgraphNames(hasContextual.map((c, i) => c ? this.names[i] : undefined).filter(isDefined))} but not in ${printSubgraphNames(hasContextual.map((c, i) => c ? undefined : this.names[i]).filter(isDefined))}`,
+            { nodes: sourceASTs(...sources.map((s) => s?.argument(argName))) },
+          ));
+        }
+        arg.remove();
+        continue;
+      }
+
       // If all the sources that have the field have the argument, we do merge it
       // and we're good, but otherwise ...
       if (sources.some((s) => s && !s.argument(argName))) {
@@ -2590,11 +2651,15 @@ class Merger {
       return;
     }
 
+    const directiveInSupergraph = this.mergedFederationDirectiveInSupergraph.get(name);
+
     if (dest.schema().directive(name)?.repeatable) {
       // For repeatable directives, we simply include each application found but with exact duplicates removed
       while (perSource.length > 0) {
         const directive = this.pickNextDirective(perSource);
-        dest.applyDirective(directive.name, directive.arguments(false));
+
+        const transformedArgs = directiveInSupergraph && directiveInSupergraph.staticArgumentTransform && directiveInSupergraph.staticArgumentTransform(this.subgraphs.values()[0], directive.arguments(false));
+        dest.applyDirective(directive.name, transformedArgs ?? directive.arguments(false));
         // We remove every instances of this particular application. That is we remove any other applicaiton with
         // the same arguments. Note that when doing so, we include default values. This allows "merging" 2 applications
         // when one rely on the default value while another don't but explicitely uses that exact default value.
@@ -2980,7 +3045,44 @@ class Merger {
         }
       }
     }
+    this.validateContextUsages();
+  }
+
+  // private traverseSelectionSetForType(
+  //   selection: string,
+  //   type: ObjectType,
+  // ) {
+  //   const selectionSet = new SelectionSet(type, )
+  // }
 
+  private validateContextUsages() {
+    // For each usage of a context, we need to validate that all set contexts could fulfill the selection of the context
+    this.contextToTypeMap.forEach(({ usages, types }, context) => {
+      for (const { usage, argumentDefinition } of usages) {
+        if (types.size === 0) {
+          this.errors.push(ERRORS.CONTEXT_NOT_SET.err(
+            `Context "${context}" is used in "${argumentDefinition.coordinate}" but is never set in any subgraph.`,
+            { nodes: sourceASTs(argumentDefinition) }
+          ));
+        }
+        // const resolvedTypes = [];
+        for (const type of types) {
+          // now ensure that for each type, the selection is satisfiable and collect the resolved type
+          try {
+            parseSelectionSet({ parentType: type, source: usage });
+          } catch (error) {
+            if (error instanceof GraphQLError) {
+              this.errors.push(ERRORS.CONTEXT_INVALID_SELECTION.err(
+                `Context "${context}" is used in "${argumentDefinition.coordinate}" but the selection is invalid: ${error.message}`,
+                { nodes: sourceASTs(argumentDefinition) }
+              ));
+            } else {
+              throw error;
+            }
+          }
+        }
+      }
+    });
   }
 
   private updateInaccessibleErrorsWithLinkToSubgraphs(

gateway-js/src/__tests__/gateway/lifecycle-hooks.test.ts

@@ -149,7 +149,7 @@ describe('lifecycle hooks', () => {
     // the supergraph (even just formatting differences), this ID will change
     // and this test will have to updated.
     expect(secondCall[0]!.compositionId).toMatchInlineSnapshot(
-      `"208492fefbbc8f62458b3f698b047a7c119f5169d4ccaef4c24b81a1ba82f87c"`,
+      `"6dc1bde2b9818fabec62208c5d8825abaa1bae89635fa6f3a5ffea7b78fc6d82"`,
     );
     // second call should have previous info in the second arg
     expect(secondCall[1]!.compositionId).toEqual(expectedFirstId);