Improve error message and retry behaviour (#111)

Be more restrictive on retries: only retry on 429 and 503 StatusCodes or
connection errors.

Return better error messages regarding retry reasons and invalid user
input.

Signed-off-by: Andre Furlan <[email protected]>

Author: andrefurlan-db

driver_e2e_test.go

@@ -266,7 +266,7 @@ func TestContextTimeoutExample(t *testing.T) {
 
 	defer ts.Close()
 
-	db, err := sql.Open("databricks", ts.URL)
+	db, err := sql.Open("databricks", ts.URL+"/path")
 	require.NoError(t, err)
 	defer db.Close()
 
@@ -363,7 +363,7 @@ func TestRetries(t *testing.T) {
 
 		connector, err := NewConnector(
 			WithServerHostname("localhost"),
-			WithHTTPPath("/500-5-retries"),
+			WithHTTPPath("/503-5-retries"),
 			WithPort(port),
 			WithRetries(2, 10*time.Millisecond, 1*time.Second),
 		)
@@ -397,7 +397,7 @@ func TestRetries(t *testing.T) {
 
 		connector, err := NewConnector(
 			WithServerHostname("localhost"),
-			WithHTTPPath("/500-5-retries"),
+			WithHTTPPath("/429-2-retries"),
 			WithPort(port),
 			WithRetries(-1, 0, 0),
 		)

internal/client/client.go

@@ -8,7 +8,9 @@ import (
 	"net"
 	"net/http"
 	"net/http/httptrace"
+	"net/url"
 	"os"
+	"regexp"
 	"time"
 
 	"github.com/apache/thrift/lib/go/thrift"
@@ -175,7 +177,11 @@ func (tsc *ThriftServiceClient) CancelOperation(ctx context.Context, req *cli_se
 // InitThriftClient is a wrapper of the http transport, so we can have access to response code and headers.
 // It is important to know the code and headers to know if we need to retry or not
 func InitThriftClient(cfg *config.Config, httpclient *http.Client) (*ThriftServiceClient, error) {
-	endpoint := cfg.ToEndpointURL()
+	var err error
+	endpoint, err := cfg.ToEndpointURL()
+	if err != nil {
+		return nil, err
+	}
 	tcfg := &thrift.TConfiguration{
 		TLSConfig: cfg.TLSConfig,
 	}
@@ -200,7 +206,6 @@ func InitThriftClient(cfg *config.Config, httpclient *http.Client) (*ThriftServi
 	}
 
 	var tTrans thrift.TTransport
-	var err error
 
 	switch cfg.ThriftTransport {
 	case "http":
@@ -270,6 +275,8 @@ func SprintGuid(bts []byte) string {
 	return fmt.Sprintf("%x", bts)
 }
 
+var retryableStatusCode = []int{http.StatusTooManyRequests, http.StatusServiceUnavailable}
+
 type Transport struct {
 	Base  *http.Transport
 	Authr auth.Authenticator
@@ -309,6 +316,27 @@ func (t *Transport) RoundTrip(req *http.Request) (*http.Response, error) {
 	if err != nil {
 		return nil, err
 	}
+	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 reason != "" {
+			logger.Err(fmt.Errorf(reason)).Msg("non retryable error")
+			return nil, errors.New(reason)
+		}
+		if terrmsg != "" {
+			logger.Err(fmt.Errorf(terrmsg)).Msg("non retryable error")
+			return nil, errors.New(terrmsg)
+		}
+		return nil, errors.New(resp.Status)
+	}
 
 	return resp, nil
 }
@@ -322,7 +350,7 @@ func RetryableClient(cfg *config.Config) *http.Client {
 		RetryWaitMax: cfg.RetryWaitMax,
 		RetryMax:     cfg.RetryMax,
 		ErrorHandler: errorHandler,
-		CheckRetry:   retryablehttp.DefaultRetryPolicy,
+		CheckRetry:   RetryPolicy,
 		Backoff:      retryablehttp.DefaultBackoff,
 	}
 	return retryableClient.StandardClient()
@@ -412,3 +440,34 @@ func errorHandler(resp *http.Response, err error, numTries int) (*http.Response,
 
 	return resp, werr
 }
+
+func RetryPolicy(ctx context.Context, resp *http.Response, err error) (bool, error) {
+	var lostConn = regexp.MustCompile(`EOF`)
+
+	// do not retry on context.Canceled or context.DeadlineExceeded
+	if ctx.Err() != nil {
+		return false, ctx.Err()
+	}
+
+	if err != nil {
+		if v, ok := err.(*url.Error); ok {
+			if lostConn.MatchString(v.Error()) {
+				return true, v
+			}
+		}
+		return false, nil
+	}
+
+	// 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
+		}
+	}
+
+	return false, nil
+
+}

internal/config/config.go

@@ -36,10 +36,19 @@ type Config struct {
 }
 
 // ToEndpointURL generates the endpoint URL from Config that a Thrift client will connect to
-func (c *Config) ToEndpointURL() string {
+func (c *Config) ToEndpointURL() (string, error) {
 	var userInfo string
 	endpointUrl := fmt.Sprintf("%s://%s%s:%d%s", c.Protocol, userInfo, c.Host, c.Port, c.HTTPPath)
-	return endpointUrl
+	if c.Host == "" {
+		return endpointUrl, errors.New("databricks: missing Hostname")
+	}
+	if c.Port == 0 {
+		return endpointUrl, errors.New("databricks: missing Port")
+	}
+	if c.HTTPPath == "" && c.Host != "localhost" {
+		return endpointUrl, errors.New("databricks: missing HTTP Path")
+	}
+	return endpointUrl, nil
 }
 
 // DeepCopy returns a true deep copy of Config

internal/config/config_test.go

@@ -17,11 +17,12 @@ func TestParseConfig(t *testing.T) {
 	}
 	tz, _ := time.LoadLocation("America/Vancouver")
 	tests := []struct {
-		name    string
-		args    args
-		wantCfg UserConfig
-		wantURL string
-		wantErr bool
+		name            string
+		args            args
+		wantCfg         UserConfig
+		wantURL         string
+		wantErr         bool
+		wantEndpointErr bool
 	}{
 		{
 			name: "base case",
@@ -253,11 +254,12 @@ func TestParseConfig(t *testing.T) {
 				RetryWaitMin:  1 * time.Second,
 				RetryWaitMax:  30 * time.Second,
 			},
-			wantURL: "https://example.cloud.databricks.com:443",
-			wantErr: false,
+			wantURL:         "https://example.cloud.databricks.com:443",
+			wantErr:         false,
+			wantEndpointErr: true,
 		},
 		{
-			name: "missing http path",
+			name: "missing http path 2",
 			args: args{dsn: "token:[email protected]:443?catalog=default&schema=system&timeout=100&maxRows=1000"},
 			wantCfg: UserConfig{
 				Protocol:      "https",
@@ -274,8 +276,9 @@ func TestParseConfig(t *testing.T) {
 				RetryWaitMin:  1 * time.Second,
 				RetryWaitMax:  30 * time.Second,
 			},
-			wantURL: "https://example.cloud.databricks.com:443",
-			wantErr: false,
+			wantURL:         "https://example.cloud.databricks.com:443",
+			wantErr:         false,
+			wantEndpointErr: true,
 		},
 		{
 			name:    "with wrong port",
@@ -331,8 +334,9 @@ func TestParseConfig(t *testing.T) {
 				RetryWaitMin:  1 * time.Second,
 				RetryWaitMax:  30 * time.Second,
 			},
-			wantURL: "https://:443",
-			wantErr: false,
+			wantURL:         "https://:443",
+			wantErr:         false,
+			wantEndpointErr: true,
 		},
 	}
 	for _, tt := range tests {
@@ -348,7 +352,11 @@ func TestParseConfig(t *testing.T) {
 			}
 			if err == nil {
 				cfg := &Config{UserConfig: got}
-				gotUrl := cfg.ToEndpointURL()
+				gotUrl, err := cfg.ToEndpointURL()
+				if (err != nil) != tt.wantEndpointErr {
+					t.Errorf("ParseConfig() error = %v, wantErr %v", err, tt.wantErr)
+					return
+				}
 				if gotUrl != tt.wantURL {
 					t.Errorf("ToEndpointURL() = %v, want %v", got, tt.wantErr)
 					return

testserver.go

@@ -12,7 +12,7 @@ type thriftHandler struct {
 	processor               thrift.TProcessor
 	inPfactory, outPfactory thrift.TProtocolFactory
 	count503_2_retries      int
-	count500_5_retries      int
+	count503_5_retries      int
 	count429_2_retries      int
 }
 
@@ -36,13 +36,13 @@ func (h *thriftHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		} else {
 			h.count429_2_retries = 0
 		}
-	case "/500-5-retries":
-		if h.count500_5_retries <= 5 {
-			w.WriteHeader(http.StatusInternalServerError)
-			h.count500_5_retries++
+	case "/503-5-retries":
+		if h.count503_5_retries <= 5 {
+			w.WriteHeader(http.StatusServiceUnavailable)
+			h.count503_5_retries++
 			return
 		} else {
-			h.count500_5_retries = 0
+			h.count503_5_retries = 0
 		}
 	}