From 6915b13d254531d826054e76227f850de0f633f9 Mon Sep 17 00:00:00 2001
From: Ariel Mashraki <mashraki@fb.com>
Date: Wed, 11 Sep 2019 12:16:03 -0700
Subject: [PATCH] ent/gen: add support for external tempaltes execution

Summary: Pull Request resolved: https://github.com/facebookincubator/ent/pull/30

Reviewed By: dlvhdr

Differential Revision: D17318127

fbshipit-source-id: eccc2963f648296e222de9d002ab3caa8fac1389
---
 dialect/sql/schema/migrate.go |  2 +-
 entc/gen/graph.go             | 24 +++++++++++++++++++++++-
 entc/gen/graph_test.go        |  6 +++++-
 entc/gen/template.go          | 27 ++++++++++++++++++---------
 4 files changed, 47 insertions(+), 12 deletions(-)

diff --git a/dialect/sql/schema/migrate.go b/dialect/sql/schema/migrate.go
index 5ffe3a2e2..f11317e03 100644
--- a/dialect/sql/schema/migrate.go
+++ b/dialect/sql/schema/migrate.go
@@ -17,7 +17,7 @@ import (
 )
 
 const (
-	// TypeTable holds the table name holding the type information.
+	// TypeTable defines the table name holding the type information.
 	TypeTable = "ent_types"
 	// MaxTypes defines the max number of types can be created when
 	// defining universal ids. The left 16-bits are reserved.
diff --git a/entc/gen/graph.go b/entc/gen/graph.go
index f47bf9bfa..08ee0e9cc 100644
--- a/entc/gen/graph.go
+++ b/entc/gen/graph.go
@@ -13,6 +13,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"text/template"
 
 	"github.com/facebookincubator/ent/dialect/sql/schema"
 	"github.com/facebookincubator/ent/entc/load"
@@ -35,6 +36,12 @@ type (
 		// IDType specifies the type of the id field in the codegen.
 		// The supported types are string and int, which also the default.
 		IDType field.Type
+		// Template specifies an alternative template to execute or to override
+		// the default. If nil, the default template is used.
+		//
+		// Note that, additional templates are executed on the Graph object and
+		// the execution output is stored in a file derived by the template name.
+		Template *template.Template
 	}
 	// Graph holds the nodes/entities of the loaded graph schema. Note that, it doesn't
 	// hold the edges of the graph. Instead, each Type holds the edges for other Types.
@@ -70,6 +77,21 @@ func NewGraph(c Config, schemas ...*load.Schema) (g *Graph, err error) {
 // Gen generates the artifacts for the graph.
 func (g *Graph) Gen() (err error) {
 	defer catch(&err)
+	var (
+		external  []GraphTemplate
+		templates = template.Must(templates.Clone())
+	)
+	if g.Template != nil {
+		for _, tmpl := range g.Template.Templates() {
+			if name := tmpl.Name(); templates.Lookup(name) == nil {
+				external = append(external, GraphTemplate{
+					Name:   name,
+					Format: snake(name) + ".go",
+				})
+			}
+			templates = template.Must(templates.AddParseTree(tmpl.Name(), tmpl.Tree))
+		}
+	}
 	for _, n := range g.Nodes {
 		path := filepath.Join(g.Config.Target, n.Package())
 		check(os.MkdirAll(path, os.ModePerm), "create dir %q", path)
@@ -80,7 +102,7 @@ func (g *Graph) Gen() (err error) {
 			check(ioutil.WriteFile(target, b.Bytes(), 0644), "create file %q", target)
 		}
 	}
-	for _, tmpl := range GraphTemplates {
+	for _, tmpl := range append(GraphTemplates[:], external...) {
 		if tmpl.Skip != nil && tmpl.Skip(g) {
 			continue
 		}
diff --git a/entc/gen/graph_test.go b/entc/gen/graph_test.go
index 74fe09cec..28f786733 100644
--- a/entc/gen/graph_test.go
+++ b/entc/gen/graph_test.go
@@ -9,6 +9,7 @@ import (
 	"os"
 	"path/filepath"
 	"testing"
+	"text/template"
 
 	"github.com/facebookincubator/ent/entc/load"
 	"github.com/facebookincubator/ent/schema/field"
@@ -207,7 +208,8 @@ func TestGraph_Gen(t *testing.T) {
 	target := filepath.Join(os.TempDir(), "ent")
 	require.NoError(os.MkdirAll(target, os.ModePerm), "creating tmpdir")
 	defer os.Remove(target)
-	graph, err := NewGraph(Config{Package: "entc/gen", Target: target, Storage: drivers}, &load.Schema{
+	external := template.Must(template.New("external").Parse("external"))
+	graph, err := NewGraph(Config{Package: "entc/gen", Target: target, Storage: drivers, Template: external}, &load.Schema{
 		Name: "T1",
 		Fields: []*load.Field{
 			{Name: "age", Type: field.TypeInt, Optional: true},
@@ -231,4 +233,6 @@ func TestGraph_Gen(t *testing.T) {
 		_, err := os.Stat(fmt.Sprintf(fmt.Sprintf("%s/%s.go", target, format), "t1"))
 		require.NoError(err)
 	}
+	_, err = os.Stat(target + "/external.go")
+	require.NoError(err)
 }
diff --git a/entc/gen/template.go b/entc/gen/template.go
index a1c3245ca..2d3796c94 100644
--- a/entc/gen/template.go
+++ b/entc/gen/template.go
@@ -18,12 +18,25 @@ import (
 
 //go:generate go-bindata -o=internal/bindata.go -pkg=internal ./template/...
 
+type (
+	// TypeTemplate specifies a template that is executed with
+	// each Type object of the graph.
+	TypeTemplate struct {
+		Name   string             // template name.
+		Format func(*Type) string // file name format.
+	}
+	// GraphTemplate specifies a template that is executed with
+	// the Graph object.
+	GraphTemplate struct {
+		Name   string            // template name.
+		Format string            // file name format.
+		Skip   func(*Graph) bool // skip condition.
+	}
+)
+
 var (
 	// Templates holds the template information for a file that the graph is generating.
-	Templates = []struct {
-		Name   string
-		Format func(*Type) string
-	}{
+	Templates = []TypeTemplate{
 		{
 			Name:   "create",
 			Format: pkgf("%s_create.go"),
@@ -56,11 +69,7 @@ var (
 		},
 	}
 	// GraphTemplates holds the templates applied on the graph.
-	GraphTemplates = []struct {
-		Name   string
-		Format string
-		Skip   func(*Graph) bool
-	}{
+	GraphTemplates = []GraphTemplate{
 		{
 			Name:   "base",
 			Format: "ent.go",