Correctly infer argument types in compound selects

cashapp · Apr 16, 2022 · 4acfcaf · 4acfcaf
1 parent c13aa6d
commit 4acfcaf
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 2 deletions.
diff --git a/sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/lang/util/Arguments.kt b/sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/lang/util/Arguments.kt
@@ -169,8 +169,10 @@ private fun SqlSelectStmt.argumentType(result: SqlResultColumn): IntermediateTyp
     else -> {
       // Check if this is part of an inner expression of a resulit column.
       val parentResult = PsiTreeUtil.getParentOfType(parentRule, SqlResultColumn::class.java)
-        ?: return null
-      (parentResult.parent as SqlSelectStmt).argumentType(parentResult)
+
+      if (parentResult == null) NamedQuery("temp", SelectQueryable(compoundSelect, compoundSelect))
+        .resultColumns[resultColumnList.indexOf(result)]
+      else (parentResult.parent as SqlSelectStmt).argumentType(parentResult)
     }
   }
 }

diff --git a/...ight-compiler/src/test/kotlin/app/cash/sqldelight/core/queries/SelectQueryFunctionTest.kt b/...ight-compiler/src/test/kotlin/app/cash/sqldelight/core/queries/SelectQueryFunctionTest.kt
@@ -1718,4 +1718,27 @@ class SelectQueryFunctionTest {
         |""".trimMargin()
     )
   }
+
+  @Test fun `parameter correctly generated for union`() {
+    val file = FixtureCompiler.parseSql(
+      """
+      |query:
+      |SELECT 'foo'
+      |UNION
+      |SELECT ?;
+      |""".trimMargin(),
+      tempFolder
+    )
+
+    val query = file.namedQueries.first()
+    val generator = SelectQueryGenerator(query)
+
+    assertThat(generator.customResultTypeFunction().toString()).isEqualTo(
+      """
+        |public fun query(expr: kotlin.String): app.cash.sqldelight.Query<kotlin.String> = QueryQuery(expr) { cursor ->
+        |  cursor.getString(0)!!
+        |}
+        |""".trimMargin()
+    )
+  }
 }