Enable @QueryProjection annotation support for KSP. (#997)

IceBlizz6 · web-flow · commit be160696ee7c · 2025-03-20T22:14:52.000+01:00
diff --git a/querydsl-examples/querydsl-example-ksp-codegen/src/main/kotlin/com/querydsl/example/ksp/PersonDTO.kt b/querydsl-examples/querydsl-example-ksp-codegen/src/main/kotlin/com/querydsl/example/ksp/PersonDTO.kt
@@ -0,0 +1,11 @@
+package com.querydsl.example.ksp
+
+import com.querydsl.core.annotations.QueryProjection
+
+data class PersonClassConstructorDTO @QueryProjection constructor(
+    val id: Int,
+    val name: String,
+)
+
+@QueryProjection
+data class PersonClassDTO (val id: Int, val name: String)
diff --git a/querydsl-examples/querydsl-example-ksp-codegen/src/test/kotlin/Tests.kt b/querydsl-examples/querydsl-example-ksp-codegen/src/test/kotlin/Tests.kt
@@ -3,6 +3,8 @@ import com.querydsl.example.ksp.CatType
 import com.querydsl.example.ksp.Person
 import com.querydsl.example.ksp.QCat
 import com.querydsl.example.ksp.QPerson
+import com.querydsl.example.ksp.QPersonClassDTO
+import com.querydsl.example.ksp.QPersonClassConstructorDTO
 import com.querydsl.jpa.impl.JPAQueryFactory
 import jakarta.persistence.EntityManagerFactory
 import org.hibernate.cfg.AvailableSettings
@@ -111,6 +113,55 @@ class Tests {
         }
     }
 
+    @Test
+    fun `select dto`() {
+        val emf = initialize()
+
+        run {
+            val em = emf.createEntityManager()
+            em.transaction.begin()
+            em.persist(Person(424, "John Smith"))
+            em.transaction.commit()
+            em.close()
+        }
+
+        run {
+            val em = emf.createEntityManager()
+            val queryFactory = JPAQueryFactory(em)
+            val q = QPerson.person
+            val personDTO = queryFactory
+                .select(QPersonClassConstructorDTO(q.id, q.name))
+                .from(q)
+                .where(q.name.eq("John Smith"))
+                .fetchOne()
+            if (personDTO == null) {
+                fail<Any>("No personDTO was returned")
+            } else {
+                assertThat(personDTO.id).isEqualTo(424)
+                assertThat(personDTO.name).isEqualTo("John Smith")
+            }
+            em.close()
+        }
+
+        run {
+            val em = emf.createEntityManager()
+            val queryFactory = JPAQueryFactory(em)
+            val q = QPerson.person
+            val personDTO = queryFactory
+                .select(QPersonClassDTO(q.id, q.name))
+                .from(q)
+                .where(q.name.eq("John Smith"))
+                .fetchOne()
+            if (personDTO == null) {
+                fail<Any>("No personDTO was returned")
+            } else {
+                assertThat(personDTO.id).isEqualTo(424)
+                assertThat(personDTO.name).isEqualTo("John Smith")
+            }
+            em.close()
+        }
+    }
+
     private fun initialize(): EntityManagerFactory {
         val configuration = Configuration()
             .setProperty(AvailableSettings.JAKARTA_JDBC_DRIVER, org.h2.Driver::class.qualifiedName!!)
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
@@ -1,5 +1,6 @@
 package com.querydsl.ksp.codegen
 
+import com.google.devtools.ksp.isConstructor
 import com.google.devtools.ksp.processing.CodeGenerator
 import com.google.devtools.ksp.processing.Resolver
 import com.google.devtools.ksp.processing.SymbolProcessor
@@ -17,9 +18,32 @@ class QueryDslProcessor(
         if (settings.enable) {
             QueryModelType.entries.forEach { type ->
                 resolver.getSymbolsWithAnnotation(type.associatedAnnotation)
-                    .map { it as KSClassDeclaration }
-                    .filter { isIncluded(it) }
-                    .forEach { declaration -> typeProcessor.add(declaration, type) }
+                    .map { 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)
+                                        declaration.parent as? KSClassDeclaration
+                                            ?: error(errorMessage)
+                                    }
+
+                                    is KSClassDeclaration -> declaration
+                                    else -> error(errorMessage)
+                                }
+                            }
+
+                            else -> declaration as KSClassDeclaration
+                        }
+                    }
+                    .filter {
+                        isIncluded(it)
+                    }
+                    .forEach { declaration ->
+                        typeProcessor.add(declaration, type)
+                    }
             }
         }
         return emptyList()
diff --git a/querydsl-tooling/querydsl-ksp-codegen/src/main/kotlin/com/querydsl/ksp/codegen/QueryModelRenderer.kt b/querydsl-tooling/querydsl-ksp-codegen/src/main/kotlin/com/querydsl/ksp/codegen/QueryModelRenderer.kt
@@ -1,25 +1,35 @@
 package com.querydsl.ksp.codegen
 
+import com.querydsl.core.types.ConstructorExpression
 import com.querydsl.core.types.Path
 import com.querydsl.core.types.PathMetadata
+import com.querydsl.core.types.Expression
 import com.querydsl.core.types.dsl.*
 import com.querydsl.ksp.codegen.Naming.toCamelCase
 import com.squareup.kotlinpoet.*
 import com.squareup.kotlinpoet.ParameterizedTypeName.Companion.parameterizedBy
 
 object QueryModelRenderer {
     fun render(model: QueryModel): TypeSpec {
-        return TypeSpec.classBuilder(model.className)
-            .setEntitySuperclass(model)
-            .addSuperProperty(model)
-            .addProperties(model)
-            .constructorForPath(model)
-            .constructorForMetadata(model)
-            .constructorForVariable(model)
-            .constructorForTypeMetadata(model)
-            .addInitializerCompanionObject(model)
-            .addInheritedProperties(model)
-            .build()
+        return when (model.type) {
+            QueryModelType.QUERY_PROJECTION -> TypeSpec.classBuilder(model.className)
+                .setPrimaryConstructor(model)
+                .setEntitySuperclass(model)
+                .addSuperConstructorParameter(model)
+                .build()
+
+            else -> TypeSpec.classBuilder(model.className)
+                .setEntitySuperclass(model)
+                .addSuperProperty(model)
+                .addProperties(model)
+                .constructorForPath(model)
+                .constructorForMetadata(model)
+                .constructorForVariable(model)
+                .constructorForTypeMetadata(model)
+                .addInitializerCompanionObject(model)
+                .addInheritedProperties(model)
+                .build()
+        }
     }
 
     private fun TypeSpec.Builder.setEntitySuperclass(model: QueryModel): TypeSpec.Builder {
@@ -33,6 +43,7 @@ object QueryModelRenderer {
             when (model.type) {
                 QueryModelType.ENTITY, QueryModelType.SUPERCLASS -> EntityPathBase::class.asClassName().parameterizedBy(constraint)
                 QueryModelType.EMBEDDABLE -> BeanPath::class.asClassName().parameterizedBy(constraint)
+                QueryModelType.QUERY_PROJECTION -> ConstructorExpression::class.asClassName().parameterizedBy(constraint)
             }
         )
         return this
@@ -221,4 +232,28 @@ object QueryModelRenderer {
         addType(companionObject)
         return this
     }
+
+    private fun TypeSpec.Builder.setPrimaryConstructor(model: QueryModel): TypeSpec.Builder {
+        val constructorSpec = FunSpec.constructorBuilder().apply {
+            model.properties.forEach {
+                addParameter(
+                    it.name,
+                    Expression::class.asClassName().parameterizedBy(it.type.originalTypeName)
+                )
+            }
+        }.build()
+        primaryConstructor(constructorSpec)
+        return this
+    }
+
+    private fun TypeSpec.Builder.addSuperConstructorParameter(model: QueryModel): TypeSpec.Builder {
+        val paramTypes = model.properties.joinToString(", ", prefix = "arrayOf(", postfix = ")") {
+            "${it.type.originalClassName}::class.java"
+        }
+        val paramNames = model.properties.joinToString(", ") { it.name }
+        addSuperclassConstructorParameter(
+            "${model.originalClassName}::class.java, $paramTypes, $paramNames"
+        )
+        return this
+    }
 }
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
@@ -2,6 +2,7 @@ package com.querydsl.ksp.codegen
 
 import com.google.devtools.ksp.symbol.KSAnnotation
 import com.google.devtools.ksp.symbol.KSClassDeclaration
+import com.querydsl.core.annotations.QueryProjection
 import com.squareup.kotlinpoet.ksp.toTypeName
 import jakarta.persistence.Embeddable
 import jakarta.persistence.Entity
@@ -12,7 +13,8 @@ enum class QueryModelType(
 ) {
     ENTITY(Entity::class.qualifiedName!!),
     EMBEDDABLE(Embeddable::class.qualifiedName!!),
-    SUPERCLASS(MappedSuperclass::class.qualifiedName!!);
+    SUPERCLASS(MappedSuperclass::class.qualifiedName!!),
+    QUERY_PROJECTION(QueryProjection::class.qualifiedName!!);
 
     companion object {
         fun autodetect(classDeclaration: KSClassDeclaration): QueryModelType? {
diff --git a/querydsl-tooling/querydsl-ksp-codegen/src/test/kotlin/RenderTest.kt b/querydsl-tooling/querydsl-ksp-codegen/src/test/kotlin/RenderTest.kt
@@ -176,6 +176,31 @@ class RenderTest {
             val features: com.querydsl.core.types.dsl.SetPath<kotlin.String, com.querydsl.core.types.dsl.StringPath> = createSet("features", kotlin.String::class.java, com.querydsl.core.types.dsl.StringPath::class.java, null)
         """.trimIndent())
     }
+
+    @Test
+    fun queryProjection() {
+        val model = QueryModel(
+            originalClassName = ClassName("", "CatDTO"),
+            typeParameterCount = 0,
+            className = ClassName("", "QCatDTO"),
+            type = QueryModelType.QUERY_PROJECTION,
+            mockk()
+        )
+        val properties = listOf(
+            QProperty("id", QPropertyType.Simple(SimpleType.QNumber(Int::class.asClassName()))),
+            QProperty("name", QPropertyType.Simple(SimpleType.QString)),
+        )
+        model.properties.addAll(properties)
+        val typeSpec = QueryModelRenderer.render(model)
+        val code = typeSpec.toString()
+        code.assertCompiles()
+        code.assertContainAll("""
+            public class QPersonDTO(
+              id: com.querydsl.core.types.Expression<kotlin.Int>,
+              name: com.querydsl.core.types.Expression<kotlin.String>,
+            ) : com.querydsl.core.types.ConstructorExpression<CatDTO>(CatDTO::class.java, arrayOf(kotlin.Int::class.java, kotlin.String::class.java), id, name)
+        """.trimIndent())
+    }
 }
 
 private fun String.assertLines(expected: String) {
@@ -234,3 +259,5 @@ class QAnimal(
 }
 
 class Cat
+
+class CatDTO
