Apply Formatting to Injected SQL

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/common/util/StringUtil.kt

@@ -16,6 +16,8 @@
 package org.domaframework.doma.intellij.common.util
 
 object StringUtil {
+    const val LINE_SEPARATE: String = "\n"
+
     fun getSqlElClassText(text: String): String =
         text
             .substringAfter("@")

src/main/kotlin/org/domaframework/doma/intellij/document/ForItemElementDocumentationProvider.kt

@@ -18,6 +18,7 @@ package org.domaframework.doma.intellij.document
 import com.intellij.lang.documentation.AbstractDocumentationProvider
 import com.intellij.psi.PsiElement
 import com.intellij.psi.util.PsiTreeUtil
+import org.domaframework.doma.intellij.common.util.StringUtil
 import org.domaframework.doma.intellij.document.generator.DocumentDaoParameterGenerator
 import org.domaframework.doma.intellij.document.generator.DocumentStaticFieldGenerator
 import org.domaframework.doma.intellij.psi.SqlElIdExpr
@@ -62,7 +63,7 @@ class ForItemElementDocumentationProvider : AbstractDocumentationProvider() {
 
         generator.generateDocument()
 
-        return result.joinToString("\n")
+        return result.joinToString(StringUtil.LINE_SEPARATE)
     }
 
     override fun generateHoverDoc(
@@ -77,6 +78,6 @@ class ForItemElementDocumentationProvider : AbstractDocumentationProvider() {
         val result: MutableList<String?> = LinkedList<String?>()
         val typeDocument = generateDoc(element, originalElement)
         result.add(typeDocument)
-        return result.joinToString("\n")
+        return result.joinToString(StringUtil.LINE_SEPARATE)
     }
 }

src/main/kotlin/org/domaframework/doma/intellij/formatter/block/comment/SqlBlockCommentBlock.kt

@@ -17,6 +17,7 @@ package org.domaframework.doma.intellij.formatter.block.comment
 
 import com.intellij.lang.ASTNode
 import com.intellij.psi.util.PsiTreeUtil
+import org.domaframework.doma.intellij.common.util.StringUtil
 import org.domaframework.doma.intellij.formatter.block.SqlBlock
 import org.domaframework.doma.intellij.formatter.util.SqlBlockFormattingContext
 
@@ -27,5 +28,6 @@ open class SqlBlockCommentBlock(
         node,
         context,
     ) {
-    override fun isSaveSpace(lastGroup: SqlBlock?): Boolean = PsiTreeUtil.prevLeaf(node.psi)?.text?.contains("\n") == true
+    override fun isSaveSpace(lastGroup: SqlBlock?): Boolean =
+        PsiTreeUtil.prevLeaf(node.psi)?.text?.contains(StringUtil.LINE_SEPARATE) == true
 }

src/main/kotlin/org/domaframework/doma/intellij/formatter/block/comment/SqlDefaultCommentBlock.kt

@@ -18,6 +18,7 @@ package org.domaframework.doma.intellij.formatter.block.comment
 import com.intellij.lang.ASTNode
 import com.intellij.psi.formatter.common.AbstractBlock
 import com.intellij.psi.util.PsiTreeUtil
+import org.domaframework.doma.intellij.common.util.StringUtil
 import org.domaframework.doma.intellij.formatter.block.SqlBlock
 import org.domaframework.doma.intellij.formatter.util.IndentType
 import org.domaframework.doma.intellij.formatter.util.SqlBlockFormattingContext
@@ -60,5 +61,5 @@ abstract class SqlDefaultCommentBlock(
             }
     }
 
-    override fun isSaveSpace(lastGroup: SqlBlock?) = PsiTreeUtil.prevLeaf(node.psi)?.text?.contains("\n") == true
+    override fun isSaveSpace(lastGroup: SqlBlock?) = PsiTreeUtil.prevLeaf(node.psi)?.text?.contains(StringUtil.LINE_SEPARATE) == true
 }

src/main/kotlin/org/domaframework/doma/intellij/formatter/block/comment/SqlElConditionLoopCommentBlock.kt

@@ -22,6 +22,7 @@ import com.intellij.psi.PsiWhiteSpace
 import com.intellij.psi.formatter.common.AbstractBlock
 import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.psi.util.elementType
+import org.domaframework.doma.intellij.common.util.StringUtil
 import org.domaframework.doma.intellij.common.util.TypeUtil
 import org.domaframework.doma.intellij.extension.expr.isConditionOrLoopDirective
 import org.domaframework.doma.intellij.formatter.block.SqlBlock
@@ -108,7 +109,7 @@ class SqlElConditionLoopCommentBlock(
                 // If the child is a condition loop directive, align its indentation with the parent directive
                 child.indent.indentLen = indent.indentLen.plus(2)
             } else if (child is SqlLineCommentBlock) {
-                if (PsiTreeUtil.prevLeaf(child.node.psi, false)?.text?.contains("\n") == true) {
+                if (PsiTreeUtil.prevLeaf(child.node.psi, false)?.text?.contains(StringUtil.LINE_SEPARATE) == true) {
                     child.indent.indentLen = indent.groupIndentLen
                 } else {
                     child.indent.indentLen = 1

src/main/kotlin/org/domaframework/doma/intellij/formatter/builder/SqlCustomSpacingBuilder.kt

@@ -19,6 +19,7 @@ import com.intellij.formatting.ASTBlock
 import com.intellij.formatting.Block
 import com.intellij.formatting.Spacing
 import com.intellij.psi.tree.IElementType
+import org.domaframework.doma.intellij.common.util.StringUtil
 import org.domaframework.doma.intellij.formatter.block.SqlBlock
 import org.domaframework.doma.intellij.formatter.block.SqlRightPatternBlock
 import org.domaframework.doma.intellij.formatter.block.SqlWhitespaceBlock
@@ -79,8 +80,8 @@ class SqlCustomSpacingBuilder {
             null -> return nonSpacing
             is SqlWhitespaceBlock -> {
                 val indentLen: Int = child2.indent.indentLen
-                val afterNewLine = child1.getNodeText().substringAfterLast("\n", "")
-                if (child1.getNodeText().contains("\n")) {
+                val afterNewLine = child1.getNodeText().substringAfterLast(StringUtil.LINE_SEPARATE, "")
+                if (child1.getNodeText().contains(StringUtil.LINE_SEPARATE)) {
                     val currentIndent = afterNewLine.length
                     val newIndent =
                         if (currentIndent != indentLen) {

src/main/kotlin/org/domaframework/doma/intellij/formatter/builder/SqlFormattingModelBuilder.kt

@@ -38,25 +38,35 @@ import org.domaframework.doma.intellij.setting.state.DomaToolsFormatEnableSettin
 class SqlFormattingModelBuilder : FormattingModelBuilder {
     override fun createModel(formattingContext: FormattingContext): FormattingModel {
         val codeStyleSettings = formattingContext.codeStyleSettings
+
+        return createRegularSQLModel(formattingContext, codeStyleSettings)
+    }
+
+    private fun createRegularSQLModel(
+        formattingContext: FormattingContext,
+        settings: CodeStyleSettings,
+    ): FormattingModel {
         val setting = DomaToolsFormatEnableSettings.getInstance()
         val isEnableFormat = setting.state.isEnableSqlFormat == true
         val formatMode = formattingContext.formattingMode
-        val spacingBuilder = createSpaceBuilder(codeStyleSettings)
+        val spacingBuilder = createSpaceBuilder(settings)
         val customSpacingBuilder = createCustomSpacingBuilder()
 
+        val block =
+            SqlFileBlock(
+                formattingContext.node,
+                Wrap.createWrap(WrapType.NONE, false),
+                Alignment.createAlignment(),
+                customSpacingBuilder,
+                spacingBuilder,
+                isEnableFormat,
+                formatMode,
+            )
         return FormattingModelProvider
             .createFormattingModelForPsiFile(
                 formattingContext.containingFile,
-                SqlFileBlock(
-                    formattingContext.node,
-                    Wrap.createWrap(WrapType.NONE, false),
-                    Alignment.createAlignment(),
-                    customSpacingBuilder,
-                    spacingBuilder,
-                    isEnableFormat,
-                    formatMode,
-                ),
-                codeStyleSettings,
+                block,
+                settings,
             )
     }
 

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

@@ -0,0 +1,77 @@
+/*
+ * 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.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 org.domaframework.doma.intellij.setting.SqlLanguage
+
+class SqlFormatPostProcessor : SqlPostProcessor() {
+    override fun processElement(
+        source: PsiElement,
+        settings: CodeStyleSettings,
+    ): PsiElement = source
+
+    override fun processText(
+        source: PsiFile,
+        rangeToReformat: TextRange,
+        settings: CodeStyleSettings,
+    ): TextRange {
+        if (!isSqlFile(source) || isInjectedSqlFile(source)) {
+            return rangeToReformat
+        }
+
+        val document = getDocument(source) ?: return rangeToReformat
+        val processedText = processDocumentText(document.text, true)
+
+        if (document.text == processedText) {
+            return rangeToReformat
+        }
+
+        updateDocument(source.project, document, processedText)
+        return TextRange(0, processedText.length)
+    }
+
+    private fun isSqlFile(source: PsiFile): Boolean = source.language == SqlLanguage.INSTANCE
+
+    private fun isInjectedSqlFile(source: PsiFile): Boolean {
+        val injectedLanguageManager = InjectedLanguageManager.getInstance(source.project)
+        return injectedLanguageManager.isInjectedFragment(source)
+    }
+
+    private fun getDocument(source: PsiFile) = PsiDocumentManager.getInstance(source.project).getDocument(source)
+
+    private fun updateDocument(
+        project: Project,
+        document: Document,
+        newText: String,
+    ) {
+        ApplicationManager.getApplication().invokeAndWait {
+            WriteCommandAction.runWriteCommandAction(project) {
+                document.setText(newText)
+                PsiDocumentManager.getInstance(project).commitDocument(document)
+            }
+        }
+    }
+}

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,10 @@ class SqlFormatPreProcessor : PreFormatProcessor {
         prevElement: PsiElement?,
         text: String,
     ): String =
-        if (prevElement?.text?.contains("\n") == false) {
-            "\n$text"
+        if (prevElement?.text?.contains(StringUtil.LINE_SEPARATE) == false &&
+            PsiTreeUtil.prevLeaf(prevElement) != null
+        ) {
+            "${StringUtil.LINE_SEPARATE}$text"
         } else {
             text
         }