Fix return type checks for batch/update with Immutable Entity

Author: xterao

src/main/kotlin/org/domaframework/doma/intellij/inspection/dao/processor/returntype/BatchAnnotationReturnTypeCheckProcessor.kt

@@ -21,6 +21,7 @@ import com.intellij.psi.impl.source.PsiClassReferenceType
 import org.domaframework.doma.intellij.common.psi.PsiDaoMethod
 import org.domaframework.doma.intellij.common.sql.PsiClassTypeUtil
 import org.domaframework.doma.intellij.common.validation.result.ValidationResult
+import org.domaframework.doma.intellij.extension.getJavaClazz
 
 /**
  * Processor for checking the return type of batch annotations.
@@ -39,11 +40,6 @@ class BatchAnnotationReturnTypeCheckProcessor(
      */
     override fun checkReturnType(): ValidationResult? {
         val methodOtherReturnType = PsiTypes.intType().createArrayType()
-        if (psiDaoMethod.useSqlAnnotation() || psiDaoMethod.sqlFileOption) {
-            return generatePsiTypeReturnTypeResult(methodOtherReturnType)
-        }
-
-        // Check if it has an immutable entity parameter
         val parameters = method.parameterList.parameters
         val immutableEntityParam = parameters.firstOrNull() ?: return null
 
@@ -52,6 +48,19 @@ class BatchAnnotationReturnTypeCheckProcessor(
         val nestPsiType = parameterType.reference.typeParameters.firstOrNull() ?: return null
         val nestClass: PsiType = PsiClassTypeUtil.convertOptionalType(nestPsiType, project)
 
+        if (psiDaoMethod.useSqlAnnotation() || psiDaoMethod.sqlFileOption) {
+            nestClass.let { methodParam ->
+                val paramTypeName = methodParam.canonicalText
+                project.getJavaClazz(paramTypeName)?.let {
+                    if (isImmutableEntity(paramTypeName)) {
+                        return checkReturnTypeImmutableEntity(nestClass)
+                    }
+                }
+            }
+            return generatePsiTypeReturnTypeResult(methodOtherReturnType)
+        }
+
+        // Check if it has an immutable entity parameter
         if (isImmutableEntity(nestClass.canonicalText)) {
             return checkReturnTypeImmutableEntity(nestClass)
         }

src/main/kotlin/org/domaframework/doma/intellij/inspection/dao/processor/returntype/UpdateAnnotationReturnTypeCheckProcessor.kt

@@ -15,14 +15,18 @@
  */
 package org.domaframework.doma.intellij.inspection.dao.processor.returntype
 
+import com.intellij.psi.PsiClass
+import com.intellij.psi.PsiClassType
 import com.intellij.psi.PsiParameter
 import com.intellij.psi.PsiTypes
 import org.domaframework.doma.intellij.common.psi.PsiDaoMethod
 import org.domaframework.doma.intellij.common.sql.PsiClassTypeUtil
+import org.domaframework.doma.intellij.common.util.DomaClassName
 import org.domaframework.doma.intellij.common.validation.result.ValidationResult
 import org.domaframework.doma.intellij.common.validation.result.ValidationReturnTypeUpdateReturningResult
 import org.domaframework.doma.intellij.extension.getJavaClazz
 import org.domaframework.doma.intellij.extension.psi.isEntity
+import org.domaframework.doma.intellij.extension.psi.psiClassType
 
 /**
  * Processor for checking the return type of update-related annotations in DAO methods.
@@ -41,14 +45,23 @@ class UpdateAnnotationReturnTypeCheckProcessor(
      */
     override fun checkReturnType(): ValidationResult? {
         val methodOtherReturnType = PsiTypes.intType()
+        val parameters = method.parameterList.parameters
+        val immutableEntityParam =
+            parameters.firstOrNull()
+                ?: return generatePsiTypeReturnTypeResult(methodOtherReturnType)
+
         if (psiDaoMethod.useSqlAnnotation() || psiDaoMethod.sqlFileOption) {
+            immutableEntityParam.let { methodParam ->
+                val paramTypeName = methodParam.type.canonicalText
+                project.getJavaClazz(paramTypeName)?.let {
+                    if (isImmutableEntity(paramTypeName)) {
+                        return checkReturnTypeImmutableEntity(immutableEntityParam)
+                    }
+                }
+            }
             return generatePsiTypeReturnTypeResult(methodOtherReturnType)
         }
 
-        val parameters = method.parameterList.parameters
-        val immutableEntityParam =
-            parameters.firstOrNull() ?: return null
-
         // Check if the method is annotated with @Returning
         if (hasReturingOption()) {
             return checkReturnTypeWithReturning(immutableEntityParam)
@@ -82,7 +95,7 @@ class UpdateAnnotationReturnTypeCheckProcessor(
             PsiClassTypeUtil.convertOptionalType(returnType, project)
         val returnTypeClass = project.getJavaClazz(checkReturnType.canonicalText)
 
-        return if (returnTypeClass?.isEntity() != true || returnType.canonicalText != paramClass.type.canonicalText) {
+        return if (!validateReturnType(returnTypeClass, paramClass)) {
             ValidationReturnTypeUpdateReturningResult(
                 paramClass.type.presentableText,
                 method.nameIdentifier,
@@ -93,6 +106,23 @@ class UpdateAnnotationReturnTypeCheckProcessor(
         }
     }
 
+    private fun validateReturnType(
+        returnTypeClass: PsiClass?,
+        paramClass: PsiParameter,
+    ): Boolean {
+        if (returnTypeClass?.isEntity() != true) return false
+
+        if (DomaClassName.OPTIONAL.isTargetClassNameStartsWith(paramClass.type.canonicalText)) {
+            val optionalType = paramClass.type as? PsiClassType
+            val optionalParam =
+                optionalType?.parameters?.firstOrNull()
+                    ?: return false
+            return optionalParam.canonicalText == paramClass.type.canonicalText
+        }
+
+        return returnTypeClass.psiClassType.canonicalText == paramClass.type.canonicalText
+    }
+
     /**
      * Checks the return type when an immutable entity parameter is present.
      *

src/test/testData/src/main/java/doma/example/dao/inspection/returntype/BatchReturnTypeTestDao.java

@@ -13,6 +13,9 @@ public interface BatchReturnTypeTestDao {
     @BatchUpdate
     String <error descr="The return type must be \"int[]\"">batchUpdateReturnsString</error>(List<Packet> e);
 
+    @BatchUpdate(sqlFile = true)
+    int[] <error descr="If a method annotated with @BatchUpdate targets immutable entities for insertion, the return type must be BatchResult<Pckt>">batchUpdateImmutableReturnsIntArray</error>(List<Pckt> e);
+
     @BatchDelete
     int[] <error descr="If a method annotated with @BatchDelete targets immutable entities for insertion, the return type must be BatchResult<Pckt>">batchDeleteReturnsIntWithImmutable</error>(List<Pckt> e);
 

src/test/testData/src/main/java/doma/example/dao/inspection/returntype/UpdateReturnTypeTestDao.java

@@ -1,5 +1,6 @@
 package doma.example.dao.inspection.returntype;
 
+import java.util.Optional;
 import org.seasar.doma.*;
 import org.seasar.doma.jdbc.Result;
 import doma.example.entity.*;
@@ -16,12 +17,18 @@ public interface UpdateReturnTypeTestDao {
     @Delete(returning = @Returning)
     Packet deleteReturningEntity(Packet e);
 
+    @Delete(returning = @Returning)
+    Optional<Packet> deleteReturningOptionalEntity(Packet e);
+
     @Delete(returning = @Returning)
     int <error descr="When \"returning = @Returning\" is specified, the return type must be \"Packet\" or Optional<Packet>">deleteReturningInt</error>(Packet e);
 
     @Update
     Result<Pckt> updateReturnsResultWithImmutable(Pckt e);
 
+    @Update(sqlFile = true)
+    int <error descr="If a method annotated with @Update targets immutable entities for insertion, the return type must be Result<Pckt>">deleteImmutableInt</error>(Pckt e);
+
     @Update
     int <error descr="If a method annotated with @Update targets immutable entities for insertion, the return type must be Result<Pckt>">updateReturnsIntWithImmutable</error>(Pckt e);
 }