From 81eeefa6a2a5bc4bb24684b4c67ef70d35ce7f89 Mon Sep 17 00:00:00 2001
From: Ariel Mashraki <ariel@mashraki.co.il>
Date: Mon, 19 Jul 2021 15:33:35 +0300
Subject: [PATCH] entc/gen: add support for helper templates

Closed #1655
---
 entc/gen/graph.go                             | 36 +++++++++++++++----
 .../template/ent/template/helper.tmpl         | 21 +++++++++++
 .../template/ent/template/node.tmpl           | 18 ++++++----
 3 files changed, 63 insertions(+), 12 deletions(-)
 create mode 100644 entc/integration/template/ent/template/helper.tmpl

diff --git a/entc/gen/graph.go b/entc/gen/graph.go
index a1055bb45..1599d7bc2 100644
--- a/entc/gen/graph.go
+++ b/entc/gen/graph.go
@@ -15,6 +15,7 @@ import (
 	"os"
 	"path/filepath"
 	"runtime/debug"
+	"strings"
 	"text/template/parse"
 
 	"entgo.io/ent/dialect/sql/schema"
@@ -598,11 +599,15 @@ func (g *Graph) typ(name string) (*Type, bool) {
 	return nil, false
 }
 
-// templates returns the template.Template for the code and external templates
-// to execute on the Graph object if provided.
+// templates returns the Template to execute on the Graph,
+// and a list of optional external templates if provided.
 func (g *Graph) templates() (*Template, []GraphTemplate) {
 	initTemplates()
-	external := make([]GraphTemplate, 0, len(g.Templates))
+	var (
+		roots    = make(map[string]struct{})
+		helpers  = make(map[string]struct{})
+		external = make([]GraphTemplate, 0, len(g.Templates))
+	)
 	for _, rootT := range g.Templates {
 		templates.Funcs(rootT.FuncMap)
 		for _, tmpl := range rootT.Templates() {
@@ -610,17 +615,36 @@ func (g *Graph) templates() (*Template, []GraphTemplate) {
 				continue
 			}
 			name := tmpl.Name()
-			// If the template does not override or extend one of
-			// the builtin templates, generate it in a new file.
-			if templates.Lookup(name) == nil && !extendExisting(name) {
+			switch {
+			// Helper templates can be either global (prefixed with "helper/"),
+			// or local, where their names follow the format: <root-tmpl>/helper/.+).
+			case strings.HasPrefix(name, "helper/"):
+			case strings.Contains(name, "/helper/"):
+				helpers[name] = struct{}{}
+			case templates.Lookup(name) == nil && !extendExisting(name):
+				// If the template does not override or extend one of
+				// the builtin templates, generate it in a new file.
 				external = append(external, GraphTemplate{
 					Name:   name,
 					Format: snake(name) + ".go",
 				})
+				roots[name] = struct{}{}
 			}
 			templates = MustParse(templates.AddParseTree(name, tmpl.Tree))
 		}
 	}
+	for name := range helpers {
+		root := name[:strings.Index(name, "/helper/")]
+		// If the name is prefixed with a name of a root
+		// template, we treat it as a local helper template.
+		if _, ok := roots[root]; ok {
+			continue
+		}
+		external = append(external, GraphTemplate{
+			Name:   name,
+			Format: snake(name) + ".go",
+		})
+	}
 	for _, f := range g.Features {
 		external = append(external, f.GraphTemplates...)
 	}
diff --git a/entc/integration/template/ent/template/helper.tmpl b/entc/integration/template/ent/template/helper.tmpl
new file mode 100644
index 000000000..133ea49b1
--- /dev/null
+++ b/entc/integration/template/ent/template/helper.tmpl
@@ -0,0 +1,21 @@
+{{/*
+Copyright 2019-present Facebook Inc. All rights reserved.
+This source code is licensed under the Apache 2.0 license found
+in the LICENSE file in the root directory of this source tree.
+*/}}
+
+{{/* The line below tells Intellij/GoLand to enable the autocompletion based *gen.Graph type. */}}
+{{/* gotype: entgo.io/ent/entc/gen.Graph */}}
+
+{{ define "helper/one" }}
+    {{ range $n := $.Nodes }}
+        {{/* Non empty global helper template. */}}
+    {{ end }}
+{{ end }}
+
+
+{{ define "helper/two" }}
+    {{ range $n := $.Nodes }}
+        {{/* Non empty global helper template. */}}
+    {{ end }}
+{{ end }}
\ No newline at end of file
diff --git a/entc/integration/template/ent/template/node.tmpl b/entc/integration/template/ent/template/node.tmpl
index cf2310f3d..83b0ccad2 100644
--- a/entc/integration/template/ent/template/node.tmpl
+++ b/entc/integration/template/ent/template/node.tmpl
@@ -164,6 +164,17 @@ func (t *tables) Load(ctx context.Context, querier querier) ([]string, error) {
 	return tables, err
 }
 
+{{ template "node/helper/loadtable" $ }}
+
+{{ end }}
+
+{{ define "client/fields/additional" }}
+	// additional fields for node api
+	tables tables
+{{ end }}
+
+{{/* An example of a local helper template. */}}
+{{ define "node/helper/loadtable" }}
 func (tables) load(ctx context.Context, querier querier) ([]string, error) {
 	rows := &sql.Rows{}
 	query, args := sql.Select("type").
@@ -177,9 +188,4 @@ func (tables) load(ctx context.Context, querier querier) ([]string, error) {
 	var tables []string
 	return tables, sql.ScanSlice(rows, &tables)
 }
-{{ end }}
-
-{{ define "client/fields/additional" }}
-	// additional fields for node api
-	tables tables
-{{ end }}
+{{ end }}
\ No newline at end of file