From 5ab582bd3f719cb97c752f22d27a057b9fa20e07 Mon Sep 17 00:00:00 2001
From: Davide <davide@hibernate.org>
Date: Mon, 13 Jun 2022 17:30:00 +0100
Subject: [PATCH] [#929] Refactor id generation for MySQL

MySQL always returns an id of type Long.
This commit changes an hack that was affecting all dbs
and preventing H2 to read identity id of type Integer or Short
---
 .../java/org/hibernate/reactive/logging/impl/Log.java |  4 ++--
 .../entity/impl/ReactiveAbstractEntityPersister.java  |  5 -----
 .../reactive/pool/impl/SqlClientConnection.java       | 11 ++++++++++-
 3 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/logging/impl/Log.java b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/logging/impl/Log.java
index d3aa642786..eeec4ef9df 100644
--- a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/logging/impl/Log.java
+++ b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/logging/impl/Log.java
@@ -104,8 +104,8 @@ public interface Log extends BasicLogger {
 	@Message(id = 34, value = "The database returned no natively generated identity value")
 	HibernateException noNativelyGeneratedValueReturned();
 
-	@Message(id = 35, value = "The database can only generate identifiers of type Long")
-	HibernateException nativelyGeneratedValueMustBeLong();
+	@Message(id = 35, value = "The database can only generate identifiers of type Long: the id %1$d cannot be cast to %2$s ")
+	HibernateException nativelyGeneratedValueMustBeLong(Long id, @FormatWith(ClassFormatter.class) Class<?> idClass);
 
 	@Message(id = 356, value = "Wrong entity type!")
 	HibernateException wrongEntityType();
diff --git a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/persister/entity/impl/ReactiveAbstractEntityPersister.java b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/persister/entity/impl/ReactiveAbstractEntityPersister.java
index 6d094e443c..05420138ec 100644
--- a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/persister/entity/impl/ReactiveAbstractEntityPersister.java
+++ b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/persister/entity/impl/ReactiveAbstractEntityPersister.java
@@ -393,11 +393,6 @@ default CompletionStage<Serializable> insertReactive(
 		} );
 
 		Class<?> idClass = delegate().getIdentifierType().getReturnedClass();
-		if ( idClass.equals(Integer.class) || idClass.equals(Short.class) ) {
-			// since on MySQL we can only retrieve Long values, adjust to Long
-			// id will be cast back to the right type by castToIdentifierType()
-			idClass = Long.class;
-		}
 		return getReactiveConnection( session )
 				//Note: in ORM core there are other ways to fetch the generated identity:
 				//      getGeneratedKeys(), or an extra round select statement. But we
diff --git a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/pool/impl/SqlClientConnection.java b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/pool/impl/SqlClientConnection.java
index e5108d7239..12e17a2a3c 100644
--- a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/pool/impl/SqlClientConnection.java
+++ b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/pool/impl/SqlClientConnection.java
@@ -309,10 +309,19 @@ public CompletionStage<Void> close() {
 	private static <T> T getLastInsertedId(RowSet<Row> rows, Class<T> idClass, String idColumnName) {
 		final Long mySqlId = rows.property( MYSQL_LAST_INSERTED_ID );
 		if ( mySqlId != null ) {
+			// MySQL will return a Long for any of these types
 			if ( Long.class.equals( idClass ) ) {
 				return (T) mySqlId;
 			}
-			throw LOG.nativelyGeneratedValueMustBeLong();
+			if ( Integer.class.equals( idClass )
+					&& mySqlId <= Integer.MAX_VALUE ) {
+				return (T) mySqlId;
+			}
+			if ( Short.class.equals( idClass )
+					&& mySqlId <= Short.MAX_VALUE) {
+				return (T) mySqlId;
+			}
+			throw LOG.nativelyGeneratedValueMustBeLong( mySqlId, idClass );
 		}
 		final Row oracleKeys = rows.property( ORACLE_GENERATED_KEYS );
 		if ( oracleKeys != null ) {