Skip to content

Commit 0f792c8

Browse files
committed
[#1548] Fix identity generation for CockroachDB
1 parent c674ea3 commit 0f792c8

File tree

4 files changed

+53
-15
lines changed

4 files changed

+53
-15
lines changed

hibernate-reactive-core/src/main/java/org/hibernate/reactive/dialect/ReactiveDialectWrapper.java

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,20 @@
1717
*/
1818
public final class ReactiveDialectWrapper extends DialectDelegateWrapper {
1919

20+
/**
21+
* A utility method to help checking the actual dialect class.
22+
* @return true, if the dialect is an instance of one of the classes
23+
*/
24+
public static boolean instanceOf(Dialect dialect, Class<?>... dialectClasses) {
25+
Dialect realDialect = extractRealDialect( dialect );
26+
for ( Class<?> dialectClass : dialectClasses ) {
27+
if ( dialectClass.isInstance( realDialect ) ) {
28+
return true;
29+
}
30+
}
31+
return false;
32+
}
33+
2034
public ReactiveDialectWrapper(Dialect wrapped) {
2135
super( wrapped );
2236
}

hibernate-reactive-core/src/main/java/org/hibernate/reactive/id/insert/ReactiveAbstractReturningDelegate.java

Lines changed: 27 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import java.lang.invoke.MethodHandles;
99
import java.util.concurrent.CompletionStage;
1010

11+
import org.hibernate.dialect.CockroachDialect;
1112
import org.hibernate.dialect.Dialect;
1213
import org.hibernate.dialect.PostgreSQLDialect;
1314
import org.hibernate.dialect.SQLServerDialect;
@@ -23,9 +24,10 @@
2324
import org.hibernate.reactive.logging.impl.LoggerFactory;
2425
import org.hibernate.reactive.pool.ReactiveConnection;
2526
import org.hibernate.reactive.session.ReactiveConnectionSupplier;
27+
import org.hibernate.type.Type;
2628

27-
import static java.util.function.Function.identity;
28-
import static org.hibernate.dialect.DialectDelegateWrapper.extractRealDialect;
29+
30+
import static org.hibernate.reactive.dialect.ReactiveDialectWrapper.instanceOf;
2931

3032
public interface ReactiveAbstractReturningDelegate extends ReactiveInsertGeneratedIdentifierDelegate {
3133

@@ -48,16 +50,36 @@ default CompletionStage<Object> reactivePerformInsert(PreparedStatementDetails i
4850
ReactiveConnection reactiveConnection = ( (ReactiveConnectionSupplier) session ).getReactiveConnection();
4951
return reactiveConnection
5052
.insertAndSelectIdentifier( insertSql, params, idType, identifierColumnName )
51-
.thenApply( identity() );
53+
.thenApply( this::validateGeneratedIdentityId );
54+
}
55+
56+
private Object validateGeneratedIdentityId(Object generatedId) {
57+
if ( generatedId == null ) {
58+
throw LOG.noNativelyGeneratedValueReturned();
59+
}
60+
61+
// CockroachDB might generate an identifier that fits an integer (and maybe a short) from time to time.
62+
// Users should not rely on it, or they might have random, hard to debug failures.
63+
Type identifierType = getPersister().getIdentifierType();
64+
if ( ( identifierType.getReturnedClass().equals( Short.class ) || identifierType.getReturnedClass().equals( Integer.class ) )
65+
&& instanceOf( getPersister().getFactory().getJdbcServices().getDialect(), CockroachDialect.class ) ) {
66+
throw LOG.invalidIdentifierTypeForCockroachDB( identifierType.getReturnedClass(), getPersister().getEntityName() );
67+
}
68+
return generatedId;
5269
}
5370

5471
private static String createInsert(PreparedStatementDetails insertStatementDetails, String identifierColumnName, Dialect dialect) {
72+
final String sqlEnd = " returning " + identifierColumnName;
5573
if ( instanceOf( dialect, PostgreSQLDialect.class ) ) {
56-
return insertStatementDetails.getSqlString() + " returning " + identifierColumnName;
74+
return insertStatementDetails.getSqlString() + sqlEnd;
75+
}
76+
if ( instanceOf( dialect, CockroachDialect.class )
77+
&& !insertStatementDetails.getSqlString().endsWith( sqlEnd ) ) {
78+
return insertStatementDetails.getSqlString() + sqlEnd;
5779
}
5880
if ( instanceOf( dialect, SQLServerDialect.class ) ) {
5981
String sql = insertStatementDetails.getSqlString();
60-
int index = sql.lastIndexOf( " returning " + identifierColumnName );
82+
int index = sql.lastIndexOf( sqlEnd );
6183
// FIXME: this is a hack for HHH-16365
6284
if ( index > -1 ) {
6385
sql = sql.substring( 0, index );
@@ -75,10 +97,6 @@ private static String createInsert(PreparedStatementDetails insertStatementDetai
7597
return insertStatementDetails.getSqlString();
7698
}
7799

78-
private static boolean instanceOf(Dialect dialect, Class<?> dialectClass) {
79-
return dialectClass.isInstance( extractRealDialect( dialect ) );
80-
}
81-
82100
@Override
83101
default CompletionStage<Object> reactivePerformInsert(String insertSQL, SharedSessionContractImplementor session, Binder binder) {
84102
throw LOG.notYetImplemented();

hibernate-reactive-core/src/main/java/org/hibernate/reactive/id/insert/ReactiveAbstractSelectingDelegate.java

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,6 @@ default CompletionStage<Object> reactivePerformInsert(
4646
final JdbcCoordinator jdbcCoordinator = session.getJdbcCoordinator();
4747
final JdbcServices jdbcServices = session.getJdbcServices();
4848

49-
jdbcServices.getSqlStatementLogger().logStatement( insertStatementDetails.getSqlString() );
50-
5149
Object[] updateParams = bind( statement -> {
5250
PreparedStatementDetails details = new PrepareStatementDetailsAdaptor( insertStatementDetails, statement, session.getJdbcServices() );
5351
jdbcValueBindings.beforeStatement( details );

hibernate-reactive-core/src/main/java/org/hibernate/reactive/persister/entity/impl/ReactiveIdentityGenerator.java

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import org.hibernate.dialect.Dialect;
1515
import org.hibernate.dialect.PostgreSQLDialect;
1616
import org.hibernate.dialect.SQLServerDialect;
17+
import org.hibernate.dialect.identity.CockroachDBIdentityColumnSupport;
1718
import org.hibernate.engine.spi.SessionFactoryImplementor;
1819
import org.hibernate.id.IdentityGenerator;
1920
import org.hibernate.id.PostInsertIdentityPersister;
@@ -24,23 +25,30 @@
2425
import org.hibernate.reactive.id.insert.ReactiveInsertReturningDelegate;
2526
import org.hibernate.sql.Insert;
2627

28+
import static org.hibernate.reactive.dialect.ReactiveDialectWrapper.instanceOf;
29+
2730
/**
2831
* Fix the insert and select id queries generated by Hibernate ORM
2932
*/
3033
public class ReactiveIdentityGenerator extends IdentityGenerator {
3134

35+
/**
36+
* @see CockroachDBIdentityColumnSupport#supportsIdentityColumns() for some limitations related to CockraochDB
37+
*/
3238
@Override
3339
public InsertGeneratedIdentifierDelegate getGeneratedIdentifierDelegate(PostInsertIdentityPersister persister) {
3440
Dialect dialect = persister.getFactory().getJdbcServices().getDialect();
3541
if ( persister.getFactory().getSessionFactoryOptions().isGetGeneratedKeysEnabled() ) {
3642
return dialect.getIdentityColumnSupport().buildGetGeneratedKeysDelegate( persister, dialect );
3743
}
38-
else if ( dialect.getIdentityColumnSupport().supportsInsertSelectIdentity() ) {
44+
if ( dialect.getIdentityColumnSupport().supportsInsertSelectIdentity()
45+
// Basically, this is the only approach that works for Cockroach at the moment.
46+
// That said, there are some limitations about identity generation with Cockroach.
47+
// Check CockroachDBIdentityColumnSupport#supportsIdentityColumns for more details
48+
|| instanceOf( dialect, CockroachDialect.class ) ) {
3949
return new ReactiveInsertReturningDelegate( persister, dialect );
4050
}
41-
else {
42-
return new ReactiveBasicSelectingDelegate( persister, dialect );
43-
}
51+
return new ReactiveBasicSelectingDelegate( persister, dialect );
4452
}
4553

4654
public static class ReactiveInsertAndSelectDelegate extends InsertReturningDelegate {

0 commit comments

Comments
 (0)