fix batch get bailing early on empty results

fixes #161

batch_test.go

@@ -90,3 +90,47 @@ func TestBatchGetWrite(t *testing.T) {
 		t.Error("expected 0 results, got", len(results))
 	}
 }
+
+func TestBatchGetEmptySets(t *testing.T) {
+	table := testDB.Table(testTable)
+	now := time.Now().UnixNano() / 1000000000
+	id := int(now)
+	entry := widget{UserID: id, Time: time.Now()}
+	if err := table.Put(entry).Run(); err != nil {
+		panic(err)
+	}
+	entry2 := widget{UserID: id + batchSize*2, Time: entry.Time}
+	if err := table.Put(entry2).Run(); err != nil {
+		panic(err)
+	}
+
+	keysToCheck := []Keyed{}
+	for i := entry.UserID; i <= entry2.UserID; i += 1 {
+		keysToCheck = append(keysToCheck, Keys{i, entry.Time})
+	}
+
+	results := []widget{}
+	err := table.Batch("UserID", "Time").Get(keysToCheck...).Consistent(true).All(&results)
+	if err != nil {
+		t.Error(err)
+	}
+	if len(results) != 2 {
+		t.Error("batch get empty set, unexpected length:", len(results), "want:", 2)
+	}
+
+	if err := table.Delete("UserID", entry.UserID).Range("Time", entry.Time).Run(); err != nil {
+		panic(err)
+	}
+
+	results = []widget{}
+	err = table.Batch("UserID", "Time").Get(keysToCheck...).Consistent(true).All(&results)
+	if len(results) != 1 {
+		t.Error("batch get empty set, unexpected length:", len(results), "want:", 1)
+	}
+
+	results = []widget{}
+	err = table.Batch("UserID", "Time").Get(keysToCheck[:len(keysToCheck)-1]...).Consistent(true).All(&results)
+	if len(results) != 0 {
+		t.Error("batch get empty set, unexpected length:", len(results), "want:", 0)
+	}
+}

batchget.go

@@ -205,6 +205,7 @@ func (itr *bgIter) NextWithContext(ctx aws.Context, out interface{}) bool {
 
 	tableName := itr.bg.batch.table.Name()
 
+redo:
 	// can we use results we already have?
 	if itr.output != nil && itr.idx < len(itr.output.Responses[tableName]) {
 		items := itr.output.Responses[tableName]
@@ -263,21 +264,8 @@ func (itr *bgIter) NextWithContext(ctx aws.Context, out interface{}) bool {
 		}
 	}
 
-	items := itr.output.Responses[tableName]
-	if len(items) == 0 {
-		if len(itr.output.UnprocessedKeys) == 0 {
-			if itr.total == 0 {
-				itr.err = ErrNotFound
-			}
-			return false
-		}
-		// need to retry to get more keys
-		return itr.NextWithContext(ctx, out)
-	}
-	itr.err = itr.unmarshal(items[itr.idx], out)
-	itr.idx++
-	itr.total++
-	return itr.err == nil
+	// we've got unprocessed results, marshal one
+	goto redo
 }
 
 // Err returns the error encountered, if any.