HHH-18979 allow Hibernate Restriction as a query method parameter

Author: gavinking

tooling/metamodel-generator/src/main/java/org/hibernate/processor/annotation/AbstractQueryMethod.java

@@ -19,6 +19,7 @@
 import static org.hibernate.processor.util.Constants.HIB_ORDER;
 import static org.hibernate.processor.util.Constants.HIB_PAGE;
 import static org.hibernate.processor.util.Constants.HIB_QUERY;
+import static org.hibernate.processor.util.Constants.HIB_RESTRICTION;
 import static org.hibernate.processor.util.Constants.HIB_SELECTION_QUERY;
 import static org.hibernate.processor.util.Constants.HIB_SORT_DIRECTION;
 import static org.hibernate.processor.util.Constants.JD_CURSORED_PAGE;
@@ -286,6 +287,21 @@ void handlePageParameters(
 		}
 	}
 
+	void handleRestrictionParameters(
+			StringBuilder declaration, List<String> paramTypes,
+			@Nullable String containerType) {
+		for ( int i = 0; i < paramNames.size(); i ++ ) {
+			final String paramName = paramNames.get(i);
+			final String paramType = paramTypes.get(i);
+			if ( isRestrictionParam(paramType) ) {
+				declaration
+						.append("\t\t\t.addRestriction(")
+						.append(paramName)
+						.append(")\n");
+			}
+		}
+	}
+
 	void convertExceptions(StringBuilder declaration) {
 		if ( dataRepository ) {
 			if ( !isReactive() ) {
@@ -329,6 +345,7 @@ void unwrapQuery(StringBuilder declaration, boolean unwrapped) {
 
 	static boolean isSpecialParam(String parameterType) {
 		return isPageParam(parameterType)
+			|| isRestrictionParam(parameterType)
 			|| isOrderParam(parameterType)
 			|| isKeyedPageParam(parameterType)
 			|| isSessionParameter(parameterType);
@@ -351,6 +368,10 @@ static boolean isOrderParam(String parameterType) {
 			|| parameterType.startsWith(JD_ORDER);
 	}
 
+	static boolean isRestrictionParam(String parameterType) {
+		return parameterType.startsWith(HIB_RESTRICTION);
+	}
+
 	static boolean isJakartaCursoredPage(@Nullable String containerType) {
 		return JD_CURSORED_PAGE.equals(containerType);
 	}

tooling/metamodel-generator/src/main/java/org/hibernate/processor/annotation/AnnotationMetaEntity.java

@@ -75,6 +75,7 @@
 import static org.hibernate.grammars.hql.HqlLexer.WHERE;
 import static org.hibernate.internal.util.StringHelper.qualify;
 import static org.hibernate.internal.util.StringHelper.unqualify;
+import static org.hibernate.processor.annotation.AbstractQueryMethod.isRestrictionParam;
 import static org.hibernate.processor.annotation.AbstractQueryMethod.isSessionParameter;
 import static org.hibernate.processor.annotation.AbstractQueryMethod.isSpecialParam;
 import static org.hibernate.processor.annotation.QueryMethod.isOrderParam;
@@ -1669,7 +1670,8 @@ private void createCriteriaFinder(
 	private void checkFinderParameter(TypeElement entity, VariableElement parameter) {
 		final Types types = context.getTypeUtils();
 		final TypeMirror parameterType = parameterType(parameter);
-		if ( isOrderParam( typeName(parameterType) ) ) {
+		final String typeName = typeName( parameterType );
+		if ( isOrderParam( typeName ) || isRestrictionParam( typeName ) ) {
 			final TypeMirror typeArgument = getTypeArgument( parameterType, entity );
 			if ( typeArgument == null ) {
 				missingTypeArgError( entity.getSimpleName().toString(), parameter );
@@ -1790,6 +1792,7 @@ private OrderBy orderByExpression(AnnotationMirror orderBy, TypeElement entityTy
 						}
 						return null;
 					case HIB_ORDER:
+					case HIB_RESTRICTION:
 					case JD_SORT:
 					case JD_ORDER:
 						for ( TypeMirror arg : type.getTypeArguments() ) {

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

@@ -58,6 +58,7 @@ void executeQuery(StringBuilder declaration, List<String> paramTypes) {
 		tryReturn( declaration, paramTypes, containerType );
 		castResult( declaration );
 		createQuery( declaration );
+		handleRestrictionParameters( declaration, paramTypes, containerType );
 		handlePageParameters( declaration, paramTypes, containerType );
 		boolean unwrapped = !isUsingEntityManager();
 		unwrapped = enableFetchProfile( declaration, unwrapped );

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

@@ -97,6 +97,7 @@ public String getAttributeDeclarationString() {
 		tryReturn( declaration, paramTypes, containerType );
 		castResult( declaration );
 		createQuery( declaration );
+		handleRestrictionParameters( declaration, paramTypes, containerType );
 		setParameters( declaration, paramTypes, "");
 		handlePageParameters( declaration, paramTypes, containerType );
 		boolean unwrapped = !isUsingEntityManager();

tooling/metamodel-generator/src/main/java/org/hibernate/processor/util/Constants.java

@@ -91,6 +91,7 @@ public final class Constants {
 	public static final String HIB_KEYED_PAGE = "org.hibernate.query.KeyedPage";
 	public static final String HIB_KEYED_RESULT_LIST = "org.hibernate.query.KeyedResultList";
 	public static final String HIB_SORT_DIRECTION = "org.hibernate.query.SortDirection";
+	public static final String HIB_RESTRICTION = "org.hibernate.query.Restriction";
 
 	public static final String CHECK_HQL = "org.hibernate.annotations.processing.CheckHQL";