From 8c694231b15379b512271d5597cf51769a9d21b5 Mon Sep 17 00:00:00 2001
From: Ariel Mashraki <ariel@mashraki.co.il>
Date: Tue, 29 Mar 2022 14:10:16 +0300
Subject: [PATCH] entc/gen: allow adding extra fields to the generated edges

---
 entc/gen/template.go                 |  3 +++
 entc/gen/template/ent.tmpl           | 21 +++++++++++++++++++--
 entc/integration/gremlin/ent/card.go |  1 -
 examples/entcpkg/ent/user.go         |  1 -
 4 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/entc/gen/template.go b/entc/gen/template.go
index c5ce3ea7a3..d91a757051 100644
--- a/entc/gen/template.go
+++ b/entc/gen/template.go
@@ -182,6 +182,8 @@ var (
 		"dialect/sql/create/additional/*",
 		"dialect/sql/create_bulk/additional/*",
 		"dialect/sql/model/additional/*",
+		"dialect/sql/model/edges/*",
+		"dialect/sql/model/edges/fields/additional/*",
 		"dialect/sql/model/fields/*",
 		"dialect/sql/select/additional/*",
 		"dialect/sql/predicate/edge/*/*",
@@ -191,6 +193,7 @@ var (
 		"import/additional/*",
 		"model/additional/*",
 		"model/comment/additional/*",
+		"model/edges/fields/additional/*",
 		"update/additional/*",
 		"query/additional/*",
 	}
diff --git a/entc/gen/template/ent.tmpl b/entc/gen/template/ent.tmpl
index cd00d473ab..239146d0d9 100644
--- a/entc/gen/template/ent.tmpl
+++ b/entc/gen/template/ent.tmpl
@@ -50,7 +50,7 @@ type {{ $.Name }} struct {
 		{{- xtemplate $tmpl . }}
 	{{- end }}
 	{{- /* Additional fields to add by the user. */}}
-	{{ template "model/fields/additional" $ }}
+	{{- template "model/fields/additional" $ }}
 }
 
 {{- with $.Edges }}
@@ -63,6 +63,8 @@ type {{ $.Name }}Edges struct {
 	// loadedTypes holds the information for reporting if a
 	// type was loaded (or requested) in eager-loading or not.
 	loadedTypes [{{ len . }}]bool
+	{{- /* Additional fields to add by the user. */}}
+	{{- template "model/edges/fields/additional" $ }}
 }
 
 {{- range $i, $e := . }}
@@ -206,7 +208,7 @@ func ({{ $receiver }} {{ $slice }}) config(cfg config) {
 	{{- end }}
 {{- end }}
 
-{{/* A template for adding additional methods or helpers for the generated model.*/}}
+{{/* A template for adding additional methods or helpers for the generated model. */}}
 {{ define "model/additional" }}
 	{{- with $tmpls := matchTemplate "model/additional/*"  }}
 		{{- range $tmpl := $tmpls }}
@@ -219,4 +221,19 @@ func ({{ $receiver }} {{ $slice }}) config(cfg config) {
 			{{- xtemplate $tmpl $ }}
 		{{- end }}
 	{{- end }}
+{{ end }}
+
+{{/* A template for adding additional fields for the Edges object of the generated model. */}}
+{{ define "model/edges/fields/additional" }}
+	{{- with $tmpls := matchTemplate "model/edges/fields/additional/*"  }}
+		{{- range $tmpl := $tmpls }}
+			{{- xtemplate $tmpl $ }}
+		{{- end }}
+	{{- end }}
+
+	{{- with $tmpls := matchTemplate (printf "dialect/%s/model/edges/fields/additional/*" $.Storage) }}
+		{{- range $tmpl := $tmpls }}
+			{{- xtemplate $tmpl $ }}
+		{{- end }}
+	{{- end }}
 {{ end }}
\ No newline at end of file
diff --git a/entc/integration/gremlin/ent/card.go b/entc/integration/gremlin/ent/card.go
index 2e67c139ed..2d0a7d0775 100644
--- a/entc/integration/gremlin/ent/card.go
+++ b/entc/integration/gremlin/ent/card.go
@@ -34,7 +34,6 @@ type Card struct {
 	// Edges holds the relations/edges for other nodes in the graph.
 	// The values are being populated by the CardQuery when eager-loading is set.
 	Edges CardEdges `json:"edges" mashraki:"edges"`
-
 	// StaticField defined by templates.
 	StaticField string `json:"boring,omitempty"`
 }
diff --git a/examples/entcpkg/ent/user.go b/examples/entcpkg/ent/user.go
index ab3f5dfaea..92b753f1a8 100644
--- a/examples/entcpkg/ent/user.go
+++ b/examples/entcpkg/ent/user.go
@@ -24,7 +24,6 @@ type User struct {
 	Name string `json:"name"`
 	// Age holds the value of the "age" field.
 	Age int `json:"age"`
-
 	// StaticField defined by templates (titled STATICFIELD).
 	StaticField string `json:"static_field,omitempty" rql:"static_field"`
 }