From 1422ff258ab7d37531c3023fcdb6c545241bc7b6 Mon Sep 17 00:00:00 2001
From: Ciaran Liedeman <3578740+cliedeman@users.noreply.github.com>
Date: Tue, 18 Jun 2024 13:23:40 +0200
Subject: [PATCH] feat: Change plugin signatures (#2011)

---
 api/generate.go                      | 14 ++++++++++++++
 plugin/federation/federation.go      | 17 +++++++++--------
 plugin/federation/federation_test.go | 10 +++++++---
 plugin/plugin.go                     | 12 ++++++++++++
 4 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/api/generate.go b/api/generate.go
index 9e7b418812f..7a0860dbbc1 100644
--- a/api/generate.go
+++ b/api/generate.go
@@ -58,6 +58,13 @@ func Generate(cfg *config.Config, option ...Option) error {
 				cfg.Sources = append(cfg.Sources, s)
 			}
 		}
+		if inj, ok := p.(plugin.EarlySourcesInjector); ok {
+			s, err := inj.InjectSourcesEarly()
+			if err != nil {
+				return fmt.Errorf("%s: %w", p.Name(), err)
+			}
+			cfg.Sources = append(cfg.Sources, s...)
+		}
 	}
 
 	if err := cfg.LoadSchema(); err != nil {
@@ -70,6 +77,13 @@ func Generate(cfg *config.Config, option ...Option) error {
 				cfg.Sources = append(cfg.Sources, s)
 			}
 		}
+		if inj, ok := p.(plugin.LateSourcesInjector); ok {
+			s, err := inj.InjectSourcesLate(cfg.Schema)
+			if err != nil {
+				return fmt.Errorf("%s: %w", p.Name(), err)
+			}
+			cfg.Sources = append(cfg.Sources, s...)
+		}
 	}
 
 	// LoadSchema again now we have everything
diff --git a/plugin/federation/federation.go b/plugin/federation/federation.go
index fd33255b61b..2975687b7a6 100644
--- a/plugin/federation/federation.go
+++ b/plugin/federation/federation.go
@@ -101,7 +101,7 @@ func (f *federation) MutateConfig(cfg *config.Config) error {
 	return nil
 }
 
-func (f *federation) InjectSourceEarly() *ast.Source {
+func (f *federation) InjectSourcesEarly() ([]*ast.Source, error) {
 	input := ``
 
 	// add version-specific changes on key directive, as well as adding the new directives for federation 2
@@ -136,7 +136,7 @@ func (f *federation) InjectSourceEarly() *ast.Source {
 	directive @interfaceObject on OBJECT
 	directive @link(import: [String!], url: String!) repeatable on SCHEMA
 	directive @override(from: String!, label: String) on FIELD_DEFINITION
-	directive @policy(policies: [[federation__Policy!]!]!) on 
+	directive @policy(policies: [[federation__Policy!]!]!) on
 	  | FIELD_DEFINITION
 	  | OBJECT
 	  | INTERFACE
@@ -144,7 +144,7 @@ func (f *federation) InjectSourceEarly() *ast.Source {
 	  | ENUM
 	directive @provides(fields: FieldSet!) on FIELD_DEFINITION
 	directive @requires(fields: FieldSet!) on FIELD_DEFINITION
-	directive @requiresScopes(scopes: [[federation__Scope!]!]!) on 
+	directive @requiresScopes(scopes: [[federation__Scope!]!]!) on
 	  | FIELD_DEFINITION
 	  | OBJECT
 	  | INTERFACE
@@ -168,16 +168,17 @@ func (f *federation) InjectSourceEarly() *ast.Source {
 	scalar federation__Scope
 `
 	}
-	return &ast.Source{
+
+	return []*ast.Source{{
 		Name:    "federation/directives.graphql",
 		Input:   input,
 		BuiltIn: true,
-	}
+	}}, nil
 }
 
 // InjectSourceLate creates a GraphQL Entity type with all
 // the fields that had the @key directive
-func (f *federation) InjectSourceLate(schema *ast.Schema) *ast.Source {
+func (f *federation) InjectSourcesLate(schema *ast.Schema) ([]*ast.Source, error) {
 	f.setEntities(schema)
 
 	var entities, resolvers, entityResolverInputDefinitions string
@@ -259,11 +260,11 @@ type Entity {
 }`
 	blocks = append(blocks, extendTypeQueryDef)
 
-	return &ast.Source{
+	return []*ast.Source{{
 		Name:    "federation/entity.graphql",
 		BuiltIn: true,
 		Input:   "\n" + strings.Join(blocks, "\n\n") + "\n",
-	}
+	}}, nil
 }
 
 func (f *federation) GenerateCode(data *codegen.Data) error {
diff --git a/plugin/federation/federation_test.go b/plugin/federation/federation_test.go
index aff0d8f689f..28670268b09 100644
--- a/plugin/federation/federation_test.go
+++ b/plugin/federation/federation_test.go
@@ -284,11 +284,15 @@ func load(t *testing.T, name string) (*federation, *config.Config) {
 	}
 
 	f := &federation{Version: cfg.Federation.Version}
-	cfg.Sources = append(cfg.Sources, f.InjectSourceEarly())
+	s, err := f.InjectSourcesEarly()
+	require.NoError(t, err)
+	cfg.Sources = append(cfg.Sources, s...)
 	require.NoError(t, cfg.LoadSchema())
 
-	if src := f.InjectSourceLate(cfg.Schema); src != nil {
-		cfg.Sources = append(cfg.Sources, src)
+	l, err := f.InjectSourcesLate(cfg.Schema)
+	require.NoError(t, err)
+	if l != nil {
+		cfg.Sources = append(cfg.Sources, l...)
 	}
 	require.NoError(t, cfg.LoadSchema())
 
diff --git a/plugin/plugin.go b/plugin/plugin.go
index a5e1ba848d7..6cc4f6dad98 100644
--- a/plugin/plugin.go
+++ b/plugin/plugin.go
@@ -22,11 +22,18 @@ type CodeGenerator interface {
 }
 
 // EarlySourceInjector is used to inject things that are required for user schema files to compile.
+// Deprecated: Use EarlySourcesInjector instead
 type EarlySourceInjector interface {
 	InjectSourceEarly() *ast.Source
 }
 
+// EarlySourcesInjector is used to inject things that are required for user schema files to compile.
+type EarlySourcesInjector interface {
+	InjectSourcesEarly() ([]*ast.Source, error)
+}
+
 // LateSourceInjector is used to inject more sources, after we have loaded the users schema.
+// Deprecated: Use LateSourcesInjector instead
 type LateSourceInjector interface {
 	InjectSourceLate(schema *ast.Schema) *ast.Source
 }
@@ -35,3 +42,8 @@ type LateSourceInjector interface {
 type ResolverImplementer interface {
 	Implement(prevImplementation string, field *codegen.Field) string
 }
+
+// LateSourcesInjector is used to inject more sources, after we have loaded the users schema.
+type LateSourcesInjector interface {
+	InjectSourcesLate(schema *ast.Schema) ([]*ast.Source, error)
+}