From 7e19b4f927283cafa668ccd91a0e1d7e11a36e0b Mon Sep 17 00:00:00 2001
From: Ariel Mashraki <7413593+a8m@users.noreply.github.com>
Date: Sat, 2 Apr 2022 22:01:28 +0300
Subject: [PATCH] gen/entc/template: add option to process nodes after query
 using external templates (#2444)

---
 entc/gen/template.go                     | 1 +
 entc/gen/template/dialect/sql/query.tmpl | 8 +++++++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/entc/gen/template.go b/entc/gen/template.go
index d91a757051..2f8f14cafd 100644
--- a/entc/gen/template.go
+++ b/entc/gen/template.go
@@ -188,6 +188,7 @@ var (
 		"dialect/sql/select/additional/*",
 		"dialect/sql/predicate/edge/*/*",
 		"dialect/sql/query/additional/*",
+		"dialect/sql/query/all/nodes/*",
 		"dialect/sql/query/from/*",
 		"dialect/sql/query/path/*",
 		"import/additional/*",
diff --git a/entc/gen/template/dialect/sql/query.tmpl b/entc/gen/template/dialect/sql/query.tmpl
index 6c3bd3f6a2..4ba95f14f6 100644
--- a/entc/gen/template/dialect/sql/query.tmpl
+++ b/entc/gen/template/dialect/sql/query.tmpl
@@ -79,12 +79,18 @@ func ({{ $receiver }} *{{ $builder }}) sqlAll(ctx context.Context, hooks ...quer
 			{{ template "dialect/sql/query/eagerloading" . }}
 		{{- end }}
 	{{- end }}
+	{{- /* Allow extensions to inject code using templates to process nodes before they are returned. */}}
+	{{- with $tmpls := matchTemplate "dialect/sql/query/all/nodes/*" }}
+		{{- range $tmpl := $tmpls }}
+			{{- xtemplate $tmpl $ }}
+		{{- end }}
+	{{- end }}
 	return nodes, nil
 }
 
 func ({{ $receiver }} *{{ $builder }}) sqlCount(ctx context.Context) (int, error) {
 	_spec := {{ $receiver }}.querySpec()
-	{{- /* Allow mutating the sqlgraph.QuerySpec by ent extensions or user templates.*/}}
+	{{- /* Allow mutating the sqlgraph.QuerySpec by ent extensions or user templates. */}}
 	{{- with $tmpls := matchTemplate "dialect/sql/query/spec/*" }}
 		{{- range $tmpl := $tmpls }}
 			{{- xtemplate $tmpl $ }}