Implement retryable flag for databricks error

Added IsRetryable() and RetryableAfter() methods to DatabricksError.
Updated client to insert an internal retryable error instance into the error chain when failing with a 503/429.  DatabricksError uses the presence of this error to set the retryable flag and retryAfter value.
Added test to driver_e2e_test.go to check that DatabricksError has the correct values for IsRetryable() and RetryAfter()

Signed-off-by: Raymond Cypher <[email protected]>

Author: rcypher-databricks

driver_e2e_test.go

@@ -421,6 +421,74 @@ func TestRetries(t *testing.T) {
 		require.ErrorContains(t, err, "after 1 attempt(s)")
 	})
 
+	t.Run("a 429 or 503 should result in a retryable error", func(t *testing.T) {
+
+		_ = logger.SetLogLevel("debug")
+		state := &callState{}
+		// load basic responses
+		loadTestData(t, "OpenSessionSuccess.json", &state.openSessionResp)
+		loadTestData(t, "CloseSessionSuccess.json", &state.closeSessionResp)
+		loadTestData(t, "CloseOperationSuccess.json", &state.closeOperationResp)
+
+		ts := getServer(state)
+
+		defer ts.Close()
+		r, err := url.Parse(ts.URL)
+		require.NoError(t, err)
+		port, err := strconv.Atoi(r.Port())
+		require.NoError(t, err)
+
+		connector, err := NewConnector(
+			WithServerHostname("localhost"),
+			WithHTTPPath("/429-5-retries"),
+			WithPort(port),
+			WithRetries(2, 10*time.Millisecond, 1*time.Second),
+		)
+		require.NoError(t, err)
+		db := sql.OpenDB(connector)
+		defer db.Close()
+
+		state.executeStatementResp = cli_service.TExecuteStatementResp{}
+		loadTestData(t, "ExecuteStatement1.json", &state.executeStatementResp)
+
+		err = db.Ping()
+		require.ErrorContains(t, err, "after 3 attempt(s)")
+
+		// The error chain should contain a databricks request error
+		b := errors.Is(err, dbsqlerr.RequestError)
+		require.True(t, b)
+		var re dbsqlerr.DBRequestError
+		b = errors.As(err, &re)
+		require.True(t, b)
+		require.NotNil(t, re)
+		require.True(t, re.IsRetryable())
+		require.Equal(t, "retry after header value", re.RetryAfter())
+
+		connector2, err := NewConnector(
+			WithServerHostname("localhost"),
+			WithHTTPPath("/503-5-retries"),
+			WithPort(port),
+			WithRetries(2, 10*time.Millisecond, 1*time.Second),
+		)
+		require.NoError(t, err)
+		db2 := sql.OpenDB(connector2)
+		defer db.Close()
+
+		state.executeStatementResp = cli_service.TExecuteStatementResp{}
+		loadTestData(t, "ExecuteStatement1.json", &state.executeStatementResp)
+
+		err = db2.Ping()
+		require.ErrorContains(t, err, "after 3 attempt(s)")
+
+		// The error chain should contain a databricks request error
+		b = errors.Is(err, dbsqlerr.RequestError)
+		require.True(t, b)
+		b = errors.As(err, &re)
+		require.True(t, b)
+		require.NotNil(t, re)
+		require.True(t, re.IsRetryable())
+	})
+
 }
 
 // TODO: add tests for x-databricks headers

errors/errors.go

@@ -59,6 +59,10 @@ type DBError interface {
 
 	// Underlying causative error. May be nil.
 	Cause() error
+
+	IsRetryable() bool
+
+	RetryAfter() string
 }
 
 // An error that is caused by an invalid request.
@@ -70,8 +74,6 @@ type DBRequestError interface {
 // A fault that is caused by Databricks services
 type DBDriverError interface {
 	DBError
-
-	IsRetryable() bool
 }
 
 // Any error that occurs after the SQL statement has been accepted (e.g. SQL syntax error).

internal/client/client.go

@@ -279,6 +279,15 @@ func SprintGuid(bts []byte) string {
 
 var retryableStatusCode = []int{http.StatusTooManyRequests, http.StatusServiceUnavailable}
 
+func isRetryable(statusCode int) bool {
+	for _, c := range retryableStatusCode {
+		if c == statusCode {
+			return true
+		}
+	}
+	return false
+}
+
 type Transport struct {
 	Base  *http.Transport
 	Authr auth.Authenticator
@@ -321,14 +330,13 @@ func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
 	if resp.StatusCode != http.StatusOK {
 		reason := resp.Header.Get("X-Databricks-Reason-Phrase")
 		terrmsg := resp.Header.Get("X-Thriftserver-Error-Message")
-		for _, c := range retryableStatusCode {
-			if c == resp.StatusCode {
-				if terrmsg != "" {
-					logger.Warn().Msg(terrmsg)
-				}
-				return resp, nil
+		if isRetryable(resp.StatusCode) {
+			if terrmsg != "" {
+				logger.Warn().Msg(terrmsg)
 			}
+			return resp, nil
 		}
+
 		if reason != "" {
 			logger.Err(fmt.Errorf(reason)).Msg("non retryable error")
 			return nil, errors.New(reason)
@@ -426,17 +434,25 @@ func errorHandler(resp *http.Response, err error, numTries int) (*http.Response,
 	if err == nil {
 		err = errors.New(fmt.Sprintf("request error after %d attempt(s)", numTries))
 	}
-	if resp != nil && resp.Header != nil {
 
+	if resp != nil {
+		var orgid, reason, terrmsg, errmsg, retryAfter 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)
 
-		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
+		if isRetryable(resp.StatusCode) {
+			err = dbsqlerrint.NewRetryableError(err, retryAfter)
+		}
 
-		werr = errors.Wrapf(err, fmt.Sprintf("orgId: %s, reason: %s, thriftErr: %s, err: %s", orgid, reason, terrmsg, errmsg))
+		werr = dbsqlerrint.WrapErr(err, msg)
 	} else {
 		werr = err
 	}
@@ -464,11 +480,8 @@ 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.
-
-	for _, c := range retryableStatusCode {
-		if c == resp.StatusCode {
-			return true, nil
-		}
+	if isRetryable(resp.StatusCode) {
+		return true, nil
 	}
 
 	return false, nil

internal/errors/err.go

@@ -10,11 +10,18 @@ import (
 	"github.com/pkg/errors"
 )
 
+// value to use with errors.Is() to determine if an error
+// chain contains a retryable error
+var RetryableError error = errors.New("Retryable Error")
+
+// base databricks error
 type databricksError struct {
 	err           error
 	correlationId string
 	connectionId  string
 	errType       string
+	isRetryable   bool
+	retryAfter    string
 }
 
 var _ error = (*databricksError)(nil)
@@ -38,11 +45,25 @@ func newDatabricksError(ctx context.Context, msg string, err error) databricksEr
 		err = errors.WithStack(err)
 	}
 
+	// If the error chain contains an instance of retryableError
+	// set the flag and retryAfter value.
+	var retryable bool = false
+	var retryAfter string
+	if errors.Is(err, RetryableError) {
+		retryable = true
+		var re retryableError
+		if ok := errors.As(err, &re); ok {
+			retryAfter = re.retryAfter
+		}
+	}
+
 	return databricksError{
 		err:           err,
 		correlationId: driverctx.CorrelationIdFromContext(ctx),
 		connectionId:  driverctx.ConnIdFromContext(ctx),
 		errType:       "unknown",
+		isRetryable:   retryable,
+		retryAfter:    retryAfter,
 	}
 }
 
@@ -75,10 +96,17 @@ func (e databricksError) Is(err error) bool {
 	return err == dbsqlerr.DatabricksError
 }
 
+func (e databricksError) IsRetryable() bool {
+	return e.isRetryable
+}
+
+func (e databricksError) RetryAfter() string {
+	return e.retryAfter
+}
+
 // driverError are issues with the driver or server, e.g. not supported operations, driver specific non-recoverable failures
 type driverError struct {
 	databricksError
-	isRetryable bool
 }
 
 var _ dbsqlerr.DBDriverError = (*driverError)(nil)
@@ -91,14 +119,10 @@ func (e driverError) Unwrap() error {
 	return e.err
 }
 
-func (e driverError) IsRetryable() bool {
-	return e.isRetryable
-}
-
 func NewDriverError(ctx context.Context, msg string, err error) *driverError {
 	dbErr := newDatabricksError(ctx, msg, err)
 	dbErr.errType = "driver error"
-	return &driverError{databricksError: dbErr, isRetryable: false}
+	return &driverError{databricksError: dbErr}
 }
 
 // requestError are errors caused by invalid requests, e.g. permission denied, warehouse not found
@@ -181,3 +205,37 @@ func WrapErrf(err error, format string, args ...interface{}) error {
 	// wrap passed in error in errors with the formatted message and a stack trace
 	return errors.Wrapf(err, format, args...)
 }
+
+type retryableError struct {
+	err        error
+	retryAfter string
+}
+
+func (e retryableError) Is(err error) bool {
+	return err == RetryableError
+}
+
+func (e retryableError) Unwrap() error {
+	return e.err
+}
+
+func (e retryableError) Error() string {
+	return fmt.Sprintf("databricks: retryableError: %s", e.err.Error())
+}
+
+func (e retryableError) RetryAfter() string {
+	return e.retryAfter
+}
+
+func NewRetryableError(err error, retryAfter string) error {
+	if err == nil {
+		err = errors.New("")
+	}
+
+	var st stackTracer
+	if ok := errors.As(err, &st); !ok {
+		err = errors.WithStack(err)
+	}
+
+	return retryableError{err: err, retryAfter: retryAfter}
+}

internal/errors/err_test.go

@@ -95,10 +95,13 @@ func TestDbSqlErrors(t *testing.T) {
 		st := requestError.StackTrace()
 		assert.NotNil(t, st)
 
+		// Get the underlying stackTracer instance, it should be
+		// the original cause
 		var str stackTracer
 		ok := errors.As(requestError.Cause(), &str)
 		assert.True(t, ok)
-		assert.NotEqual(t, requestError, str)
+		ss := str.StackTrace()
+		assert.NotNil(t, ss)
 		assert.Equal(t, cause, str)
 
 		cause = &boringError{}
@@ -107,9 +110,12 @@ func TestDbSqlErrors(t *testing.T) {
 		st = requestError.StackTrace()
 		assert.NotNil(t, st)
 
+		// Get the underlying stackTracer instance, it should not be
+		// the original cause
 		ok = errors.As(requestError.Cause(), &str)
 		assert.True(t, ok)
-		assert.NotEqual(t, requestError, str)
+		ss = str.StackTrace()
+		assert.NotNil(t, ss)
 		assert.NotEqual(t, cause, str)
 	})
 

testserver.go

@@ -44,6 +44,15 @@ func (h *thriftHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		} else {
 			h.count503_5_retries = 0
 		}
+	case "/429-5-retries":
+		if h.count503_5_retries <= 5 {
+			w.Header().Set("Retry-After", "retry after header value")
+			w.WriteHeader(http.StatusServiceUnavailable)
+			h.count503_5_retries++
+			return
+		} else {
+			h.count503_5_retries = 0
+		}
 	}
 
 	thriftHandler := thrift.NewThriftHandlerFunc(h.processor, h.inPfactory, h.outPfactory)