Skip directive extraction for custom directives which happen to confl…

…ict with default federation directive names (#3108)

When we extract demand control directives from the supergraph to the
query planner's subgraphs, we look for well-known directive names and
copy instances of them over. When looking up the directive name from the
imports, we would use the default expected name. If a custom directive
with the default name was imported, we would wrongly try to copy it over
as the federation directive.

This change makes it so we only extract directives which are imported
from `specs.apollo.dev`, such that we don't confuse a custom `@cost`
directive with the one defined in federation.

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

@@ -206,6 +206,36 @@ const subgraphWithRenamedListSizeFromFederationSpec = {
     `,
 };
 
+const subgraphWithUnimportedCost = {
+  name: 'subgraphWithCost',
+  typeDefs: asFed2SubgraphDocument(gql`
+    enum AorB @federation__cost(weight: 15) {
+      A
+      B
+    }
+
+    input InputTypeWithCost {
+      somethingWithCost: Int @federation__cost(weight: 20)
+    }
+
+    type Query {
+      fieldWithCost: Int @federation__cost(weight: 5)
+      argWithCost(arg: Int @federation__cost(weight: 10)): Int
+      enumWithCost: AorB
+      inputWithCost(someInput: InputTypeWithCost): Int
+    }
+  `),
+};
+
+const subgraphWithUnimportedListSize = {
+  name: 'subgraphWithListSize',
+  typeDefs: asFed2SubgraphDocument(gql`
+    type Query {
+      fieldWithListSize: [String!] @federation__listSize(assumedSize: 2000, requireOneSlicingArgument: false)
+    }
+  `),
+};
+
 // Used to test @cost applications on FIELD_DEFINITION
 function fieldWithCost(result: CompositionResult): FieldDefinition<ObjectType> | undefined {
   return result
@@ -482,18 +512,42 @@ describe('demand control directive composition', () => {
       `);
     });
   });
+
+  describe('when using fully qualified names instead of importing', () => {
+    it('propagates @federation__cost and @federation__listSize to the supergraph', () => {
+      const result = composeServices([subgraphWithUnimportedCost, subgraphWithUnimportedListSize]);
+      assertCompositionSuccess(result);
+      expect(result.hints).toEqual([]);
+
+      const costDirectiveApplications = fieldWithCost(result)?.appliedDirectivesOf('federation__cost');
+      expect(costDirectiveApplications?.toString()).toMatchString(`@federation__cost(weight: 5)`);
+
+      const argCostDirectiveApplications = argumentWithCost(result)?.appliedDirectivesOf('federation__cost');
+      expect(argCostDirectiveApplications?.toString()).toMatchString(`@federation__cost(weight: 10)`);
+
+      const enumCostDirectiveApplications = enumWithCost(result)?.appliedDirectivesOf('federation__cost');
+      expect(enumCostDirectiveApplications?.toString()).toMatchString(`@federation__cost(weight: 15)`);
+
+      const inputCostDirectiveApplications = inputWithCost(result)?.field('somethingWithCost')?.appliedDirectivesOf('federation__cost');
+      expect(inputCostDirectiveApplications?.toString()).toMatchString(`@federation__cost(weight: 20)`);
+
+      const listSizeDirectiveApplications = fieldWithListSize(result)?.appliedDirectivesOf('federation__listSize');
+      expect(listSizeDirectiveApplications?.toString()).toMatchString(`@federation__listSize(assumedSize: 2000, requireOneSlicingArgument: false)`);
+    });
+  })
 });
 
 describe('demand control directive extraction', () => {
   it.each([
     subgraphWithCost,
     subgraphWithRenamedCost,
     subgraphWithCostFromFederationSpec,
-    subgraphWithRenamedCostFromFederationSpec
+    subgraphWithRenamedCostFromFederationSpec,
+    subgraphWithUnimportedCost,
   ])('extracts @cost from the supergraph', (subgraph: ServiceDefinition) => {
     const result = composeServices([subgraph]);
     assertCompositionSuccess(result);
-    const extracted = Supergraph.build(result.supergraphSdl).subgraphs().get(subgraphWithCost.name);
+    const extracted = Supergraph.build(result.supergraphSdl).subgraphs().get(subgraph.name);
 
     expect(extracted?.toString()).toMatchString(`
       schema
@@ -537,11 +591,12 @@ describe('demand control directive extraction', () => {
     subgraphWithListSize,
     subgraphWithRenamedListSize,
     subgraphWithListSizeFromFederationSpec,
-    subgraphWithRenamedListSizeFromFederationSpec
+    subgraphWithRenamedListSizeFromFederationSpec,
+    subgraphWithUnimportedListSize,
   ])('extracts @listSize from the supergraph', (subgraph: ServiceDefinition) => {
     const result = composeServices([subgraph]);
     assertCompositionSuccess(result);
-    const extracted = Supergraph.build(result.supergraphSdl).subgraphs().get(subgraphWithListSize.name);
+    const extracted = Supergraph.build(result.supergraphSdl).subgraphs().get(subgraph.name);
 
     expect(extracted?.toString()).toMatchString(`
       schema
@@ -646,4 +701,64 @@ describe('demand control directive extraction', () => {
       expect(supergraph.subgraphs().get(subgraphB.name)?.toString()).toMatchString(expectedSubgraph);
     });
   });
+
+  describe('when the supergraph uses custom directives with the same name', () => {
+    it('does not attempt to extract them to the subgraphs', () => {
+      const subgraphA = {
+        name: 'subgraph-a',
+        typeDefs: asFed2SubgraphDocument(gql`
+          extend schema
+            @link(url: "https://example.com/myCustomDirective/v1.0", import: ["@cost"])
+            @composeDirective(name: "@cost")
+
+          directive @cost(name: String!) on FIELD_DEFINITION
+
+          type Query {
+            a: Int @cost(name: "cost")
+          }
+        `)
+      };
+      const subgraphB = {
+        name: 'subgraph-b',
+        typeDefs: asFed2SubgraphDocument(gql`
+          extend schema
+            @link(url: "https://example.com/myOtherCustomDirective/v1.0", import: ["@listSize"])
+            @composeDirective(name: "@listSize")
+          
+          directive @listSize(name: String!) on FIELD_DEFINITION
+
+          type Query {
+            b: [Int] @listSize(name: "listSize")
+          }
+        `)
+      };
+
+      const result = composeServices([subgraphA, subgraphB]);
+      assertCompositionSuccess(result);
+      const supergraph = Supergraph.build(result.supergraphSdl);
+
+      expect(supergraph.subgraphs().get(subgraphA.name)?.toString()).toMatchString(`
+        schema
+          ${FEDERATION2_LINK_WITH_AUTO_EXPANDED_IMPORTS}
+        {
+          query: Query
+        }
+
+        type Query {
+          a: Int
+        }  
+      `);
+      expect(supergraph.subgraphs().get(subgraphB.name)?.toString()).toMatchString(`
+        schema
+          ${FEDERATION2_LINK_WITH_AUTO_EXPANDED_IMPORTS}
+        {
+          query: Query
+        }
+
+        type Query {
+          b: [Int]
+        }  
+      `);
+    });
+  });
 });

internals-js/src/extractSubgraphsFromSupergraph.ts

@@ -40,7 +40,7 @@ import { parseSelectionSet } from "./operations";
 import fs from 'fs';
 import path from 'path';
 import { validateStringContainsBoolean } from "./utils";
-import { CONTEXT_VERSIONS, ContextSpecDefinition, DirectiveDefinition, FederationDirectiveName, SchemaElement, errorCauses, isFederationDirectiveDefinedInSchema, printErrors } from ".";
+import { CONTEXT_VERSIONS, ContextSpecDefinition, DirectiveDefinition, FeatureUrl, FederationDirectiveName, SchemaElement, errorCauses, isFederationDirectiveDefinedInSchema, printErrors } from ".";
 
 function filteredTypes(
   supergraph: Schema,
@@ -224,7 +224,7 @@ export function extractSubgraphsFromSupergraph(supergraph: Schema, validateExtra
     }
 
     const types = filteredTypes(supergraph, joinSpec, coreFeatures.coreDefinition);
-    const originalDirectiveNames = getOriginalDirectiveNames(supergraph);
+    const originalDirectiveNames = getApolloDirectiveNames(supergraph);
     const args: ExtractArguments = {
        supergraph,
        subgraphs,
@@ -485,13 +485,40 @@ function extractObjOrItfContent(args: ExtractArguments, info: TypeInfo<ObjectTyp
   }
 }
 
-function getOriginalDirectiveNames(supergraph: Schema): Record<string, string> {
+/**
+ * Builds a map of original name to new name for Apollo feature directives. This is
+ * used to handle cases where a directive is renamed via an import statement. For
+ * example, importing a directive with a custom name like
+ * ```graphql
+ * @link(url: "https://specs.apollo.dev/cost/v0.1", import: [{ name: "@cost", as: "@renamedCost" }])
+ * ```
+ * results in a map entry of `cost -> renamedCost` with the `@` prefix removed.
+ * 
+ * If the directive is imported under its default name, that also results in an entry. So,
+ * ```graphql
+ * @link(url: "https://specs.apollo.dev/cost/v0.1", import: ["@cost"])
+ * ```
+ * results in a map entry of `cost -> cost`. This duals as a way to check if a directive
+ * is included in the supergraph schema.
+ * 
+ * **Important:** This map does _not_ include directives imported from identities other
+ * than `specs.apollo.dev`. This helps us avoid extracting directives to subgraphs
+ * when a custom directive's name conflicts with that of a default one.
+ */
+function getApolloDirectiveNames(supergraph: Schema): Record<string, string> {
   const originalDirectiveNames: Record<string, string> = {};
   for (const linkDirective of supergraph.schemaDefinition.appliedDirectivesOf("link")) {
     if (linkDirective.arguments().url && linkDirective.arguments().import) {
+      const url = FeatureUrl.maybeParse(linkDirective.arguments().url);
+      if (!url?.identity.includes("specs.apollo.dev")) {
+        continue;
+      }
+
       for (const importedDirective of linkDirective.arguments().import) {
         if (importedDirective.name && importedDirective.as) {
           originalDirectiveNames[importedDirective.name.replace('@', '')] = importedDirective.as.replace('@', '');
+        } else if (typeof importedDirective === 'string') {
+          originalDirectiveNames[importedDirective.replace('@', '')] = importedDirective.replace('@', '');
         }
       }
     }
@@ -652,16 +679,20 @@ function maybeDumpSubgraphSchema(subgraph: Subgraph): string {
 }
 
 function propagateDemandControlDirectives(source: SchemaElement<any, any>, dest: SchemaElement<any, any>, subgraph: Subgraph, originalDirectiveNames?: Record<string, string>) {
-  const costDirectiveName = originalDirectiveNames?.[FederationDirectiveName.COST] ?? FederationDirectiveName.COST;
-  const costDirective = source.appliedDirectivesOf(costDirectiveName).pop();
-  if (costDirective) {
-    dest.applyDirective(subgraph.metadata().costDirective().name, costDirective.arguments());
+  const costDirectiveName = originalDirectiveNames?.[FederationDirectiveName.COST];
+  if (costDirectiveName) {
+    const costDirective = source.appliedDirectivesOf(costDirectiveName).pop();
+    if (costDirective) {
+      dest.applyDirective(subgraph.metadata().costDirective().name, costDirective.arguments());
+    }
   }
 
-  const listSizeDirectiveName = originalDirectiveNames?.[FederationDirectiveName.LIST_SIZE] ?? FederationDirectiveName.LIST_SIZE;
-  const listSizeDirective = source.appliedDirectivesOf(listSizeDirectiveName).pop();
-  if (listSizeDirective) {
-    dest.applyDirective(subgraph.metadata().listSizeDirective().name, listSizeDirective.arguments());
+  const listSizeDirectiveName = originalDirectiveNames?.[FederationDirectiveName.LIST_SIZE];
+  if (listSizeDirectiveName) {
+    const listSizeDirective = source.appliedDirectivesOf(listSizeDirectiveName).pop();
+    if (listSizeDirective) {
+      dest.applyDirective(subgraph.metadata().listSizeDirective().name, listSizeDirective.arguments());
+    }
   }
 }