From e5bceccf1dd3ccbcfe7ac59cba8e0ccae0e74cdc Mon Sep 17 00:00:00 2001 From: Ariel Mashraki Date: Tue, 1 Jun 2021 16:05:15 +0300 Subject: [PATCH] entc/integration: add query modifier example --- entc/integration/template/ent/group_query.go | 13 ++++++++--- entc/integration/template/ent/pet_query.go | 13 ++++++++--- .../template/ent/template/modifier.tmpl | 23 +++++++++++++++++++ .../template/ent/template/query.tmpl | 4 ++-- entc/integration/template/ent/user_query.go | 13 ++++++++--- entc/integration/template/template_test.go | 14 +++++++++++ 6 files changed, 69 insertions(+), 11 deletions(-) create mode 100644 entc/integration/template/ent/template/modifier.tmpl diff --git a/entc/integration/template/ent/group_query.go b/entc/integration/template/ent/group_query.go index 0cc131ddb..d547051cb 100644 --- a/entc/integration/template/ent/group_query.go +++ b/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 diff --git a/entc/integration/template/ent/pet_query.go b/entc/integration/template/ent/pet_query.go index ce66bc135..86bd3fb84 100644 --- a/entc/integration/template/ent/pet_query.go +++ b/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 diff --git a/entc/integration/template/ent/template/modifier.tmpl b/entc/integration/template/ent/template/modifier.tmpl new file mode 100644 index 000000000..b3ea867ce --- /dev/null +++ b/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 -}} diff --git a/entc/integration/template/ent/template/query.tmpl b/entc/integration/template/ent/template/query.tmpl index 12a37ce49..3e2ec8c43 100644 --- a/entc/integration/template/ent/template/query.tmpl +++ b/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 }}", diff --git a/entc/integration/template/ent/user_query.go b/entc/integration/template/ent/user_query.go index 56f8db307..a70cc935b 100644 --- a/entc/integration/template/ent/user_query.go +++ b/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 diff --git a/entc/integration/template/template_test.go b/entc/integration/template/template_test.go index f43717c5d..652dbc703 100644 --- a/entc/integration/template/template_test.go +++ b/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) }