From 1e6282e64ba00ee470fe67ca23d83e623ffc76fb Mon Sep 17 00:00:00 2001
From: IceBlizz6 <iceblizz6@proton.me>
Date: Tue, 9 Sep 2025 10:53:27 +0200
Subject: [PATCH] QueryEntity annotation support

---
 .../querydsl/ksp/codegen/QueryDslProcessor.kt | 69 ++++++++++---------
 .../querydsl/ksp/codegen/QueryModelType.kt    | 22 +++---
 2 files changed, 50 insertions(+), 41 deletions(-)

diff --git a/querydsl-tooling/querydsl-ksp-codegen/src/main/kotlin/com/querydsl/ksp/codegen/QueryDslProcessor.kt b/querydsl-tooling/querydsl-ksp-codegen/src/main/kotlin/com/querydsl/ksp/codegen/QueryDslProcessor.kt
index 1308d82519..8933601ef7 100644
--- a/querydsl-tooling/querydsl-ksp-codegen/src/main/kotlin/com/querydsl/ksp/codegen/QueryDslProcessor.kt
+++ b/querydsl-tooling/querydsl-ksp-codegen/src/main/kotlin/com/querydsl/ksp/codegen/QueryDslProcessor.kt
@@ -17,39 +17,42 @@ class QueryDslProcessor(
     override fun process(resolver: Resolver): List<KSAnnotated> {
         if (settings.enable) {
             QueryModelType.entries.forEach { type ->
-                resolver.getSymbolsWithAnnotation(type.associatedAnnotation)
-                    .forEach { declaration ->
-                        when {
-                            type == QueryModelType.QUERY_PROJECTION -> {
-                                val errorMessage = "${type.associatedAnnotation} annotation" +
-                                        " must be declared on a constructor function or class"
-                                when (declaration) {
-                                    is KSFunctionDeclaration -> {
-                                        if (!declaration.isConstructor()) error(errorMessage)
-                                        val parentDeclaration = declaration.parent as? KSClassDeclaration
-                                            ?: error(errorMessage)
-                                        if (isIncluded(parentDeclaration)) {
-                                            typeProcessor.addConstructor(parentDeclaration, declaration)
-                                        }
-                                    }
-                                    is KSClassDeclaration -> {
-                                        if (isIncluded(declaration)) {
-                                            typeProcessor.addClass(declaration, type)
-                                        }
-                                    }
-                                    else -> error(errorMessage)
-                                }
-                            }
-                            declaration is KSClassDeclaration -> {
-                                if (isIncluded(declaration)) {
-                                    typeProcessor.addClass(declaration, type)
-                                }
-                            }
-                            else -> {
-                                error("Annotated element was expected to be class or constructor, instead got ${declaration}")
-                            }
-                        }
-                    }
+				type.associatedAnnotations.forEach { associatedAnnotation ->
+					resolver.getSymbolsWithAnnotation(associatedAnnotation)
+						.forEach { declaration ->
+							when {
+								type == QueryModelType.QUERY_PROJECTION -> {
+									val errorMessage = "$associatedAnnotation annotation" +
+											" must be declared on a constructor function or class"
+									when (declaration) {
+										is KSFunctionDeclaration -> {
+											if (!declaration.isConstructor()) error(errorMessage)
+											val parentDeclaration = declaration.parent as? KSClassDeclaration
+												?: error(errorMessage)
+											if (isIncluded(parentDeclaration)) {
+												typeProcessor.addConstructor(parentDeclaration, declaration)
+											}
+										}
+										is KSClassDeclaration -> {
+											if (isIncluded(declaration)) {
+												typeProcessor.addClass(declaration, type)
+											}
+										}
+										else -> error(errorMessage)
+									}
+								}
+								declaration is KSClassDeclaration -> {
+									if (isIncluded(declaration)) {
+										typeProcessor.addClass(declaration, type)
+									}
+								}
+								else -> {
+									error("Annotated element was expected to be class or constructor, instead got ${declaration}")
+								}
+							}
+						}
+				}
+
             }
         }
         return emptyList()
diff --git a/querydsl-tooling/querydsl-ksp-codegen/src/main/kotlin/com/querydsl/ksp/codegen/QueryModelType.kt b/querydsl-tooling/querydsl-ksp-codegen/src/main/kotlin/com/querydsl/ksp/codegen/QueryModelType.kt
index 7fe0313374..8eed2f4c8d 100644
--- a/querydsl-tooling/querydsl-ksp-codegen/src/main/kotlin/com/querydsl/ksp/codegen/QueryModelType.kt
+++ b/querydsl-tooling/querydsl-ksp-codegen/src/main/kotlin/com/querydsl/ksp/codegen/QueryModelType.kt
@@ -9,20 +9,26 @@ import jakarta.persistence.Entity
 import jakarta.persistence.MappedSuperclass
 
 enum class QueryModelType(
-    val associatedAnnotation: String
+    val associatedAnnotations: List<String>
 ) {
-    ENTITY(Entity::class.qualifiedName!!),
-    EMBEDDABLE(Embeddable::class.qualifiedName!!),
-    SUPERCLASS(MappedSuperclass::class.qualifiedName!!),
-    QUERY_PROJECTION(QueryProjection::class.qualifiedName!!);
+    ENTITY(
+		listOf(
+			Entity::class.qualifiedName!!,
+			"com.querydsl.core.annotations.QueryEntity",
+			"org.springframework.data.mongodb.core.mapping.Document"
+		)
+	),
+    EMBEDDABLE(listOf(Embeddable::class.qualifiedName!!)),
+    SUPERCLASS(listOf(MappedSuperclass::class.qualifiedName!!)),
+    QUERY_PROJECTION(listOf(QueryProjection::class.qualifiedName!!));
 
     companion object {
         fun autodetect(classDeclaration: KSClassDeclaration): QueryModelType? {
             for (annotation in classDeclaration.annotations) {
                 for (type in QueryModelType.entries) {
-                    if (annotation.isEqualTo(type.associatedAnnotation)) {
-                        return type
-                    }
+					if (type.associatedAnnotations.any { ann -> annotation.isEqualTo(ann) }) {
+						return type
+					}
                 }
             }
             return null