add offset step to query builder (#1049)

Summary:
Pull Request resolved: https://github.com/facebookexternal/fbc/pull/1049

Useful for real paging

Reviewed By: noamsch

Differential Revision: D16003607

fbshipit-source-id: 6a85d0e3d71a2582bc3cd8f1d66748dda4f2a10e

entc/gen/template/builder/query.tmpl

@@ -10,10 +10,11 @@
 // {{ $builder }} is the builder for querying {{ pascal $.Name }} entities.
 type {{ $builder }} struct {
 	config
-	limit *int
+	limit		*int
-	order  []Order
+	offset		*int
-	unique []string
+	order		[]Order
-	predicates []ent.Predicate
+	unique		[]string
+	predicates 	[]ent.Predicate
 	// intermediate queries.
 	sql     *sql.Selector
 	gremlin *dsl.Traversal
@@ -31,6 +32,12 @@ func ({{ $receiver }} *{{ $builder }}) Limit(limit int) *{{ $builder }} {
 	return {{ $receiver }}
 }
 
+// Offset adds an offset step to the query.
+func ({{ $receiver }} *{{ $builder }}) Offset(offset int) *{{ $builder }} {
+	{{ $receiver }}.offset = &offset
+	return {{ $receiver }}
+}
+
 // Order adds an order step to the query.
 func ({{ $receiver }} *{{ $builder }}) Order(o ...Order) *{{ $builder }} {
 	{{ $receiver }}.order = append({{ $receiver }}.order, o...)

entc/gen/template/dialect/gremlin/query.tmpl

@@ -65,7 +65,12 @@ func ({{ $receiver }} *{{ $builder }}) gremlinQuery() *dsl.Traversal {
 			p.Gremlin(v)
 		}
 	}
-	if limit := {{ $receiver }}.limit; limit != nil {
+	switch limit, offset := {{ $receiver }}.limit, {{ $receiver }}.offset; {
+	case limit != nil && offset != nil:
+		v.Range(*offset, *offset + *limit)
+	case offset != nil:
+		v.Range(*offset, math.MaxInt64)
+	case limit != nil:
 		v.Limit(*limit)
 	}
 	if unique := {{ $receiver }}.unique; len(unique) == 0 {

entc/gen/template/dialect/sql/query.tmpl

@@ -80,6 +80,11 @@ func ({{ $receiver }} *{{ $builder }}) sqlQuery() *sql.Selector {
 	for _, p := range {{ $receiver }}.order {
 		p.SQL(selector)
 	}
+	if offset := {{ $receiver }}.offset; offset != nil {
+		// limit is mandatory for offset clause. We start
+		// with default value, and override it below if needed.
+		selector.Offset(*offset).Limit(math.MaxInt64)
+	}
 	if limit := {{ $receiver }}.limit; limit != nil {
 		selector.Limit(*limit)
 	}

entc/gen/template/import.tmpl

@@ -3,6 +3,7 @@ import (
 	"fmt"
 	"context"
 	"errors"
+	"math"
 	"strconv"
 	"strings"
 	"time"

entc/integration/ent/card_query.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 
 	"fbc/ent/entc/integration/ent/card"
 	"fbc/ent/entc/integration/ent/user"
@@ -24,6 +25,7 @@ import (
 type CardQuery struct {
 	config
 	limit      *int
+	offset     *int
 	order      []Order
 	unique     []string
 	predicates []ent.Predicate
@@ -44,6 +46,12 @@ func (cq *CardQuery) Limit(limit int) *CardQuery {
 	return cq
 }
 
+// Offset adds an offset step to the query.
+func (cq *CardQuery) Offset(offset int) *CardQuery {
+	cq.offset = &offset
+	return cq
+}
+
 // Order adds an order step to the query.
 func (cq *CardQuery) Order(o ...Order) *CardQuery {
 	cq.order = append(cq.order, o...)
@@ -363,6 +371,11 @@ func (cq *CardQuery) sqlQuery() *sql.Selector {
 	for _, p := range cq.order {
 		p.SQL(selector)
 	}
+	if offset := cq.offset; offset != nil {
+		// limit is mandatory for offset clause. We start
+		// with default value, and override it below if needed.
+		selector.Offset(*offset).Limit(math.MaxInt64)
+	}
 	if limit := cq.limit; limit != nil {
 		selector.Limit(*limit)
 	}
@@ -432,7 +445,12 @@ func (cq *CardQuery) gremlinQuery() *dsl.Traversal {
 			p.Gremlin(v)
 		}
 	}
-	if limit := cq.limit; limit != nil {
+	switch limit, offset := cq.limit, cq.offset; {
+	case limit != nil && offset != nil:
+		v.Range(*offset, *offset+*limit)
+	case offset != nil:
+		v.Range(*offset, math.MaxInt64)
+	case limit != nil:
 		v.Limit(*limit)
 	}
 	if unique := cq.unique; len(unique) == 0 {

entc/integration/ent/comment_query.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 
 	"fbc/ent/entc/integration/ent/comment"
 
@@ -23,6 +24,7 @@ import (
 type CommentQuery struct {
 	config
 	limit      *int
+	offset     *int
 	order      []Order
 	unique     []string
 	predicates []ent.Predicate
@@ -43,6 +45,12 @@ func (cq *CommentQuery) Limit(limit int) *CommentQuery {
 	return cq
 }
 
+// Offset adds an offset step to the query.
+func (cq *CommentQuery) Offset(offset int) *CommentQuery {
+	cq.offset = &offset
+	return cq
+}
+
 // Order adds an order step to the query.
 func (cq *CommentQuery) Order(o ...Order) *CommentQuery {
 	cq.order = append(cq.order, o...)
@@ -330,6 +338,11 @@ func (cq *CommentQuery) sqlQuery() *sql.Selector {
 	for _, p := range cq.order {
 		p.SQL(selector)
 	}
+	if offset := cq.offset; offset != nil {
+		// limit is mandatory for offset clause. We start
+		// with default value, and override it below if needed.
+		selector.Offset(*offset).Limit(math.MaxInt64)
+	}
 	if limit := cq.limit; limit != nil {
 		selector.Limit(*limit)
 	}
@@ -399,7 +412,12 @@ func (cq *CommentQuery) gremlinQuery() *dsl.Traversal {
 			p.Gremlin(v)
 		}
 	}
-	if limit := cq.limit; limit != nil {
+	switch limit, offset := cq.limit, cq.offset; {
+	case limit != nil && offset != nil:
+		v.Range(*offset, *offset+*limit)
+	case offset != nil:
+		v.Range(*offset, math.MaxInt64)
+	case limit != nil:
 		v.Limit(*limit)
 	}
 	if unique := cq.unique; len(unique) == 0 {

entc/integration/ent/file_query.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 
 	"fbc/ent/entc/integration/ent/file"
 
@@ -23,6 +24,7 @@ import (
 type FileQuery struct {
 	config
 	limit      *int
+	offset     *int
 	order      []Order
 	unique     []string
 	predicates []ent.Predicate
@@ -43,6 +45,12 @@ func (fq *FileQuery) Limit(limit int) *FileQuery {
 	return fq
 }
 
+// Offset adds an offset step to the query.
+func (fq *FileQuery) Offset(offset int) *FileQuery {
+	fq.offset = &offset
+	return fq
+}
+
 // Order adds an order step to the query.
 func (fq *FileQuery) Order(o ...Order) *FileQuery {
 	fq.order = append(fq.order, o...)
@@ -343,6 +351,11 @@ func (fq *FileQuery) sqlQuery() *sql.Selector {
 	for _, p := range fq.order {
 		p.SQL(selector)
 	}
+	if offset := fq.offset; offset != nil {
+		// limit is mandatory for offset clause. We start
+		// with default value, and override it below if needed.
+		selector.Offset(*offset).Limit(math.MaxInt64)
+	}
 	if limit := fq.limit; limit != nil {
 		selector.Limit(*limit)
 	}
@@ -412,7 +425,12 @@ func (fq *FileQuery) gremlinQuery() *dsl.Traversal {
 			p.Gremlin(v)
 		}
 	}
-	if limit := fq.limit; limit != nil {
+	switch limit, offset := fq.limit, fq.offset; {
+	case limit != nil && offset != nil:
+		v.Range(*offset, *offset+*limit)
+	case offset != nil:
+		v.Range(*offset, math.MaxInt64)
+	case limit != nil:
 		v.Limit(*limit)
 	}
 	if unique := fq.unique; len(unique) == 0 {

entc/integration/ent/group_query.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 
 	"fbc/ent/entc/integration/ent/file"
 	"fbc/ent/entc/integration/ent/group"
@@ -26,6 +27,7 @@ import (
 type GroupQuery struct {
 	config
 	limit      *int
+	offset     *int
 	order      []Order
 	unique     []string
 	predicates []ent.Predicate
@@ -46,6 +48,12 @@ func (gq *GroupQuery) Limit(limit int) *GroupQuery {
 	return gq
 }
 
+// Offset adds an offset step to the query.
+func (gq *GroupQuery) Offset(offset int) *GroupQuery {
+	gq.offset = &offset
+	return gq
+}
+
 // Order adds an order step to the query.
 func (gq *GroupQuery) Order(o ...Order) *GroupQuery {
 	gq.order = append(gq.order, o...)
@@ -427,6 +435,11 @@ func (gq *GroupQuery) sqlQuery() *sql.Selector {
 	for _, p := range gq.order {
 		p.SQL(selector)
 	}
+	if offset := gq.offset; offset != nil {
+		// limit is mandatory for offset clause. We start
+		// with default value, and override it below if needed.
+		selector.Offset(*offset).Limit(math.MaxInt64)
+	}
 	if limit := gq.limit; limit != nil {
 		selector.Limit(*limit)
 	}
@@ -496,7 +509,12 @@ func (gq *GroupQuery) gremlinQuery() *dsl.Traversal {
 			p.Gremlin(v)
 		}
 	}
-	if limit := gq.limit; limit != nil {
+	switch limit, offset := gq.limit, gq.offset; {
+	case limit != nil && offset != nil:
+		v.Range(*offset, *offset+*limit)
+	case offset != nil:
+		v.Range(*offset, math.MaxInt64)
+	case limit != nil:
 		v.Limit(*limit)
 	}
 	if unique := gq.unique; len(unique) == 0 {

entc/integration/ent/groupinfo_query.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 
 	"fbc/ent/entc/integration/ent/group"
 	"fbc/ent/entc/integration/ent/groupinfo"
@@ -24,6 +25,7 @@ import (
 type GroupInfoQuery struct {
 	config
 	limit      *int
+	offset     *int
 	order      []Order
 	unique     []string
 	predicates []ent.Predicate
@@ -44,6 +46,12 @@ func (giq *GroupInfoQuery) Limit(limit int) *GroupInfoQuery {
 	return giq
 }
 
+// Offset adds an offset step to the query.
+func (giq *GroupInfoQuery) Offset(offset int) *GroupInfoQuery {
+	giq.offset = &offset
+	return giq
+}
+
 // Order adds an order step to the query.
 func (giq *GroupInfoQuery) Order(o ...Order) *GroupInfoQuery {
 	giq.order = append(giq.order, o...)
@@ -363,6 +371,11 @@ func (giq *GroupInfoQuery) sqlQuery() *sql.Selector {
 	for _, p := range giq.order {
 		p.SQL(selector)
 	}
+	if offset := giq.offset; offset != nil {
+		// limit is mandatory for offset clause. We start
+		// with default value, and override it below if needed.
+		selector.Offset(*offset).Limit(math.MaxInt64)
+	}
 	if limit := giq.limit; limit != nil {
 		selector.Limit(*limit)
 	}
@@ -432,7 +445,12 @@ func (giq *GroupInfoQuery) gremlinQuery() *dsl.Traversal {
 			p.Gremlin(v)
 		}
 	}
-	if limit := giq.limit; limit != nil {
+	switch limit, offset := giq.limit, giq.offset; {
+	case limit != nil && offset != nil:
+		v.Range(*offset, *offset+*limit)
+	case offset != nil:
+		v.Range(*offset, math.MaxInt64)
+	case limit != nil:
 		v.Limit(*limit)
 	}
 	if unique := giq.unique; len(unique) == 0 {

entc/integration/ent/node_query.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 
 	"fbc/ent/entc/integration/ent/node"
 
@@ -23,6 +24,7 @@ import (
 type NodeQuery struct {
 	config
 	limit      *int
+	offset     *int
 	order      []Order
 	unique     []string
 	predicates []ent.Predicate
@@ -43,6 +45,12 @@ func (nq *NodeQuery) Limit(limit int) *NodeQuery {
 	return nq
 }
 
+// Offset adds an offset step to the query.
+func (nq *NodeQuery) Offset(offset int) *NodeQuery {
+	nq.offset = &offset
+	return nq
+}
+
 // Order adds an order step to the query.
 func (nq *NodeQuery) Order(o ...Order) *NodeQuery {
 	nq.order = append(nq.order, o...)
@@ -381,6 +389,11 @@ func (nq *NodeQuery) sqlQuery() *sql.Selector {
 	for _, p := range nq.order {
 		p.SQL(selector)
 	}
+	if offset := nq.offset; offset != nil {
+		// limit is mandatory for offset clause. We start
+		// with default value, and override it below if needed.
+		selector.Offset(*offset).Limit(math.MaxInt64)
+	}
 	if limit := nq.limit; limit != nil {
 		selector.Limit(*limit)
 	}
@@ -450,7 +463,12 @@ func (nq *NodeQuery) gremlinQuery() *dsl.Traversal {
 			p.Gremlin(v)
 		}
 	}
-	if limit := nq.limit; limit != nil {
+	switch limit, offset := nq.limit, nq.offset; {
+	case limit != nil && offset != nil:
+		v.Range(*offset, *offset+*limit)
+	case offset != nil:
+		v.Range(*offset, math.MaxInt64)
+	case limit != nil:
 		v.Limit(*limit)
 	}
 	if unique := nq.unique; len(unique) == 0 {

entc/integration/ent/pet_query.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 
 	"fbc/ent/entc/integration/ent/pet"
 	"fbc/ent/entc/integration/ent/user"
@@ -24,6 +25,7 @@ import (
 type PetQuery struct {
 	config
 	limit      *int
+	offset     *int
 	order      []Order
 	unique     []string
 	predicates []ent.Predicate
@@ -44,6 +46,12 @@ func (pq *PetQuery) Limit(limit int) *PetQuery {
 	return pq
 }
 
+// Offset adds an offset step to the query.
+func (pq *PetQuery) Offset(offset int) *PetQuery {
+	pq.offset = &offset
+	return pq
+}
+
 // Order adds an order step to the query.
 func (pq *PetQuery) Order(o ...Order) *PetQuery {
 	pq.order = append(pq.order, o...)
@@ -382,6 +390,11 @@ func (pq *PetQuery) sqlQuery() *sql.Selector {
 	for _, p := range pq.order {
 		p.SQL(selector)
 	}
+	if offset := pq.offset; offset != nil {
+		// limit is mandatory for offset clause. We start
+		// with default value, and override it below if needed.
+		selector.Offset(*offset).Limit(math.MaxInt64)
+	}
 	if limit := pq.limit; limit != nil {
 		selector.Limit(*limit)
 	}
@@ -451,7 +464,12 @@ func (pq *PetQuery) gremlinQuery() *dsl.Traversal {
 			p.Gremlin(v)
 		}
 	}
-	if limit := pq.limit; limit != nil {
+	switch limit, offset := pq.limit, pq.offset; {
+	case limit != nil && offset != nil:
+		v.Range(*offset, *offset+*limit)
+	case offset != nil:
+		v.Range(*offset, math.MaxInt64)
+	case limit != nil:
 		v.Limit(*limit)
 	}
 	if unique := pq.unique; len(unique) == 0 {

entc/integration/ent/user_query.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"math"
 
 	"fbc/ent/entc/integration/ent/card"
 	"fbc/ent/entc/integration/ent/file"
@@ -27,6 +28,7 @@ import (
 type UserQuery struct {
 	config
 	limit      *int
+	offset     *int
 	order      []Order
 	unique     []string
 	predicates []ent.Predicate
@@ -47,6 +49,12 @@ func (uq *UserQuery) Limit(limit int) *UserQuery {
 	return uq
 }
 
+// Offset adds an offset step to the query.
+func (uq *UserQuery) Offset(offset int) *UserQuery {
+	uq.offset = &offset
+	return uq
+}
+
 // Order adds an order step to the query.
 func (uq *UserQuery) Order(o ...Order) *UserQuery {
 	uq.order = append(uq.order, o...)
@@ -576,6 +584,11 @@ func (uq *UserQuery) sqlQuery() *sql.Selector {
 	for _, p := range uq.order {
 		p.SQL(selector)
 	}
+	if offset := uq.offset; offset != nil {
+		// limit is mandatory for offset clause. We start
+		// with default value, and override it below if needed.
+		selector.Offset(*offset).Limit(math.MaxInt64)
+	}
 	if limit := uq.limit; limit != nil {
 		selector.Limit(*limit)
 	}
@@ -645,7 +658,12 @@ func (uq *UserQuery) gremlinQuery() *dsl.Traversal {
 			p.Gremlin(v)
 		}
 	}
-	if limit := uq.limit; limit != nil {
+	switch limit, offset := uq.limit, uq.offset; {
+	case limit != nil && offset != nil:
+		v.Range(*offset, *offset+*limit)
+	case offset != nil:
+		v.Range(*offset, math.MaxInt64)
+	case limit != nil:
 		v.Limit(*limit)
 	}
 	if unique := uq.unique; len(unique) == 0 {

entc/integration/integration_test.go

@@ -96,6 +96,7 @@ func TestGremlin(t *testing.T) {
 var tests = []func(*testing.T, *ent.Client){
 	Tx,
 	Sanity,
+	Paging,
 	Relation,
 	UniqueConstraint,
 	O2OTwoTypes,
@@ -194,6 +195,30 @@ func Sanity(t *testing.T, client *ent.Client) {
 	client.User.Delete().Where(user.IDIn(ids...)).ExecX(ctx)
 }
 
+func Paging(t *testing.T, client *ent.Client) {
+	require := require.New(t)
+	ctx := context.Background()
+	for i := 1; i <= 10; i++ {
+		client.User.Create().SetName(fmt.Sprintf("name-%d", i)).SetAge(i).SaveX(ctx)
+	}
+
+	require.Equal(10, client.User.Query().CountX(ctx))
+	require.Len(client.User.Query().Offset(5).AllX(ctx), 5)
+	require.Len(client.User.Query().Offset(6).AllX(ctx), 4)
+	require.Equal(
+		[]int{7, 8},
+		client.User.Query().
+			Offset(6).
+			Limit(2).
+			Order(ent.Asc(user.FieldAge)).
+			GroupBy(user.FieldAge).
+			IntsX(ctx),
+	)
+	for i := 0; i < 10; i++ {
+		require.Equal(i+1, client.User.Query().Order(ent.Asc(user.FieldAge)).Offset(i).Limit(1).AllX(ctx)[0].Age)
+	}
+}
+
 func Relation(t *testing.T, client *ent.Client) {
 	require := require.New(t)
 	ctx := context.Background()