Skip to content

Commit

Permalink
fix to handle error codes instead of error message
Browse files Browse the repository at this point in the history
  • Loading branch information
yas-okadatech committed May 30, 2024
1 parent 9c2219d commit 364ecd6
Showing 1 changed file with 29 additions and 11 deletions.
40 changes: 29 additions & 11 deletions src/main/java/org/embulk/output/SnowflakeOutputPlugin.java
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,16 @@ public static MatchByColumnName fromString(String value) {
}
}

// error codes which need reauthenticate
// ref:
// https://github.com/snowflakedb/snowflake-jdbc/blob/v3.13.26/src/main/java/net/snowflake/client/jdbc/SnowflakeUtil.java#L42
private static final int ID_TOKEN_EXPIRED_GS_CODE = 390110;
private static final int SESSION_NOT_EXIST_GS_CODE = 390111;
private static final int MASTER_TOKEN_NOTFOUND = 390113;
private static final int MASTER_EXPIRED_GS_CODE = 390114;
private static final int MASTER_TOKEN_INVALID_GS_CODE = 390115;
private static final int ID_TOKEN_INVALID_LOGIN_REQUEST_GS_CODE = 390195;

@Override
protected Class<? extends PluginTask> getTaskClass() {
return SnowflakePluginTask.class;
Expand Down Expand Up @@ -227,17 +237,25 @@ private void runDropStage(
try {
snowflakeCon.runDropStage(stageIdentifier);
} catch (SnowflakeSQLException ex) {
logger.info("SnowflakeSQLException was caught: {}", ex.getMessage());

if (ex.getMessage().startsWith("Authentication token has expired.")
|| ex.getMessage().startsWith("Session no longer exists.")) {

// INFO: If runCreateStage consumed a lot of time, authentication might be expired.
// In this case, retry to drop stage.
snowflakeCon = (SnowflakeOutputConnection) getConnector(task, true).connect(true);
snowflakeCon.runDropStage(stageIdentifier);
} else {
throw ex;
// INFO: Don't handle only SnowflakeReauthenticationRequest here
// because SnowflakeSQLException with following error codes may be thrown in some cases.

logger.info("SnowflakeSQLException was caught: ({}) {}", ex.getErrorCode(), ex.getMessage());

switch (ex.getErrorCode()) {
case ID_TOKEN_EXPIRED_GS_CODE:
case SESSION_NOT_EXIST_GS_CODE:
case MASTER_TOKEN_NOTFOUND:
case MASTER_EXPIRED_GS_CODE:
case MASTER_TOKEN_INVALID_GS_CODE:
case ID_TOKEN_INVALID_LOGIN_REQUEST_GS_CODE:
// INFO: If runCreateStage consumed a lot of time, authentication might be expired.
// In this case, retry to drop stage.
snowflakeCon = (SnowflakeOutputConnection) getConnector(task, true).connect(true);
snowflakeCon.runDropStage(stageIdentifier);
break;
default:
throw ex;
}
}
}
Expand Down

0 comments on commit 364ecd6

Please sign in to comment.