HHH-16991 Fix array based id restriction with EnhanceUserType

Author: beikov

hibernate-core/src/main/java/org/hibernate/loader/ast/internal/CollectionBatchLoaderArrayParam.java

@@ -109,12 +109,14 @@ public CollectionBatchLoaderArrayParam(
 				.buildSelectTranslator( getSessionFactory(), sqlSelect )
 				.translate( JdbcParameterBindings.NO_BINDINGS, QueryOptions.NONE );
 	}
+
 	@Override
 	public PersistentCollection<?> load(Object keyBeingLoaded, SharedSessionContractImplementor session) {
 		final ForeignKeyDescriptor keyDescriptor = getLoadable().getKeyDescriptor();
-		if ( keyDescriptor.isEmbedded() ) {
+		if ( keyDescriptor.isEmbedded()
+			|| keyDescriptor.getKeyPart().getSingleJdbcMapping().getValueConverter() != null ) {
 			assert keyDescriptor.getJdbcTypeCount() == 1;
-			return loadEmbeddable( keyBeingLoaded, session, keyDescriptor );
+			return loadWithConversion( keyBeingLoaded, session, keyDescriptor );
 		}
 		else {
 			return super.load( keyBeingLoaded, session );
@@ -123,13 +125,13 @@ public PersistentCollection<?> load(Object keyBeingLoaded, SharedSessionContract
 	}
 
 	@AllowReflection
-	private PersistentCollection<?> loadEmbeddable(
+	private PersistentCollection<?> loadWithConversion(
 			Object keyBeingLoaded,
 			SharedSessionContractImplementor session,
 			ForeignKeyDescriptor keyDescriptor) {
 		if ( MULTI_KEY_LOAD_LOGGER.isTraceEnabled() ) {
 			MULTI_KEY_LOAD_LOGGER.trace( "Batch fetching collection: "
-					+ collectionInfoString( getLoadable(), keyBeingLoaded ) );
+										+ collectionInfoString( getLoadable(), keyBeingLoaded ) );
 		}
 
 		final int length = getDomainBatchSize();
@@ -141,14 +143,14 @@ private PersistentCollection<?> loadEmbeddable(
 						.getComponentType(),
 				length
 		);
-		final Object[] embeddedKeys = (Object[]) newInstance( keyDomainType, length );
+		final Object[] domainKeys = (Object[]) newInstance( keyDomainType, length );
 		session.getPersistenceContextInternal().getBatchFetchQueue()
 				.collectBatchLoadableCollectionKeys(
 						length,
 						(index, key) ->
 								keyDescriptor.forEachJdbcValue( key, (i, value, jdbcMapping) -> {
 									keysToInitialize[index] = value;
-									embeddedKeys[index] = key;
+									domainKeys[index] = key;
 								}, session )
 						,
 						keyBeingLoaded,
@@ -163,7 +165,7 @@ private PersistentCollection<?> loadEmbeddable(
 
 		initializeKeys( keyBeingLoaded, keys, session );
 
-		for ( Object initializedKey : embeddedKeys ) {
+		for ( Object initializedKey : domainKeys ) {
 			if ( initializedKey != null ) {
 				finishInitializingKey( initializedKey, session );
 			}
@@ -186,7 +188,7 @@ void initializeKeys(Object key, Object[] keysToInitialize, SharedSessionContract
 		assert jdbcSelectOperation != null;
 		assert jdbcParameter != null;
 
-		final JdbcParameterBindings jdbcParameterBindings = new JdbcParameterBindingsImpl(1);
+		final JdbcParameterBindings jdbcParameterBindings = new JdbcParameterBindingsImpl( 1 );
 		jdbcParameterBindings.addBinding(
 				jdbcParameter,
 				new JdbcParameterBindingImpl( arraySqlTypedMapping.getJdbcMapping(), keysToInitialize )
@@ -219,27 +221,10 @@ void finishInitializingKeys(Object[] keys, SharedSessionContractImplementor sess
 	}
 
 	@Override
-	@AllowReflection
 	Object[] resolveKeysToInitialize(Object keyBeingLoaded, SharedSessionContractImplementor session) {
-		final ForeignKeyDescriptor keyDescriptor = getLoadable().getKeyDescriptor();
-		if( keyDescriptor.isEmbedded()){
-			assert keyDescriptor.getJdbcTypeCount() == 1;
-			final int length = getDomainBatchSize();
-			final Object[] keysToInitialize = (Object[]) newInstance( keyDescriptor.getSingleJdbcMapping().getJdbcJavaType().getJavaTypeClass(), length );
-			session.getPersistenceContextInternal().getBatchFetchQueue()
-					.collectBatchLoadableCollectionKeys(
-							length,
-							(index, key) ->
-								keyDescriptor.forEachJdbcValue( key, (i, value, jdbcMapping) -> {
-									keysToInitialize[index] = value;
-								}, session )
-							,
-							keyBeingLoaded,
-							getLoadable()
-					);
-			// now trim down the array to the number of keys we found
-			return trimIdBatch( length, keysToInitialize );
-		}
+		assert !getLoadable().getKeyDescriptor().isEmbedded()
+			&& getLoadable().getKeyDescriptor().getKeyPart().getSingleJdbcMapping().getValueConverter() == null
+				: "Should use loadWithConversion() instead";
 		return super.resolveKeysToInitialize( keyBeingLoaded, session );
 	}
 }