feat: add pagination support when retrieving partition values

mgjules · Oct 6, 2021 · 9cd1568 · 9cd1568
1 parent c59357d
commit 9cd1568
Showing 1 changed file with 25 additions and 13 deletions.
diff --git a/cmd/backup.go b/cmd/backup.go
@@ -179,21 +179,33 @@ var Backup = &cli.Command{
 
 		sugar.Debugw("retrieving partition values...", "sql", sql)
 
-		partitionOutput, err := querySvc.QueryWithContext(c.Context, &timestreamquery.QueryInput{
-			QueryString: aws.String(sql),
-			MaxRows:     aws.Int64(100),
-		})
-		if err != nil {
-			sugar.Errorw("retrieve partition values", "error", err, "sql", sql)
-			return fmt.Errorf("retrieve partition values: %v", err)
-		}
+		var (
+			nextToken       *string = nil
+			partitionValues []string
+		)
+
+		for {
+			partitionOutput, err := querySvc.QueryWithContext(c.Context, &timestreamquery.QueryInput{
+				QueryString: aws.String(sql),
+				NextToken:   nextToken,
+				MaxRows:     aws.Int64(100),
+			})
+			if err != nil {
+				sugar.Errorw("retrieve partition values", "error", err, "sql", sql)
+				return fmt.Errorf("retrieve partition values: %v", err)
+			}
+
+			for _, row := range partitionOutput.Rows {
+				if row.Data[0].ScalarValue == nil {
+					continue
+				}
+				partitionValues = append(partitionValues, *row.Data[0].ScalarValue)
+			}
 
-		var partitionValues []string
-		for _, row := range partitionOutput.Rows {
-			if row.Data[0].ScalarValue == nil {
-				continue
+			nextToken = partitionOutput.NextToken
+			if nextToken == nil {
+				break
 			}
-			partitionValues = append(partitionValues, *row.Data[0].ScalarValue)
 		}
 
 		sugar.Debugw("partition", "column", partitionColumn, "values", partitionValues)