Retry FetchResults and error handling

Allow retries on calls to fetch results.  The drivers handling of fetching result pages is robust enough to handle unexpected pagination due to retries.
Updated the error handling to make sure that the client error handler will be called so that information in databricks response headers will be added to the error info.
Signed-off-by: Raymond Cypher <raymond.cypher@databricks.com>

Author: rcypher-databricks

internal/client/client.go

@@ -62,6 +62,9 @@ var idempotentClientMethods map[clientMethod]any = map[clientMethod]any{
 	getOperationStatus:   struct{}{},
 	closeOperation:       struct{}{},
 	cancelOperation:      struct{}{},
+	// fetchResults is treated as idempotent because the drivers fetching logic is robust enoug
+	// to fetch results in correct order
+	fetchResults: struct{}{},
 }
 
 // OpenSession is a wrapper around the thrift operation OpenSession
@@ -360,23 +363,8 @@ func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
 	// req.Body is assumed to be closed by the base RoundTripper.
 	reqBodyClosed = true
 	resp, err := t.Base.RoundTrip(req2)
-	if err != nil {
-		return resp, err
-	}
-
-	if resp.StatusCode != http.StatusOK {
-		reason := resp.Header.Get("X-Databricks-Reason-Phrase")
-		terrmsg := resp.Header.Get("X-Thriftserver-Error-Message")
 
-		if reason != "" {
-			logger.Err(fmt.Errorf(reason)).Msg(resp.Status)
-		}
-		if terrmsg != "" {
-			logger.Err(fmt.Errorf(terrmsg)).Msg(resp.Status)
-		}
-	}
-
-	return resp, nil
+	return resp, err
 }
 
 func RetryableClient(cfg *config.Config) *http.Client {
@@ -459,28 +447,27 @@ func (l *leveledLogger) Warn(msg string, keysAndValues ...interface{}) {
 
 func errorHandler(resp *http.Response, err error, numTries int) (*http.Response, error) {
 	var werr error
+	msg := fmt.Sprintf("request error after %d attempt(s)", numTries)
 	if err == nil {
-		err = errors.New(fmt.Sprintf("request error after %d attempt(s)", numTries))
+		err = errors.New(msg)
+	} else {
+		err = errors.Wrap(err, msg)
 	}
 
 	if resp != nil {
-		var orgid, reason, terrmsg, errmsg, retryAfter string
+		var orgid, reason, terrmsg, errmsg string
 		// TODO @mattdeekay: convert these to specific error types
 		if resp.Header != nil {
 			orgid = resp.Header.Get("X-Databricks-Org-Id")
 			reason = resp.Header.Get("X-Databricks-Reason-Phrase") // TODO note: shown on notebook
 			terrmsg = resp.Header.Get("X-Thriftserver-Error-Message")
 			errmsg = resp.Header.Get("x-databricks-error-or-redirect-message")
-			retryAfter = resp.Header.Get("Retry-After")
 			// TODO note: need to see if there's other headers
 		}
 		msg := fmt.Sprintf("orgId: %s, reason: %s, thriftErr: %s, err: %s", orgid, reason, terrmsg, errmsg)
 
-		if isRetryableServerResponse(resp) {
-			err = dbsqlerrint.NewRetryableError(err, retryAfter)
-		}
-
 		werr = dbsqlerrint.WrapErr(err, msg)
+		logger.Err(werr).Msg(resp.Status)
 	} else {
 		werr = err
 	}
@@ -508,7 +495,6 @@ var (
 )
 
 func RetryPolicy(ctx context.Context, resp *http.Response, err error) (bool, error) {
-
 	// do not retry on context.Canceled or context.DeadlineExceeded
 	if ctx.Err() != nil {
 		return false, ctx.Err()
@@ -532,24 +518,35 @@ func RetryPolicy(ctx context.Context, resp *http.Response, err error) (bool, err
 		return true, nil
 	}
 
+	var checkErr error
+	if resp.StatusCode != http.StatusOK {
+		checkErr = fmt.Errorf("unexpected HTTP status %s", resp.Status)
+	}
+
 	// 429 Too Many Requests or 503 service unavailable is recoverable. Sometimes the server puts
 	// a Retry-After response header to indicate when the server is
 	// available to start processing request from client.
 	if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode == http.StatusServiceUnavailable {
-		return true, nil
+		var retryAfter string
+		if resp.Header != nil {
+			retryAfter = resp.Header.Get("Retry-After")
+		}
+
+		return true, dbsqlerrint.NewRetryableError(checkErr, retryAfter)
 	}
 
 	if resp.StatusCode == 0 || (resp.StatusCode >= 500 && resp.StatusCode != http.StatusNotImplemented) {
 		callerAny := ctx.Value(ClientMethod)
 		if caller, ok := callerAny.(clientMethod); ok {
 			if _, ok := idempotentClientMethods[caller]; ok {
-				return true, fmt.Errorf("unexpected HTTP status %s", resp.Status)
+				return true, checkErr
 			}
 		}
 	}
 
-	return false, nil
-
+	// checkErr will be non-nil if the response code was not StatusOK.
+	// Returning it here ensures that the error handler will be called.
+	return false, checkErr
 }
 
 func backoff(min, max time.Duration, attemptNum int, resp *http.Response) time.Duration {

internal/client/client_test.go

@@ -84,8 +84,8 @@ func TestRetryPolicy(t *testing.T) {
 
 		nonRetryableCodes := []int{200, 300, 400, 501}
 
-		idempotentOps := []clientMethod{closeSession, getResultSetMetadata, getOperationStatus, closeOperation, cancelOperation}
-		nonIdempotentOps := []clientMethod{openSession, fetchResults, executeStatement}
+		idempotentOps := []clientMethod{closeSession, getResultSetMetadata, getOperationStatus, closeOperation, cancelOperation, fetchResults}
+		nonIdempotentOps := []clientMethod{openSession, executeStatement}
 
 		cancelled, cancel := context.WithCancel(context.Background())
 		cancel()

internal/rows/rows_test.go

@@ -832,6 +832,40 @@ func TestRowsCloseOptimization(t *testing.T) {
 	assert.Equal(t, 0, closeCount)
 }
 
+type fetch struct {
+	direction      cli_service.TFetchOrientation
+	resultStartRec int
+}
+
+func TestFetchResultsWithRetries(t *testing.T) {
+	t.Parallel()
+
+	// Simulate a scenario where network issues and retries cause unexpected jumps
+	// across multiple result pages.
+	fetches := []fetch{}
+
+	client := getRowsTestSimpleClient2(&fetches)
+	rowSet := &rows{client: client, hasMoreRows: true}
+
+	// next row number is zero so should fetch first result page
+	err := rowSet.fetchResultPage()
+	assert.Nil(t, err)
+	assert.Len(t, fetches, 1)
+	assert.Equal(t, fetches[0].direction, cli_service.TFetchOrientation_FETCH_NEXT)
+
+	// row number four is still in the first page so there should be no calls to fetch
+	rowSet.nextRowNumber = 4
+	err = rowSet.fetchResultPage()
+	assert.Nil(t, err)
+	assert.Len(t, fetches, 1)
+
+	rowSet.nextRowNumber = 5
+	err = rowSet.fetchResultPage()
+	assert.Nil(t, err)
+	assert.Len(t, fetches, 5)
+
+}
+
 type rowTestPagingResult struct {
 	getMetadataCount  int
 	fetchResultsCount int
@@ -1245,3 +1279,102 @@ func getRowsTestSimpleClient(getMetadataCount, fetchResultsCount *int) cli_servi
 
 	return client
 }
+
+// Build a simple test client
+func getRowsTestSimpleClient2(fetches *[]fetch) cli_service.TCLIService {
+	// Metadata for the different types is based on the results returned when querying a table with
+	// all the different types which was created in a test shard.
+	metadata := &cli_service.TGetResultSetMetadataResp{
+		Status: &cli_service.TStatus{
+			StatusCode: cli_service.TStatusCode_SUCCESS_STATUS,
+		},
+		Schema: &cli_service.TTableSchema{
+			Columns: []*cli_service.TColumnDesc{
+				{
+					ColumnName: "bool_col",
+					TypeDesc: &cli_service.TTypeDesc{
+						Types: []*cli_service.TTypeEntry{
+							{
+								PrimitiveEntry: &cli_service.TPrimitiveTypeEntry{
+									Type: cli_service.TTypeId_BOOLEAN_TYPE,
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	getMetadata := func(ctx context.Context, req *cli_service.TGetResultSetMetadataReq) (_r *cli_service.TGetResultSetMetadataResp, _err error) {
+		return metadata, nil
+	}
+
+	moreRows := true
+	noMoreRows := false
+	colVals := []*cli_service.TColumn{{BoolVal: &cli_service.TBoolColumn{Values: []bool{true, false, true, false, true}}}}
+
+	pages := []*cli_service.TFetchResultsResp{
+		{
+			Status: &cli_service.TStatus{
+				StatusCode: cli_service.TStatusCode_SUCCESS_STATUS,
+			},
+			HasMoreRows: &moreRows,
+			Results: &cli_service.TRowSet{
+				StartRowOffset: 0,
+				Columns:        colVals,
+			},
+		},
+		{
+			Status: &cli_service.TStatus{
+				StatusCode: cli_service.TStatusCode_SUCCESS_STATUS,
+			},
+			HasMoreRows: &moreRows,
+			Results: &cli_service.TRowSet{
+				StartRowOffset: 5,
+				Columns:        colVals,
+			},
+		},
+		{
+			Status: &cli_service.TStatus{
+				StatusCode: cli_service.TStatusCode_SUCCESS_STATUS,
+			},
+			HasMoreRows: &noMoreRows,
+			Results: &cli_service.TRowSet{
+				StartRowOffset: 10,
+				Columns:        colVals,
+			},
+		},
+		{
+			Status: &cli_service.TStatus{
+				StatusCode: cli_service.TStatusCode_SUCCESS_STATUS,
+			},
+			HasMoreRows: &noMoreRows,
+			Results: &cli_service.TRowSet{
+				StartRowOffset: 15,
+				Columns:        []*cli_service.TColumn{},
+			},
+		},
+	}
+
+	// We are simulating the scenario where network errors and retry behaviour cause the fetch
+	// result request to be sent multiple times, resulting in jumping past the next/previous result
+	// page. Behaviour should be robust enough to handle this by changing the fetch orientation.
+	pageSequence := []int{0, 3, 2, 0, 1, 2}
+	pageIndex := -1
+
+	fetchResults := func(ctx context.Context, req *cli_service.TFetchResultsReq) (_r *cli_service.TFetchResultsResp, _err error) {
+		pageIndex++
+
+		p := pages[pageSequence[pageIndex]]
+		*fetches = append(*fetches, fetch{direction: req.Orientation, resultStartRec: int(p.Results.StartRowOffset)})
+		return p, nil
+	}
+
+	client := &client.TestClient{
+		FnGetResultSetMetadata: getMetadata,
+		FnFetchResults:         fetchResults,
+	}
+
+	return client
+}