Merge pull request #109 from erizocosmico/bugfix/batcher-limit

fix: set limit correctly in batcher queries

Author: Miguel Molina

batcher.go

@@ -57,9 +57,17 @@ func (r *batchQueryRunner) next() (Record, error) {
 	}
 
 	if len(r.records) == 0 {
-		records, err := r.loadNextBatch()
-		if err != nil {
-			return nil, err
+		var (
+			records []Record
+			err     error
+		)
+
+		limit := r.q.GetLimit()
+		if limit <= 0 || limit > uint64(r.total) {
+			records, err = r.loadNextBatch()
+			if err != nil {
+				return nil, err
+			}
 		}
 
 		if len(records) == 0 {

batcher_test.go

@@ -31,3 +31,38 @@ func TestOneToManyWithFilterNotWritable(t *testing.T) {
 	r.NoError(err)
 	r.False(record.IsWritable())
 }
+
+func TestBatcherLimit(t *testing.T) {
+	r := require.New(t)
+	db, err := openTestDB()
+	r.NoError(err)
+	setupTables(t, db)
+	defer db.Close()
+	defer teardownTables(t, db)
+
+	store := NewStore(db)
+	for i := 0; i < 10; i++ {
+		m := newModel("foo", "bar", 1)
+		r.NoError(store.Insert(ModelSchema, m))
+
+		for i := 0; i < 4; i++ {
+			r.NoError(store.Insert(RelSchema, newRel(m.GetID(), fmt.Sprint(i))))
+		}
+	}
+
+	q := NewBaseQuery(ModelSchema)
+	q.BatchSize(2)
+	q.Limit(5)
+	r.NoError(q.AddRelation(RelSchema, "rels", OneToMany, Eq(f("foo"), "1")))
+	runner := newBatchQueryRunner(ModelSchema, squirrel.NewStmtCacher(db), q)
+	rs := NewBatchingResultSet(runner)
+
+	var count int
+	for rs.Next() {
+		_, err := rs.Get(nil)
+		r.NoError(err)
+		count++
+	}
+	r.NoError(err)
+	r.Equal(5, count)
+}