8000 Make the emulated "on conflict" upserts race-free on SQL Server and Oracle by ahus1 · Pull Request #54825 · keycloak/keycloak · GitHub
[go: up one dir, main page]

Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -24,12 +24,18 @@
import org.hibernate.sql.ast.SqlAstTranslatorFactory;
import org.hibernate.sql.ast.spi.StandardSqlAstTranslatorFactory;
import org.hibernate.sql.ast.tree.Statement;
import org.hibernate.sql.ast.tree.insert.ConflictClause;
import org.hibernate.sql.ast.tree.insert.InsertSelectStatement;
import org.hibernate.sql.exec.spi.JdbcOperation;

/**
* Fixes Hibernate's Oracle dialect to correctly handle HQL {@code insert ... on conflict ... do nothing}
* by emitting a MERGE statement instead of a plain INSERT that silently drops the conflict clause.
* Fixes Hibernate's Oracle dialect to correctly handle HQL {@code insert ... on conflict (...) do nothing}, which it
* would otherwise render as a plain INSERT that silently drops the conflict clause.
* <p>
* It renders an INSERT with the {@code IGNORE_ROW_ON_DUPKEY_INDEX} hint, so that a duplicate key is skipped and
* counted as 0 rows. Unlike a MERGE, this also holds when a concurrent transaction inserts the same key. The conflict
* columns must match a unique index. {@code do update} clauses are still rendered as MERGE, which fails with a unique
* constraint violation in that case.
* <p>
* See <a href="https://hibernate.atlassian.net/browse/HHH-20826">HHH-20826</a> for the upstream tracker.
*/
Expand All @@ -44,13 +50,31 @@ protected <T extends JdbcOperation> SqlAstTranslator<T> buildTranslator(
return new OracleSqlAstTranslator<>(sessionFactory, statement) {
@Override
protected void visitInsertStatementOnly(InsertSelectStatement statement) {
if (statement.getConflictClause() != null
&& !statement.getConflictClause().getConstraintColumnNames().isEmpty()) {
ConflictClause conflictClause = statement.getConflictClause();
if (conflictClause != null && !conflictClause.isDoNothing()
&& !conflictClause.getConstraintColumnNames().isEmpty()) {
visitInsertStatementEmulateMerge(statement);
} else {
super.visitInsertStatementOnly(statement);
}
}

@Override
protected void renderInsertCommand(InsertSelectStatement statement) {
ConflictClause conflictClause = statement.getConflictClause();
if (conflictClause != null && conflictClause.isDoNothing()
&& !conflictClause.getConstraintColumnNames().isEmpty()) {
// The hint names the table as it appears in the statement, without a schema
String table = statement.getTargetTable().getTableExpression();
appendSql("insert /*+ IGNORE_ROW_ON_DUPKEY_INDEX(");
appendSql(table.substring(table.lastIndexOf('.') + 1));
appendSql('(');
appendSql(String.join(",", conflictClause.getConstraintColumnNames()));
appendSql(")) */ into ");
} else {
super.renderInsertCommand(statement);
}
}
};
}
};
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,20 +17,25 @@

package org.keycloak.connections.jpa.dialect;

import org.hibernate.LockMode;
import org.hibernate.dialect.SQLServerDialect;
import org.hibernate.dialect.sql.ast.SQLServerSqlAstTranslator;
import org.hibernate.engine.spi.SessionFactoryImplementor;
import org.hibernate.sql.ast.SqlAstTranslator;
import org.hibernate.sql.ast.SqlAstTranslatorFactory;
import org.hibernate.sql.ast.spi.StandardSqlAstTranslatorFactory;
import org.hibernate.sql.ast.tree.Statement;
import org.hibernate.sql.ast.tree.from.NamedTableReference;
import org.hibernate.sql.ast.tree.insert.InsertSelectStatement;
import org.hibernate.sql.exec.spi.JdbcOperation;

/**
* Fixes Hibernate's SQL Server dialect to correctly handle HQL {@code insert ... on conflict ... do nothing}
* by emitting a MERGE statement instead of a plain INSERT that silently drops the conflict clause.
* <p>
* The MERGE target gets a {@code HOLDLOCK} hint for all conflict clauses. Without it, two concurrent upserts of the
* same new key can both take the "not matched" branch, and one fails with a unique constraint violation.
* <p>
* See <a href="https://hibernate.atlassian.net/browse/HHH-20826">HHH-20826</a> for the upstream tracker.
*/
public class KeycloakSQLServerDialect extends SQLServerDialect {
Expand All @@ -42,16 +47,33 @@ public SqlAstTranslatorFactory getSqlAstTranslatorFactory() {
protected <T extends JdbcOperation> SqlAstTranslator<T> buildTranslator(
SessionFactoryImplementor sessionFactory, Statement statement) {
return new SQLServerSqlAstTranslator<>(sessionFactory, statement) {
private boolean renderingMergeTarget;

@Override
protected void visitInsertStatementOnly(InsertSelectStatement statement) {
if (statement.getConflictClause() != null
&& !statement.getConflictClause().getConstraintColumnNames().isEmpty()) {
renderingMergeTarget = true;
visitInsertStatementEmulateMerge(statement);
appendSql(';');
} else {
super.visitInsertStatementOnly(statement);
}
}

@Override
protected boolean renderNamedTableReference(NamedTableReference tableReference, LockMode lockMode) {
if (renderingMergeTarget) {
renderingMergeTarget = false;
// In a MERGE the table hint goes between the table name and the alias
appendSql(tableReference.getTableExpression());
registerAffectedTable(tableReference);
appendSql(" with (holdlock)");
renderTableReferenceIdentificationVariable(tableReference);
return true;
}
return super.renderNamedTableReference(tableReference, lockMode);
}
};
}
};
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@

import org.hibernate.dialect.Dialect;
import org.hibernate.dialect.MySQLDialect;
import org.hibernate.dialect.OracleDialect;
import org.hibernate.engine.jdbc.env.spi.IdentifierHelper;
import org.hibernate.engine.spi.SessionFactoryImplementor;
import org.hibernate.internal.SessionFactoryImpl;
Expand Down Expand Up @@ -70,6 +71,16 @@ public static boolean isUpsertRowCountUnreliable(EntityManager em) {
.getJdbcServices().getDialect() instanceof MySQLDialect;
}

/**
* Oracle emulates {@code on conflict ... do update} with a MERGE, which fails with a unique constraint violation
* when a concurrent transaction inserts the same key. {@code on conflict ... do nothing} is safe there, see
* {@link org.keycloak.connections.jpa.dialect.KeycloakOracleDialect}.
*/
public static boolean isConflictUpdateRacy(EntityManager em) {
return em.getEntityManagerFactory().unwrap(SessionFactoryImplementor.class)
.getJdbcServices().getDialect() instanceof OracleDialect;
}

public static String getTableNameForNativeQuery(String tableName, EntityManager em) {
final Dialect dialect = em.getEntityManagerFactory().unwrap(SessionFactoryImpl.class).getJdbcServices().getDialect();
IdentifierHelper identifierHelper = em.getEntityManagerFactory().unwrap(SessionFactoryImpl.class).getJdbcServices().getJdbcEnvironment().getIdentifierHelper();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,15 @@
" on conflict (id) do update set expire = :expire" +
" where expire <= :currentTime"
),
@NamedQuery(
name = "deleteExpiredRevokeToken",
query = "delete from RevokedTokenEntity t where t.id = :id and t.expire <= :currentTime"
),
@NamedQuery(
name = "insertRevokeTokenIfAbsentNoUpdate",
query = "insert into RevokedTokenEntity (id, expire) values (:id, :expire)" +
" on conflict (id) do nothing"
),
})
@Table(name="REVOKED_TOKEN")
@Entity
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,15 @@ public boolean put(String id, long lifespanSeconds) {
.setParameter(1, id)
.setParameter(2, expire)
.executeUpdate();
} else if (JpaUtils.isConflictUpdateRacy(em)) {
em.createNamedQuery("deleteExpiredRevokeToken")
.setParameter("id", id)
.setParameter("currentTime", currentTime)
.executeUpdate();
rows = em.createNamedQuery("insertRevokeTokenIfAbsentNoUpdate")
.setParameter("id", id)
.setParameter("expire", expire)
.executeUpdate();
} else {
rows = em.createNamedQuery("insertRevokeTokenIfAbsent")
.setParameter("id", id)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,16 @@ public boolean putIfAbsent(String key, long lifespanInSeconds) {
.setParameter(2, SingleUseObjectSerialization.notesToString(key, Map.of()))
.setParameter(3, currentTime + lifespanInSeconds)
.executeUpdate();
} else if (JpaUtils.isConflictUpdateRacy(em)) {
em.createNamedQuery("deleteExpiredSingleUseObject")
.setParameter("id", key)
.setParameter("currentTime", currentTime)
.executeUpdate();
rows = em.createNamedQuery("insertIfAbsentSingleUseObject")
.setParameter("id", key)
.setParameter("notes", SingleUseObjectSerialization.notesToString(key, Map.of()))
.setParameter("expire", currentTime + lifespanInSeconds)
.executeUpdate();
} else {
rows = em.createNamedQuery("insertIfAbsentOrExpiredSingleUseObject")
.setParameter("id", key)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,15 @@
" on conflict (id) do update set notes = :notes, expire = :expire" +
" where expire <= :currentTime"
),
@NamedQuery(
name = "deleteExpiredSingleUseObject",
query = "delete from SingleUseObjectEntity e where e.id = :id and e.expire <= :currentTime"
),
@NamedQuery(
name = "insertIfAbsentSingleUseObject",
query = "insert into SingleUseObjectEntity (id, notes, expire) values (:id, :notes, :expire)" +
" on conflict (id) do nothing"
),
@NamedQuery(
name = "updateIfNotExpiredSingleUseObject",
query = "update SingleUseObjectEntity set notes = :notes" +
Expand Down
Loading
Loading
0