Fix compilation of grouped statements with returns

cashapp · Oct 6, 2023 · 125ab42 · 125ab42
1 parent 92c0591
commit 125ab42
Show file tree

Hide file tree

Showing 5 changed files with 114 additions and 2 deletions.
diff --git a/...elight-compiler/src/main/kotlin/app/cash/sqldelight/core/compiler/QueriesTypeGenerator.kt b/...elight-compiler/src/main/kotlin/app/cash/sqldelight/core/compiler/QueriesTypeGenerator.kt
@@ -61,7 +61,7 @@ class QueriesTypeGenerator(
           type.addFunction(generator.defaultResultTypeFunction())
         }
 
-        if (query.arguments.isNotEmpty()) {
+        if (query.needsQuerySubType()) {
           type.addType(generator.querySubtype())
         }
       }

diff --git a/...elight-compiler/src/main/kotlin/app/cash/sqldelight/core/compiler/SelectQueryGenerator.kt b/...elight-compiler/src/main/kotlin/app/cash/sqldelight/core/compiler/SelectQueryGenerator.kt
@@ -223,7 +223,7 @@ class SelectQueryGenerator(
     }
     mapperLambda.unindent().add("}\n")
 
-    if (query.arguments.isEmpty()) {
+    if (!query.needsQuerySubType()) {
       // No need for a custom query type, return an instance of Query:
       // return Query(statement, selectForId) { resultSet -> ... }
       val tablesObserved = query.tablesObserved

diff --git a/sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/compiler/model/NamedQuery.kt b/sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/compiler/model/NamedQuery.kt
@@ -27,6 +27,7 @@ import app.cash.sqldelight.core.lang.util.name
 import app.cash.sqldelight.core.lang.util.sqFile
 import app.cash.sqldelight.core.lang.util.tablesObserved
 import app.cash.sqldelight.core.lang.util.type
+import app.cash.sqldelight.core.psi.SqlDelightStmtClojureStmtList
 import app.cash.sqldelight.dialect.api.IntermediateType
 import app.cash.sqldelight.dialect.api.PrimitiveType.ARGUMENT
 import app.cash.sqldelight.dialect.api.PrimitiveType.BLOB
@@ -141,6 +142,8 @@ data class NamedQuery(
 
   internal fun needsWrapper() = (resultColumns.size > 1 || resultColumns[0].javaType.isNullable)
 
+  internal fun needsQuerySubType() = arguments.isNotEmpty() || statement is SqlDelightStmtClojureStmtList
+
   internal val tablesObserved: List<TableNameElement>? by lazy {
     if (queryable is SelectQueryable && queryable.select == queryable.statement) {
       queryable.select.tablesObserved()

diff --git a/sqldelight-compiler/src/test/kotlin/app/cash/sqldelight/core/QueriesTypeTest.kt b/sqldelight-compiler/src/test/kotlin/app/cash/sqldelight/core/QueriesTypeTest.kt
@@ -846,4 +846,68 @@ class QueriesTypeTest {
       """.trimMargin(),
     )
   }
+
+  @Test fun `grouped statement with return and no arguments gets a query type`() {
+    val result = FixtureCompiler.compileSql(
+      """
+      |CREATE TABLE data (
+      |  id INTEGER PRIMARY KEY,
+      |  value TEXT
+      |);
+      |
+      |insertAndReturn {
+      |  INSERT INTO data (value)
+      |  VALUES (NULL)
+      |  ;
+      |  SELECT last_insert_rowid();
+      |}
+      """.trimMargin(),
+      temporaryFolder,
+      fileName = "Data.sq",
+    )
+
+    val query = result.compiledFile.namedQueries.first()
+    assertThat(result.errors).isEmpty()
+
+    val dataQueries = File(result.outputDirectory, "com/example/DataQueries.kt")
+    assertThat(result.compilerOutput).containsKey(dataQueries)
+    assertThat(result.compilerOutput[dataQueries].toString()).isEqualTo(
+      """
+      |package com.example
+      |
+      |import app.cash.sqldelight.ExecutableQuery
+      |import app.cash.sqldelight.TransacterImpl
+      |import app.cash.sqldelight.db.QueryResult
+      |import app.cash.sqldelight.db.SqlCursor
+      |import app.cash.sqldelight.db.SqlDriver
+      |import kotlin.Any
+      |import kotlin.Long
+      |import kotlin.String
+      |
+      |public class DataQueries(
+      |  driver: SqlDriver,
+      |) : TransacterImpl(driver) {
+      |  public fun insertAndReturn(): ExecutableQuery<Long> = InsertAndReturnQuery() { cursor ->
+      |    cursor.getLong(0)!!
+      |  }
+      |
+      |  private inner class InsertAndReturnQuery<out T : Any>(
+      |    mapper: (SqlCursor) -> T,
+      |  ) : ExecutableQuery<T>(mapper) {
+      |    override fun <R> execute(mapper: (SqlCursor) -> QueryResult<R>): QueryResult<R> =
+      |        transactionWithResult {
+      |      driver.execute(${query.idForIndex(0).withUnderscores}, ""${'"'}
+      |          |INSERT INTO data (value)
+      |          |  VALUES (NULL)
+      |          ""${'"'}.trimMargin(), 0)
+      |      driver.executeQuery(${query.idForIndex(1).withUnderscores}, ""${'"'}SELECT last_insert_rowid()""${'"'}, mapper, 0)
+      |    }
+      |
+      |    override fun toString(): String = "Data.sq:insertAndReturn"
+      |  }
+      |}
+      |
+      """.trimMargin(),
+    )
+  }
 }
diff --git a/sqldelight-compiler/src/test/kotlin/app/cash/sqldelight/core/queries/SelectQueryTypeTest.kt b/sqldelight-compiler/src/test/kotlin/app/cash/sqldelight/core/queries/SelectQueryTypeTest.kt
@@ -1692,6 +1692,51 @@ class SelectQueryTypeTest {
     )
   }
 
+  @Test
+  fun `grouped statements with return and no parameters`() {
+    val file = FixtureCompiler.parseSql(
+      """
+      |CREATE TABLE data (
+      |  id INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
+      |  value INTEGER NOT NULL
+      |);
+      |
+      |insertAndReturn {
+      |  INSERT INTO data (value)
+      |  VALUES (NULL)
+      |  ;
+      |  SELECT last_insert_rowid();
+      |}
+      |
+      """.trimMargin(),
+      tempFolder,
+    )
+
+    val query = file.namedQueries.first()
+    val generator = SelectQueryGenerator(query)
+
+    println(generator.defaultResultTypeFunction().toString())
+
+    assertThat(generator.querySubtype().toString()).isEqualTo(
+      """
+      |private inner class InsertAndReturnQuery<out T : kotlin.Any>(
+      |  mapper: (app.cash.sqldelight.db.SqlCursor) -> T,
+      |) : app.cash.sqldelight.ExecutableQuery<T>(mapper) {
+      |  override fun <R> execute(mapper: (app.cash.sqldelight.db.SqlCursor) -> app.cash.sqldelight.db.QueryResult<R>): app.cash.sqldelight.db.QueryResult<R> = transactionWithResult {
+      |    driver.execute(${query.idForIndex(0).withUnderscores}, ""${'"'}
+      |        |INSERT INTO data (value)
+      |        |  VALUES (NULL)
+      |        ""${'"'}.trimMargin(), 0)
+      |    driver.executeQuery(${query.idForIndex(1).withUnderscores}, ""${'"'}SELECT last_insert_rowid()""${'"'}, mapper, 0)
+      |  }
+      |
+      |  override fun toString(): kotlin.String = "Test.sq:insertAndReturn"
+      |}
+      |
+      """.trimMargin(),
+    )
+  }
+
   @Test
   fun `proper exposure of case arguments function`() {
     val file = FixtureCompiler.parseSql(