Preserves source of union members and enum values in supergraph

composition-js/CHANGELOG.md

@@ -4,6 +4,8 @@ This CHANGELOG pertains only to Apollo Federation packages in the 2.x range. The
 
 ## vNext
 
+- Preserves source of union members and enum values in supergraph [PR #2288](https://github.com/apollographql/federation/pull/2288).
+
 ## 2.2.0
 
 - __BREAKING__: composition now rejects `@shareable` on interface fields. The `@shareable` directive is about

composition-js/src/__tests__/__snapshots__/compose.composeDirective.test.ts.snap

@@ -20,6 +20,10 @@ directive @join__field(graph: join__Graph, requires: join__FieldSet, provides: j
 
 directive @join__implements(graph: join__Graph!, interface: String!) repeatable on OBJECT | INTERFACE
 
+directive @join__unionMember(graph: join__Graph!, member: String!) repeatable on UNION
+
+directive @join__enumValue(graph: join__Graph!) repeatable on ENUM_VALUE
+
 directive @mytag(name: String!) repeatable on FIELD_DEFINITION | OBJECT | INTERFACE | UNION | ARGUMENT_DEFINITION | SCALAR | ENUM | ENUM_VALUE | INPUT_OBJECT | INPUT_FIELD_DEFINITION
 
 directive @tag(name: String!, prop: String!) on FIELD_DEFINITION | OBJECT

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

@@ -365,7 +365,7 @@ describe('composing custom core directives', () => {
   });
 
   it.each([
-    '@join__field', '@join__graph', '@join__implements', '@join__type',
+    '@join__field', '@join__graph', '@join__implements', '@join__type', '@join__unionMember', '@join__enumValue'
   ])('join spec directives should result in an error', (directive) => {
     const subgraphA = generateSubgraph({
       name: 'subgraphA',
@@ -376,6 +376,8 @@ describe('composing custom core directives', () => {
         directive @join__graph(name: String!, url: String!) on ENUM_VALUE
         directive @join__implements(graph: join__Graph!, interface: String!) repeatable on OBJECT | INTERFACE
         directive @join__type(graph: join__Graph!, key: join__FieldSet, extension: Boolean! = false, resolvable: Boolean! = true) repeatable on OBJECT | INTERFACE | UNION | ENUM | INPUT_OBJECT | SCALAR
+        directive @join__unionMember(graph: join__Graph!, member: String!) repeatable on UNION
+        directive @join__enumValue(graph: join__Graph!) repeatable on ENUM_VALUE
 
         scalar join__FieldSet
 
@@ -735,7 +737,7 @@ describe('composing custom core directives', () => {
   });
 
   it.each([
-    '@join__field', '@join__graph', '@join__implements', '@join__type'
+    '@join__field', '@join__graph', '@join__implements', '@join__type', '@join__unionMember', '@join__enumValue'
   ])('naming conflict with join spec directives', (directive) => {
     const subgraphA = generateSubgraph({
       name: 'subgraphA',

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

@@ -35,6 +35,12 @@ describe('composition', () => {
         type T @key(fields: "k") {
           k: ID
         }
+
+        type S {
+          x: Int
+        }
+
+        union U = S | T
       `
     }
 
@@ -47,6 +53,11 @@ describe('composition', () => {
           a: Int
           b: String
         }
+
+        enum E {
+          V1
+          V2
+        }
       `
     }
 
@@ -61,6 +72,8 @@ describe('composition', () => {
         query: Query
       }
 
+      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) repeatable on FIELD_DEFINITION | INPUT_FIELD_DEFINITION
 
       directive @join__graph(name: String!, url: String!) on ENUM_VALUE
@@ -69,8 +82,17 @@ describe('composition', () => {
 
       directive @join__type(graph: join__Graph!, key: join__FieldSet, extension: Boolean! = false, resolvable: Boolean! = true, isInterfaceObject: Boolean! = false) repeatable on OBJECT | INTERFACE | UNION | ENUM | INPUT_OBJECT | SCALAR
 
+      directive @join__unionMember(graph: join__Graph!, member: String!) repeatable on UNION
+
       directive @link(url: String, as: String, for: link__Purpose, import: [link__Import]) repeatable on SCHEMA
 
+      enum E
+        @join__type(graph: SUBGRAPH2)
+      {
+        V1 @join__enumValue(graph: SUBGRAPH2)
+        V2 @join__enumValue(graph: SUBGRAPH2)
+      }
+
       scalar join__FieldSet
 
       enum join__Graph {
@@ -99,6 +121,12 @@ describe('composition', () => {
         t: T @join__field(graph: SUBGRAPH1)
       }
 
+      type S
+        @join__type(graph: SUBGRAPH1)
+      {
+        x: Int
+      }
+
       type T
         @join__type(graph: SUBGRAPH1, key: "k")
         @join__type(graph: SUBGRAPH2, key: "k")
@@ -107,19 +135,36 @@ describe('composition', () => {
         a: Int @join__field(graph: SUBGRAPH2)
         b: String @join__field(graph: SUBGRAPH2)
       }
+
+      union U
+        @join__type(graph: SUBGRAPH1)
+        @join__unionMember(graph: SUBGRAPH1, member: "S")
+        @join__unionMember(graph: SUBGRAPH1, member: "T")
+       = S | T
     `);
 
     const [_, api] = schemas(result);
     expect(printSchema(api)).toMatchString(`
+      enum E {
+        V1
+        V2
+      }
+
       type Query {
         t: T
       }
 
+      type S {
+        x: Int
+      }
+
       type T {
         k: ID
         a: Int
         b: String
       }
+
+      union U = S | T
     `);
   })
 
@@ -2140,6 +2185,8 @@ describe('composition', () => {
         query: Query
       }
 
+      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) repeatable on FIELD_DEFINITION | INPUT_FIELD_DEFINITION
 
       directive @join__graph(name: String!, url: String!) on ENUM_VALUE
@@ -2148,6 +2195,8 @@ describe('composition', () => {
 
       directive @join__type(graph: join__Graph!, key: join__FieldSet, extension: Boolean! = false, resolvable: Boolean! = true, isInterfaceObject: Boolean! = false) repeatable on OBJECT | INTERFACE | UNION | ENUM | INPUT_OBJECT | SCALAR
 
+      directive @join__unionMember(graph: join__Graph!, member: String!) repeatable on UNION
+
       directive @link(url: String, as: String, for: link__Purpose, import: [link__Import]) repeatable on SCHEMA
 
       scalar join__FieldSet

composition-js/src/__tests__/matchers/index.ts

@@ -1 +1,2 @@
 import './toMatchString';
+import './toMatchSubgraph';

composition-js/src/__tests__/matchers/toMatchSubgraph.ts

@@ -0,0 +1,30 @@
+import { defaultPrintOptions, orderPrintedDefinitions, Subgraph } from "@apollo/federation-internals";
+
+// Make this file a module (See: https://github.com/microsoft/TypeScript/issues/17736)
+export {};
+
+declare global {
+  namespace jest {
+    interface Matchers<R> {
+      toMatchSubgraph(actual: Subgraph): R;
+    }
+  }
+}
+
+expect.extend({
+  toMatchSubgraph(expected: Subgraph, actual: Subgraph) {
+    // Note: we use `Subgraph.toString`, not `printSchema()` because 1) it's simpler and 2) the former skips federation
+    // specific definitions, making errors diffs more readable.
+    const printOptions = orderPrintedDefinitions(defaultPrintOptions);
+    const expectedString = expected.toString(printOptions);
+    const actualString = actual.toString(printOptions);
+    const pass = this.equals(expectedString, actualString);
+    const msgBase = `For subgraph ${expected.name}\n`
+      + this.utils.matcherHint('toMatchSubgraph', undefined, undefined)
+      + '\n\n'
+    const message = pass
+      ? () => msgBase + `Expected: not ${this.printExpected(expectedString)}`
+      : () =>  msgBase + this.utils.printDiffOrStringify(expectedString, actualString, 'Expected', 'Received', true);
+    return {actual, expected, message, name: 'toMatchString', pass};
+  }
+});

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

@@ -0,0 +1,92 @@
+import { assertCompositionSuccess, composeAsFed2Subgraphs } from "./testHelper";
+import gql from 'graphql-tag';
+import { asFed2SubgraphDocument, buildSubgraph, buildSupergraphSchema, extractSubgraphsFromSupergraph, ServiceDefinition } from "@apollo/federation-internals";
+import './matchers';
+
+function composeAndTestReversibility(subgraphs: ServiceDefinition[]) {
+  const result = composeAsFed2Subgraphs(subgraphs);
+  assertCompositionSuccess(result);
+
+  const extracted = extractSubgraphsFromSupergraph(buildSupergraphSchema(result.supergraphSdl)[0]);
+  for (const expectedSubgraph of subgraphs) {
+    const actual = extracted.get(expectedSubgraph.name)!;
+    // Note: the subgraph extracted from the supergraph are created with their `@link` on the schema definition, not as an extension (no
+    // strong reason for that, it's how the code was written), so let's match that so our follwoing `toMatchSubgraph` don't fail for that.
+    const expected = buildSubgraph(expectedSubgraph.name, '', asFed2SubgraphDocument(expectedSubgraph.typeDefs, { addAsSchemaExtension: false }));
+    expect(actual).toMatchSubgraph(expected);
+  }
+}
+
+it('preserves the source of union members', () => {
+  const s1 = {
+    typeDefs: gql`
+      type Query {
+        uFromS1: U
+      }
+
+      union U = A | B
+
+      type A {
+        a: Int
+      }
+
+      type B {
+        b: Int @shareable
+      }
+    `,
+    name: 'S1',
+  };
+
+  const s2 = {
+    typeDefs: gql`
+      type Query {
+        uFromS2: U
+      }
+
+      union U = B | C
+
+      type B {
+        b: Int @shareable
+      }
+
+      type C {
+        c: Int
+      }
+    `,
+    name: 'S2',
+  };
+
+  composeAndTestReversibility([s1, s2]);
+});
+
+it('preserves the source of enum values', () => {
+  const s1 = {
+    typeDefs: gql`
+      type Query {
+        eFromS1: E
+      }
+
+      enum E {
+        A,
+        B
+      }
+    `,
+    name: 'S1',
+  };
+
+  const s2 = {
+    typeDefs: gql`
+      type Query {
+        eFromS2: E
+      }
+
+      enum E {
+        B,
+        C
+      }
+    `,
+    name: 'S2',
+  };
+
+  composeAndTestReversibility([s1, s2]);
+});

composition-js/src/merging/merge.ts

@@ -1689,10 +1689,32 @@ class Merger {
       }
     }
     for (const type of dest.types()) {
+      this.addJoinUnionMember(sources, dest, type);
       this.hintOnInconsistentUnionMember(sources, dest, type.name);
     }
   }
 
+  private addJoinUnionMember(sources: (UnionType | undefined)[], dest: UnionType, member: ObjectType) {
+    const joinUnionMemberDirective = joinSpec.unionMemberDirective(this.merged);
+    // We should always be merging with the latest join spec, so this should exists, but well, in prior versions where
+    // the directive didn't existed, we simply did had any replacement so ...
+    if (!joinUnionMemberDirective) {
+      return;
+    }
+
+    for (const [idx, source] of sources.entries()) {
+      if (!source?.hasTypeMember(member.name)) {
+        continue;
+      }
+
+      const name = this.joinSpecName(idx);
+      dest.applyDirective(joinUnionMemberDirective, {
+        graph: name,
+        member: member.name,
+      });
+    }
+  }
+
   private hintOnInconsistentUnionMember(
     sources: (UnionType | undefined)[],
     dest: UnionType,
@@ -1768,6 +1790,7 @@ class Merger {
     const valueSources = sources.map(s => s?.value(value.name));
     this.mergeDescription(valueSources, value);
     this.recordAppliedDirectivesToMerge(valueSources, value);
+    this.addJoinEnumValue(valueSources, value);
 
     const inaccessibleInSupergraph = this.mergedFederationDirectiveInSupergraph.get(inaccessibleSpec.inaccessibleDirectiveSpec.name);
     const isInaccessible = inaccessibleInSupergraph && value.hasAppliedDirective(inaccessibleInSupergraph);
@@ -1817,6 +1840,26 @@ class Merger {
     }
   }
 
+  private addJoinEnumValue(sources: (EnumValue | undefined)[], dest: EnumValue) {
+    const joinEnumValueDirective = joinSpec.enumValueDirective(this.merged);
+    // We should always be merging with the latest join spec, so this should exists, but well, in prior versions where
+    // the directive didn't existed, we simply did had any replacement so ...
+    if (!joinEnumValueDirective) {
+      return;
+    }
+
+    for (const [idx, source] of sources.entries()) {
+      if (!source) {
+        continue;
+      }
+
+      const name = this.joinSpecName(idx);
+      dest.applyDirective(joinEnumValueDirective, {
+        graph: name,
+      });
+    }
+  }
+
   private hintOnInconsistentOutputEnumValue(
     sources: (EnumType | undefined)[],
     dest: EnumType,

gateway-js/CHANGELOG.md

@@ -5,7 +5,7 @@ This CHANGELOG pertains only to Apollo Federation packages in the 2.x range. The
 ## vNext
 
 - Adds support for `@interfaceObject` and keys on interfaces [PR #2279](https://github.com/apollographql/federation/pull/2279).
-
+- Preserves source of union members and enum values in supergraph [PR #2288](https://github.com/apollographql/federation/pull/2288).
 
 ## 2.2.2
 

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

@@ -147,7 +147,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).toEqual(
-      '97f11725be210d703ced94cd941ef00d72ab26a9e04bdb724207b9e89e87359e',
+      'ed8cb418d55e7cd069f11d093b2ea29316e1a913a5757f383cc78ed399414104',
     );
     // second call should have previous info in the second arg
     expect(secondCall[1]!.compositionId).toEqual(expectedFirstId);

internals-js/src/__tests__/extractSubgraphsFromSupergraph.test.ts

@@ -596,7 +596,6 @@ test('preserves default values of input object fields', () => {
   expect(inputFieldA?.defaultValue).toBe(1234)
 })
 
-
 test('throw meaningful error for invalid federation directive fieldSet', () => {
   const supergraph = `
     schema