Simplified error message

Simplified error message  to use either X-Databricks-Reason-Phrase or X-Thriftserver-Error-Message header.
Always honour Retry-After header when retrying.  Not just on 429 or 503.
Signed-off-by: Raymond Cypher <raymond.cypher@databricks.com>

Author: rcypher-databricks

internal/client/client.go

@@ -56,16 +56,7 @@ const (
 	cancelOperation
 )
 
-var idempotentClientMethods map[clientMethod]any = map[clientMethod]any{
-	closeSession:         struct{}{},
-	getResultSetMetadata: struct{}{},
-	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{}{},
-}
+var nonRetryableClientMethods map[clientMethod]any = map[clientMethod]any{executeStatement: struct{}{}}
 
 // OpenSession is a wrapper around the thrift operation OpenSession
 // If RecordResults is true, the results will be marshalled to JSON format and written to OpenSession<index>.json
@@ -449,27 +440,24 @@ func errorHandler(resp *http.Response, err error, numTries int) (*http.Response,
 	var werr error
 	msg := fmt.Sprintf("request error after %d attempt(s)", numTries)
 	if err == nil {
-		err = errors.New(msg)
+		werr = errors.New(msg)
 	} else {
-		err = errors.Wrap(err, msg)
+		werr = errors.Wrap(err, msg)
 	}
 
 	if resp != nil {
-		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")
-			// TODO note: need to see if there's other headers
+			reason := resp.Header.Get("X-Databricks-Reason-Phrase")
+			terrmsg := resp.Header.Get("X-Thriftserver-Error-Message")
+
+			if reason != "" {
+				werr = dbsqlerrint.WrapErr(werr, reason)
+			} else if terrmsg != "" {
+				werr = dbsqlerrint.WrapErr(werr, terrmsg)
+			}
 		}
-		msg := fmt.Sprintf("orgId: %s, reason: %s, thriftErr: %s, err: %s", orgid, reason, terrmsg, errmsg)
 
-		werr = dbsqlerrint.WrapErr(err, msg)
 		logger.Err(werr).Msg(resp.Status)
-	} else {
-		werr = err
 	}
 
 	return resp, werr
@@ -526,7 +514,7 @@ func RetryPolicy(ctx context.Context, resp *http.Response, err error) (bool, err
 	// 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 {
+	if isRetryableServerResponse(resp) {
 		var retryAfter string
 		if resp.Header != nil {
 			retryAfter = resp.Header.Get("Retry-After")
@@ -538,7 +526,7 @@ func RetryPolicy(ctx context.Context, resp *http.Response, err error) (bool, err
 	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 {
+			if _, noRetry := nonRetryableClientMethods[caller]; !noRetry {
 				return true, checkErr
 			}
 		}
@@ -550,14 +538,16 @@ func RetryPolicy(ctx context.Context, resp *http.Response, err error) (bool, err
 }
 
 func backoff(min, max time.Duration, attemptNum int, resp *http.Response) time.Duration {
-	if resp != nil && isRetryableServerResponse(resp) {
+	// honour the Retry-After header
+	if resp != nil && resp.Header != nil {
 		if s, ok := resp.Header["Retry-After"]; ok {
 			if sleep, err := strconv.ParseInt(s[0], 10, 64); err == nil {
 				return time.Second * time.Duration(sleep)
 			}
 		}
 	}
 
+	// exponential backoff
 	mult := math.Pow(2, float64(attemptNum)) * float64(min)
 	sleep := time.Duration(mult)
 	if float64(sleep) != mult || sleep > max {

internal/client/client_test.go

@@ -84,14 +84,23 @@ func TestRetryPolicy(t *testing.T) {
 
 		nonRetryableCodes := []int{200, 300, 400, 501}
 
-		idempotentOps := []clientMethod{closeSession, getResultSetMetadata, getOperationStatus, closeOperation, cancelOperation, fetchResults}
-		nonIdempotentOps := []clientMethod{openSession, executeStatement}
+		retryableOps := []clientMethod{
+			closeSession,
+			getResultSetMetadata,
+			getOperationStatus,
+			closeOperation,
+			cancelOperation,
+			fetchResults,
+			openSession,
+		}
+
+		nonRetryableOps := []clientMethod{executeStatement}
 
 		cancelled, cancel := context.WithCancel(context.Background())
 		cancel()
 
 		for _, code := range retryableCodes {
-			for _, op := range idempotentOps {
+			for _, op := range retryableOps {
 				resp.StatusCode = code
 				ctx := context.WithValue(context.Background(), ClientMethod, op)
 				retry, _ := RetryPolicy(ctx, resp, nil)
@@ -102,7 +111,7 @@ func TestRetryPolicy(t *testing.T) {
 				require.False(t, retry)
 			}
 
-			for _, op := range nonIdempotentOps {
+			for _, op := range nonRetryableOps {
 				resp.StatusCode = code
 				ctx := context.WithValue(context.Background(), ClientMethod, op)
 				retry, _ := RetryPolicy(ctx, resp, nil)
@@ -115,7 +124,7 @@ func TestRetryPolicy(t *testing.T) {
 		}
 
 		for _, code := range maybeRetryableCodes {
-			for _, op := range idempotentOps {
+			for _, op := range retryableOps {
 				resp.StatusCode = code
 				ctx := context.WithValue(context.Background(), ClientMethod, op)
 				retry, _ := RetryPolicy(ctx, resp, nil)
@@ -126,7 +135,7 @@ func TestRetryPolicy(t *testing.T) {
 				require.False(t, retry)
 			}
 
-			for _, op := range nonIdempotentOps {
+			for _, op := range nonRetryableOps {
 				resp.StatusCode = code
 				ctx := context.WithValue(context.Background(), ClientMethod, op)
 				retry, _ := RetryPolicy(ctx, resp, nil)
@@ -139,7 +148,7 @@ func TestRetryPolicy(t *testing.T) {
 		}
 
 		for _, code := range nonRetryableCodes {
-			for _, op := range idempotentOps {
+			for _, op := range retryableOps {
 				resp.StatusCode = code
 				ctx := context.WithValue(context.Background(), ClientMethod, op)
 				retry, _ := RetryPolicy(ctx, resp, nil)
@@ -150,7 +159,7 @@ func TestRetryPolicy(t *testing.T) {
 				require.False(t, retry)
 			}
 
-			for _, op := range nonIdempotentOps {
+			for _, op := range nonRetryableOps {
 				resp.StatusCode = code
 				ctx := context.WithValue(context.Background(), ClientMethod, op)
 				retry, _ := RetryPolicy(ctx, resp, nil)