Add support for new OKTA OIE and retry 429 error code

authokta.go

@@ -21,7 +21,8 @@ type authOKTARequest struct {
 }
 
 type authOKTAResponse struct {
-	CookieToken string `json:"cookieToken"`
+	CookieToken  string `json:"cookieToken"`
+	SessionToken string `json:"sessionToken"`
 }
 
 /*
@@ -137,7 +138,13 @@ func authenticateBySAML(
 	logger.WithContext(ctx).Info("step 4: query IDP URL snowflake app to get SAML response")
 	params = &url.Values{}
 	params.Add("RelayState", "/some/deep/link")
-	params.Add("onetimetoken", respa.CookieToken)
+	var oneTimeToken string
+	if respa.SessionToken != "" {
+		oneTimeToken = respa.SessionToken
+	} else {
+		oneTimeToken = respa.CookieToken
+	}
+	params.Add("onetimetoken", oneTimeToken)
 
 	headers = make(map[string]string)
 	headers[httpHeaderAccept] = "*/*"

retry.go

@@ -239,7 +239,7 @@ func (r *retryHTTP) execute() (res *http.Response, err error) {
 			logger.WithContext(r.ctx).Warningf(
 				"failed http connection. no response is returned. err: %v. retrying...\n", err)
 		} else {
-			if res.StatusCode == http.StatusOK || r.raise4XX && res != nil && res.StatusCode >= 400 && res.StatusCode < 500 {
+			if res.StatusCode == http.StatusOK || r.raise4XX && res != nil && res.StatusCode >= 400 && res.StatusCode < 500 && res.StatusCode != 429 {
 				// exit if success
 				// or
 				// abort connection if raise4XX flag is enabled and the range of HTTP status code are 4XX.

retry_test.go

@@ -46,10 +46,11 @@ func (b *fakeResponseBody) Close() error {
 }
 
 type fakeHTTPClient struct {
-	cnt     int    // number of retry
-	success bool   // return success after retry in cnt times
-	timeout bool   // timeout
-	body    []byte // return body
+	cnt        int    // number of retry
+	success    bool   // return success after retry in cnt times
+	timeout    bool   // timeout
+	body       []byte // return body
+	statusCode int    // status code
 }
 
 func (c *fakeHTTPClient) Do(req *http.Request) (*http.Response, error) {
@@ -71,7 +72,11 @@ func (c *fakeHTTPClient) Do(req *http.Request) (*http.Response, error) {
 				timeout: true,
 			}
 		}
-		retcode = 0
+		if c.statusCode != 0 {
+			retcode = c.statusCode
+		} else {
+			retcode = 0
+		}
 	}
 
 	ret := &http.Response{
@@ -263,3 +268,29 @@ func TestRetryLoginRequest(t *testing.T) {
 		t.Fatalf("no retry counter should be attached: %v", retryCounterKey)
 	}
 }
+
+func TestLoginRetry429(t *testing.T) {
+	client := &fakeHTTPClient{
+		cnt:        3,
+		success:    true,
+		statusCode: 429,
+	}
+	urlPtr, err := url.Parse("https://fakeaccountretrylogin.snowflakecomputing.com:443/login-request?request_id=testid")
+	if err != nil {
+		t.Fatal("failed to parse the test URL")
+	}
+	_, err = newRetryHTTP(context.TODO(),
+		client,
+		fakeRequestFunc, urlPtr, make(map[string]string), 60*time.Second).doRaise4XX(true).doPost().setBody([]byte{0}).execute() // enable doRaise4XXX
+	if err != nil {
+		t.Fatal("failed to run retry")
+	}
+	var values url.Values
+	values, err = url.ParseQuery(urlPtr.RawQuery)
+	if err != nil {
+		t.Fatalf("failed to fail to parse the URL: %v", err)
+	}
+	if values.Get(retryCounterKey) != "" {
+		t.Fatalf("no retry counter should be attached: %v", retryCounterKey)
+	}
+}