Optionally render entity requires populator function for advanced @re…

…quires use cases (#2884)

* Adding generation of new functions to populate requires representations. WIP.

* Something.

* Adding config option for Package to allow for enabling flags.  Added flag to render explicit requires function.

* Adding fixed by @ldebruijn.

* Adding explicit requires testsing and make requires follow federation package.

* Fix test failure.

* Using embeded template like federation gotpl.  Fix rewriter not using correct dir.

* Update generated code.

* Docs (#1)

* Adding initial docs for explicit requires

* Add example docs for explicit requires

* Adding ordering fix.

* Regenerate.

---------

Co-authored-by: Jesse Lovelace <jesse.lovelace@gmail.com>
Co-authored-by: Brandon Him <brandonhim@live.com>

codegen/config/package.go

@@ -10,10 +10,11 @@ import (
 )
 
 type PackageConfig struct {
-	Filename      string `yaml:"filename,omitempty"`
-	Package       string `yaml:"package,omitempty"`
-	Version       int    `yaml:"version,omitempty"`
-	ModelTemplate string `yaml:"model_template,omitempty"`
+	Filename      string          `yaml:"filename,omitempty"`
+	Package       string          `yaml:"package,omitempty"`
+	Version       int             `yaml:"version,omitempty"`
+	ModelTemplate string          `yaml:"model_template,omitempty"`
+	Options       map[string]bool `yaml:"options,omitempty"`
 }
 
 func (c *PackageConfig) ImportPath() string {

docs/content/recipes/federation.md

@@ -201,3 +201,48 @@ should return
   }
 }
 ```
+
+## Explicit `@requires` Directive
+If you need to support **nested** or **array** fields in the `@requires` directive, this can be enabled in the configuration by setting `federation.options.explicit_requires` to true.
+
+```yml
+federation:
+  filename: graph/federation.go
+  package: graph
+  version: 2
+  options:
+    explicit_requires: true
+```
+
+Enabling this will generate corresponding functions with the entity representations received in the request. This allows for the entity model to be explicitly populated with the required data provided.
+
+### Example
+Take a simple todo app schema that needs to provide a formatted status text to be used across all clients by referencing the assignee's name.
+
+```graphql
+type Todo @key(fields:"id") {
+  id: ID!
+  text: String!
+  statusText: String! @requires(fields: "assignee { name }")
+  status: String!
+  owner: User!
+  assignee: User!
+}
+
+type User @key(fields:"id") {
+  id: ID!
+  name: String! @external
+}
+```
+
+A `PopulateTodoRequires` function is generated, and can be modified accordingly to use the todo representation with the assignee name.
+
+```golang
+// PopulateTodoRequires is the requires populator for the Todo entity.
+func (ec *executionContext) PopulateTodoRequires(ctx context.Context, entity *model.Todo, reps map[string]interface{}) error {
+	if reps["assignee"] != nil {
+		entity.StatusText = entity.Status + " by " + reps["assignee"].(map[string]interface{})["name"].(string)
+	}
+	return nil
+}
+```

plugin/federation/entity.go

@@ -2,6 +2,7 @@ package federation
 
 import (
 	"go/types"
+	"strings"
 
 	"github.com/vektah/gqlparser/v2/ast"
 
@@ -18,6 +19,7 @@ type Entity struct {
 	Resolvers []*EntityResolver
 	Requires  []*Requires
 	Multi     bool
+	Type      types.Type
 }
 
 type EntityResolver struct {
@@ -116,3 +118,9 @@ func (e *Entity) keyFields() []string {
 	}
 	return keyFields
 }
+
+// GetTypeInfo - get the imported package & type name combo.  package.TypeName
+func (e Entity) GetTypeInfo() string {
+	typeParts := strings.Split(e.Type.String(), "/")
+	return typeParts[len(typeParts)-1]
+}

plugin/federation/federation.go

@@ -11,16 +11,21 @@ import (
 	"github.com/99designs/gqlgen/codegen"
 	"github.com/99designs/gqlgen/codegen/config"
 	"github.com/99designs/gqlgen/codegen/templates"
+	"github.com/99designs/gqlgen/internal/rewrite"
 	"github.com/99designs/gqlgen/plugin"
 	"github.com/99designs/gqlgen/plugin/federation/fieldset"
 )
 
 //go:embed federation.gotpl
 var federationTemplate string
 
+//go:embed requires.gotpl
+var explicitRequiresTemplate string
+
 type federation struct {
-	Entities []*Entity
-	Version  int
+	Entities       []*Entity
+	Version        int
+	PackageOptions map[string]bool
 }
 
 // New returns a federation plugin that injects
@@ -263,6 +268,16 @@ type Entity {
 }
 
 func (f *federation) GenerateCode(data *codegen.Data) error {
+	// requires imports
+	requiresImports := make(map[string]bool, 0)
+	requiresImports["context"] = true
+	requiresImports["fmt"] = true
+
+	requiresEntities := make(map[string]*Entity, 0)
+
+	// Save package options on f for template use
+	f.PackageOptions = data.Config.Federation.Options
+
 	if len(f.Entities) > 0 {
 		if data.Objects.ByName("Entity") != nil {
 			data.Objects.ByName("Entity").Root = true
@@ -302,9 +317,19 @@ func (f *federation) GenerateCode(data *codegen.Data) error {
 					fmt.Println("skipping @requires field " + reqField.Name + " in " + e.Def.Name)
 					continue
 				}
+				// keep track of which entities have requires
+				requiresEntities[e.Def.Name] = e
+				// make a proper import path
+				typeString := strings.Split(obj.Type.String(), ".")
+				requiresImports[strings.Join(typeString[:len(typeString)-1], ".")] = true
+
 				cgField := reqField.Field.TypeReference(obj, data.Objects)
 				reqField.Type = cgField.TypeReference
 			}
+
+			// add type info to entity
+			e.Type = obj.Type
+
 		}
 	}
 
@@ -325,6 +350,75 @@ func (f *federation) GenerateCode(data *codegen.Data) error {
 		}
 	}
 
+	if data.Config.Federation.Options["explicit_requires"] && len(requiresEntities) > 0 {
+		// check for existing requires functions
+		type Populator struct {
+			FuncName       string
+			Exists         bool
+			Comment        string
+			Implementation string
+			Entity         *Entity
+		}
+		populators := make([]Populator, 0)
+
+		rewriter, err := rewrite.New(data.Config.Federation.Dir())
+		if err != nil {
+			return err
+		}
+
+		for name, entity := range requiresEntities {
+			populator := Populator{
+				FuncName: fmt.Sprintf("Populate%sRequires", name),
+				Entity:   entity,
+			}
+
+			populator.Comment = strings.TrimSpace(strings.TrimLeft(rewriter.GetMethodComment("executionContext", populator.FuncName), `\`))
+			populator.Implementation = strings.TrimSpace(rewriter.GetMethodBody("executionContext", populator.FuncName))
+
+			if populator.Implementation == "" {
+				populator.Exists = false
+				populator.Implementation = fmt.Sprintf("panic(fmt.Errorf(\"not implemented: %v\"))", populator.FuncName)
+			}
+			populators = append(populators, populator)
+		}
+
+		sort.Slice(populators, func(i, j int) bool {
+			return populators[i].FuncName < populators[j].FuncName
+		})
+
+		requiresFile := data.Config.Federation.Dir() + "/federation.requires.go"
+		existingImports := rewriter.ExistingImports(requiresFile)
+		for _, imp := range existingImports {
+			if imp.Alias == "" {
+				// import exists in both places, remove
+				delete(requiresImports, imp.ImportPath)
+			}
+		}
+
+		for k := range requiresImports {
+			existingImports = append(existingImports, rewrite.Import{ImportPath: k})
+		}
+
+		// render requires populators
+		err = templates.Render(templates.Options{
+			PackageName: data.Config.Federation.Package,
+			Filename:    requiresFile,
+			Data: struct {
+				federation
+				ExistingImports []rewrite.Import
+				Populators      []Populator
+				OriginalSource  string
+			}{*f, existingImports, populators, ""},
+			GeneratedHeader: false,
+			Packages:        data.Config.Packages,
+			Template:        explicitRequiresTemplate,
+		})
+		if err != nil {
+			return err
+		}
+
+	}
+
 	return templates.Render(templates.Options{
 		PackageName:     data.Config.Federation.Package,
 		Filename:        data.Config.Federation.Filename,

plugin/federation/federation.gotpl

@@ -5,6 +5,7 @@
 {{ reserveImport "sync"  }}
 
 {{ reserveImport "github.com/99designs/gqlgen/plugin/federation/fedruntime" }}
+{{ $options := .PackageOptions }}
 
 var (
 	ErrUnknownType = errors.New("unknown type")
@@ -103,11 +104,18 @@ func (ec *executionContext) __resolve_entities(ctx context.Context, representati
 						if err != nil {
 							return fmt.Errorf(`resolving Entity "{{$entity.Def.Name}}": %w`, err)
 						}
-						{{ range $entity.Requires }}
-							entity.{{.Field.JoinGo `.`}}, err = ec.{{.Type.UnmarshalFunc}}(ctx, rep["{{.Field.Join `"].(map[string]interface{})["`}}"])
+						{{ if and (index $options "explicit_requires") $entity.Requires }}
+							err = ec.Populate{{$entity.Def.Name}}Requires(ctx, entity, rep)
 							if err != nil {
-								return err
+								return fmt.Errorf(`populating requires for Entity "{{$entity.Def.Name}}": %w`, err)
 							}
+						{{- else }}
+							{{ range $entity.Requires }}
+								entity.{{.Field.JoinGo `.`}}, err = ec.{{.Type.UnmarshalFunc}}(ctx, rep["{{.Field.Join `"].(map[string]interface{})["`}}"])
+								if err != nil {
+									return err
+								}
+							{{- end }}
 						{{- end }}
 						list[idx[i]] = entity
 						return nil