Merge pull request #364 from domaframework/fix/sql-format-one-line-injection

Formatting Support for Single-Line Injected SQL in @SQL Annotations

Author: xterao

src/main/kotlin/org/domaframework/doma/intellij/common/FileTypeCheck.kt

@@ -18,8 +18,6 @@ package org.domaframework.doma.intellij.common
 import com.intellij.openapi.fileTypes.FileTypeManager
 import com.intellij.psi.PsiFile
 
-val sourceExtensionNames: List<String> = listOf("JAVA", "Kotlin", "CLASS")
-
 /**
  * Get extension by file type identifier
  */
@@ -56,13 +54,3 @@ fun isSupportFileType(file: PsiFile): Boolean {
         else -> false
     }
 }
-
-fun isInjectionSqlFile(file: PsiFile): Boolean {
-    val extension = file.fileType.defaultExtension
-    val filePath = file.virtualFile?.path ?: return false
-    return when (extension) {
-        "sql" -> true
-        else -> false
-    } &&
-        !(filePath.endsWith(".sql") || filePath.endsWith(".script"))
-}

src/main/kotlin/org/domaframework/doma/intellij/common/dao/DaoMethodUtil.kt

@@ -32,8 +32,8 @@ import org.domaframework.doma.intellij.common.RESOURCES_META_INF_PATH
 import org.domaframework.doma.intellij.common.getExtension
 import org.domaframework.doma.intellij.common.getJarRoot
 import org.domaframework.doma.intellij.common.getMethodDaoFilePath
-import org.domaframework.doma.intellij.common.isInjectionSqlFile
 import org.domaframework.doma.intellij.common.isSupportFileType
+import org.domaframework.doma.intellij.common.util.InjectionSqlUtil.isInjectedSqlFile
 import org.domaframework.doma.intellij.extension.getContentRoot
 import org.domaframework.doma.intellij.extension.getJavaClazz
 import org.domaframework.doma.intellij.extension.getModule
@@ -51,7 +51,7 @@ fun findDaoMethod(
     val virtualFile = originalFile.virtualFile ?: return null
     val module = project.getModule(virtualFile)
 
-    if (isInjectionSqlFile(originalFile)) {
+    if (isInjectedSqlFile(originalFile)) {
         originalFile.let {
             return PsiTreeUtil.getParentOfType(originalFile.context, PsiMethod::class.java)
         }
@@ -135,7 +135,7 @@ private fun getDaoPathFromSqlFilePath(
     sqlFile: PsiFile,
     contentRootPath: String,
 ): PsiClass? {
-    if (isInjectionSqlFile(sqlFile)) {
+    if (isInjectedSqlFile(sqlFile)) {
         return null
     }
     val module = sqlFile.module ?: return null

src/main/kotlin/org/domaframework/doma/intellij/common/util/InjectionSqlUtil.kt

@@ -41,4 +41,9 @@ object InjectionSqlUtil {
         val injectedLanguageManager = InjectedLanguageManager.getInstance(source.project)
         return injectedLanguageManager.isInjectedFragment(source)
     }
+
+    fun getLiteralExpressionHost(source: PsiFile): PsiLiteralExpression? {
+        val injectedLanguageManager = InjectedLanguageManager.getInstance(source.project)
+        return injectedLanguageManager.getInjectionHost(source) as? PsiLiteralExpression
+    }
 }

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

@@ -16,11 +16,13 @@
 package org.domaframework.doma.intellij.formatter.processor
 
 import com.intellij.lang.ASTNode
+import com.intellij.lang.injection.InjectedLanguageManager
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiLiteralExpression
 import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.TokenType
 import com.intellij.psi.impl.source.codeStyle.PreFormatProcessor
@@ -50,6 +52,11 @@ class SqlFormatPreProcessor : PreFormatProcessor {
             SqlTypes.OTHER,
         )
 
+    data class ProcessResult(
+        val document: Document?,
+        val range: TextRange,
+    )
+
     override fun process(
         node: ASTNode,
         rangeToReformat: TextRange,
@@ -65,13 +72,26 @@ class SqlFormatPreProcessor : PreFormatProcessor {
             return rangeToReformat
         }
 
+        // Do not execute processor processing in single-line text state
+        if (isInjectedSqlFile(source)) {
+            val host = InjectedLanguageManager.getInstance(source.project).getInjectionHost(source) as? PsiLiteralExpression
+            if (host?.isTextBlock != true) return rangeToReformat
+        }
+        val result = updateDocument(source, rangeToReformat)
+        return result.range
+    }
+
+    fun updateDocument(
+        source: PsiFile,
+        rangeToReformat: TextRange,
+    ): ProcessResult {
         logging()
 
         val visitor = SqlFormatVisitor()
         source.accept(visitor)
 
         val docManager = PsiDocumentManager.getInstance(source.project)
-        val document = docManager.getDocument(source) ?: return rangeToReformat
+        val document = docManager.getDocument(source) ?: return ProcessResult(null, rangeToReformat)
 
         val keywordList = visitor.replaces.filter { it.elementType != TokenType.WHITE_SPACE }
         val replaceKeywordList = visitor.replaces.filter { it.elementType == SqlTypes.KEYWORD }
@@ -122,7 +142,7 @@ class SqlFormatPreProcessor : PreFormatProcessor {
 
         docManager.commitDocument(document)
 
-        return rangeToReformat.grown(visitor.replaces.size)
+        return ProcessResult(document, rangeToReformat.grown(visitor.replaces.size))
     }
 
     private fun removeSpacesAroundNewline(

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

@@ -19,13 +19,27 @@ 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.common.util.InjectionSqlUtil
 import org.domaframework.doma.intellij.formatter.visitor.DaoInjectionSqlVisitor
+import org.domaframework.doma.intellij.formatter.visitor.FormattingTask
 
+/**
+ * Post-processor for SQL injection formatting.
+ *
+ * This processor handles SQL formatting in two contexts:
+ * 1. **File formatting**: When formatting entire DAO files (Java/Kotlin) containing SQL annotations
+ * 2. **Code formatting**: When formatting injected SQL fragments within string literals
+ *
+ * The context is determined by checking:
+ * - If the source is an injected fragment (`InjectedLanguageManager.isInjectedFragment()` returns true),
+ *   it's being called from code formatting for a specific SQL string literal
+ * - If the source is a regular DAO file (Java/Kotlin with @Dao annotation),
+ *   it's being called from file formatting to process all SQL strings in the file
+ */
 class SqlInjectionPostProcessor : SqlPostProcessor() {
     override fun processElement(
         element: PsiElement,
@@ -37,49 +51,55 @@ class SqlInjectionPostProcessor : SqlPostProcessor() {
         rangeToReformat: TextRange,
         settings: CodeStyleSettings,
     ): TextRange {
-        if (!shouldProcessFile(source)) {
+        if (!isProcessFile(source)) {
             return rangeToReformat
         }
 
         val manager = InjectedLanguageManager.getInstance(source.project)
         if (manager.isInjectedFragment(source)) {
-            processInjectedFragment(source, manager)
+            processInjectedFragment(source)
         } else {
             processRegularFile(source)
         }
 
         return rangeToReformat
     }
 
-    private fun shouldProcessFile(source: PsiFile): Boolean {
+    private fun isProcessFile(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
+    /**
+     * Processes all SQL injections in a DAO file during file formatting.
+     * This is called when formatting an entire DAO file containing SQL annotations.
+     */
+    private fun processInjectedFragment(source: PsiFile) {
+        val host = InjectionSqlUtil.getLiteralExpressionHost(source) ?: return
         val originalText = host.value?.toString() ?: return
 
-        val visitor = DaoInjectionSqlVisitor(hostDaoFile, source.project)
-        val formattingTask = DaoInjectionSqlVisitor.FormattingTask(host, originalText)
+        val visitor = DaoInjectionSqlVisitor(source.project)
+        val formattingTask = FormattingTask(host, originalText, host.isTextBlock)
 
-        visitor.replaceHostStringLiteral(formattingTask) { text ->
+        visitor.convertExpressionToTextBlock(formattingTask.expression)
+        visitor.processFormattingTask(formattingTask) { text ->
             processDocumentText(text)
         }
     }
 
+    /**
+     * Processes injected SQL fragments during code formatting.
+     * This is called when formatting a specific SQL string literal within a DAO file.
+     */
     private fun processRegularFile(source: PsiFile) {
-        val visitor = DaoInjectionSqlVisitor(source, source.project)
+        val visitor = DaoInjectionSqlVisitor(source.project)
         source.accept(visitor)
-
-        visitor.processAll { text ->
+        visitor.processAllTextBlock()
+        source.accept(visitor)
+        visitor.processAllReFormat { text ->
             processDocumentText(text)
         }
     }