gen/template: allow overriding of client struct and initialization (#2503)

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.

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.

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

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())
 }