Give Rows a context for fetchNextPage

Author: njvrzm

pkg/redshift/driver/db.go

@@ -65,8 +65,8 @@ func (d *db) CancelQuery(_ context.Context, queryID string) error {
 	return d.api.Stop(&sqlAPI.ExecuteQueryOutput{ID: queryID})
 }
 
-func (d *db) GetRows(_ context.Context, queryID string) (driver.Rows, error) {
-	return newRows(d.api.DataClient, queryID)
+func (d *db) GetRows(ctx context.Context, queryID string) (driver.Rows, error) {
+	return newRows(ctx, d.api.DataClient, queryID)
 }
 
 func (d *db) Ping(ctx context.Context) error {

pkg/redshift/driver/rows.go

@@ -19,15 +19,17 @@ import (
 type Rows struct {
 	service redshiftdata.GetStatementResultAPIClient
 	queryID string
+	context context.Context
 
 	done   bool
 	result *redshiftdata.GetStatementResultOutput
 }
 
-func newRows(service redshiftdata.GetStatementResultAPIClient, queryId string) (*Rows, error) {
+func newRows(ctx context.Context, service redshiftdata.GetStatementResultAPIClient, queryId string) (*Rows, error) {
 	r := Rows{
 		service: service,
 		queryID: queryId,
+		context: ctx,
 	}
 
 	if err := r.fetchNextPage(nil); err != nil {
@@ -178,7 +180,7 @@ func (r *Rows) Close() error {
 func (r *Rows) fetchNextPage(token *string) error {
 	var err error
 
-	r.result, err = r.service.GetStatementResult(context.TODO(), &redshiftdata.GetStatementResultInput{
+	r.result, err = r.service.GetStatementResult(r.context, &redshiftdata.GetStatementResultInput{
 		Id:        aws.String(r.queryID),
 		NextToken: token,
 	})

pkg/redshift/driver/rows_test.go

@@ -1,6 +1,7 @@
 package driver
 
 import (
+	"context"
 	"database/sql/driver"
 	"fmt"
 	"io"
@@ -17,7 +18,7 @@ import (
 func TestOnePageSuccess(t *testing.T) {
 	redshiftServiceMock := &mock.RedshiftService{}
 	redshiftServiceMock.CalledTimesCountDown = 1
-	rows, rowErr := newRows(redshiftServiceMock, mock.SinglePageResponseQueryId)
+	rows, rowErr := newRows(context.Background(), redshiftServiceMock, mock.SinglePageResponseQueryId)
 	require.NoError(t, rowErr)
 	cnt := 0
 	for {
@@ -39,7 +40,7 @@ func TestOnePageSuccess(t *testing.T) {
 func TestMultiPageSuccess(t *testing.T) {
 	redshiftServiceMock := &mock.RedshiftService{}
 	redshiftServiceMock.CalledTimesCountDown = 5
-	rows, rowErr := newRows(redshiftServiceMock, mock.MultiPageResponseQueryId)
+	rows, rowErr := newRows(context.Background(), redshiftServiceMock, mock.MultiPageResponseQueryId)
 	require.NoError(t, rowErr)
 	cnt := 0
 	for {