HHH-16633 use unwrap() in preference to typecasts

perhaps it would be better to preemptively unwrap the
EntityManager to Session, because this approach results
in raw query types (but we don't really care since we
already know it's sound)

tooling/metamodel-generator/src/main/java/org/hibernate/jpamodelgen/annotation/CriteriaFinderMethod.java

@@ -13,8 +13,6 @@ import org.hibernate.jpamodelgen.util.Constants;
 
 import java.util.List;
 
-import static org.hibernate.internal.util.StringHelper.join;
-
 /**
  * @author Gavin King
  */
@@ -60,6 +58,8 @@ public class CriteriaFinderMethod implements MetaAttribute {
 
 	@Override
 	public String getAttributeDeclarationString() {
+		final boolean usingEntityManager = Constants.ENTITY_MANAGER.equals(sessionType);
+
 		StringBuilder declaration = new StringBuilder();
 		declaration
 				.append("\n/**\n * @see ")
@@ -67,7 +67,7 @@ public class CriteriaFinderMethod implements MetaAttribute {
 				.append("#")
 				.append(methodName)
 				.append("(")
-				.append(join(",", paramTypes.stream().map(this::strip).map(annotationMetaEntity::importType).toArray()))
+				.append(parameterList())
 				.append(")")
 				.append("\n **/\n");
 		if ( belongsToDao ) {
@@ -109,7 +109,7 @@ public class CriteriaFinderMethod implements MetaAttribute {
 		declaration
 				.append(") {")
 				.append("\n\tvar builder = entityManager")
-				.append(Constants.ENTITY_MANAGER.equals(sessionType)
+				.append(usingEntityManager
 						? ".getEntityManagerFactory()"
 						: ".getFactory()")
 				.append(".getCriteriaBuilder();")
@@ -135,27 +135,19 @@ public class CriteriaFinderMethod implements MetaAttribute {
 					.append(paramNames.get(i))
 					.append(")");
 		}
-
 		declaration
 				.append("\n\t);")
-				.append("\n\treturn ");
+				.append("\n\treturn entityManager.createQuery(query)");
-		if ( containerType != null
+		boolean hasEnabledFetchProfiles = !fetchProfiles.isEmpty();
-				&& Constants.ENTITY_MANAGER.equals(sessionType)
+		boolean hasNativeReturnType = containerType != null && containerType.startsWith("org.hibernate");
-				&& containerType.startsWith("org.hibernate") ) {
+		boolean unwrap =
+				( hasEnabledFetchProfiles || hasNativeReturnType )
+						&& usingEntityManager;
+		if ( unwrap ) {
 			declaration
-					.append("(")
+					.append("\n\t\t\t.unwrap(")
-					.append(type)
+					.append(annotationMetaEntity.importType(Constants.HIB_SELECTION_QUERY))
-					.append(") ");
+					.append(".class)");
-		}
-		declaration
-				.append("entityManager.createQuery(query)");
-		if ( !fetchProfiles.isEmpty() ) {
-			if ( Constants.ENTITY_MANAGER.equals(sessionType) ) {
-				declaration
-						.append("\n\t\t\t.unwrap(")
-						.append(annotationMetaEntity.importType(Constants.HIB_SELECTION_QUERY))
-						.append(".class)");
-			}
 		}
 		for ( String profile : fetchProfiles ) {
 			declaration
@@ -163,14 +155,17 @@ public class CriteriaFinderMethod implements MetaAttribute {
 					.append(profile)
 					.append(")");
 		}
-		if ( !fetchProfiles.isEmpty() ) {
-			declaration.append("\n\t\t\t");
-		}
 		if ( containerType == null) {
+			if ( unwrap || hasEnabledFetchProfiles) {
+				declaration.append("\n\t\t\t");
+			}
 			declaration
 					.append(".getSingleResult()");
 		}
 		else if ( containerType.equals(Constants.LIST) ) {
+			if ( unwrap || hasEnabledFetchProfiles) {
+				declaration.append("\n\t\t\t");
+			}
 			declaration
 					.append(".getResultList()");
 		}
@@ -179,6 +174,14 @@ public class CriteriaFinderMethod implements MetaAttribute {
 		return declaration.toString();
 	}
 
+	private String parameterList() {
+		return paramTypes.stream()
+				.map(this::strip)
+				.map(annotationMetaEntity::importType)
+				.reduce((x, y) -> x + y)
+				.orElse("");
+	}
+
 	private String strip(String type) {
 		int index = type.indexOf("<");
 		String stripped = index > 0 ? type.substring(0, index) : type;

tooling/metamodel-generator/src/main/java/org/hibernate/jpamodelgen/annotation/IdFinderMethod.java

@@ -54,6 +54,8 @@ public class IdFinderMethod implements MetaAttribute {
 
 	@Override
 	public String getAttributeDeclarationString() {
+		final boolean usingEntityManager = Constants.ENTITY_MANAGER.equals(sessionType);
+
 		StringBuilder declaration = new StringBuilder();
 		declaration
 				.append("\n/**\n * @see ")
@@ -99,7 +101,7 @@ public class IdFinderMethod implements MetaAttribute {
 					.append("\n}");
 		}
 		else {
-			if ( Constants.ENTITY_MANAGER.equals(sessionType) ) {
+			if ( usingEntityManager ) {
 				declaration
 						.append(".unwrap(")
 						.append(annotationMetaEntity.importType(Constants.HIB_SESSION))

tooling/metamodel-generator/src/main/java/org/hibernate/jpamodelgen/annotation/QueryMethod.java

@@ -71,6 +71,8 @@ public class QueryMethod implements MetaAttribute {
 
 	@Override
 	public String getAttributeDeclarationString() {
+		final boolean usingEntityManager = Constants.ENTITY_MANAGER.equals(sessionType);
+
 		List<String> paramTypes = this.paramTypes.stream()
 				.map(ptype -> isOrderParam(ptype) && ptype.endsWith("[]")
 						? ptype.substring(0, ptype.length()-2) + "..."
@@ -146,11 +148,13 @@ public class QueryMethod implements MetaAttribute {
 				.append(")")
 				.append(" {")
 				.append("\n\treturn ");
-		if ( isNative && returnTypeName != null
+		if ( isNative && returnTypeName != null && containerTypeName == null
-				|| containerTypeName != null
+				&& usingEntityManager) {
-					&& Constants.ENTITY_MANAGER.equals(sessionType)
+			// EntityManager.createNativeQuery() does not return TypedQuery,
-					&& containerTypeName.startsWith("org.hibernate") ) {
+			// so we need to cast to the entity type
-			declaration.append("(").append(type).append(") ");
+			declaration.append("(")
+					.append(type)
+					.append(") ");
 		}
 		declaration
 				.append("entityManager.")
@@ -164,7 +168,7 @@ public class QueryMethod implements MetaAttribute {
 					.append(".class");
 		}
 		declaration.append(")");
-		boolean unwrapped = !Constants.ENTITY_MANAGER.equals(sessionType);
+		boolean unwrapped = !usingEntityManager;
 		for (int i = 1; i <= paramNames.size(); i++) {
 			String param = paramNames.get(i-1);
 			String ptype = paramTypes.get(i-1);
@@ -219,6 +223,17 @@ public class QueryMethod implements MetaAttribute {
 			declaration
 					.append("\n\t\t\t.getResultList()");
 		}
+		else {
+			if ( usingEntityManager && !unwrapped
+					&& ( containerTypeName.startsWith("org.hibernate")
+						|| isNative && returnTypeName != null ) ) {
+				declaration
+						.append("\n\t\t\t.unwrap(")
+						.append(annotationMetaEntity.importType(containerTypeName))
+						.append(".class)");
+
+			}
+		}
 		declaration.append(";\n}");
 		return declaration.toString();
 	}

tooling/metamodel-generator/src/test/java/org/hibernate/jpamodelgen/test/hqlsql/QueryMethodTest.java

@@ -19,7 +19,7 @@ import static org.hibernate.jpamodelgen.test.util.TestUtil.assertMetamodelClassG
 public class QueryMethodTest extends CompilationTest {
 	@Test
 	@WithClasses({ Book.class, Dao.class })
-	public void testGeneratedAnnotationNotGenerated() {
+	public void testQueryMethod() {
 		System.out.println( TestUtil.getMetaModelSourceAsString( Dao.class ) );
 		assertMetamodelClassGeneratedFor( Book.class );
 		assertMetamodelClassGeneratedFor( Dao.class );