From 0debcce6b01e2973f8216bf6c454c714b2261280 Mon Sep 17 00:00:00 2001
From: Michael Hayes <michael@hayes.io>
Date: Fri, 1 Mar 2024 14:15:52 -0800
Subject: [PATCH] [federation] explude non-resolvable entities from Entity
 union

---
 .changeset/popular-spiders-yawn.md            |  5 ++++
 .../plugin-federation/src/global-types.ts     | 25 ++++++++++++-----
 .../plugin-federation/src/schema-builder.ts   | 28 ++++++++++++++++---
 packages/plugin-federation/src/util.ts        |  2 +-
 .../tests/__snapshots__/index.test.ts.snap    |  2 +-
 .../tests/example/accounts/schema.ts          |  1 +
 6 files changed, 50 insertions(+), 13 deletions(-)
 create mode 100644 .changeset/popular-spiders-yawn.md

diff --git a/.changeset/popular-spiders-yawn.md b/.changeset/popular-spiders-yawn.md
new file mode 100644
index 000000000..0c10ac98c
--- /dev/null
+++ b/.changeset/popular-spiders-yawn.md
@@ -0,0 +1,5 @@
+---
+'@pothos/plugin-federation': patch
+---
+
+Exclude non resolveAble entities from Entity union
diff --git a/packages/plugin-federation/src/global-types.ts b/packages/plugin-federation/src/global-types.ts
index e0c0ff37d..f9572c296 100644
--- a/packages/plugin-federation/src/global-types.ts
+++ b/packages/plugin-federation/src/global-types.ts
@@ -98,14 +98,25 @@ declare global {
         param: Param,
         options: {
           key: KeySelection | KeySelection[];
-          resolvable?: boolean;
           interfaceObject?: Param extends ObjectRef<unknown> ? boolean : never;
-          resolveReference: (
-            parent: KeySelection[typeof selectionShapeKey],
-            context: Types['Context'],
-            info: GraphQLResolveInfo,
-          ) => MaybePromise<ShapeFromTypeParam<Types, Param, true>>;
-        },
+        } & (
+          | {
+              resolvable: false;
+              resolveReference?: (
+                parent: KeySelection[typeof selectionShapeKey],
+                context: Types['Context'],
+                info: GraphQLResolveInfo,
+              ) => MaybePromise<ShapeFromTypeParam<Types, Param, true>>;
+            }
+          | {
+              resolvable?: true;
+              resolveReference: (
+                parent: KeySelection[typeof selectionShapeKey],
+                context: Types['Context'],
+                info: GraphQLResolveInfo,
+              ) => MaybePromise<ShapeFromTypeParam<Types, Param, true>>;
+            }
+        ),
       ) => void;
     }
 
diff --git a/packages/plugin-federation/src/schema-builder.ts b/packages/plugin-federation/src/schema-builder.ts
index ba0aac37e..fdec9182e 100644
--- a/packages/plugin-federation/src/schema-builder.ts
+++ b/packages/plugin-federation/src/schema-builder.ts
@@ -18,12 +18,32 @@ import { entityMapping, getUsedDirectives, mergeDirectives } from './util';
 
 const schemaBuilderProto = SchemaBuilder.prototype as PothosSchemaTypes.SchemaBuilder<SchemaTypes>;
 
-export function hasDirective(type: GraphQLNamedType, directive: string) {
+function hasResolvableKey(type: GraphQLNamedType) {
   if (Array.isArray(type.extensions?.directives)) {
-    return type.extensions?.directives.some((d) => (d as { name: string }).name === directive);
+    return type.extensions?.directives.some(
+      (d: { name: string; args: Record<string, unknown> }) =>
+        d.name === 'key' && d.args.resolvable !== false,
+    );
   }
 
-  return directive in ((type.extensions?.directives ?? {}) as {});
+  const directives = (type.extensions?.directives ?? {}) as {
+    key?:
+      | {
+          resolvable?: boolean;
+        }[]
+      | {
+          resolvable?: boolean;
+        };
+  };
+  if (!('key' in directives)) {
+    return false;
+  }
+
+  if (Array.isArray(directives.key)) {
+    return directives.key.some((d) => d.resolvable !== false);
+  }
+
+  return directives.key?.resolvable !== false;
 }
 
 schemaBuilderProto.selection = <Shape extends object>(selection: SelectionFromShape<Shape>) => ({
@@ -89,7 +109,7 @@ schemaBuilderProto.toSubGraphSchema = function toSubGraphSchema(
   queryType?.toConfig();
 
   const entityTypes = Object.values(types).filter(
-    (type) => isObjectType(type) && hasDirective(type, 'key'),
+    (type) => isObjectType(type) && hasResolvableKey(type),
   );
 
   const hasEntities = entityTypes.length > 0;
diff --git a/packages/plugin-federation/src/util.ts b/packages/plugin-federation/src/util.ts
index 67f925c25..45cc7ff77 100644
--- a/packages/plugin-federation/src/util.ts
+++ b/packages/plugin-federation/src/util.ts
@@ -51,7 +51,7 @@ export const entityMapping = new WeakMap<
       key: Selection<object> | Selection<object>[];
       resolvable?: boolean;
       interfaceObject?: boolean;
-      resolveReference: (val: object, context: {}, info: GraphQLResolveInfo) => unknown;
+      resolveReference?: (val: object, context: {}, info: GraphQLResolveInfo) => unknown;
     }
   >
 >();
diff --git a/packages/plugin-federation/tests/__snapshots__/index.test.ts.snap b/packages/plugin-federation/tests/__snapshots__/index.test.ts.snap
index cb2b1d42b..4aec7b1d8 100644
--- a/packages/plugin-federation/tests/__snapshots__/index.test.ts.snap
+++ b/packages/plugin-federation/tests/__snapshots__/index.test.ts.snap
@@ -17,7 +17,7 @@ type Query {
 }
 
 type User
-  @key(fields: "id")
+  @key(fields: "id", resolvable: true)
 {
   id: ID
   name: String
diff --git a/packages/plugin-federation/tests/example/accounts/schema.ts b/packages/plugin-federation/tests/example/accounts/schema.ts
index e6aca5d61..1aff77538 100644
--- a/packages/plugin-federation/tests/example/accounts/schema.ts
+++ b/packages/plugin-federation/tests/example/accounts/schema.ts
@@ -45,6 +45,7 @@ const UserType = builder.objectRef<User>('User').implement({
 
 builder.asEntity(UserType, {
   key: builder.selection<{ id: string }>('id'),
+  resolvable: true,
   resolveReference: (user) => users.find(({ id }) => user.id === id),
 });