[#1548] Fix identity generation for CockroachDB

Author: DavideD

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

@@ -17,6 +17,20 @@
  */
 public final class ReactiveDialectWrapper extends DialectDelegateWrapper {
 
+	/**
+	 * A utility method to help checking the actual dialect class.
+	 * @return true, if the dialect is an instance of one of the classes
+	 */
+	public static boolean instanceOf(Dialect dialect, Class<?>... dialectClasses) {
+		Dialect realDialect = extractRealDialect( dialect );
+		for ( Class<?> dialectClass : dialectClasses ) {
+			if ( dialectClass.isInstance( realDialect ) ) {
+				return true;
+			}
+		}
+		return false;
+	}
+
 	public ReactiveDialectWrapper(Dialect wrapped) {
 		super( wrapped );
 	}

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

@@ -8,6 +8,7 @@
 import java.lang.invoke.MethodHandles;
 import java.util.concurrent.CompletionStage;
 
+import org.hibernate.dialect.CockroachDialect;
 import org.hibernate.dialect.Dialect;
 import org.hibernate.dialect.PostgreSQLDialect;
 import org.hibernate.dialect.SQLServerDialect;
@@ -23,9 +24,10 @@
 import org.hibernate.reactive.logging.impl.LoggerFactory;
 import org.hibernate.reactive.pool.ReactiveConnection;
 import org.hibernate.reactive.session.ReactiveConnectionSupplier;
+import org.hibernate.type.Type;
 
-import static java.util.function.Function.identity;
-import static org.hibernate.dialect.DialectDelegateWrapper.extractRealDialect;
+
+import static org.hibernate.reactive.dialect.ReactiveDialectWrapper.instanceOf;
 
 public interface ReactiveAbstractReturningDelegate extends ReactiveInsertGeneratedIdentifierDelegate {
 
@@ -48,16 +50,36 @@ default CompletionStage<Object> reactivePerformInsert(PreparedStatementDetails i
 		ReactiveConnection reactiveConnection = ( (ReactiveConnectionSupplier) session ).getReactiveConnection();
 		return reactiveConnection
 				.insertAndSelectIdentifier( insertSql, params, idType, identifierColumnName )
-				.thenApply( identity() );
+				.thenApply( this::validateGeneratedIdentityId );
+	}
+
+	private Object validateGeneratedIdentityId(Object generatedId) {
+		if ( generatedId == null ) {
+			throw LOG.noNativelyGeneratedValueReturned();
+		}
+
+		// CockroachDB might generate an identifier that fits an integer (and maybe a short) from time to time.
+		// Users should not rely on it, or they might have random, hard to debug failures.
+		Type identifierType = getPersister().getIdentifierType();
+		if ( ( identifierType.getReturnedClass().equals( Short.class ) || identifierType.getReturnedClass().equals( Integer.class ) )
+				&& instanceOf( getPersister().getFactory().getJdbcServices().getDialect(), CockroachDialect.class ) ) {
+			throw LOG.invalidIdentifierTypeForCockroachDB( identifierType.getReturnedClass(), getPersister().getEntityName() );
+		}
+		return generatedId;
 	}
 
 	private static String createInsert(PreparedStatementDetails insertStatementDetails, String identifierColumnName, Dialect dialect) {
+		final String sqlEnd = " returning " + identifierColumnName;
 		if ( instanceOf( dialect, PostgreSQLDialect.class ) ) {
-			return insertStatementDetails.getSqlString() + " returning " + identifierColumnName;
+			return insertStatementDetails.getSqlString() + sqlEnd;
+		}
+		if ( instanceOf( dialect, CockroachDialect.class )
+				&& !insertStatementDetails.getSqlString().endsWith( sqlEnd ) ) {
+			return insertStatementDetails.getSqlString() + sqlEnd;
 		}
 		if ( instanceOf( dialect, SQLServerDialect.class ) ) {
 			String sql = insertStatementDetails.getSqlString();
-			int index = sql.lastIndexOf( " returning " + identifierColumnName );
+			int index = sql.lastIndexOf( sqlEnd );
 			// FIXME: this is a hack for HHH-16365
 			if ( index > -1 ) {
 				sql = sql.substring( 0, index );
@@ -75,10 +97,6 @@ private static String createInsert(PreparedStatementDetails insertStatementDetai
 		return insertStatementDetails.getSqlString();
 	}
 
-	private static boolean instanceOf(Dialect dialect, Class<?> dialectClass) {
-		return dialectClass.isInstance( extractRealDialect( dialect ) );
-	}
-
 	@Override
 	default CompletionStage<Object> reactivePerformInsert(String insertSQL, SharedSessionContractImplementor session, Binder binder) {
 		throw LOG.notYetImplemented();

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

@@ -46,8 +46,6 @@ default CompletionStage<Object> reactivePerformInsert(
 		final JdbcCoordinator jdbcCoordinator = session.getJdbcCoordinator();
 		final JdbcServices jdbcServices = session.getJdbcServices();
 
-		jdbcServices.getSqlStatementLogger().logStatement( insertStatementDetails.getSqlString() );
-
 		Object[] updateParams = bind( statement -> {
 			PreparedStatementDetails details = new PrepareStatementDetailsAdaptor( insertStatementDetails, statement, session.getJdbcServices() );
 			jdbcValueBindings.beforeStatement( details );

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

@@ -14,6 +14,7 @@
 import org.hibernate.dialect.Dialect;
 import org.hibernate.dialect.PostgreSQLDialect;
 import org.hibernate.dialect.SQLServerDialect;
+import org.hibernate.dialect.identity.CockroachDBIdentityColumnSupport;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.id.IdentityGenerator;
 import org.hibernate.id.PostInsertIdentityPersister;
@@ -24,23 +25,30 @@
 import org.hibernate.reactive.id.insert.ReactiveInsertReturningDelegate;
 import org.hibernate.sql.Insert;
 
+import static org.hibernate.reactive.dialect.ReactiveDialectWrapper.instanceOf;
+
 /**
  * Fix the insert and select id queries generated by Hibernate ORM
  */
 public class ReactiveIdentityGenerator extends IdentityGenerator {
 
+	/**
+	 * @see CockroachDBIdentityColumnSupport#supportsIdentityColumns() for some limitations related to CockraochDB
+	 */
 	@Override
 	public InsertGeneratedIdentifierDelegate getGeneratedIdentifierDelegate(PostInsertIdentityPersister persister) {
 		Dialect dialect = persister.getFactory().getJdbcServices().getDialect();
 		if ( persister.getFactory().getSessionFactoryOptions().isGetGeneratedKeysEnabled() ) {
 			return dialect.getIdentityColumnSupport().buildGetGeneratedKeysDelegate( persister, dialect );
 		}
-		else if ( dialect.getIdentityColumnSupport().supportsInsertSelectIdentity() ) {
+		if ( dialect.getIdentityColumnSupport().supportsInsertSelectIdentity()
+				// Basically, this is the only approach that works for Cockroach at the moment.
+				// That said, there are some limitations about identity generation with Cockroach.
+				// Check CockroachDBIdentityColumnSupport#supportsIdentityColumns for more details
+				|| instanceOf( dialect, CockroachDialect.class ) ) {
 			return new ReactiveInsertReturningDelegate( persister, dialect );
 		}
-		else {
-			return new ReactiveBasicSelectingDelegate( persister, dialect );
-		}
+		return new ReactiveBasicSelectingDelegate( persister, dialect );
 	}
 
 	public static class ReactiveInsertAndSelectDelegate extends InsertReturningDelegate {