From 8c1d07b3715c049a8bd81c293851e32573cd5eb4 Mon Sep 17 00:00:00 2001
From: Trevor Scheer <trevor@apollographql.com>
Date: Tue, 23 Jul 2019 14:10:27 -0700
Subject: [PATCH] Types cleanup

---
 .../src/composition/compose.ts                |  8 +--
 .../src/composition/types.ts                  | 61 +++++++++++++++----
 .../composition/validate/sdl/matchingEnums.ts | 16 ++---
 3 files changed, 58 insertions(+), 27 deletions(-)

diff --git a/packages/apollo-federation/src/composition/compose.ts b/packages/apollo-federation/src/composition/compose.ts
index 20144c0f317..d46b442a730 100644
--- a/packages/apollo-federation/src/composition/compose.ts
+++ b/packages/apollo-federation/src/composition/compose.ts
@@ -14,6 +14,8 @@ import {
   DocumentNode,
   GraphQLObjectType,
   specifiedDirectives,
+  TypeDefinitionNode,
+  TypeExtensionNode,
 } from 'graphql';
 import { mapValues } from 'apollo-env';
 import { transformSchema } from 'apollo-graphql';
@@ -31,8 +33,6 @@ import {
   ServiceName,
   ExternalFieldDefinition,
   ServiceNameToKeyDirectivesMap,
-  FederatedTypeDefinitionNode,
-  FederatedTypeExtensionNode,
 } from './types';
 import { validateSDL } from 'graphql/validation/validate';
 import { compositionRules } from './rules';
@@ -52,11 +52,11 @@ const EmptyMutationDefinition = {
 
 // Map of all definitions to eventually be passed to extendSchema
 interface DefinitionsMap {
-  [name: string]: FederatedTypeDefinitionNode[];
+  [name: string]: TypeDefinitionNode[];
 }
 // Map of all extensions to eventually be passed to extendSchema
 interface ExtensionsMap {
-  [name: string]: FederatedTypeExtensionNode[];
+  [name: string]: TypeExtensionNode[];
 }
 
 /**
diff --git a/packages/apollo-federation/src/composition/types.ts b/packages/apollo-federation/src/composition/types.ts
index 9334ee9ecda..50c5b44e62d 100644
--- a/packages/apollo-federation/src/composition/types.ts
+++ b/packages/apollo-federation/src/composition/types.ts
@@ -1,10 +1,4 @@
-import {
-  SelectionNode,
-  DocumentNode,
-  FieldDefinitionNode,
-  TypeDefinitionNode,
-  TypeExtensionNode,
-} from 'graphql';
+import { SelectionNode, DocumentNode, FieldDefinitionNode } from 'graphql';
 
 export type ServiceName = string | null;
 
@@ -81,10 +75,51 @@ declare module 'graphql/type/definition' {
   }
 }
 
-export type FederatedTypeDefinitionNode = TypeDefinitionNode & {
-  serviceName: string | null;
-};
+declare module 'graphql/language/ast' {
+  interface UnionTypeDefinitionNode {
+    serviceName?: string | null;
+  }
+  interface UnionTypeExtensionNode {
+    serviceName?: string | null;
+  }
+
+  interface EnumTypeDefinitionNode {
+    serviceName?: string | null;
+  }
+
+  interface EnumTypeExtensionNode {
+    serviceName?: string | null;
+  }
+
+  interface ScalarTypeDefinitionNode {
+    serviceName?: string | null;
+  }
+
+  interface ScalarTypeExtensionNode {
+    serviceName?: string | null;
+  }
+
+  interface ObjectTypeDefinitionNode {
+    serviceName?: string | null;
+  }
+
+  interface ObjectTypeExtensionNode {
+    serviceName?: string | null;
+  }
 
-export type FederatedTypeExtensionNode = TypeExtensionNode & {
-  serviceName: string | null;
-};
+  interface InterfaceTypeDefinitionNode {
+    serviceName?: string | null;
+  }
+
+  interface InterfaceTypeExtensionNode {
+    serviceName?: string | null;
+  }
+
+  interface InputObjectTypeDefinitionNode {
+    serviceName?: string | null;
+  }
+
+  interface InputObjectTypeExtensionNode {
+    serviceName?: string | null;
+  }
+}
diff --git a/packages/apollo-federation/src/composition/validate/sdl/matchingEnums.ts b/packages/apollo-federation/src/composition/validate/sdl/matchingEnums.ts
index 86038e711bf..bb162963a3e 100644
--- a/packages/apollo-federation/src/composition/validate/sdl/matchingEnums.ts
+++ b/packages/apollo-federation/src/composition/validate/sdl/matchingEnums.ts
@@ -4,21 +4,17 @@ import {
   Kind,
   EnumTypeDefinitionNode,
   EnumValueDefinitionNode,
+  TypeDefinitionNode,
 } from 'graphql';
-import { FederatedTypeDefinitionNode } from '../../types';
 import { errorWithCode, logServiceAndType } from '../../utils';
 import { isString } from 'util';
 
-function isEnumDefinition(node: FederatedTypeDefinitionNode) {
+function isEnumDefinition(node: TypeDefinitionNode) {
   return node.kind === Kind.ENUM_TYPE_DEFINITION;
 }
 
 type TypeToDefinitionsMap = {
-  [typeNems: string]: FederatedTypeDefinitionNode[];
-};
-
-type FederatedEnumDefinition = EnumTypeDefinitionNode & {
-  serviceName: string;
+  [typeNems: string]: TypeDefinitionNode[];
 };
 
 export function MatchingEnums(context: SDLValidationContext): ASTVisitor {
@@ -27,8 +23,8 @@ export function MatchingEnums(context: SDLValidationContext): ASTVisitor {
   // group all definitions by name
   // { MyTypeName: [{ serviceName: "A", name: {...}}]}
   let definitionsByName: {
-    [typeName: string]: FederatedTypeDefinitionNode[];
-  } = (definitions as FederatedTypeDefinitionNode[]).reduce(
+    [typeName: string]: TypeDefinitionNode[];
+  } = (definitions as TypeDefinitionNode[]).reduce(
     (typeToDefinitionsMap: TypeToDefinitionsMap, node) => {
       const name = node.name.value;
       if (typeToDefinitionsMap[name]) {
@@ -54,7 +50,7 @@ export function MatchingEnums(context: SDLValidationContext): ASTVisitor {
       for (const {
         values,
         serviceName,
-      } of definitions as FederatedEnumDefinition[]) {
+      } of definitions as EnumTypeDefinitionNode[]) {
         if (serviceName && values)
           simpleEnumDefs.push({
             serviceName,