Move the logic for retrieving injected files to a utility class.

Author: xterao

src/main/kotlin/org/domaframework/doma/intellij/common/psi/PsiDaoMethod.kt

@@ -16,7 +16,6 @@
 package org.domaframework.doma.intellij.common.psi
 
 import com.intellij.lang.Language
-import com.intellij.lang.injection.InjectedLanguageManager
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.fileEditor.FileEditorManager
@@ -150,7 +149,6 @@ class PsiDaoMethod(
         }
         // the injection part as a custom language file
         getSqlAnnotation()?.let { annotation ->
-            InjectedLanguageManager.getInstance(psiProject)
             annotation.parameterList.children
                 .firstOrNull { it is PsiNameValuePair }
                 ?.let { sql ->

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

@@ -0,0 +1,44 @@
+/*
+ * Copyright Doma Tools Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.domaframework.doma.intellij.common.util
+
+import com.intellij.lang.injection.InjectedLanguageManager
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiFile
+import com.intellij.psi.PsiLiteralExpression
+import org.domaframework.doma.intellij.common.isJavaOrKotlinFileType
+
+object InjectionSqlUtil {
+    fun initInjectionElement(
+        basePsiFile: PsiFile,
+        project: Project,
+        literal: PsiLiteralExpression,
+    ): PsiFile? =
+        if (isJavaOrKotlinFileType(basePsiFile)) {
+            val injectedLanguageManager = InjectedLanguageManager.getInstance(project)
+            injectedLanguageManager
+                .getInjectedPsiFiles(literal)
+                ?.firstOrNull()
+                ?.first as? PsiFile
+        } else {
+            null
+        }
+
+    fun isInjectedSqlFile(source: PsiFile): Boolean {
+        val injectedLanguageManager = InjectedLanguageManager.getInstance(source.project)
+        return injectedLanguageManager.isInjectedFragment(source)
+    }
+}

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

@@ -26,7 +26,9 @@ import com.intellij.psi.TokenType
 import com.intellij.psi.impl.source.codeStyle.PreFormatProcessor
 import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.psi.util.elementType
+import org.domaframework.doma.intellij.common.util.InjectionSqlUtil.isInjectedSqlFile
 import org.domaframework.doma.intellij.common.util.PluginLoggerUtil
+import org.domaframework.doma.intellij.common.util.StringUtil
 import org.domaframework.doma.intellij.formatter.util.CreateQueryType
 import org.domaframework.doma.intellij.formatter.util.SqlKeywordUtil
 import org.domaframework.doma.intellij.formatter.visitor.SqlFormatVisitor
@@ -58,7 +60,10 @@ class SqlFormatPreProcessor : PreFormatProcessor {
         rangeToReformat: TextRange,
     ): TextRange {
         // Turn on by default the code formatter that only runs when explicitly invoked by the user.
-        if (source.language != SqlLanguage.INSTANCE) return rangeToReformat
+        // Handle both direct SQL files and injected SQL in Java files
+        if (source.language != SqlLanguage.INSTANCE && !isInjectedSqlFile(source)) {
+            return rangeToReformat
+        }
 
         logging()
 
@@ -151,7 +156,6 @@ class SqlFormatPreProcessor : PreFormatProcessor {
         element: PsiWhiteSpace,
     ) {
         val singleSpace = " "
-        val newLine = "\n"
         val range = element.textRange
         val originalText = document.getText(range)
         val nextElement = element.nextSibling
@@ -162,24 +166,28 @@ class SqlFormatPreProcessor : PreFormatProcessor {
             newText = originalText.replace(originalText, singleSpace)
         } else {
             newText =
-                when (nextElement.elementType) {
-                    SqlTypes.LINE_COMMENT -> {
-                        if (nextElementText.startsWith(newLine)) {
-                            originalText.replace(originalText, singleSpace)
-                        } else if (originalText.contains(newLine)) {
-                            originalText.replace(Regex("\\s*\\n\\s*"), newLine)
-                        } else {
-                            originalText.replace(originalText, singleSpace)
+                if (element.prevSibling == null) {
+                    ""
+                } else {
+                    when (nextElement.elementType) {
+                        SqlTypes.LINE_COMMENT -> {
+                            if (nextElementText.startsWith(StringUtil.LINE_SEPARATE)) {
+                                originalText.replace(originalText, singleSpace)
+                            } else if (originalText.contains(StringUtil.LINE_SEPARATE)) {
+                                originalText.replace(Regex("\\s*\\n\\s*"), StringUtil.LINE_SEPARATE)
+                            } else {
+                                originalText.replace(originalText, singleSpace)
+                            }
                         }
-                    }
 
-                    else -> {
-                        if (nextElementText.contains(newLine) == true) {
-                            originalText.replace(originalText, singleSpace)
-                        } else if (originalText.contains(newLine)) {
-                            originalText.replace(Regex("\\s*\\n\\s*"), newLine)
-                        } else {
-                            originalText.replace(originalText, newLine)
+                        else -> {
+                            if (nextElementText.contains(StringUtil.LINE_SEPARATE) == true) {
+                                originalText.replace(originalText, singleSpace)
+                            } else if (originalText.contains(StringUtil.LINE_SEPARATE)) {
+                                originalText.replace(Regex("\\s*\\n\\s*"), StringUtil.LINE_SEPARATE)
+                            } else {
+                                originalText.replace(originalText, StringUtil.LINE_SEPARATE)
+                            }
                         }
                     }
                 }
@@ -264,8 +272,8 @@ class SqlFormatPreProcessor : PreFormatProcessor {
         prevElement: PsiElement?,
         text: String,
     ): String =
-        if (prevElement?.text?.contains("\n") == false) {
-            "\n$text"
+        if (prevElement?.text?.contains(StringUtil.LINE_SEPARATE) == false) {
+            "${StringUtil.LINE_SEPARATE}$text"
         } else {
             text
         }

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

@@ -15,20 +15,24 @@
  */
 package org.domaframework.doma.intellij.formatter.processor
 
+import com.intellij.lang.injection.InjectedLanguageManager
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.command.WriteCommandAction
+import com.intellij.openapi.editor.Document
 import com.intellij.openapi.project.Project
 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.codeStyle.CodeStyleSettings
 import com.intellij.psi.impl.source.codeStyle.PostFormatProcessor
+import org.domaframework.doma.intellij.common.util.InjectionSqlUtil.isInjectedSqlFile
+import org.domaframework.doma.intellij.common.util.StringUtil
 import org.domaframework.doma.intellij.setting.SqlLanguage
 
 class SqlPostProcessor : PostFormatProcessor {
     companion object {
-        private const val FILE_END_PADDING = " \n"
+        private const val FILE_END_PADDING = " ${StringUtil.LINE_SEPARATE}"
     }
 
     private val trailingSpacesRegex = Regex(" +(\r?\n)")
@@ -43,12 +47,13 @@ class SqlPostProcessor : PostFormatProcessor {
         rangeToReformat: TextRange,
         settings: CodeStyleSettings,
     ): TextRange {
-        if (!isSqlFile(source)) {
+        if (!isSqlFile(source) || isInjectedSqlFile(source)) {
             return rangeToReformat
         }
 
-        val document = getDocument(source) ?: return rangeToReformat
-        val processedText = processDocumentText(document.text)
+        val originalDocument = getDocument(source)
+        val document = originalDocument ?: source.fileDocument
+        val processedText = processDocumentText(document.text, originalDocument != null)
 
         if (document.text == processedText) {
             return rangeToReformat
@@ -62,18 +67,26 @@ class SqlPostProcessor : PostFormatProcessor {
 
     private fun getDocument(source: PsiFile) = PsiDocumentManager.getInstance(source.project).getDocument(source)
 
-    private fun processDocumentText(originalText: String): String {
+    private fun processDocumentText(
+        originalText: String,
+        existsOriginalDocument: Boolean,
+    ): String {
         val withoutTrailingSpaces = removeTrailingSpaces(originalText)
-        return ensureProperFileEnding(withoutTrailingSpaces)
+        return ensureProperFileEnding(withoutTrailingSpaces, existsOriginalDocument)
     }
 
     private fun removeTrailingSpaces(text: String): String = text.replace(trailingSpacesRegex, "$1")
 
-    private fun ensureProperFileEnding(text: String): String = text.trimEnd() + FILE_END_PADDING
+    private fun ensureProperFileEnding(
+        text: String,
+        isEndSpace: Boolean,
+    ): String =
+        text.trimEnd() +
+            if (isEndSpace) FILE_END_PADDING else ""
 
     private fun updateDocument(
         project: Project,
-        document: com.intellij.openapi.editor.Document,
+        document: Document,
         newText: String,
     ) {
         ApplicationManager.getApplication().invokeAndWait {

src/main/kotlin/org/domaframework/doma/intellij/inspection/sql/visitor/SqlVisitorBase.kt

@@ -15,11 +15,10 @@
  */
 package org.domaframework.doma.intellij.inspection.sql.visitor
 
-import com.intellij.lang.injection.InjectedLanguageManager
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiLiteralExpression
-import org.domaframework.doma.intellij.common.isJavaOrKotlinFileType
+import org.domaframework.doma.intellij.common.util.InjectionSqlUtil
 import org.domaframework.doma.intellij.psi.SqlVisitor
 
 open class SqlVisitorBase : SqlVisitor() {
@@ -31,16 +30,9 @@ open class SqlVisitorBase : SqlVisitor() {
         project: Project,
         literal: PsiLiteralExpression,
     ): PsiFile? =
-        when (isJavaOrKotlinFileType(basePsiFile)) {
-            true -> {
-                val injectedLanguageManager =
-                    InjectedLanguageManager.getInstance(project)
-                injectedLanguageManager
-                    .getInjectedPsiFiles(literal)
-                    ?.firstOrNull()
-                    ?.first as? PsiFile
-            }
-
-            false -> null
-        }
+        InjectionSqlUtil.initInjectionElement(
+            basePsiFile,
+            project,
+            literal,
+        )
 }