HSEARCH-4869 Fix potential bug where we would use the entity ID type …

…instead of the document ID type in Hibernate ORM loading
hibernate · Sep 25, 2023 · 7265a45 · 7265a45
1 parent 4ad4fc0
commit 7265a45
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 12 deletions.
diff --git a/.../java/org/hibernate/search/mapper/orm/loading/impl/ConditionalExpressionQueryFactory.java b/.../java/org/hibernate/search/mapper/orm/loading/impl/ConditionalExpressionQueryFactory.java
@@ -16,9 +16,11 @@
 public abstract class ConditionalExpressionQueryFactory<E, I> implements TypeQueryFactory<E, I> {
 
 	private static final String TYPES_PARAM_NAME = "HIBERNATE_SEARCH_INCLUDED_TYPES_FILTER";
+	protected final Class<I> uniquePropertyType;
 	protected final String uniquePropertyName;
 
-	public ConditionalExpressionQueryFactory(String uniquePropertyName) {
+	public ConditionalExpressionQueryFactory(Class<I> uniquePropertyType, String uniquePropertyName) {
+		this.uniquePropertyType = uniquePropertyType;
 		this.uniquePropertyName = uniquePropertyName;
 	}
 
@@ -32,12 +34,11 @@ public Query<Long> createQueryForCount(SharedSessionContractImplementor session,
 	}
 
 	@Override
-	@SuppressWarnings("unchecked") // Can't do better here: EntityMappingType has no generics
 	public Query<I> createQueryForIdentifierListing(SharedSessionContractImplementor session, EntityMappingType entityMappingType,
 			Set<? extends Class<? extends E>> includedTypesFilter, ConditionalExpression conditionalExpression) {
 		return createQueryWithConditionalExpression( session,
 				"select e. " + uniquePropertyName + " from " + entityMappingType.getEntityName() + " e",
-				(Class<I>) entityMappingType.getIdentifierMapping().getJavaType().getJavaTypeClass(), "e",
+				uniquePropertyType, "e",
 				includedTypesFilter, conditionalExpression
 		);
 	}

diff --git a/.../src/main/java/org/hibernate/search/mapper/orm/loading/impl/CriteriaTypeQueryFactory.java b/.../src/main/java/org/hibernate/search/mapper/orm/loading/impl/CriteriaTypeQueryFactory.java
@@ -26,16 +26,14 @@ class CriteriaTypeQueryFactory<E, I> extends ConditionalExpressionQueryFactory<E
 
 	public static <E> CriteriaTypeQueryFactory<E, ?> create(EntityDomainType<E> type,
 			String uniquePropertyName) {
-		return new CriteriaTypeQueryFactory<>( type, uniquePropertyName,
-				type.getSingularAttribute( uniquePropertyName ) );
+		return new CriteriaTypeQueryFactory<>( type, type.getSingularAttribute( uniquePropertyName ) );
 	}
 
 	private final EntityDomainType<E> type;
 	private final SingularAttribute<? super E, I> uniqueProperty;
 
-	private CriteriaTypeQueryFactory(EntityDomainType<E> type,
-			String uniquePropertyName, SingularAttribute<? super E, I> uniqueProperty) {
-		super( uniquePropertyName );
+	private CriteriaTypeQueryFactory(EntityDomainType<E> type, SingularAttribute<? super E, I> uniqueProperty) {
+		super( uniqueProperty.getJavaType(), uniqueProperty.getName() );
 		this.type = type;
 		this.uniqueProperty = uniqueProperty;
 	}

diff --git a/...r/orm/src/main/java/org/hibernate/search/mapper/orm/loading/impl/HqlTypeQueryFactory.java b/...r/orm/src/main/java/org/hibernate/search/mapper/orm/loading/impl/HqlTypeQueryFactory.java
@@ -18,8 +18,12 @@ class HqlTypeQueryFactory<E, I> extends ConditionalExpressionQueryFactory<E, I>
 
 	private final EntityMappingType entityMappingType;
 
+	@SuppressWarnings("unchecked") // Can't do better here: EntityMappingType has no generics
 	HqlTypeQueryFactory(EntityMappingType entityMappingType, String uniquePropertyName) {
-		super( uniquePropertyName );
+		super( uniquePropertyName.equals( entityMappingType.getIdentifierMapping().getAttributeName() )
+				? (Class<I>) entityMappingType.getIdentifierMapping().getJavaType().getJavaTypeClass()
+				: (Class<I>) entityMappingType.findAttributeMapping( uniquePropertyName ).getJavaType().getJavaTypeClass(),
+				uniquePropertyName );
 		this.entityMappingType = entityMappingType;
 	}
 
@@ -32,14 +36,12 @@ public Query<Long> createQueryForCount(SharedSessionContractImplementor session,
 				"e", includedTypesFilter );
 	}
 
-	@SuppressWarnings("unchecked")
 	@Override
 	public Query<I> createQueryForIdentifierListing(SharedSessionContractImplementor session,
 			Set<? extends Class<? extends E>> includedTypesFilter) {
 		return createQueryWithTypesFilter( session,
 				"select e. " + uniquePropertyName + " from " + entityMappingType.getEntityName() + " e",
-				(Class<I>) entityMappingType.getIdentifierMapping().getJavaType().getJavaTypeClass(),
-				"e", includedTypesFilter );
+				uniquePropertyType, "e", includedTypesFilter );
 	}
 
 	@SuppressWarnings("unchecked")