Skip to content

Commit

Permalink
Optimize cleanupTokenByTokenValue logic.
Browse files Browse the repository at this point in the history
  • Loading branch information
@mpmadhavig
mpmadhavig committed Jun 27, 2024
1 parent 444ed6b commit 63f7f9f
Show file tree
Hide file tree
Showing 2 changed files with 70 additions and 37 deletions.
104 changes: 67 additions & 37 deletions ...n.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/dao/OldTokensCleanDAO.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,48 +64,57 @@ public void cleanupTokenByTokenId(String tokenId, Connection connection) throws
}
}

public void cleanupTokenByTokenValue(String token, Connection connection) throws SQLException {
OldAccessTokenDO oldAccessTokenObject = new OldAccessTokenDO();

String sql;
if (OAuth2ServiceComponentHolder.isIDPIdColumnEnabled()) {
sql = SQLQueries.RETRIEVE_OLD_TOKEN_BY_TOKEN_HASH_WITH_IDP_NAME;
} else {
sql = SQLQueries.RETRIEVE_OLD_TOKEN_BY_TOKEN_HASH;
}
/**
* This method is used to clean-up old tokens by token hash value.
*
* @param tokenHash Token hash value.
* @param connection Database connection.
* @throws SQLException If an error occurs while cleaning up the token.
*/
public void cleanupTokenByTokenValue(String tokenHash, Connection connection) throws SQLException {

PreparedStatement prepStmt = connection.prepareStatement(sql);
prepStmt.setString(1, token);
ResultSet resultSet = prepStmt.executeQuery();
//iterate result set and insert to AccessTokenDO object.
if (resultSet.next()) {
oldAccessTokenObject.setTokenId(resultSet.getString(1));
oldAccessTokenObject.setAccessToken(resultSet.getString(2));
oldAccessTokenObject.setRefreshToken(resultSet.getString(3));
oldAccessTokenObject.setConsumerKeyId(resultSet.getInt(4));
oldAccessTokenObject.setAuthzUser(resultSet.getString(5));
oldAccessTokenObject.setTenantId(resultSet.getInt(6));
oldAccessTokenObject.setUserDomain(resultSet.getString(7));
oldAccessTokenObject.setUserType(resultSet.getString(8));
oldAccessTokenObject.setGrantType(resultSet.getString(9));
oldAccessTokenObject.setTimeCreated(resultSet.getTimestamp(10));
oldAccessTokenObject.setRefreshTokenTimeCreated(resultSet.getTimestamp(11));
oldAccessTokenObject.setValdityPeriod(resultSet.getLong(12));
oldAccessTokenObject.setRefreshTokenValidityPeriod(resultSet.getLong(13));
oldAccessTokenObject.setTokenScopeHash(resultSet.getString(14));
oldAccessTokenObject.setTokenState(resultSet.getString(15));
oldAccessTokenObject.setTokenStateId(resultSet.getString(16));
oldAccessTokenObject.setSubjectIdentifier(resultSet.getString(17));
oldAccessTokenObject.setAccessTokenHash(resultSet.getString(18));
oldAccessTokenObject.setRefreshTokenHash(resultSet.getString(19));
OldAccessTokenDO oldAccessTokenObject = new OldAccessTokenDO();
if (OAuthServerConfiguration.getInstance().useRetainOldAccessTokens()) {
String sql;
if (OAuth2ServiceComponentHolder.isIDPIdColumnEnabled()) {
oldAccessTokenObject.setIdpId(resultSet.getInt(20));
sql = SQLQueries.RETRIEVE_OLD_TOKEN_BY_TOKEN_HASH_WITH_IDP_NAME;
} else {
sql = SQLQueries.RETRIEVE_OLD_TOKEN_BY_TOKEN_HASH;
}

PreparedStatement prepStmt = connection.prepareStatement(sql);
prepStmt.setString(1, tokenHash);
ResultSet resultSet = prepStmt.executeQuery();
//iterate result set and insert to AccessTokenDO object.
if (resultSet.next()) {
oldAccessTokenObject.setTokenId(resultSet.getString(1));
oldAccessTokenObject.setAccessToken(resultSet.getString(2));
oldAccessTokenObject.setRefreshToken(resultSet.getString(3));
oldAccessTokenObject.setConsumerKeyId(resultSet.getInt(4));
oldAccessTokenObject.setAuthzUser(resultSet.getString(5));
oldAccessTokenObject.setTenantId(resultSet.getInt(6));
oldAccessTokenObject.setUserDomain(resultSet.getString(7));
oldAccessTokenObject.setUserType(resultSet.getString(8));
oldAccessTokenObject.setGrantType(resultSet.getString(9));
oldAccessTokenObject.setTimeCreated(resultSet.getTimestamp(10));
oldAccessTokenObject.setRefreshTokenTimeCreated(resultSet.getTimestamp(11));
oldAccessTokenObject.setValdityPeriod(resultSet.getLong(12));
oldAccessTokenObject.setRefreshTokenValidityPeriod(resultSet.getLong(13));
oldAccessTokenObject.setTokenScopeHash(resultSet.getString(14));
oldAccessTokenObject.setTokenState(resultSet.getString(15));
oldAccessTokenObject.setTokenStateId(resultSet.getString(16));
oldAccessTokenObject.setSubjectIdentifier(resultSet.getString(17));
oldAccessTokenObject.setAccessTokenHash(resultSet.getString(18));
oldAccessTokenObject.setRefreshTokenHash(resultSet.getString(19));
if (OAuth2ServiceComponentHolder.isIDPIdColumnEnabled()) {
oldAccessTokenObject.setIdpId(resultSet.getInt(20));
}
}
}
if (OAuthServerConfiguration.getInstance().useRetainOldAccessTokens()) {
saveTokenInAuditTable(oldAccessTokenObject, connection);
removeTokenFromMainTable(oldAccessTokenObject.getTokenId(), connection);
} else {
removeTokenFromMainTableByTokenHash(tokenHash, connection);
}
removeTokenFromMainTable(oldAccessTokenObject.getTokenId(), connection);
}

private void saveTokenInAuditTable(OldAccessTokenDO oldAccessTokenDAO, Connection connection) throws SQLException {
Expand Down Expand Up @@ -168,6 +177,27 @@ private void removeTokenFromMainTable(String oldAccessTokenID, Connection connec
}
}

private void removeTokenFromMainTableByTokenHash(String accessTokenHash, Connection connection)
throws SQLException {

connection.setAutoCommit(false);
try {
PreparedStatement deletefromaccesstokentable =
connection.prepareStatement(SQLQueries.DELETE_OLD_TOKEN_BY_TOKEN_HASH);
deletefromaccesstokentable.setString(1, accessTokenHash);
deletefromaccesstokentable.executeUpdate();
if (log.isDebugEnabled()) {
log.debug(
"Successfully old access token deleted from access token table. Token Hash: " +
accessTokenHash);
}
connection.commit();
} catch (SQLException e) {
connection.rollback();
log.error("SQL error occurred while remove token from main table", e);
}
}

public void cleanupTokensInBatch(List<String> oldTokens, Connection connection) throws SQLException {
for (String token : oldTokens) {
cleanupTokenByTokenValue(token, connection);
Expand Down
3 changes: 3 additions & 0 deletions ...2.carbon.identity.oauth/src/main/java/org/wso2/carbon/identity/oauth2/dao/SQLQueries.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -963,6 +963,9 @@ public class SQLQueries {

public static final String DELETE_OLD_TOKEN_BY_ID = "DELETE FROM IDN_OAUTH2_ACCESS_TOKEN WHERE TOKEN_ID = ?";

public static final String DELETE_OLD_TOKEN_BY_TOKEN_HASH =
"DELETE FROM IDN_OAUTH2_ACCESS_TOKEN WHERE ACCESS_TOKEN_HASH = ?";

public static final String DELETE_ACCESS_TOKEN = "DELETE FROM IDN_OAUTH2_ACCESS_TOKEN WHERE ACCESS_TOKEN = ? ";

public static final String REVOKE_SAAS_TOKENS_OF_OTHER_TENANTS = "UPDATE IDN_OAUTH2_ACCESS_TOKEN " +
Expand Down

0 comments on commit 63f7f9f

Please sign in to comment.