Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[SNOW-1853185] Okta 429 retry PoC #1997

Draft
wants to merge 1 commit into
base: master
Choose a base branch
from
+27 −4
Draft

[SNOW-1853185] Okta 429 retry PoC #1997

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 24 additions & 3 deletions src/main/java/net/snowflake/client/core/SessionUtil.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
import net.snowflake.client.log.ArgSupplier;
import net.snowflake.client.log.SFLogger;
import net.snowflake.client.log.SFLoggerFactory;
import net.snowflake.client.util.DecorrelatedJitterBackoff;
import net.snowflake.client.util.SecretDetector;
import net.snowflake.client.util.Stopwatch;
import net.snowflake.common.core.SqlState;
Expand Down Expand Up @@ -69,6 +70,7 @@ public class SessionUtil {
private static final String SF_PATH_TOKEN_REQUEST = "/session/token-request";
public static final String SF_PATH_AUTHENTICATOR_REQUEST = "/session/authenticator-request";
public static final String SF_PATH_CONSOLE_LOGIN_REQUEST = "/console/login";
public static final String OKTA_PATH_AUTH = "/api/v1/authn";

public static final String SF_QUERY_SESSION_DELETE = "delete";

Expand Down Expand Up @@ -1456,22 +1458,41 @@ private static void handleFederatedFlowError(SFLoginInput loginInput, Exception
*/
private static String getSamlResponseUsingOkta(SFLoginInput loginInput)
throws SnowflakeSQLException {
DecorrelatedJitterBackoff backoff = new DecorrelatedJitterBackoff(2000, 16000);
long i = 0;
int retryCount = 1;
long backoffMilliSeconds = 0;
String tokenUrl="";
String ssoUrl="";
while (true) {
try {
JsonNode dataNode = federatedFlowStep1(loginInput);
String tokenUrl = dataNode.path("tokenUrl").asText();
String ssoUrl = dataNode.path("ssoUrl").asText();
Thread.sleep(backoffMilliSeconds);
// We don't need to keep sending HTTP requests to the authenticator-request endpoint on every retry
// because it's unlikely that the token and SSO URLs are going to change.
if (retryCount == 1) {
JsonNode dataNode = federatedFlowStep1(loginInput);
tokenUrl = dataNode.path("tokenUrl").asText();
ssoUrl = dataNode.path("ssoUrl").asText();
}
federatedFlowStep2(loginInput, tokenUrl, ssoUrl);
final String oneTimeToken = federatedFlowStep3(loginInput, tokenUrl);
return federatedFlowStep4(loginInput, ssoUrl, oneTimeToken);
} catch (SnowflakeSQLException ex) {
i += 1000; // increment by 1000 ms
// This error gets thrown if the okta request encountered a retry-able error that
// requires getting a new one-time token.
if (ex.getErrorCode() == ErrorCode.AUTHENTICATOR_REQUEST_TIMEOUT.getMessageCode()) {
logger.debug("Failed to get Okta SAML response. Retrying without changing retry count.");
backoffMilliSeconds = backoff.nextSleepTime(i);
logger.debug(
"Going to backoff for " + backoffMilliSeconds + " ms. Retry Count: " + retryCount);
retryCount++;
} else {
throw ex;
}
} catch (InterruptedException ex) {
throw new SnowflakeSQLException(
"Thread was interrupted before retrying Okta request for SAML response.", null);
}
}
}
Expand Down
4 changes: 3 additions & 1 deletion src/main/java/net/snowflake/client/jdbc/RestRequest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -466,7 +466,9 @@ public static CloseableHttpResponse execute(

// If this was a request for an Okta one-time token that failed with a retry-able error,
// throw exception to renew the token before trying again.
if (String.valueOf(httpRequest.getURI()).contains("okta.com/api/v1/authn")) {
// TODO: The URI won't necessrily contain okta.com with customers using vanity URLs, i.e
// somecompany.com/api/v1/authn
if (String.valueOf(httpRequest.getURI()).contains("/api/v1/authn")) {
throw new SnowflakeSQLException(
ErrorCode.AUTHENTICATOR_REQUEST_TIMEOUT,
retryCount,
Expand Down
Loading