From 7d80b8a3db20a4e74f4f83b4bd8644d7fc793020 Mon Sep 17 00:00:00 2001
From: Gavin King <gavin@hibernate.org>
Date: Sun, 31 Mar 2024 20:20:08 +0200
Subject: [PATCH] fix to @OrderBy for @Query methods with primary entity type

don't think this is even strictly-speaking required by Jakarta Data
but we will allow it

Signed-off-by: Gavin King <gavin@hibernate.org>
---
 .../annotation/AnnotationMetaEntity.java      | 47 ++++++++++++-------
 1 file changed, 31 insertions(+), 16 deletions(-)

diff --git a/tooling/metamodel-generator/src/main/java/org/hibernate/processor/annotation/AnnotationMetaEntity.java b/tooling/metamodel-generator/src/main/java/org/hibernate/processor/annotation/AnnotationMetaEntity.java
index 317d4cb2bc..6b877d1e8c 100644
--- a/tooling/metamodel-generator/src/main/java/org/hibernate/processor/annotation/AnnotationMetaEntity.java
+++ b/tooling/metamodel-generator/src/main/java/org/hibernate/processor/annotation/AnnotationMetaEntity.java
@@ -1541,23 +1541,26 @@ public class AnnotationMetaEntity extends AnnotationMeta {
 				Diagnostic.Kind.ERROR );
 	}
 
-	private List<OrderBy> orderByList(ExecutableElement method, TypeElement entityType) {
-		final AnnotationMirror orderByList =
-				getAnnotationMirror( method, "jakarta.data.repository.OrderBy.List" );
-		if ( orderByList != null ) {
-			final List<OrderBy> result = new ArrayList<>();
-			@SuppressWarnings("unchecked")
-			final List<AnnotationValue> list = (List<AnnotationValue>)
-					castNonNull( getAnnotationValue( orderByList, "value" ) ).getValue();
-			for ( AnnotationValue element : list ) {
-				result.add( orderByExpression( castNonNull( (AnnotationMirror) element.getValue() ), entityType, method ) );
+	private List<OrderBy> orderByList(ExecutableElement method, TypeElement returnType) {
+		final TypeElement entityType = implicitEntityType( returnType );
+		if ( entityType != null ) {
+			final AnnotationMirror orderByList =
+					getAnnotationMirror( method, "jakarta.data.repository.OrderBy.List" );
+			if ( orderByList != null ) {
+				final List<OrderBy> result = new ArrayList<>();
+				@SuppressWarnings("unchecked")
+				final List<AnnotationValue> list = (List<AnnotationValue>)
+						castNonNull( getAnnotationValue( orderByList, "value" ) ).getValue();
+				for ( AnnotationValue element : list ) {
+					result.add( orderByExpression( castNonNull( (AnnotationMirror) element.getValue() ), entityType, method ) );
+				}
+				return result;
+			}
+			final AnnotationMirror orderBy =
+					getAnnotationMirror( method, "jakarta.data.repository.OrderBy" );
+			if ( orderBy != null ) {
+				return List.of( orderByExpression(orderBy, entityType, method) );
 			}
-			return result;
-		}
-		final AnnotationMirror orderBy =
-				getAnnotationMirror( method, "jakarta.data.repository.OrderBy" );
-		if ( orderBy != null ) {
-			return List.of( orderByExpression(orderBy, entityType, method) );
 		}
 		return emptyList();
 	}
@@ -2214,6 +2217,18 @@ public class AnnotationMetaEntity extends AnnotationMeta {
 		}
 	}
 
+	private @Nullable TypeElement implicitEntityType(@Nullable TypeElement resultType) {
+		if ( resultType != null && hasAnnotation(resultType, ENTITY) ) {
+			return resultType;
+		}
+		else if ( primaryEntity != null ) {
+			return primaryEntity;
+		}
+		else {
+			return null;
+		}
+	}
+
 	private static String entityName(DeclaredType resultType, AnnotationMirror annotation) {
 		final AnnotationValue name = getAnnotationValue(annotation, "name");
 		if (name != null) {