Support directives on enum type definitions and extensions (#70)

.changeset/thin-dolphins-allow.md

@@ -0,0 +1,5 @@
+---
+"@theguild/federation-composition": minor
+---
+
+Support directives on enum type definitions and extensions

__tests__/ast.spec.ts

@@ -988,6 +988,30 @@ describe('enum object type', () => {
       }
     `);
   });
+
+  test('directives', () => {
+    expect(
+      createEnumTypeNode({
+        name: 'Media',
+        values: [
+          {
+            name: 'BOOK',
+          },
+          {
+            name: 'MOVIE',
+          },
+        ],
+        ast: {
+          directives: [createDirective('custom')],
+        },
+      }),
+    ).toEqualGraphQL(/* GraphQL */ `
+      enum Media @custom {
+        BOOK
+        MOVIE
+      }
+    `);
+  });
 });
 
 describe('schema', () => {

__tests__/composition.spec.ts

@@ -2326,6 +2326,58 @@ testImplementations(api => {
       `);
     });
 
+    test('preserve directive on enums 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
+
+              enum UserType @whatever {
+                ADMIN
+                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
+      `);
+
+      expect(result.supergraphSdl).toContainGraphQL(/* GraphQL */ `
+        enum UserType @whatever @join__type(graph: A) {
+          ADMIN @join__enumValue(graph: A)
+          REGULAR @join__enumValue(graph: A)
+        }
+      `);
+    });
+
     test('preserve directive on interface its field and argument if included in @composeDirective', () => {
       const result = composeServices([
         {
@@ -6916,16 +6968,22 @@ testImplementations(api => {
           extend schema
             @link(
               url: "https://specs.apollo.dev/federation/v2.3"
-              import: ["@key", "@composeDirective"]
+              import: ["@key", "@shareable", "@composeDirective"]
             )
             @link(url: "https://myspecs.dev/lowercase/v1.0", import: ["@lowercase"])
             @composeDirective(name: "@lowercase")
 
-          directive @lowercase on FIELD_DEFINITION
+          directive @lowercase on FIELD_DEFINITION | ENUM
 
           type User @key(fields: "id") {
             id: ID! @lowercase
             age: Int!
+            type: UserType! @shareable
+          }
+
+          enum UserType @lowercase {
+            REGULAR
+            ADMIN
           }
 
           type Query {
@@ -6939,17 +6997,23 @@ testImplementations(api => {
           extend schema
             @link(
               url: "https://specs.apollo.dev/federation/v2.3"
-              import: ["@key", "@external", "@requires", "@composeDirective"]
+              import: ["@key", "@external", "@requires", "@shareable", "@composeDirective"]
             )
             @link(url: "https://myspecs.dev/lowercase/v1.0", import: ["@lowercase"])
             @composeDirective(name: "@lowercase")
 
-          directive @lowercase on FIELD_DEFINITION
+          directive @lowercase on FIELD_DEFINITION | ENUM
 
           type User @key(fields: "id") {
             id: ID! @lowercase
             age: Int! @external
             birthday: String @requires(fields: "age")
+            type: UserType! @lowercase @shareable
+          }
+
+          enum UserType @lowercase {
+            REGULAR
+            ADMIN
           }
 
           type Query {
@@ -6967,6 +7031,14 @@ testImplementations(api => {
         id: ID! @lowercase
         age: Int! @join__field(external: true, graph: B) @join__field(graph: A)
         birthday: String @join__field(graph: B, requires: "age")
+        type: UserType! @lowercase
+      }
+    `);
+
+    expect(result.supergraphSdl).toContainGraphQL(/* GraphQL */ `
+      enum UserType @join__type(graph: A) @join__type(graph: B) @lowercase {
+        ADMIN @join__enumValue(graph: A) @join__enumValue(graph: B)
+        REGULAR @join__enumValue(graph: A) @join__enumValue(graph: B)
       }
     `);
   });

src/subgraph/state.ts

@@ -159,6 +159,9 @@ export interface EnumType {
   referencedByOutputType: boolean;
   inputTypeReferences: Set<string>;
   outputTypeReferences: Set<string>;
+  ast: {
+    directives: DirectiveNode[];
+  };
 }
 
 export interface Field {
@@ -790,6 +793,11 @@ export function createSubgraphStateBuilder(
                 }
                 break;
               }
+              case Kind.ENUM_TYPE_DEFINITION:
+              case Kind.ENUM_TYPE_EXTENSION: {
+                enumTypeBuilder.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`);
@@ -1812,6 +1820,9 @@ function enumTypeFactory(state: SubgraphState) {
       getOrCreateEnumType(state, typeName).referencedByOutputType = true;
       getOrCreateEnumType(state, typeName).outputTypeReferences.add(schemaCoordinate);
     },
+    setDirective(typeName: string, directive: DirectiveNode) {
+      getOrCreateEnumType(state, typeName).ast.directives.push(directive);
+    },
     value: {
       setValue(typeName: string, valueName: string) {
         getOrCreateEnumValue(state, typeName, valueName);
@@ -2050,6 +2061,9 @@ function getOrCreateEnumType(state: SubgraphState, typeName: string): EnumType {
     referencedByOutputType: false,
     inputTypeReferences: new Set(),
     outputTypeReferences: new Set(),
+    ast: {
+      directives: [],
+    },
   };
 
   state.types.set(typeName, enumType);

src/supergraph/composition/enum-type.ts

@@ -1,7 +1,8 @@
+import type { DirectiveNode } from 'graphql';
 import { FederationVersion } from '../../specifications/federation.js';
 import { Deprecated, Description, EnumType } from '../../subgraph/state.js';
 import { createEnumTypeNode } from './ast.js';
-import type { MapByGraph, TypeBuilder } from './common.js';
+import { convertToConst, type MapByGraph, type TypeBuilder } from './common.js';
 
 export function enumTypeBuilder(): TypeBuilder<EnumType, EnumTypeState> {
   return {
@@ -51,6 +52,10 @@ export function enumTypeBuilder(): TypeBuilder<EnumType, EnumTypeState> {
         });
       }
 
+      type.ast.directives.forEach(directive => {
+        enumTypeState.ast.directives.push(directive);
+      });
+
       enumTypeState.byGraph.set(graph.id, {
         inaccessible: type.inaccessible,
         version: graph.version,
@@ -121,6 +126,9 @@ export function enumTypeBuilder(): TypeBuilder<EnumType, EnumTypeState> {
             graph: graphName.toUpperCase(),
           })),
         },
+        ast: {
+          directives: convertToConst(enumType.ast.directives),
+        },
       });
     },
   };
@@ -166,6 +174,9 @@ export type EnumTypeState = {
   inputTypeReferences: Set<string>;
   outputTypeReferences: Set<string>;
   values: Map<string, EnumValueState>;
+  ast: {
+    directives: DirectiveNode[];
+  };
 };
 
 type EnumValueState = {
@@ -209,6 +220,9 @@ function getOrCreateEnumType(state: Map<string, EnumTypeState>, typeName: string
     inputTypeReferences: new Set(),
     outputTypeReferences: new Set(),
     byGraph: new Map(),
+    ast: {
+      directives: [],
+    },
   };
 
   state.set(typeName, def);