entc/gen: allow adding extra fields to the generated edges

ent · Mar 29, 2022 · b983f65 · b983f65
1 parent 317594e
commit b983f65
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git 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
@@ -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 }}