entc/integration: add query modifier example

entc/integration/template/ent/group_query.go

@@ -28,10 +28,9 @@ type GroupQuery struct {
 	order      []OrderFunc
 	fields     []string
 	predicates []predicate.Group
-
 	// additional query fields.
-	extra string
-
+	extra     string
+	modifiers []func(s *sql.Selector)
 	// intermediate query (i.e. traversal path).
 	sql  *sql.Selector
 	path func(context.Context) (*sql.Selector, error)
@@ -415,6 +414,9 @@ func (gq *GroupQuery) sqlQuery(ctx context.Context) *sql.Selector {
 		selector = gq.sql
 		selector.Select(selector.Columns(columns...)...)
 	}
+	for _, m := range gq.modifiers {
+		m(selector)
+	}
 	for _, p := range gq.predicates {
 		p(selector)
 	}
@@ -432,6 +434,11 @@ func (gq *GroupQuery) sqlQuery(ctx context.Context) *sql.Selector {
 	return selector
 }
 
+func (gq *GroupQuery) Modify(modifier func(s *sql.Selector)) *GroupQuery {
+	gq.modifiers = append(gq.modifiers, modifier)
+	return gq
+}
+
 // GroupGroupBy is the group-by builder for Group entities.
 type GroupGroupBy struct {
 	config

entc/integration/template/ent/pet_query.go

@@ -32,10 +32,9 @@ type PetQuery struct {
 	// eager-loading edges.
 	withOwner *UserQuery
 	withFKs   bool
-
 	// additional query fields.
-	extra string
-
+	extra     string
+	modifiers []func(s *sql.Selector)
 	// intermediate query (i.e. traversal path).
 	sql  *sql.Selector
 	path func(context.Context) (*sql.Selector, error)
@@ -494,6 +493,9 @@ func (pq *PetQuery) sqlQuery(ctx context.Context) *sql.Selector {
 		selector = pq.sql
 		selector.Select(selector.Columns(columns...)...)
 	}
+	for _, m := range pq.modifiers {
+		m(selector)
+	}
 	for _, p := range pq.predicates {
 		p(selector)
 	}
@@ -511,6 +513,11 @@ func (pq *PetQuery) sqlQuery(ctx context.Context) *sql.Selector {
 	return selector
 }
 
+func (pq *PetQuery) Modify(modifier func(s *sql.Selector)) *PetQuery {
+	pq.modifiers = append(pq.modifiers, modifier)
+	return pq
+}
+
 // PetGroupBy is the group-by builder for Pet entities.
 type PetGroupBy struct {
 	config

entc/integration/template/ent/template/modifier.tmpl

@@ -0,0 +1,23 @@
+{{/* The line below tells Intellij/GoLand to enable the autocompletion based *gen.Type type. */}}
+{{/* gotype: entgo.io/ent/entc/gen.Type */}}
+
+{{/* A template for adding the Modify method to the query builder. */}}
+{{ define "dialect/sql/query/additional/modify" }}
+    {{ $builder := pascal $.Scope.Builder }}
+    {{ $receiver := receiver $builder }}
+    func ({{ $receiver }} *{{ $builder }}) Modify(modifier func(s *sql.Selector)) *{{ $builder }} {
+        {{ $receiver }}.modifiers = append({{ $receiver }}.modifiers, modifier)
+        return {{ $receiver }}
+    }
+{{ end }}
+
+{{ define "dialect/sql/query/fields/additional/modify" -}}
+    modifiers []func(s *sql.Selector)
+{{- end -}}
+
+{{ define "dialect/sql/query/selector/modify" }}
+    {{- $receiver := pascal $.Scope.Builder | receiver }}
+    for _, m := range {{ $receiver }}.modifiers {
+        m(selector)
+    }
+{{- end -}}

entc/integration/template/ent/template/query.tmpl

@@ -7,10 +7,10 @@ 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 "dialect/sql/query/fields/additional/extra" }}
+{{ define "dialect/sql/query/fields/additional/extra" -}}
 	// additional query fields.
 	extra string
-{{ end }}
+{{- end -}}
 
 {{ define "dialect/sql/query/fields/init/extra" }}
 	extra: "{{ $.Name }}",

entc/integration/template/ent/user_query.go

@@ -33,10 +33,9 @@ type UserQuery struct {
 	// eager-loading edges.
 	withPets    *PetQuery
 	withFriends *UserQuery
-
 	// additional query fields.
-	extra string
-
+	extra     string
+	modifiers []func(s *sql.Selector)
 	// intermediate query (i.e. traversal path).
 	sql  *sql.Selector
 	path func(context.Context) (*sql.Selector, error)
@@ -588,6 +587,9 @@ func (uq *UserQuery) sqlQuery(ctx context.Context) *sql.Selector {
 		selector = uq.sql
 		selector.Select(selector.Columns(columns...)...)
 	}
+	for _, m := range uq.modifiers {
+		m(selector)
+	}
 	for _, p := range uq.predicates {
 		p(selector)
 	}
@@ -605,6 +607,11 @@ func (uq *UserQuery) sqlQuery(ctx context.Context) *sql.Selector {
 	return selector
 }
 
+func (uq *UserQuery) Modify(modifier func(s *sql.Selector)) *UserQuery {
+	uq.modifiers = append(uq.modifiers, modifier)
+	return uq
+}
+
 // UserGroupBy is the group-by builder for User entities.
 type UserGroupBy struct {
 	config

entc/integration/template/template_test.go

@@ -10,9 +10,11 @@ import (
 	"reflect"
 	"testing"
 
+	"entgo.io/ent/dialect/sql"
 	"entgo.io/ent/entc/integration/template/ent"
 	"entgo.io/ent/entc/integration/template/ent/hook"
 	"entgo.io/ent/entc/integration/template/ent/migrate"
+	"entgo.io/ent/entc/integration/template/ent/pet"
 	"entgo.io/ent/entc/integration/template/ent/user"
 
 	_ "github.com/mattn/go-sqlite3"
@@ -65,4 +67,16 @@ func TestCustomTemplate(t *testing.T) {
 	result := client.User.Query().Where(user.NameGlob("a8*")).
 		AllX(ctx)
 	require.Equal(t, 1, len(result))
+
+	var v []struct{ ID, Owner int }
+	client.Pet.Query().
+		Modify(func(s *sql.Selector) {
+			t := sql.Table(user.Table)
+			s.Join(t).On(s.C(pet.OwnerColumn), t.C(user.FieldID))
+			s.Select(s.C(pet.FieldID), sql.As(t.C(user.FieldID), "owner"))
+		}).
+		Select(pet.FieldID).
+		ScanX(ctx, &v)
+	require.Equal(t, p.ID, v[0].ID)
+	require.Equal(t, u.ID, v[0].Owner)
 }