Merge pull request #348 from domaframework/fix/sql-format-reformat-code

Support code reformatting on injected SQL

Author: xterao

src/main/kotlin/org/domaframework/doma/intellij/formatter/processor/SqlInjectionPostProcessor.kt

@@ -15,13 +15,15 @@
  */
 package org.domaframework.doma.intellij.formatter.processor
 
-import com.intellij.openapi.project.Project
+import com.intellij.lang.injection.InjectedLanguageManager
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiLiteralExpression
 import com.intellij.psi.codeStyle.CodeStyleSettings
 import org.domaframework.doma.intellij.common.dao.getDaoClass
 import org.domaframework.doma.intellij.common.isJavaOrKotlinFileType
+import org.domaframework.doma.intellij.common.isSupportFileType
 import org.domaframework.doma.intellij.formatter.visitor.DaoInjectionSqlVisitor
 
 class SqlInjectionPostProcessor : SqlPostProcessor() {
@@ -35,17 +37,49 @@ class SqlInjectionPostProcessor : SqlPostProcessor() {
         rangeToReformat: TextRange,
         settings: CodeStyleSettings,
     ): TextRange {
-        if (!isJavaOrKotlinFileType(source) || getDaoClass(source) == null) return rangeToReformat
+        if (!shouldProcessFile(source)) {
+            return rangeToReformat
+        }
+
+        val manager = InjectedLanguageManager.getInstance(source.project)
+        if (manager.isInjectedFragment(source)) {
+            processInjectedFragment(source, manager)
+        } else {
+            processRegularFile(source)
+        }
 
-        processInjected(source)
         return rangeToReformat
     }
 
-    private fun processInjected(element: PsiFile) {
-        val project: Project = element.project
-        val visitor = DaoInjectionSqlVisitor(element, project)
-        element.accept(visitor)
-        visitor.processAll { text, skipFinalLineBreak ->
+    private fun shouldProcessFile(source: PsiFile): Boolean {
+        val manager = InjectedLanguageManager.getInstance(source.project)
+        val isInjectedSql = if (isSupportFileType(source)) manager.isInjectedFragment(source) else false
+        val isDaoFile = isJavaOrKotlinFileType(source) && getDaoClass(source) != null
+
+        return isInjectedSql || isDaoFile
+    }
+
+    private fun processInjectedFragment(
+        source: PsiFile,
+        manager: InjectedLanguageManager,
+    ) {
+        val host = manager.getInjectionHost(source) as? PsiLiteralExpression ?: return
+        val hostDaoFile = host.containingFile
+        val originalText = host.value?.toString() ?: return
+
+        val visitor = DaoInjectionSqlVisitor(hostDaoFile, source.project)
+        val formattingTask = DaoInjectionSqlVisitor.FormattingTask(host, originalText)
+
+        visitor.replaceHostStringLiteral(formattingTask) { text ->
+            processDocumentText(text)
+        }
+    }
+
+    private fun processRegularFile(source: PsiFile) {
+        val visitor = DaoInjectionSqlVisitor(source, source.project)
+        source.accept(visitor)
+
+        visitor.processAll { text ->
             processDocumentText(text)
         }
     }

src/main/kotlin/org/domaframework/doma/intellij/formatter/visitor/DaoInjectionSqlVisitor.kt

@@ -36,7 +36,7 @@ class DaoInjectionSqlVisitor(
     private val element: PsiFile,
     private val project: Project,
 ) : JavaRecursiveElementVisitor() {
-    private data class FormattingTask(
+    data class FormattingTask(
         val expression: PsiLiteralExpression,
         val formattedText: String,
     )
@@ -58,8 +58,7 @@ class DaoInjectionSqlVisitor(
         if (injected != null) {
             // Format SQL and store the task
             val originalText = expression.value?.toString() ?: return
-            val removeIndent = removeIndentLines(originalText)
-            formattingTasks.add(FormattingTask(expression, removeIndent))
+            formattingTasks.add(FormattingTask(expression, originalText))
         }
     }
 
@@ -88,7 +87,7 @@ class DaoInjectionSqlVisitor(
      * Processes all collected formatting tasks in a single write action.
      * @param removeSpace Function to remove trailing spaces from formatted text
      */
-    fun processAll(removeSpace: (String, Boolean) -> String) {
+    fun processAll(removeSpace: (String) -> String) {
         if (formattingTasks.isEmpty()) return
 
         // Apply all formatting tasks in a single write action
@@ -130,9 +129,9 @@ class DaoInjectionSqlVisitor(
     /**
      * Replaces the host Java string literal with formatted SQL text.
      */
-    private fun replaceHostStringLiteral(
+    fun replaceHostStringLiteral(
         task: FormattingTask,
-        sqlPostProcessorProcess: (String, Boolean) -> String,
+        sqlPostProcessorProcess: (String) -> String,
     ) {
         try {
             // Keep the current top line indent
@@ -145,11 +144,12 @@ class DaoInjectionSqlVisitor(
 
     private fun createFormattedLiteral(
         task: FormattingTask,
-        sqlPostProcessorProcess: (String, Boolean) -> String,
+        sqlPostProcessorProcess: (String) -> String,
     ): String {
         // Retrieve the same formatted string as when formatting a regular SQL file.
-        val formattedSql = formatAsTemporarySqlFile(task.formattedText)
-        val cleanedText = sqlPostProcessorProcess(formattedSql, false)
+        val removeIndent = removeIndentLines(task.formattedText)
+        val formattedSql = formatAsTemporarySqlFile(removeIndent)
+        val cleanedText = sqlPostProcessorProcess(formattedSql)
         // Generate text aligned with the literal element using the formatted string.
         val newLiteralText = createFormattedLiteralText(cleanedText)
         val normalizedText = normalizeIndentation(newLiteralText)