From 4e434b2b732e7e2feb7b6678796bf9dcc4a17516 Mon Sep 17 00:00:00 2001
From: MasseElch <12862103+masseelch@users.noreply.github.com>
Date: Thu, 28 Apr 2022 09:07:44 +0200
Subject: [PATCH] gen/template: allow overriding of client struct and
 initialization (#2503)

---
 entc/gen/template/client.tmpl                 | 44 ++++++++++--------
 entc/integration/template/ent/client.go       |  7 +++
 .../template/ent/template/client.tmpl         | 46 +++++++++++++++++++
 entc/integration/template/template_test.go    |  2 +
 4 files changed, 79 insertions(+), 20 deletions(-)

diff --git a/entc/gen/template/client.tmpl b/entc/gen/template/client.tmpl
index ebcf907c6..f6416f939 100644
--- a/entc/gen/template/client.tmpl
+++ b/entc/gen/template/client.tmpl
@@ -6,25 +6,7 @@ in the LICENSE file in the root directory of this source tree.
 
 {{/* gotype: entgo.io/ent/entc/gen.Graph */}}
 
-{{ define "client" }}
-
-{{ $pkg := base $.Config.Package }}
-{{ template "header" $ }}
-
-import (
-	"log"
-
-	"{{ $.Config.Package }}/migrate"
-	{{ range $n := $.Nodes }}
-		{{ $n.PackageAlias }} "{{ $n.Config.Package }}/{{ $n.PackageDir }}"
-	{{- end }}
-
-	"entgo.io/ent/dialect"
-	{{ range $import := $.Storage.Imports -}}
-		"{{ $import }}"
-	{{ end -}}
-)
-
+{{ define "client/init" }}
 // Client is the client that holds all ent builders.
 type Client struct {
 	config
@@ -58,9 +40,31 @@ func (c *Client) init() {
 		c.Schema = migrate.NewSchema(c.driver)
 	{{- end }}
 	{{- range $n := $.Nodes }}
-    	c.{{ $n.Name }} =  New{{ $n.Name }}Client(c.config)
+    	c.{{ $n.Name }} = New{{ $n.Name }}Client(c.config)
 	{{- end }}
 }
+{{ end }}
+
+{{ define "client" }}
+
+{{ $pkg := base $.Config.Package }}
+{{ template "header" $ }}
+
+import (
+	"log"
+
+	"{{ $.Config.Package }}/migrate"
+	{{ range $n := $.Nodes }}
+		{{ $n.PackageAlias }} "{{ $n.Config.Package }}/{{ $n.PackageDir }}"
+	{{- end }}
+
+	"entgo.io/ent/dialect"
+	{{ range $import := $.Storage.Imports -}}
+		"{{ $import }}"
+	{{ end -}}
+)
+
+{{ template "client/init" $ }}
 
 // Open opens a database/sql.DB specified by the driver name and
 // the data source name, and returns a new client attached to it.
diff --git a/entc/integration/template/ent/client.go b/entc/integration/template/ent/client.go
index f68a6fbc7..0cb6ee607 100644
--- a/entc/integration/template/ent/client.go
+++ b/entc/integration/template/ent/client.go
@@ -25,6 +25,7 @@ import (
 // Client is the client that holds all ent builders.
 type Client struct {
 	config
+	extraHiddenField int
 	// Schema is the client for creating, migrating and dropping schema.
 	Schema *migrate.Schema
 	// Group is the client for interacting with the Group builders.
@@ -47,12 +48,18 @@ func NewClient(opts ...Option) *Client {
 }
 
 func (c *Client) init() {
+	c.extraHiddenField = 20
 	c.Schema = migrate.NewSchema(c.driver)
 	c.Group = NewGroupClient(c.config)
 	c.Pet = NewPetClient(c.config)
 	c.User = NewUserClient(c.config)
 }
 
+// HiddenData returns the extraHiddenField value.
+func (c *Client) HiddenData() int {
+	return c.extraHiddenField
+}
+
 // Open opens a database/sql.DB specified by the driver name and
 // the data source name, and returns a new client attached to it.
 // Optional parameters can be added for configuring the client.
diff --git a/entc/integration/template/ent/template/client.tmpl b/entc/integration/template/ent/template/client.tmpl
index 2348925e4..d2986bdb8 100644
--- a/entc/integration/template/ent/template/client.tmpl
+++ b/entc/integration/template/ent/template/client.tmpl
@@ -7,6 +7,52 @@ 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 "client/init" }}
+// Client is the client that holds all ent builders.
+type Client struct {
+	config
+	extraHiddenField int
+	{{- if $.SupportMigrate }}
+		// Schema is the client for creating, migrating and dropping schema.
+		Schema *migrate.Schema
+	{{- end }}
+	{{- range $n := $.Nodes }}
+		// {{ $n.Name }} is the client for interacting with the {{ $n.Name }} builders.
+		{{ $n.Name }} *{{ $n.Name }}Client
+	{{- end }}
+	{{- template "client/fields/additional" $ }}
+	{{- with $tmpls := matchTemplate "client/fields/additional/*" }}
+		{{- range $tmpl := $tmpls }}
+			{{- xtemplate $tmpl $ }}
+		{{- end }}
+	{{- end }}
+}
+
+// NewClient creates a new client configured with the given options.
+func NewClient(opts ...Option) *Client {
+	cfg := config{log: log.Println, hooks: &hooks{}}
+	cfg.options(opts...)
+	client := &Client{config: cfg}
+	client.init()
+	return client
+}
+
+func (c *Client) init() {
+	c.extraHiddenField = 20
+	{{- if $.SupportMigrate }}
+		c.Schema = migrate.NewSchema(c.driver)
+	{{- end }}
+	{{- range $n := $.Nodes }}
+    	c.{{ $n.Name }} =  New{{ $n.Name }}Client(c.config)
+	{{- end }}
+}
+
+// HiddenData returns the extraHiddenField value.
+func (c *Client) HiddenData() int {
+	return c.extraHiddenField
+}
+{{ end }}
+
 {{ define "client/fields/additional" }}
 	// additional fields.
 	sync.Mutex
diff --git a/entc/integration/template/template_test.go b/entc/integration/template/template_test.go
index f69cf1bc4..b0fd37f64 100644
--- a/entc/integration/template/template_test.go
+++ b/entc/integration/template/template_test.go
@@ -91,4 +91,6 @@ func TestCustomTemplate(t *testing.T) {
 		Select().
 		IntX(ctx)
 	require.Equal(t, sum, got)
+
+	require.Equal(t, 20, client.HiddenData())
 }