From dd649e13fb2a3f5f05ad6ea2a167e61c22f77821 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 d3aa64278..eeec4ef9d 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 6d094e443..05420138e 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 e5108d723..12e17a2a3 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 ) {