Ensure updates and deletes with RETURNING statements execute queries

dialects/postgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/PostgreSqlTypeResolver.kt

@@ -14,8 +14,10 @@ import app.cash.sqldelight.dialects.postgresql.PostgreSqlType.BIG_INT
 import app.cash.sqldelight.dialects.postgresql.PostgreSqlType.SMALL_INT
 import app.cash.sqldelight.dialects.postgresql.PostgreSqlType.TIMESTAMP
 import app.cash.sqldelight.dialects.postgresql.PostgreSqlType.TIMESTAMP_TIMEZONE
+import app.cash.sqldelight.dialects.postgresql.grammar.psi.PostgreSqlDeleteStmtLimited
 import app.cash.sqldelight.dialects.postgresql.grammar.psi.PostgreSqlInsertStmt
 import app.cash.sqldelight.dialects.postgresql.grammar.psi.PostgreSqlTypeName
+import app.cash.sqldelight.dialects.postgresql.grammar.psi.PostgreSqlUpdateStmtLimited
 import com.alecstrong.sql.psi.core.psi.SqlAnnotatedElement
 import com.alecstrong.sql.psi.core.psi.SqlFunctionExpr
 import com.alecstrong.sql.psi.core.psi.SqlStmt
@@ -94,6 +96,26 @@ class PostgreSqlTypeResolver(private val parentResolver: TypeResolver) : TypeRes
         }
       }
     }
+    sqlStmt.updateStmtLimited?.let { update ->
+      check(update is PostgreSqlUpdateStmtLimited)
+      update.returningClause?.let {
+        return object : QueryWithResults {
+          override var statement: SqlAnnotatedElement = update
+          override val select = it
+          override val pureTable = update.qualifiedTableName.tableName
+        }
+      }
+    }
+    sqlStmt.deleteStmtLimited?.let { delete ->
+      check(delete is PostgreSqlDeleteStmtLimited)
+      delete.returningClause?.let {
+        return object : QueryWithResults {
+          override var statement: SqlAnnotatedElement = delete
+          override val select = it
+          override val pureTable = delete.qualifiedTableName?.tableName
+        }
+      }
+    }
     return parentResolver.queryWithResults(sqlStmt)
   }
 }

dialects/sqlite-3-35/src/main/kotlin/app/cash/sqldelight/dialects/sqlite_3_35/SqliteTypeResolver.kt

@@ -2,7 +2,9 @@ package app.cash.sqldelight.dialects.sqlite_3_35
 
 import app.cash.sqldelight.dialect.api.QueryWithResults
 import app.cash.sqldelight.dialect.api.TypeResolver
+import app.cash.sqldelight.dialects.sqlite_3_35.grammar.psi.SqliteDeleteStmtLimited
 import app.cash.sqldelight.dialects.sqlite_3_35.grammar.psi.SqliteInsertStmt
+import app.cash.sqldelight.dialects.sqlite_3_35.grammar.psi.SqliteUpdateStmtLimited
 import com.alecstrong.sql.psi.core.psi.SqlAnnotatedElement
 import com.alecstrong.sql.psi.core.psi.SqlStmt
 import app.cash.sqldelight.dialects.sqlite_3_24.SqliteTypeResolver as Sqlite324TypeResolver
@@ -19,6 +21,26 @@ class SqliteTypeResolver(private val parentResolver: TypeResolver) : Sqlite324Ty
         }
       }
     }
+    sqlStmt.updateStmtLimited?.let { update ->
+      check(update is SqliteUpdateStmtLimited)
+      update.returningClause?.let {
+        return object : QueryWithResults {
+          override var statement: SqlAnnotatedElement = update
+          override val select = it
+          override val pureTable = update.qualifiedTableName.tableName
+        }
+      }
+    }
+    sqlStmt.deleteStmtLimited?.let { delete ->
+      check(delete is SqliteDeleteStmtLimited)
+      delete.returningClause?.let {
+        return object : QueryWithResults {
+          override var statement: SqlAnnotatedElement = delete
+          override val select = it
+          override val pureTable = delete.qualifiedTableName?.tableName
+        }
+      }
+    }
     return parentResolver.queryWithResults(sqlStmt)
   }
 }

sqldelight-compiler/src/test/kotlin/app/cash/sqldelight/core/queries/SelectQueryTypeTest.kt

@@ -60,6 +60,53 @@ class SelectQueryTypeTest {
     )
   }
 
+  @Test fun `returning clause in an update correctly generates a query function`(dialect: TestDialect) {
+    assumeTrue(dialect in listOf(TestDialect.POSTGRESQL, TestDialect.SQLITE_3_35))
+    val file = FixtureCompiler.parseSql(
+      """
+      |CREATE TABLE IF NOT EXISTS users(
+      |    id ${dialect.intType} PRIMARY KEY,
+      |    firstname ${dialect.textType} NOT NULL,
+      |    lastname ${dialect.textType} NOT NULL
+      |);
+      |
+      |update:
+      |UPDATE users SET
+      |    firstname = :firstname,
+      |    lastname = :lastname
+      |WHERE id = :id
+      |RETURNING id, firstname, lastname;
+      |""".trimMargin(),
+      tempFolder,
+      dialect = PostgreSqlDialect()
+    )
+
+    val query = file.namedQueries.first()
+    val generator = SelectQueryGenerator(query)
+
+    assertThat(generator.customResultTypeFunction().toString()).isEqualTo(
+      """
+      |public fun <T : kotlin.Any> update(
+      |  firstname: kotlin.String,
+      |  lastname: kotlin.String,
+      |  id: kotlin.Int,
+      |  mapper: (
+      |    id: kotlin.Int,
+      |    firstname: kotlin.String,
+      |    lastname: kotlin.String,
+      |  ) -> T,
+      |): app.cash.sqldelight.ExecutableQuery<T> = UpdateQuery(firstname, lastname, id) { cursor ->
+      |  check(cursor is app.cash.sqldelight.driver.jdbc.JdbcCursor)
+      |  mapper(
+      |    cursor.getLong(0)!!.toInt(),
+      |    cursor.getString(1)!!,
+      |    cursor.getString(2)!!
+      |  )
+      |}
+      |""".trimMargin()
+    )
+  }
+
   @Test fun `query type generates properly`() {
     val file = FixtureCompiler.parseSql(
       """