feat(federation): add FieldSet type for selections with inline fragments

.changeset/spicy-pugs-repeat.md

@@ -0,0 +1,12 @@
+---
+'@pothos/plugin-federation': minor
+---
+
+Add `FieldSet` type that can be used to define selections that can't be expressed with
+`SelectionFromShape`, like selections with inline fragments on union or interface fields:
+
+```ts
+builder.selection(
+  'media { ... on Image { url } ... on Video { url } }' as FieldSet<{ media: Media[] }>,
+)
+```

packages/plugin-federation/README.md

@@ -119,6 +119,38 @@ ProductRef.implement({
 });
 ```
 
+### Selections with inline fragments
+
+The template-literal type that checks selection strings can not express selections that use inline
+fragments (e.g. selecting through a union or interface field). For these cases, a string can be cast
+to `FieldSet<Shape>` to bypass the selection string checks. The cast replaces the generic argument
+of `builder.selection` — the shape is inferred from the cast, and is still used for the resolver's
+`parent` type:
+
+```typescript
+import { type FieldSet } from '@pothos/plugin-federation';
+
+type Media = { __typename: 'Image'; url: string } | { __typename: 'Video'; url: string };
+
+PostRef.implement({
+  externalFields: (t) => ({
+    media: t.field({ type: [MediaUnion] }),
+  }),
+  fields: (t) => ({
+    mediaUrls: t.stringList({
+      requires: builder.selection(
+        'media { ... on Image { url } ... on Video { url } }' as FieldSet<{ media: Media[] }>,
+      ),
+      resolve: (post) => post.media.map((media) => media.url),
+    }),
+  }),
+});
+```
+
+`FieldSet` is accepted anywhere a selection string is expected, including `builder.selection` and
+`ref.provides`. The selection string is not validated against the shape, so make sure the selection
+matches the fields described by the generic argument.
+
 To set the `resolvable` property of an external field to `false`, can use `builder.keyDirective`:
 
 ```ts

packages/plugin-federation/src/external-ref.ts

@@ -9,6 +9,7 @@ import {
 import type { GraphQLResolveInfo } from 'graphql';
 import type {
   ExternalEntityOptions,
+  FieldSet,
   Selection,
   SelectionFromShape,
   selectionShapeKey,
@@ -91,7 +92,7 @@ export class ExternalEntityRef<
     return this;
   }
 
-  provides<T extends object>(selection: SelectionFromShape<T>) {
+  provides<T extends object>(selection: FieldSet<T> | SelectionFromShape<T>) {
     const ref = Object.create(this) as ExternalEntityRef<Types, Shape & T, Key>;
 
     providesMap.set(ref, selection);

packages/plugin-federation/src/global-types.ts

@@ -14,6 +14,7 @@ import type {
 import type { GraphQLResolveInfo, GraphQLSchema } from 'graphql';
 import type { ExternalEntityRef } from './external-ref.js';
 import type {
+  FieldSet,
   KeyDirective,
   PothosFederationPlugin,
   Selection,
@@ -103,7 +104,9 @@ declare global {
         ) => MaybePromise<Shape | null | undefined>,
       ) => ExternalEntityRef<Types, Shape, KeySelection>;
 
-      selection: <Shape extends object>(selection: SelectionFromShape<Shape>) => Selection<Shape>;
+      selection: <Shape extends object>(
+        selection: FieldSet<Shape> | SelectionFromShape<Shape>,
+      ) => Selection<Shape>;
 
       keyDirective: <Shape extends object, Resolvable extends boolean = true>(
         key: Selection<Shape>,

packages/plugin-federation/src/schema-builder.ts

@@ -13,7 +13,12 @@ import {
   lexicographicSortSchema,
 } from 'graphql';
 import { ExternalEntityRef } from './external-ref.js';
-import { type Selection, type SelectionFromShape, selectionShapeKey } from './types.js';
+import {
+  type FieldSet,
+  type Selection,
+  type SelectionFromShape,
+  selectionShapeKey,
+} from './types.js';
 import { entityMapping, getUsedDirectives, mergeDirectives } from './util.js';
 
 const schemaBuilderProto = SchemaBuilder.prototype as PothosSchemaTypes.SchemaBuilder<SchemaTypes>;
@@ -46,7 +51,9 @@ export function hasResolvableKey(type: GraphQLNamedType) {
   return directives.key?.resolvable !== false;
 }
 
-schemaBuilderProto.selection = <Shape extends object>(selection: SelectionFromShape<Shape>) => ({
+schemaBuilderProto.selection = <Shape extends object>(
+  selection: FieldSet<Shape> | SelectionFromShape<Shape>,
+) => ({
   selection,
   [selectionShapeKey]: {} as unknown as Shape,
 });

packages/plugin-federation/src/types.ts

@@ -9,6 +9,26 @@ import type {
 
 export const selectionShapeKey = Symbol.for('Pothos.federationSelectionKey');
 
+declare const fieldSetShapeKey: unique symbol;
+
+/**
+ * A branded selection string for `Shape` that bypasses the `SelectionFromShape` checks.
+ *
+ * `SelectionFromShape` can not express every valid FieldSet (e.g. selections that
+ * require inline fragments to select through a union or interface field). Casting a
+ * string to `FieldSet<Shape>` allows it to be passed anywhere a selection for `Shape`
+ * is expected, replacing the generic argument entirely:
+ *
+ * ```ts
+ * builder.selection(
+ *   'media { ... on Image { url } ... on Video { url } }' as FieldSet<{ media: Media[] }>,
+ * )
+ * ```
+ */
+export type FieldSet<Shape extends object = object> = string & {
+  readonly [fieldSetShapeKey]: Shape;
+};
+
 export type EntityObjectFieldsShape<Types extends SchemaTypes, Shape, Fields extends FieldMap> = (
   t: PothosSchemaTypes.FieldBuilder<Types, Shape, 'EntityObject'>,
 ) => Fields;

packages/plugin-federation/tests/__snapshots__/index.test.ts.snap

@@ -101,6 +101,66 @@ type _Service {
 }"
 `;
 
+exports[`federation > media schema > generates expected schema 1`] = `
+"extend schema
+  @link(url: "https://specs.apollo.dev/federation/v2.6", import: ["@extends", "@external", "@key", "@requires"])
+
+type Image {
+  url: String
+}
+
+union Media = Image | Video
+
+type Post
+  @key(fields: "id")
+  @extends
+{
+  media: [Media!] @external
+  id: String
+  mediaUrls: [String!] @requires(fields: "media { ... on Image { url } ... on Video { url } }")
+}
+
+type Video {
+  url: String
+  previewUrl: String
+}"
+`;
+
+exports[`federation > media schema > generates expected schema 2`] = `
+"type Image {
+  url: String
+}
+
+union Media = Image | Video
+
+type Post {
+  id: String
+  media: [Media!]
+  mediaUrls: [String!]
+}
+
+type Query {
+  _entities(representations: [_Any!]!): [_Entity]!
+  _service: _Service!
+}
+
+type Video {
+  previewUrl: String
+  url: String
+}
+
+scalar _Any
+
+union _Entity = Post
+
+type _Service {
+  """
+  The sdl representing the federated service capabilities. Includes federation directives, removes federation types, and includes rest of full schema after schema directives have been applied
+  """
+  sdl: String
+}"
+`;
+
 exports[`federation > products schema > generates expected schema 1`] = `
 "extend schema
   @link(url: "https://specs.apollo.dev/federation/v2.6", import: ["@key", "@composeDirective"])

packages/plugin-federation/tests/example/media/schema.ts

@@ -0,0 +1,64 @@
+import SchemaBuilder from '@pothos/core';
+import DirectivesPlugin from '@pothos/plugin-directives';
+import FederationPlugin, { type FieldSet } from '../../../src';
+
+const builder = new SchemaBuilder({
+  plugins: [DirectivesPlugin, FederationPlugin],
+  directives: {
+    useGraphQLToolsUnorderedDirectives: true,
+  },
+});
+
+interface Image {
+  __typename: 'Image';
+  url: string;
+}
+
+interface Video {
+  __typename: 'Video';
+  url: string;
+  previewUrl: string;
+}
+
+type Media = Image | Video;
+
+const ImageRef = builder.objectRef<Image>('Image').implement({
+  fields: (t) => ({
+    url: t.exposeString('url'),
+  }),
+});
+
+const VideoRef = builder.objectRef<Video>('Video').implement({
+  fields: (t) => ({
+    url: t.exposeString('url'),
+    previewUrl: t.exposeString('previewUrl'),
+  }),
+});
+
+const MediaUnion = builder.unionType('Media', {
+  types: [ImageRef, VideoRef],
+  resolveType: (media) => media.__typename,
+});
+
+const PostRef = builder.externalRef(
+  'Post',
+  builder.selection<{ id: string }>('id'),
+  (entity) => entity,
+);
+
+PostRef.implement({
+  externalFields: (t) => ({
+    media: t.field({ type: [MediaUnion] }),
+  }),
+  fields: (t) => ({
+    id: t.exposeString('id'),
+    mediaUrls: t.stringList({
+      requires: builder.selection(
+        'media { ... on Image { url } ... on Video { url } }' as FieldSet<{ media: Media[] }>,
+      ),
+      resolve: (post) => post.media.map((media) => media.url),
+    }),
+  }),
+});
+
+export const schema = builder.toSubGraphSchema({});

packages/plugin-federation/tests/field-set-types.ts

@@ -0,0 +1,30 @@
+// Type-level assertions for the FieldSet escape hatch (checked by `pnpm type`).
+import SchemaBuilder from '@pothos/core';
+import DirectivesPlugin from '@pothos/plugin-directives';
+import FederationPlugin, { type FieldSet, type Selection } from '../src';
+
+const builder = new SchemaBuilder({
+  plugins: [DirectivesPlugin, FederationPlugin],
+});
+
+type Media = { __typename: 'Image'; url: string } | { __typename: 'Video'; url: string };
+
+// checked selection strings work exactly as before
+export const checked = builder.selection<{ upc: string; price?: number }>('upc price');
+
+// @ts-expect-error incomplete selections are still rejected
+export const missingField = builder.selection<{ upc: string; price?: number }>('upc');
+
+// a FieldSet cast allows selections SelectionFromShape can not express; the cast
+// replaces the generic entirely — the shape is inferred from the brand
+export const withFragments: Selection<{ media: Media[] }> = builder.selection(
+  'media { ... on Image { url } ... on Video { url } }' as FieldSet<{ media: Media[] }>,
+);
+
+declare const plainString: string;
+
+// @ts-expect-error un-branded plain strings are still rejected
+export const unchecked = builder.selection<{ upc: string }>(plainString);
+
+// @ts-expect-error a FieldSet branded with a different shape does not satisfy the selection
+export const wrongShape = builder.selection<{ upc: string }>('upc' as FieldSet<{ id: string }>);

packages/plugin-federation/tests/index.test.ts

@@ -6,6 +6,7 @@ import { printSchema } from 'graphql';
 import { schema as accountsSchema } from './example/accounts/schema';
 import { createGateway } from './example/gateway';
 import { schema as inventorySchema } from './example/inventory/schema';
+import { schema as mediaSchema } from './example/media/schema';
 import { schema as productsSchema } from './example/products/schema';
 import { schema as reviewsSchema } from './example/reviews/schema';
 import { startServers } from './example/servers';
@@ -25,6 +26,19 @@ describe('federation', () => {
     });
   });
 
+  describe('media schema', () => {
+    it('generates expected schema', () => {
+      expect(printSubgraphSchema(mediaSchema)).toMatchSnapshot();
+      expect(printSchema(mediaSchema)).toMatchSnapshot();
+    });
+
+    it('supports selections with inline fragments via FieldSet', () => {
+      expect(printSubgraphSchema(mediaSchema)).toContain(
+        '@requires(fields: "media { ... on Image { url } ... on Video { url } }")',
+      );
+    });
+  });
+
   describe('products schema', () => {
     it('generates expected schema', () => {
       expect(printSubgraphSchema(productsSchema)).toMatchSnapshot();

website/content/docs/plugins/federation.mdx

@@ -122,6 +122,38 @@ ProductRef.implement({
 });
 ```
 
+### Selections with inline fragments
+
+The template-literal type that checks selection strings can not express selections that use inline
+fragments (e.g. selecting through a union or interface field). For these cases, a string can be cast
+to `FieldSet<Shape>` to bypass the selection string checks. The cast replaces the generic argument
+of `builder.selection` — the shape is inferred from the cast, and is still used for the resolver's
+`parent` type:
+
+```typescript
+import { type FieldSet } from '@pothos/plugin-federation';
+
+type Media = { __typename: 'Image'; url: string } | { __typename: 'Video'; url: string };
+
+PostRef.implement({
+  externalFields: (t) => ({
+    media: t.field({ type: [MediaUnion] }),
+  }),
+  fields: (t) => ({
+    mediaUrls: t.stringList({
+      requires: builder.selection(
+        'media { ... on Image { url } ... on Video { url } }' as FieldSet<{ media: Media[] }>,
+      ),
+      resolve: (post) => post.media.map((media) => media.url),
+    }),
+  }),
+});
+```
+
+`FieldSet` is accepted anywhere a selection string is expected, including `builder.selection` and
+`ref.provides`. The selection string is not validated against the shape, so make sure the selection
+matches the fields described by the generic argument.
+
 To set the `resolvable` property of an external field to `false`, can use `builder.keyDirective`:
 
 ```ts