From 21bea79f2d8a8caf3f00973ac769a60b359c69ed Mon Sep 17 00:00:00 2001
From: Ben Newman <ben@apollographql.com>
Date: Tue, 30 Jan 2024 16:42:40 -0500
Subject: [PATCH] Allow `@sourceField` on fields of object types _and_ type
 extensions.

---
 composition-js/src/__tests__/compose.test.ts | 43 ++++++++++++++++++++
 internals-js/src/specs/sourceSpec.ts         |  6 ++-
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/composition-js/src/__tests__/compose.test.ts b/composition-js/src/__tests__/compose.test.ts
index c3770a657..c5e638afb 100644
--- a/composition-js/src/__tests__/compose.test.ts
+++ b/composition-js/src/__tests__/compose.test.ts
@@ -5206,5 +5206,48 @@ describe('@source* directives', () => {
         '[renamed] @api(name: "not an identifier") must specify name using only [a-zA-Z0-9-_] characters'
       );
     });
+
+    const extendTypeQuerySchema = gql`
+      extend schema
+        @link(url: "https://specs.apollo.dev/federation/v2.7", import: ["@key"])
+        @link(
+          url: "https://specs.apollo.dev/source/v0.1"
+          import: ["@sourceAPI", "@sourceField"]
+          for: EXECUTION
+        )
+        @sourceAPI(
+          name: "json"
+          http: { baseURL: "https://jsonplaceholder.typicode.com/" }
+        )
+
+      type Query {
+        users: [User]
+          @sourceField(api: "json", http: { GET: "/users" }, selection: "id name")
+      }
+
+      # This test is a regression test for a spurious validation error triggered
+      # by using 'extend type ' here instead of just 'type ' (as above).
+      extend type Query {
+        user(id: ID!): User
+          @sourceField(
+            api: "json"
+            http: { GET: "/users/{id}" }
+            selection: "id name"
+          )
+      }
+
+      type User {
+        id: ID!
+        name: String
+      }
+    `;
+
+    it('allows @sourceField on extended Query type', () => {
+      const result = composeServices([{
+        name: 'extendTypeQuery',
+        typeDefs: extendTypeQuerySchema,
+      }]);
+      expect(result.errors ?? []).toEqual([]);
+    });
   });
 });
diff --git a/internals-js/src/specs/sourceSpec.ts b/internals-js/src/specs/sourceSpec.ts
index 325fac1b6..b43f83731 100644
--- a/internals-js/src/specs/sourceSpec.ts
+++ b/internals-js/src/specs/sourceSpec.ts
@@ -483,7 +483,11 @@ export class SourceSpecDefinition extends FeatureDefinition {
         ));
       } else {
         const typeGrandparent = fieldParent.parent as SchemaElement<any, any>;
-        if (typeGrandparent.sourceAST?.kind !== Kind.OBJECT_TYPE_DEFINITION) {
+        const typeKind = typeGrandparent.sourceAST?.kind;
+        if (
+          typeKind !== Kind.OBJECT_TYPE_DEFINITION &&
+          typeKind !== Kind.OBJECT_TYPE_EXTENSION
+        ) {
           errors.push(ERRORS.SOURCE_FIELD_NOT_ON_ROOT_OR_ENTITY_FIELD.err(
             `${sourceField} must be applied to field of object type`,
             { nodes: application.sourceAST },