Handle multiple composable directives from the same spec (#3086)

# Overview

This PR targets the branch of an existing open PR. The main difference
is that this one combines the new `@cost` and `@listSize` directives are
combined into the same spec. This required a small update to how the
supergraph schema is prepared.

When a directive spec definition is imported by a subgraph, it was
previously included in the supergraph without an import argument. So a
subgraph with `@link(url: "https://specs.apollo.dev/cost/v0.1", import:
["@cost"])`, is translated to the supergraph as `@link(url:
"https://specs.apollo.dev/cost/v0.1")`. If multiple imports are
included, the supergraph wouldn't get any of them. With this change, the
supergraph now gets an equivalent `@link(url:
"https://specs.apollo.dev/cost/v0.1", import: ["@cost"])`.

## Renaming

The core/link directive can have an argument called "as" to rename the
directive from a spec. This would take a subgraph import `@link(url:
"https://specs.apollo.dev/cost/v0.1", import: [{ name: "@cost", as:
"@renamedCost" }])`, and translate it to `@link(url:
"https://specs.apollo.dev/cost/v0.1", as: "renamedCost")`. With this
change, the supergraph now gets an equivalent `@link(url:
"https://specs.apollo.dev/cost/v0.1", import: [{ name: "@cost", as:
"@renamedCost" }])`.

## Import aggregation

As we add directive spec features to the supergraph, we check for
existing links with the same feature version (both the link identity and
the version must match). If one is found, we append the new directive to
the existing list of imports.

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

@@ -4,7 +4,7 @@ exports[`composing custom core directives custom tag directive works when federa
 "schema
   @link(url: \\"https://specs.apollo.dev/link/v1.0\\")
   @link(url: \\"https://specs.apollo.dev/join/v0.3\\", for: EXECUTION)
-  @link(url: \\"https://specs.apollo.dev/tag/v0.3\\", as: \\"mytag\\")
+  @link(url: \\"https://specs.apollo.dev/tag/v0.3\\", import: [{name: \\"@tag\\", as: \\"@mytag\\"}])
   @link(url: \\"https://custom.dev/tag/v1.0\\", import: [\\"@tag\\"])
 {
   query: Query

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

@@ -926,8 +926,8 @@ describe('composing custom core directives', () => {
     expectCoreFeature(schema, 'https://custom.dev/tag', '1.0', [{ name: '@tag' }]);
     const feature = schema.coreFeatures?.getByIdentity('https://specs.apollo.dev/tag');
     expect(feature?.url.toString()).toBe('https://specs.apollo.dev/tag/v0.3');
-    expect(feature?.imports).toEqual([]);
-    expect(feature?.nameInSchema).toEqual('mytag');
+    expect(feature?.imports).toEqual([{ name: '@tag', as: '@mytag' }]);
+    expect(feature?.nameInSchema).toEqual('tag');
     expect(printSchema(schema)).toMatchSnapshot();
   });
 

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

@@ -39,7 +39,7 @@ const subgraphWithCost = {
 const subgraphWithListSize = {
   name: 'subgraphWithListSize',
   typeDefs: asFed2SubgraphDocument(gql`
-    extend schema @link(url: "https://specs.apollo.dev/listSize/v0.1", import: ["@listSize"])
+    extend schema @link(url: "https://specs.apollo.dev/cost/v0.1", import: ["@listSize"])
 
     type Query {
       fieldWithListSize: [String!] @listSize(assumedSize: 2000, requireOneSlicingArgument: false)
@@ -73,7 +73,7 @@ const subgraphWithRenamedCost = {
 const subgraphWithRenamedListSize = {
   name: 'subgraphWithListSize',
   typeDefs: asFed2SubgraphDocument(gql`
-    extend schema @link(url: "https://specs.apollo.dev/listSize/v0.1", import: [{ name: "@listSize", as: "@renamedListSize" }])
+    extend schema @link(url: "https://specs.apollo.dev/cost/v0.1", import: [{ name: "@listSize", as: "@renamedListSize" }])
 
     type Query {
       fieldWithListSize: [String!] @renamedListSize(assumedSize: 2000, requireOneSlicingArgument: false)
@@ -238,11 +238,9 @@ describe('demand control directive composition', () => {
 
       // Ensure the new directive names are specified in the supergraph so we can use them during extraction
       const links = result.schema.schemaDefinition.appliedDirectivesOf("link");
-      const costLink = links.find((link) => link.arguments().url === "https://specs.apollo.dev/cost/v0.1");
-      expect(costLink?.arguments().as).toBe("renamedCost");
-
-      const listSizeLink = links.find((link) => link.arguments().url === "https://specs.apollo.dev/listSize/v0.1");
-      expect(listSizeLink?.arguments().as).toBe("renamedListSize");
+      const costLinks = links.filter((link) => link.arguments().url === "https://specs.apollo.dev/cost/v0.1");
+      expect(costLinks.length).toBe(1);
+      expect(costLinks[0].toString()).toEqual(`@link(url: "https://specs.apollo.dev/cost/v0.1", import: [{name: "@cost", as: "@renamedCost"}, {name: "@listSize", as: "@renamedListSize"}])`);
 
       // Ensure the directives are applied to the expected fields with the new names
       const costDirectiveApplications = fieldWithCost(result)?.appliedDirectivesOf('renamedCost');
@@ -341,7 +339,7 @@ describe('demand control directive composition', () => {
       const subgraphA = {
         name: 'subgraph-a',
         typeDefs: asFed2SubgraphDocument(gql`
-          extend schema @link(url: "https://specs.apollo.dev/listSize/v0.1", import: ["@listSize"])
+          extend schema @link(url: "https://specs.apollo.dev/cost/v0.1", import: ["@listSize"])
   
           type Query {
             sharedWithListSize: [Int] @shareable @listSize(assumedSize: 10)
@@ -351,7 +349,7 @@ describe('demand control directive composition', () => {
       const subgraphB = {
         name: 'subgraph-b',
         typeDefs: asFed2SubgraphDocument(gql`
-          extend schema @link(url: "https://specs.apollo.dev/listSize/v0.1", import: ["@listSize"])
+          extend schema @link(url: "https://specs.apollo.dev/cost/v0.1", import: ["@listSize"])
   
           type Query {
             sharedWithListSize: [Int] @shareable @listSize(assumedSize: 20)
@@ -448,7 +446,7 @@ describe('demand control directive extraction', () => {
     const subgraphA = {
       name: 'subgraph-a',
       typeDefs: asFed2SubgraphDocument(gql`
-        extend schema @link(url: "https://specs.apollo.dev/listSize/v0.1", import: ["@listSize"])
+        extend schema @link(url: "https://specs.apollo.dev/cost/v0.1", import: ["@listSize"])
 
         type Query {
           sizedList(first: Int!): HasInts @shareable @listSize(slicingArguments: ["first"], sizedFields: ["ints"], requireOneSlicingArgument: true)
@@ -462,7 +460,7 @@ describe('demand control directive extraction', () => {
     const subgraphB = {
       name: 'subgraph-b',
       typeDefs: asFed2SubgraphDocument(gql`
-        extend schema @link(url: "https://specs.apollo.dev/listSize/v0.1", import: ["@listSize"])
+        extend schema @link(url: "https://specs.apollo.dev/cost/v0.1", import: ["@listSize"])
 
         type Query {
           sizedList(first: Int!): HasInts @shareable @listSize(slicingArguments: ["first"], sizedFields: ["ints"], requireOneSlicingArgument: false)
@@ -545,7 +543,7 @@ describe('demand control directive extraction', () => {
       const subgraphA = {
         name: 'subgraph-a',
         typeDefs: asFed2SubgraphDocument(gql`
-          extend schema @link(url: "https://specs.apollo.dev/listSize/v0.1", import: ["@listSize"])
+          extend schema @link(url: "https://specs.apollo.dev/cost/v0.1", import: ["@listSize"])
   
           type Query {
             sharedWithListSize: [Int] @shareable @listSize(assumedSize: 10)
@@ -555,7 +553,7 @@ describe('demand control directive extraction', () => {
       const subgraphB = {
         name: 'subgraph-b',
         typeDefs: asFed2SubgraphDocument(gql`
-          extend schema @link(url: "https://specs.apollo.dev/listSize/v0.1", import: ["@listSize"])
+          extend schema @link(url: "https://specs.apollo.dev/cost/v0.1", import: ["@listSize"])
   
           type Query {
             sharedWithListSize: [Int] @shareable @listSize(assumedSize: 20)

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

@@ -219,7 +219,7 @@ describe('composition of directive with non-trivial argument strategies', () =>
     const s = result.schema;
 
     expect(directiveStrings(s.schemaDefinition, name)).toStrictEqual([
-      `@link(url: "https://specs.apollo.dev/${name}/v0.1")`
+      `@link(url: "https://specs.apollo.dev/${name}/v0.1", import: ["@${name}"])`
     ]);
 
     const t = s.type('T') as ObjectType;

composition-js/src/composeDirectiveManager.ts

@@ -67,7 +67,6 @@ const DISALLOWED_IDENTITIES = [
   'https://specs.apollo.dev/source',
   'https://specs.apollo.dev/context',
   'https://specs.apollo.dev/cost',
-  'https://specs.apollo.dev/listSize',
 ];
 
 export class ComposeDirectiveManager {

composition-js/src/merging/merge.ts

@@ -449,7 +449,7 @@ class Merger {
     // don't bother adding the spec to the supergraph.
     if (nameInSupergraph) {
       const specInSupergraph = compositionSpec.supergraphSpecification(this.latestFedVersionUsed);
-      const errors = this.linkSpec.applyFeatureToSchema(this.merged, specInSupergraph, nameInSupergraph === specInSupergraph.url.name ? undefined : nameInSupergraph, specInSupergraph.defaultCorePurpose);
+      const errors = this.linkSpec.applyFeatureAsLink(this.merged, specInSupergraph, specInSupergraph.defaultCorePurpose, [{ name, as: name === nameInSupergraph ? undefined : nameInSupergraph }], );
       assert(errors.length === 0, "We shouldn't have errors adding the join spec to the (still empty) supergraph schema");
       const feature = this.merged?.coreFeatures?.getByIdentity(specInSupergraph.url.identity);
       assert(feature, 'Should have found the feature we just added');
@@ -459,7 +459,7 @@ class Merger {
         throw argumentsMerger;
       }
       this.mergedFederationDirectiveNames.add(nameInSupergraph);
-      this.mergedFederationDirectiveInSupergraph.set(specInSupergraph.url.name, {
+      this.mergedFederationDirectiveInSupergraph.set(name, {
         definition: this.merged.directive(nameInSupergraph)!,
         argumentsMerger,
         staticArgumentTransform: compositionSpec.staticArgumentTransform,

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

@@ -149,7 +149,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).toMatchInlineSnapshot(
-      `"6dc1bde2b9818fabec62208c5d8825abaa1bae89635fa6f3a5ffea7b78fc6d82"`,
+      `"4aa2278e35df345ff5959a30546d2e9ef9e997204b4ffee4a42344b578b36068"`,
     );
     // second call should have previous info in the second arg
     expect(secondCall[1]!.compositionId).toEqual(expectedFirstId);

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, FeatureUrl, FederationDirectiveName, SchemaElement, costIdentity, errorCauses, isFederationDirectiveDefinedInSchema, listSizeIdentity, printErrors } from ".";
+import { CONTEXT_VERSIONS, ContextSpecDefinition, DirectiveDefinition, FederationDirectiveName, SchemaElement, errorCauses, isFederationDirectiveDefinedInSchema, printErrors } from ".";
 
 function filteredTypes(
   supergraph: Schema,
@@ -483,17 +483,11 @@ function extractObjOrItfContent(args: ExtractArguments, info: TypeInfo<ObjectTyp
 function getOriginalDirectiveNames(supergraph: Schema): Record<string, string> {
   const originalDirectiveNames: Record<string, string> = {};
   for (const linkDirective of supergraph.schemaDefinition.appliedDirectivesOf("link")) {
-    if (linkDirective.arguments().url && linkDirective.arguments().as) {
-      const parsedUrl = FeatureUrl.maybeParse(linkDirective.arguments().url);
-      // Ideally, there's a map somewhere that can do this lookup instead of enumerating all the directives we care about,
-      // but it seems the original names are being stripped from the supergraph schema.
-      switch (parsedUrl?.identity) {
-        case costIdentity:
-          originalDirectiveNames[FederationDirectiveName.COST] = linkDirective.arguments().as;
-          break;
-        case listSizeIdentity:
-          originalDirectiveNames[FederationDirectiveName.LIST_SIZE] = linkDirective.arguments().as;
-          break;
+    if (linkDirective.arguments().url && linkDirective.arguments().import) {
+      for (const importedDirective of linkDirective.arguments().import) {
+        if (importedDirective.name && importedDirective.as) {
+          originalDirectiveNames[importedDirective.name.replace('@', '')] = importedDirective.as.replace('@', '');
+        }
       }
     }
   }

internals-js/src/federation.ts

@@ -100,8 +100,7 @@ import {
   SourceFieldDirectiveArgs,
   SourceTypeDirectiveArgs,
 } from "./specs/sourceSpec";
-import { CostDirectiveArguments } from "./specs/costSpec";
-import { ListSizeDirectiveArguments } from "./specs/listSizeSpec";
+import { CostDirectiveArguments, ListSizeDirectiveArguments } from "./specs/costSpec";
 
 const linkSpec = LINK_VERSIONS.latest();
 const tagSpec = TAG_VERSIONS.latest();

internals-js/src/index.ts

@@ -26,4 +26,3 @@ export * from './specs/requiresScopesSpec';
 export * from './specs/policySpec';
 export * from './specs/sourceSpec';
 export * from './specs/costSpec';
-export * from './specs/listSizeSpec';

internals-js/src/specs/coreSpec.ts

@@ -552,6 +552,27 @@ export class CoreSpecDefinition extends FeatureDefinition {
     return feature.addElementsToSchema(schema);
   }
 
+  applyFeatureAsLink(schema: Schema, feature: FeatureDefinition, purpose?: CorePurpose, imports?: CoreImport[]): GraphQLError[] {
+    const existing = schema.schemaDefinition.appliedDirectivesOf(linkDirectiveDefaultName).find((link) => link.arguments().url === feature.toString());
+    if (existing) {
+      existing.remove();
+    }
+
+    const coreDirective = this.coreDirective(schema);
+    const args: LinkDirectiveArgs = {
+      url: feature.toString(),
+      import: (existing?.arguments().import ?? []).concat(imports?.map((i) => i.as ? { name: `@${i.name}`, as: `@${i.as}` } : `@${i.name}`)),
+      feature: undefined,
+    };
+
+    if (this.supportPurposes() && purpose) {
+      args.for = purpose;
+    }
+
+    schema.schemaDefinition.applyDirective(coreDirective, args);
+    return feature.addElementsToSchema(schema);
+  }
+
   extractFeatureUrl(args: CoreOrLinkDirectiveArgs): FeatureUrl {
     return FeatureUrl.parse(args[this.urlArgName()]!);
   }

internals-js/src/specs/costSpec.ts

@@ -1,7 +1,7 @@
 import { DirectiveLocation } from 'graphql';
 import { createDirectiveSpecification } from '../directiveAndTypeSpecification';
 import { FeatureDefinition, FeatureDefinitions, FeatureUrl, FeatureVersion } from './coreSpec';
-import { NonNullType } from '../definitions';
+import { ListType, NonNullType } from '../definitions';
 import { registerKnownFeature } from '../knownCoreFeatures';
 import { ARGUMENT_COMPOSITION_STRATEGIES } from '../argumentCompositionStrategies';
 
@@ -26,6 +26,20 @@ export class CostSpecDefinition extends FeatureDefinition {
       repeatable: false,
       supergraphSpecification: (fedVersion) => COST_VERSIONS.getMinimumRequiredVersion(fedVersion),
     }));
+
+    this.registerDirective(createDirectiveSpecification({
+      name: 'listSize',
+      locations: [DirectiveLocation.FIELD_DEFINITION],
+      args: [
+        { name: 'assumedSize', type: (schema) => schema.intType(), compositionStrategy: ARGUMENT_COMPOSITION_STRATEGIES.NULLABLE_MAX },
+        { name: 'slicingArguments', type: (schema) => new ListType(new NonNullType(schema.stringType())), compositionStrategy: ARGUMENT_COMPOSITION_STRATEGIES.NULLABLE_UNION },
+        { name: 'sizedFields', type: (schema) => new ListType(new NonNullType(schema.stringType())), compositionStrategy: ARGUMENT_COMPOSITION_STRATEGIES.NULLABLE_UNION },
+        { name: 'requireOneSlicingArgument', type: (schema) => schema.booleanType(), defaultValue: true, compositionStrategy: ARGUMENT_COMPOSITION_STRATEGIES.NULLABLE_AND },
+      ],
+      composes: true,
+      repeatable: false,
+      supergraphSpecification: (fedVersion) => COST_VERSIONS.getMinimumRequiredVersion(fedVersion)
+    }));
   }
 }
 
@@ -37,3 +51,10 @@ registerKnownFeature(COST_VERSIONS);
 export interface CostDirectiveArguments {
   weight: number;
 }
+
+export interface ListSizeDirectiveArguments {
+  assumedSize?: number;
+  slicingArguments?: string[];
+  sizedFields?: string[];
+  requireOneSlicingArgument?: boolean;
+}

internals-js/src/specs/federationSpec.ts

@@ -21,7 +21,6 @@ import { POLICY_VERSIONS } from './policySpec';
 import { SOURCE_VERSIONS } from './sourceSpec';
 import { CONTEXT_VERSIONS } from './contextSpec';
 import { COST_VERSIONS } from "./costSpec";
-import { LIST_SIZE_VERSIONS } from "./listSizeSpec";
 
 export const federationIdentity = 'https://specs.apollo.dev/federation';
 
@@ -189,7 +188,6 @@ export class FederationSpecDefinition extends FeatureDefinition {
 
     if (version.gte(new FeatureVersion(2, 9))) {
       this.registerSubFeature(COST_VERSIONS.find(new FeatureVersion(0, 1))!);
-      this.registerSubFeature(LIST_SIZE_VERSIONS.find(new FeatureVersion(0, 1))!);
     }
   }
 }

internals-js/src/supergraphs.ts

@@ -41,7 +41,6 @@ export const ROUTER_SUPPORTED_SUPERGRAPH_FEATURES = new Set([
   'https://specs.apollo.dev/source/v0.1',
   'https://specs.apollo.dev/context/v0.1',
   'https://specs.apollo.dev/cost/v0.1',
-  'https://specs.apollo.dev/listSize/v0.1',
 ]);
 
 const coreVersionZeroDotOneUrl = FeatureUrl.parse('https://specs.apollo.dev/core/v0.1');