entc/gen: add Aggregate to <T>Select and <T>Query

This commit is contained in:
Ariel Mashraki
2022-10-22 22:39:13 +03:00
committed by Ariel Mashraki
parent 1bc4d48a51
commit 765ec09d31
165 changed files with 2629 additions and 226 deletions

View File

@@ -267,6 +267,7 @@ func IsConstraintError(err error) bool {
type selector struct {
label string
flds *[]string
fns []AggregateFunc
scan func(context.Context, any) error
}

View File

@@ -300,6 +300,11 @@ func (uq *UserQuery) Select(fields ...string) *UserSelect {
return selbuild
}
// Aggregate returns a UserSelect configured with the given aggregations.
func (uq *UserQuery) Aggregate(fns ...AggregateFunc) *UserSelect {
return uq.Select().Aggregate(fns...)
}
func (uq *UserQuery) prepareQuery(ctx context.Context) error {
for _, f := range uq.fields {
if !user.ValidColumn(f) {
@@ -499,8 +504,6 @@ func (ugb *UserGroupBy) sqlQuery() *sql.Selector {
for _, fn := range ugb.fns {
aggregation = append(aggregation, fn(selector))
}
// If no columns were selected in a custom aggregation function, the default
// selection is the fields used for "group-by", and the aggregation functions.
if len(selector.SelectedColumns()) == 0 {
columns := make([]string, 0, len(ugb.fields)+len(ugb.fns))
for _, f := range ugb.fields {
@@ -520,6 +523,12 @@ type UserSelect struct {
sql *sql.Selector
}
// Aggregate adds the given aggregation functions to the selector query.
func (us *UserSelect) Aggregate(fns ...AggregateFunc) *UserSelect {
us.fns = append(us.fns, fns...)
return us
}
// Scan applies the selector query and scans the result into the given value.
func (us *UserSelect) Scan(ctx context.Context, v any) error {
if err := us.prepareQuery(ctx); err != nil {
@@ -530,6 +539,16 @@ func (us *UserSelect) Scan(ctx context.Context, v any) error {
}
func (us *UserSelect) sqlScan(ctx context.Context, v any) error {
aggregation := make([]string, 0, len(us.fns))
for _, fn := range us.fns {
aggregation = append(aggregation, fn(us.sql))
}
switch n := len(*us.selector.flds); {
case n == 0 && len(aggregation) > 0:
us.sql.Select(aggregation...)
case n != 0 && len(aggregation) > 0:
us.sql.AppendSelect(aggregation...)
}
rows := &sql.Rows{}
query, args := us.sql.Query()
if err := us.driver.Query(ctx, query, args, rows); err != nil {