Implement line break formatting rule for Returning.

xterao · xterao · commit 9eb7f91f2f5a · 2025-07-09T12:58:00.000+09:00
diff --git a/src/main/kotlin/org/domaframework/doma/intellij/formatter/block/SqlRightPatternBlock.kt b/src/main/kotlin/org/domaframework/doma/intellij/formatter/block/SqlRightPatternBlock.kt
@@ -20,11 +20,11 @@ import com.intellij.psi.formatter.common.AbstractBlock
 import org.domaframework.doma.intellij.common.util.TypeUtil.isExpectedClassType
 import org.domaframework.doma.intellij.formatter.block.conflict.SqlConflictClauseBlock
 import org.domaframework.doma.intellij.formatter.block.group.column.SqlColumnDefinitionRawGroupBlock
+import org.domaframework.doma.intellij.formatter.block.group.column.SqlColumnRawGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlKeywordGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.create.SqlCreateTableColumnDefinitionGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.insert.SqlInsertColumnGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.update.SqlUpdateColumnGroupBlock
-import org.domaframework.doma.intellij.formatter.block.group.keyword.update.SqlUpdateSetGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.update.SqlUpdateValueGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.with.SqlWithQuerySubGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.subgroup.SqlFunctionParamBlock
diff --git a/src/main/kotlin/org/domaframework/doma/intellij/formatter/block/group/keyword/insert/SqlInsertQueryGroupBlock.kt b/src/main/kotlin/org/domaframework/doma/intellij/formatter/block/group/keyword/insert/SqlInsertQueryGroupBlock.kt
@@ -16,7 +16,7 @@
 package org.domaframework.doma.intellij.formatter.block.group.keyword.insert
 
 import com.intellij.lang.ASTNode
-import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlValuesGroupBlock
+import org.domaframework.doma.intellij.formatter.block.group.keyword.second.SqlValuesGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.top.SqlTopQueryGroupBlock
 import org.domaframework.doma.intellij.formatter.util.SqlBlockFormattingContext
 
diff --git a/src/main/kotlin/org/domaframework/doma/intellij/formatter/block/group/keyword/second/SqlReturningGroupBlock.kt b/src/main/kotlin/org/domaframework/doma/intellij/formatter/block/group/keyword/second/SqlReturningGroupBlock.kt
@@ -0,0 +1,16 @@
+package org.domaframework.doma.intellij.formatter.block.group.keyword.second
+
+import com.intellij.lang.ASTNode
+import org.domaframework.doma.intellij.formatter.util.SqlBlockFormattingContext
+
+class SqlReturningGroupBlock(
+    node: ASTNode,
+    context: SqlBlockFormattingContext,
+) : SqlSecondKeywordBlock(node, context) {
+    override fun createBlockIndentLen(): Int {
+        parentBlock?.let { parent ->
+            return parent.indent.indentLen
+        }
+        return 0
+    }
+}
diff --git a/src/main/kotlin/org/domaframework/doma/intellij/formatter/block/group/keyword/second/SqlValuesGroupBlock.kt b/src/main/kotlin/org/domaframework/doma/intellij/formatter/block/group/keyword/second/SqlValuesGroupBlock.kt
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.domaframework.doma.intellij.formatter.block.group.keyword
+package org.domaframework.doma.intellij.formatter.block.group.keyword.second
 
 import com.intellij.lang.ASTNode
 import org.domaframework.doma.intellij.formatter.block.SqlBlock
diff --git a/src/main/kotlin/org/domaframework/doma/intellij/formatter/block/group/keyword/top/SqlJoinQueriesGroupBlock.kt b/src/main/kotlin/org/domaframework/doma/intellij/formatter/block/group/keyword/top/SqlJoinQueriesGroupBlock.kt
@@ -1,22 +1,8 @@
-/*
- * 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.block.group.keyword
+package org.domaframework.doma.intellij.formatter.block.group.keyword.top
 
 import com.intellij.lang.ASTNode
 import org.domaframework.doma.intellij.formatter.block.SqlBlock
+import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlKeywordGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.with.SqlWithQuerySubGroupBlock
 import org.domaframework.doma.intellij.formatter.util.IndentType
 import org.domaframework.doma.intellij.formatter.util.SqlBlockFormattingContext
diff --git a/src/main/kotlin/org/domaframework/doma/intellij/formatter/block/group/subgroup/SqlSubGroupBlock.kt b/src/main/kotlin/org/domaframework/doma/intellij/formatter/block/group/subgroup/SqlSubGroupBlock.kt
@@ -25,8 +25,8 @@ import org.domaframework.doma.intellij.formatter.block.SqlRightPatternBlock
 import org.domaframework.doma.intellij.formatter.block.comment.SqlCommentBlock
 import org.domaframework.doma.intellij.formatter.block.conflict.SqlDoGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.SqlNewGroupBlock
-import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlJoinQueriesGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.create.SqlCreateViewGroupBlock
+import org.domaframework.doma.intellij.formatter.block.group.keyword.top.SqlJoinQueriesGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.with.SqlWithColumnGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.with.SqlWithCommonTableGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.with.SqlWithQueryGroupBlock
diff --git a/src/main/kotlin/org/domaframework/doma/intellij/formatter/processor/SqlSetParentGroupProcessor.kt b/src/main/kotlin/org/domaframework/doma/intellij/formatter/processor/SqlSetParentGroupProcessor.kt
@@ -18,14 +18,17 @@ package org.domaframework.doma.intellij.formatter.processor
 import org.domaframework.doma.intellij.common.util.TypeUtil.isExpectedClassType
 import org.domaframework.doma.intellij.formatter.block.SqlBlock
 import org.domaframework.doma.intellij.formatter.block.SqlCommaBlock
+import org.domaframework.doma.intellij.formatter.block.SqlKeywordBlock
 import org.domaframework.doma.intellij.formatter.block.SqlRightPatternBlock
+import org.domaframework.doma.intellij.formatter.block.conflict.SqlDoGroupBlock
 import org.domaframework.doma.intellij.formatter.block.expr.SqlElConditionLoopCommentBlock
 import org.domaframework.doma.intellij.formatter.block.group.column.SqlColumnDefinitionRawGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.column.SqlColumnRawGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlInlineGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlInlineSecondGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlKeywordGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.create.SqlCreateViewGroupBlock
+import org.domaframework.doma.intellij.formatter.block.group.keyword.second.SqlReturningGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.update.SqlUpdateQueryGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.with.SqlWithCommonTableGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.with.SqlWithQuerySubGroupBlock
@@ -167,6 +170,19 @@ class SqlSetParentGroupProcessor(
             }
 
             blockBuilder.removeLastGroupTopNodeIndexHistory()
+            if (childBlock is SqlReturningGroupBlock) {
+                // Since `DO UPDATE` does not include a `RETURNING` clause, it should be registered as a child of the parent `INSERT` query.
+                // The `DO` keyword should align with the `INSERT` query, and therefore it will serve as the **indentation anchor** for the following update block.
+                setParentGroups(context) { history ->
+                    val lastGroup = history.lastOrNull()
+                    return@setParentGroups if (lastGroup is SqlUpdateQueryGroupBlock && lastGroup.parentBlock is SqlDoGroupBlock) {
+                        lastGroup.parentBlock
+                    } else {
+                        history.lastOrNull()
+                    }
+                }
+                return
+            }
             updateGroupBlockLastGroupParentAddGroup(
                 lastGroupBlock,
                 childBlock,
diff --git a/src/main/kotlin/org/domaframework/doma/intellij/formatter/util/InsertClauseUtil.kt b/src/main/kotlin/org/domaframework/doma/intellij/formatter/util/InsertClauseUtil.kt
@@ -17,10 +17,10 @@ package org.domaframework.doma.intellij.formatter.util
 
 import com.intellij.lang.ASTNode
 import org.domaframework.doma.intellij.formatter.block.SqlBlock
-import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlValuesGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.insert.SqlInsertColumnGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.insert.SqlInsertQueryGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.insert.SqlInsertValueGroupBlock
+import org.domaframework.doma.intellij.formatter.block.group.keyword.second.SqlValuesGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.subgroup.SqlSubGroupBlock
 
 object InsertClauseUtil {
diff --git a/src/main/kotlin/org/domaframework/doma/intellij/formatter/util/NotQueryGroupUtil.kt b/src/main/kotlin/org/domaframework/doma/intellij/formatter/util/NotQueryGroupUtil.kt
@@ -21,6 +21,7 @@ import com.intellij.psi.util.elementType
 import org.domaframework.doma.intellij.formatter.block.SqlBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.condition.SqlConditionKeywordGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.condition.SqlConditionalExpressionGroupBlock
+import org.domaframework.doma.intellij.formatter.block.group.keyword.second.SqlReturningGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.subgroup.SqlFunctionParamBlock
 import org.domaframework.doma.intellij.formatter.block.group.subgroup.SqlParallelListBlock
 import org.domaframework.doma.intellij.psi.SqlTypes
@@ -51,4 +52,15 @@ object NotQueryGroupUtil {
 
         return null
     }
+
+    fun getKeywordGroup(
+        child: ASTNode,
+        sqlBlockFormattingCtx: SqlBlockFormattingContext,
+    ): SqlBlock? {
+        val keyword = child.text.lowercase()
+        if (keyword == "returning") {
+            return SqlReturningGroupBlock(child, sqlBlockFormattingCtx)
+        }
+        return null
+    }
 }
diff --git a/src/main/kotlin/org/domaframework/doma/intellij/formatter/util/SqlBlockUtil.kt b/src/main/kotlin/org/domaframework/doma/intellij/formatter/util/SqlBlockUtil.kt
@@ -42,18 +42,18 @@ import org.domaframework.doma.intellij.formatter.block.group.column.SqlColumnRaw
 import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlInlineGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlInlineSecondGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlJoinGroupBlock
-import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlJoinQueriesGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlKeywordGroupBlock
-import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlSecondKeywordBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlSecondOptionKeywordGroupBlock
-import org.domaframework.doma.intellij.formatter.block.group.keyword.SqlValuesGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.condition.SqlConditionKeywordGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.create.SqlCreateKeywordGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.create.SqlCreateTableColumnDefinitionGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.create.SqlCreateTableColumnDefinitionRawGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.create.SqlCreateViewGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.insert.SqlInsertQueryGroupBlock
+import org.domaframework.doma.intellij.formatter.block.group.keyword.second.SqlSecondKeywordBlock
+import org.domaframework.doma.intellij.formatter.block.group.keyword.second.SqlValuesGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.top.SqlDeleteQueryGroupBlock
+import org.domaframework.doma.intellij.formatter.block.group.keyword.top.SqlJoinQueriesGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.top.SqlSelectQueryGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.update.SqlUpdateQueryGroupBlock
 import org.domaframework.doma.intellij.formatter.block.group.keyword.update.SqlUpdateSetGroupBlock
@@ -260,6 +260,13 @@ class SqlBlockUtil(
                         WithClauseUtil
                             .getWithClauseKeywordGroup(lastGroupBlock, child, sqlBlockFormattingCtx)
                             ?.let { return it }
+
+                        NotQueryGroupUtil
+                            .getKeywordGroup(
+                                child,
+                                sqlBlockFormattingCtx,
+                            )?.let { return it }
+
                         SqlSecondKeywordBlock(
                             child,
                             sqlBlockFormattingCtx,
diff --git a/src/test/testData/sql/formatter/WithDelete_format.sql b/src/test/testData/sql/formatter/WithDelete_format.sql
@@ -2,7 +2,7 @@ WITH deleted_posts AS (
     DELETE FROM posts
      WHERE deleted = TRUE
        AND deleted_at < CURRENT_DATE - INTERVAL '30 days'
- RETURNING id
+    RETURNING id
 )
 SELECT *
   FROM deleted_posts 
diff --git a/src/test/testData/sql/formatter/WithInsert.sql b/src/test/testData/sql/formatter/WithInsert.sql
@@ -1,5 +1,5 @@
 WITH new_user AS ( INSERT INTO users (name, email)
-         VALUES ('Alice Example', 'alice@example.com') RETURNING id)
+         VALUES ('Alice Example', 'alice@example.com') RETURNING id,name)
 INSERT INTO user_profiles (user_id, bio)
-SELECT id , 'This is Alice' , profile
+SELECT id , name
   FROM new_user 
diff --git a/src/test/testData/sql/formatter/WithInsert_format.sql b/src/test/testData/sql/formatter/WithInsert_format.sql
@@ -4,12 +4,12 @@ WITH new_user AS (
                  , email)
          VALUES ('Alice Example'
                  , 'alice@example.com')
-      RETURNING id
+    RETURNING id
+              , name
 )
 INSERT INTO user_profiles
             (user_id
              , bio)
 SELECT id
-       , 'This is Alice'
-       , profile
+       , name
   FROM new_user 
diff --git a/src/test/testData/sql/formatter/WithUpdate_format.sql b/src/test/testData/sql/formatter/WithUpdate_format.sql
@@ -9,7 +9,7 @@ WITH recent_activity AS (
        SET last_login_at = ra.last_login
       FROM recent_activity ra
      WHERE u.id = ra.user_id
- RETURNING id
+    RETURNING id
 )
 SELECT *
   FROM updated_users 
