the-guild-org · kamilkisiela · Oct 2, 2024 · Oct 2, 2024
diff --git a/.changeset/short-oranges-melt.md b/.changeset/short-oranges-melt.md
@@ -0,0 +1,5 @@
+---
+"@theguild/federation-composition": minor
+---
+
+Support directives on enum values and unions
diff --git a/__tests__/ast.spec.ts b/__tests__/ast.spec.ts
@@ -786,6 +786,20 @@ describe('union type', () => {
       union Media @inaccessible = Movie | Book
     `);
   });
+
+  test('directives', () => {
+    expect(
+      createUnionTypeNode({
+        name: 'Media',
+        members: ['Book', 'Movie'],
+        ast: {
+          directives: [createDirective('custom')],
+        },
+      }),
+    ).toEqualGraphQL(/* GraphQL */ `
+      union Media @custom = Movie | Book
+    `);
+  });
 });
 
 describe('input object type', () => {
@@ -996,6 +1010,9 @@ describe('enum object type', () => {
         values: [
           {
             name: 'BOOK',
+            ast: {
+              directives: [createDirective('any')],
+            },
           },
           {
             name: 'MOVIE',
@@ -1007,7 +1024,7 @@ describe('enum object type', () => {
       }),
     ).toEqualGraphQL(/* GraphQL */ `
       enum Media @custom {
-        BOOK
+        BOOK @any
         MOVIE
       }
     `);

diff --git a/__tests__/composition.spec.ts b/__tests__/composition.spec.ts
@@ -2378,6 +2378,114 @@ testImplementations(api => {
       `);
     });
 
+    test('preserve directive on enum values if included in @composeDirective', () => {
+      const result = composeServices([
+        {
+          name: 'a',
+          typeDefs: parse(/* GraphQL */ `
+              extend schema
+                @link(
+                  url: "https://specs.apollo.dev/federation/${version}"
+                  import: ["@key", "@composeDirective"]
+                )
+                @link(url: "https://myspecs.dev/whatever/v1.0", import: ["@whatever"])
+                @composeDirective(name: "@whatever")
+
+              directive @whatever on ENUM_VALUE
+
+              enum UserType {
+                ADMIN @whatever
+                REGULAR
+              }
+
+              type User @key(fields: "id") {
+                id: ID!
+                name: String!
+                type: UserType!
+              }
+
+              type Query {
+                users: [User]
+              }
+            `),
+        },
+      ]);
+
+      if (version === 'v2.0') {
+        assertCompositionFailure(result);
+        return;
+      }
+
+      assertCompositionSuccess(result);
+
+      expect(result.supergraphSdl).toContainGraphQL(/* GraphQL */ `
+        directive @whatever on ENUM_VALUE
+      `);
+
+      expect(result.supergraphSdl).toContainGraphQL(/* GraphQL */ `
+        enum UserType @join__type(graph: A) {
+          ADMIN @join__enumValue(graph: A) @whatever
+          REGULAR @join__enumValue(graph: A)
+        }
+      `);
+    });
+
+    test('preserve directive on union types if included in @composeDirective', () => {
+      const result = composeServices([
+        {
+          name: 'a',
+          typeDefs: parse(/* GraphQL */ `
+              extend schema
+                @link(
+                  url: "https://specs.apollo.dev/federation/${version}"
+                  import: ["@key", "@composeDirective"]
+                )
+                @link(url: "https://myspecs.dev/whatever/v1.0", import: ["@whatever"])
+                @composeDirective(name: "@whatever")
+
+              directive @whatever on UNION
+
+              type User @key(fields: "id") {
+                id: ID!
+                name: String!
+              }
+
+              type Admin @key(fields: "id") {
+                id: ID!
+                name: String!
+              }
+
+              union Whoever @whatever = User | Admin
+
+              type Query {
+                whoever: [Whoever]
+              }
+            `),
+        },
+      ]);
+
+      if (version === 'v2.0') {
+        assertCompositionFailure(result);
+        return;
+      }
+
+      assertCompositionSuccess(result);
+
+      expect(result.supergraphSdl).toContainGraphQL(/* GraphQL */ `
+        directive @whatever on UNION
+      `);
+
+      expect(result.supergraphSdl).toContainGraphQL(/* GraphQL */ `
+        union Whoever
+          @join__type(graph: A)
+          @join__unionMember(graph: A, member: "Admin")
+          @join__unionMember(graph: A, member: "User")
+          @whatever =
+          | Admin
+          | User
+      `);
+    });
+
     test('preserve directive on interface its field and argument if included in @composeDirective', () => {
       const result = composeServices([
         {

diff --git a/src/subgraph/state.ts b/src/subgraph/state.ts
@@ -142,6 +142,9 @@ export interface UnionType {
   inaccessible: boolean;
   isDefinition: boolean;
   description?: Description;
+  ast: {
+    directives: DirectiveNode[];
+  };
 }
 
 export interface EnumType {
@@ -211,6 +214,9 @@ export interface EnumValue {
   tags: Set<string>;
   description?: Description;
   deprecated?: Deprecated;
+  ast: {
+    directives: DirectiveNode[];
+  };
 }
 
 export interface Argument {
@@ -724,6 +730,7 @@ export function createSubgraphStateBuilder(
           if (composedDirectives.has(node.name.value)) {
             const typeDef = typeNodeInfo.getTypeDef();
             const fieldDef = typeNodeInfo.getFieldDef();
+            const enumValueDef = typeNodeInfo.getValueDef();
             const argDef = typeNodeInfo.getArgumentDef();
 
             if (!typeDef) {
@@ -795,12 +802,27 @@ export function createSubgraphStateBuilder(
               }
               case Kind.ENUM_TYPE_DEFINITION:
               case Kind.ENUM_TYPE_EXTENSION: {
-                enumTypeBuilder.setDirective(typeDef.name.value, node);
+                if (enumValueDef) {
+                  enumTypeBuilder.value.setDirective(
+                    typeDef.name.value,
+                    enumValueDef.name.value,
+                    node,
+                  );
+                } else {
+                  enumTypeBuilder.setDirective(typeDef.name.value, node);
+                }
+                break;
+              }
+              case Kind.UNION_TYPE_DEFINITION:
+              case Kind.UNION_TYPE_EXTENSION: {
+                unionTypeBuilder.setDirective(typeDef.name.value, node);
                 break;
               }
+
               default:
-                // TODO: T07 support directives on other locations than OBJECT, FIELD_DEFINITION, ARGUMENT_DEFINITION
-                throw new Error(`Directives on "${typeDef.kind}" types are not supported yet`);
+                throw new Error(
+                  `Directives on "${typeof typeDef === 'object' && typeDef !== null && 'kind' in typeDef ? (typeDef as any).kind : typeDef}" types are not supported yet`,
+                );
             }
           } else if (node.name.value === 'specifiedBy') {
             const typeDef = typeNodeInfo.getTypeDef();
@@ -1786,6 +1808,9 @@ function unionTypeFactory(state: SubgraphState) {
     setMember(typeName: string, member: string) {
       getOrCreateUnionType(state, typeName).members.add(member);
     },
+    setDirective(typeName: string, directive: DirectiveNode) {
+      getOrCreateUnionType(state, typeName).ast.directives.push(directive);
+    },
   };
 }
 
@@ -1830,6 +1855,9 @@ function enumTypeFactory(state: SubgraphState) {
       setDescription(typeName: string, valueName: string, description: Description) {
         getOrCreateEnumValue(state, typeName, valueName).description = description;
       },
+      setDirective(typeName: string, valueName: string, directive: DirectiveNode) {
+        getOrCreateEnumValue(state, typeName, valueName).ast.directives.push(directive);
+      },
       setInaccessible(typeName: string, valueName: string) {
         getOrCreateEnumValue(state, typeName, valueName).inaccessible = true;
       },
@@ -2089,6 +2117,9 @@ function getOrCreateUnionType(state: SubgraphState, typeName: string): UnionType
     inaccessible: false,
     tags: new Set(),
     isDefinition: false,
+    ast: {
+      directives: [],
+    },
   };
 
   state.types.set(typeName, unionType);
@@ -2230,6 +2261,9 @@ function getOrCreateEnumValue(
     name: enumValueName,
     inaccessible: false,
     tags: new Set(),
+    ast: {
+      directives: [],
+    },
   };
 
   enumType.values.set(enumValueName, enumValue);

diff --git a/src/supergraph/composition/enum-type.ts b/src/supergraph/composition/enum-type.ts
@@ -80,6 +80,10 @@ export function enumTypeBuilder(): TypeBuilder<EnumType, EnumTypeState> {
           valueState.description = value.description;
         }
 
+        value.ast.directives.forEach(directive => {
+          valueState.ast.directives.push(directive);
+        });
+
         valueState.byGraph.set(graph.id, {
           inaccessible: value.inaccessible,
           version: graph.version,
@@ -114,6 +118,9 @@ export function enumTypeBuilder(): TypeBuilder<EnumType, EnumTypeState> {
           inaccessible: value.inaccessible,
           description: value.description,
           deprecated: value.deprecated,
+          ast: {
+            directives: convertToConst(value.ast.directives),
+          },
         })),
         tags: Array.from(enumType.tags),
         inaccessible: enumType.inaccessible,
@@ -185,6 +192,9 @@ type EnumValueState = {
   inaccessible: boolean;
   deprecated?: Deprecated;
   description?: Description;
+  ast: {
+    directives: DirectiveNode[];
+  };
   byGraph: MapByGraph<EnumValueStateInGraph>;
 };
 
@@ -242,6 +252,9 @@ function getOrCreateEnumValue(enumTypeState: EnumTypeState, enumValueName: strin
     tags: new Set(),
     inaccessible: false,
     byGraph: new Map(),
+    ast: {
+      directives: [],
+    },
   };
 
   enumTypeState.values.set(enumValueName, def);

diff --git a/src/supergraph/composition/union-type.ts b/src/supergraph/composition/union-type.ts
@@ -1,7 +1,8 @@
+import type { DirectiveNode } from 'graphql';
 import { FederationVersion } from '../../specifications/federation.js';
 import { Description, UnionType } from '../../subgraph/state.js';
 import { createUnionTypeNode } from './ast.js';
-import type { MapByGraph, TypeBuilder } from './common.js';
+import { convertToConst, type MapByGraph, type TypeBuilder } from './common.js';
 
 export function unionTypeBuilder(): TypeBuilder<UnionType, UnionTypeState> {
   return {
@@ -23,6 +24,10 @@ export function unionTypeBuilder(): TypeBuilder<UnionType, UnionTypeState> {
         unionTypeState.description = type.description;
       }
 
+      type.ast.directives.forEach(directive => {
+        unionTypeState.ast.directives.push(directive);
+      });
+
       unionTypeState.byGraph.set(graph.id, {
         members: type.members,
         version: graph.version,
@@ -50,6 +55,9 @@ export function unionTypeBuilder(): TypeBuilder<UnionType, UnionTypeState> {
             })
             .flat(1),
         },
+        ast: {
+          directives: convertToConst(unionType.ast.directives),
+        },
       });
     },
   };
@@ -64,6 +72,9 @@ export type UnionTypeState = {
   inaccessible: boolean;
   byGraph: MapByGraph<UnionTypeInGraph>;
   members: Set<string>;
+  ast: {
+    directives: DirectiveNode[];
+  };
 };
 
 type UnionTypeInGraph = {
@@ -86,6 +97,9 @@ function getOrCreateUnionType(state: Map<string, UnionTypeState>, typeName: stri
     inaccessible: false,
     hasDefinition: false,
     byGraph: new Map(),
+    ast: {
+      directives: [],
+    },
   };
 
   state.set(typeName, def);