Prefer Kotlin type for some functions (#4517)

sqldelight · Apr 2, 2024 · 9245d3d · 9245d3d
1 parent e76eb29
commit 9245d3d
Show file tree

Hide file tree

Showing 11 changed files with 530 additions and 18 deletions.
diff --git a/dialects/hsql/src/main/kotlin/app/cash/sqldelight/dialects/hsql/HsqlTypeResolver.kt b/dialects/hsql/src/main/kotlin/app/cash/sqldelight/dialects/hsql/HsqlTypeResolver.kt
@@ -7,7 +7,7 @@ import app.cash.sqldelight.dialect.api.PrimitiveType.INTEGER
 import app.cash.sqldelight.dialect.api.PrimitiveType.REAL
 import app.cash.sqldelight.dialect.api.PrimitiveType.TEXT
 import app.cash.sqldelight.dialect.api.TypeResolver
-import app.cash.sqldelight.dialect.api.encapsulatingType
+import app.cash.sqldelight.dialect.api.encapsulatingTypePreferringKotlin
 import app.cash.sqldelight.dialects.hsql.HsqlType.BIG_INT
 import app.cash.sqldelight.dialects.hsql.HsqlType.SMALL_INT
 import app.cash.sqldelight.dialects.hsql.HsqlType.TINY_INT
@@ -42,9 +42,31 @@ class HsqlTypeResolver(private val parentResolver: TypeResolver) : TypeResolver
   }
 
   private fun SqlFunctionExpr.hsqlFunctionType() = when (functionName.text.lowercase()) {
-    "coalesce", "ifnull" -> encapsulatingType(exprList, TINY_INT, SMALL_INT, HsqlType.INTEGER, INTEGER, BIG_INT, REAL, TEXT, BLOB)
-    "max" -> encapsulatingType(exprList, TINY_INT, SMALL_INT, HsqlType.INTEGER, INTEGER, BIG_INT, REAL, TEXT, BLOB).asNullable()
-    "min" -> encapsulatingType(exprList, BLOB, TEXT, TINY_INT, SMALL_INT, INTEGER, HsqlType.INTEGER, BIG_INT, REAL).asNullable()
+    "coalesce", "ifnull" -> encapsulatingTypePreferringKotlin(exprList, TINY_INT, SMALL_INT, HsqlType.INTEGER, INTEGER, BIG_INT, REAL, TEXT, BLOB)
+    "greatest" -> encapsulatingTypePreferringKotlin(
+      exprList,
+      TINY_INT,
+      SMALL_INT,
+      HsqlType.INTEGER,
+      INTEGER,
+      BIG_INT,
+      REAL,
+      TEXT,
+      BLOB,
+    )
+    "least" -> encapsulatingTypePreferringKotlin(
+      exprList,
+      BLOB,
+      TEXT,
+      TINY_INT,
+      SMALL_INT,
+      INTEGER,
+      HsqlType.INTEGER,
+      BIG_INT,
+      REAL,
+    )
+    "max" -> encapsulatingTypePreferringKotlin(exprList, TINY_INT, SMALL_INT, HsqlType.INTEGER, INTEGER, BIG_INT, REAL, TEXT, BLOB).asNullable()
+    "min" -> encapsulatingTypePreferringKotlin(exprList, BLOB, TEXT, TINY_INT, SMALL_INT, INTEGER, HsqlType.INTEGER, BIG_INT, REAL).asNullable()
     "length", "char_length", "character_length" -> IntermediateType(BIG_INT).nullableIf(resolvedType(exprList[0]).javaType.isNullable)
     else -> null
   }

diff --git a/dialects/mysql/src/main/kotlin/app/cash/sqldelight/dialects/mysql/MySqlTypeResolver.kt b/dialects/mysql/src/main/kotlin/app/cash/sqldelight/dialects/mysql/MySqlTypeResolver.kt
@@ -9,6 +9,7 @@ import app.cash.sqldelight.dialect.api.PrimitiveType.REAL
 import app.cash.sqldelight.dialect.api.PrimitiveType.TEXT
 import app.cash.sqldelight.dialect.api.TypeResolver
 import app.cash.sqldelight.dialect.api.encapsulatingType
+import app.cash.sqldelight.dialect.api.encapsulatingTypePreferringKotlin
 import app.cash.sqldelight.dialects.mysql.MySqlType.BIG_INT
 import app.cash.sqldelight.dialects.mysql.MySqlType.SMALL_INT
 import app.cash.sqldelight.dialects.mysql.MySqlType.TINY_INT
@@ -93,16 +94,45 @@ class MySqlTypeResolver(
   }
 
   private fun SqlFunctionExpr.mySqlFunctionType() = when (functionName.text.lowercase()) {
-    "greatest" -> encapsulatingType(exprList, INTEGER, REAL, TEXT, BLOB)
+    "greatest" -> encapsulatingTypePreferringKotlin(
+      exprList,
+      TINY_INT,
+      SMALL_INT,
+      MySqlType.INTEGER,
+      INTEGER,
+      BIG_INT,
+      REAL,
+      MySqlType.TIMESTAMP,
+      MySqlType.DATE,
+      MySqlType.DATETIME,
+      MySqlType.TIME,
+      TEXT,
+      BLOB,
+    )
+    "least" -> encapsulatingTypePreferringKotlin(
+      exprList,
+      BLOB,
+      TEXT,
+      MySqlType.TIME,
+      MySqlType.DATETIME,
+      MySqlType.DATE,
+      MySqlType.TIMESTAMP,
+      TINY_INT,
+      SMALL_INT,
+      INTEGER,
+      MySqlType.INTEGER,
+      BIG_INT,
+      REAL,
+    )
     "concat" -> encapsulatingType(exprList, TEXT)
     "last_insert_id" -> IntermediateType(INTEGER)
     "row_count" -> IntermediateType(INTEGER)
     "microsecond", "second", "minute", "hour", "day", "week", "month", "year" -> IntermediateType(
       INTEGER,
     )
     "sin", "cos", "tan" -> IntermediateType(REAL)
-    "coalesce", "ifnull" -> encapsulatingType(exprList, TINY_INT, SMALL_INT, MySqlType.INTEGER, INTEGER, BIG_INT, REAL, TEXT, BLOB)
-    "max" -> encapsulatingType(
+    "coalesce", "ifnull" -> encapsulatingTypePreferringKotlin(exprList, TINY_INT, SMALL_INT, MySqlType.INTEGER, INTEGER, BIG_INT, REAL, TEXT, BLOB)
+    "max" -> encapsulatingTypePreferringKotlin(
       exprList,
       TINY_INT,
       SMALL_INT,
@@ -117,7 +147,7 @@ class MySqlTypeResolver(
       TEXT,
       BLOB,
     ).asNullable()
-    "min" -> encapsulatingType(
+    "min" -> encapsulatingTypePreferringKotlin(
       exprList,
       BLOB,
       TEXT,

diff --git a/...tgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/PostgreSqlTypeResolver.kt b/...tgresql/src/main/kotlin/app/cash/sqldelight/dialects/postgresql/PostgreSqlTypeResolver.kt
@@ -11,6 +11,7 @@ import app.cash.sqldelight.dialect.api.QueryWithResults
 import app.cash.sqldelight.dialect.api.ReturningQueryable
 import app.cash.sqldelight.dialect.api.TypeResolver
 import app.cash.sqldelight.dialect.api.encapsulatingType
+import app.cash.sqldelight.dialect.api.encapsulatingTypePreferringKotlin
 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
@@ -90,12 +91,35 @@ class PostgreSqlTypeResolver(private val parentResolver: TypeResolver) : TypeRes
   }
 
   private fun SqlFunctionExpr.postgreSqlFunctionType() = when (functionName.text.lowercase()) {
-    "greatest" -> encapsulatingType(exprList, INTEGER, REAL, TEXT, BLOB)
+    "greatest" -> encapsulatingTypePreferringKotlin(
+      exprList,
+      SMALL_INT,
+      PostgreSqlType.INTEGER,
+      INTEGER,
+      BIG_INT,
+      REAL,
+      TEXT,
+      BLOB,
+      TIMESTAMP_TIMEZONE,
+      TIMESTAMP,
+    )
+    "least" -> encapsulatingTypePreferringKotlin(
+      exprList,
+      BLOB,
+      TEXT,
+      SMALL_INT,
+      INTEGER,
+      PostgreSqlType.INTEGER,
+      BIG_INT,
+      REAL,
+      TIMESTAMP_TIMEZONE,
+      TIMESTAMP,
+    )
     "concat" -> encapsulatingType(exprList, TEXT)
     "substring" -> IntermediateType(TEXT).nullableIf(resolvedType(exprList[0]).javaType.isNullable)
-    "coalesce", "ifnull" -> encapsulatingType(exprList, SMALL_INT, PostgreSqlType.INTEGER, INTEGER, BIG_INT, REAL, TEXT, BLOB)
-    "max" -> encapsulatingType(exprList, SMALL_INT, PostgreSqlType.INTEGER, INTEGER, BIG_INT, REAL, TEXT, BLOB, TIMESTAMP_TIMEZONE, TIMESTAMP).asNullable()
-    "min" -> encapsulatingType(exprList, BLOB, TEXT, SMALL_INT, INTEGER, PostgreSqlType.INTEGER, BIG_INT, REAL, TIMESTAMP_TIMEZONE, TIMESTAMP).asNullable()
+    "coalesce", "ifnull" -> encapsulatingTypePreferringKotlin(exprList, SMALL_INT, PostgreSqlType.INTEGER, INTEGER, BIG_INT, REAL, TEXT, BLOB)
+    "max" -> encapsulatingTypePreferringKotlin(exprList, SMALL_INT, PostgreSqlType.INTEGER, INTEGER, BIG_INT, REAL, TEXT, BLOB, TIMESTAMP_TIMEZONE, TIMESTAMP).asNullable()
+    "min" -> encapsulatingTypePreferringKotlin(exprList, BLOB, TEXT, SMALL_INT, INTEGER, PostgreSqlType.INTEGER, BIG_INT, REAL, TIMESTAMP_TIMEZONE, TIMESTAMP).asNullable()
     "sum" -> {
       val type = resolvedType(exprList.single())
       if (type.dialectType == REAL) {

diff --git a/sqldelight-compiler/dialect/src/main/kotlin/app/cash/sqldelight/dialect/api/TypeResolver.kt b/sqldelight-compiler/dialect/src/main/kotlin/app/cash/sqldelight/dialect/api/TypeResolver.kt
@@ -46,13 +46,23 @@ fun TypeResolver.encapsulatingType(
   vararg typeOrder: DialectType,
 ) = encapsulatingType(exprList = exprList, nullableIfAny = false, typeOrder = typeOrder)
 
+/**
+ * @return the type from the expr list which is the highest order in the typeOrder list
+ */
+fun TypeResolver.encapsulatingTypePreferringKotlin(
+  exprList: List<SqlExpr>,
+  vararg typeOrder: DialectType,
+  nullableIfAny: Boolean = false,
+) = encapsulatingType(exprList = exprList, nullableIfAny = nullableIfAny, preferKotlinType = true, typeOrder = typeOrder)
+
 /**
  * @return the type from the expr list which is the highest order in the typeOrder list
  */
 fun TypeResolver.encapsulatingType(
   exprList: List<SqlExpr>,
   nullableIfAny: Boolean,
   vararg typeOrder: DialectType,
+  preferKotlinType: Boolean = false,
 ): IntermediateType {
   val types = exprList.map { resolvedType(it) }
   val sqlTypes = types.map { it.dialectType }
@@ -68,12 +78,18 @@ fun TypeResolver.encapsulatingType(
     error("The Kotlin type of the argument cannot be inferred, use CAST instead.")
   }
 
-  val type = typeOrder.last { it in sqlTypes }
+  // stripping nullability because that shouldn't affect the type comparison
+  val isTypesHomogeneous = types.map { it.dialectType to it.javaType.copy(nullable = false) }.distinct().size == 1
+  val type = if (preferKotlinType && isTypesHomogeneous) {
+    types.first()
+  } else {
+    IntermediateType(typeOrder.last { it in sqlTypes })
+  }
 
   if (!nullableIfAny && types.all { it.javaType.isNullable } ||
     nullableIfAny && types.any { it.javaType.isNullable }
   ) {
-    return IntermediateType(type).asNullable()
+    return type.asNullable()
   }
-  return IntermediateType(type)
+  return type
 }
diff --git a/sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/lang/util/ExprUtil.kt b/sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/lang/util/ExprUtil.kt
@@ -31,6 +31,7 @@ import app.cash.sqldelight.dialect.api.QueryWithResults
 import app.cash.sqldelight.dialect.api.SelectQueryable
 import app.cash.sqldelight.dialect.api.TypeResolver
 import app.cash.sqldelight.dialect.api.encapsulatingType
+import app.cash.sqldelight.dialect.api.encapsulatingTypePreferringKotlin
 import com.alecstrong.sql.psi.core.psi.SqlBetweenExpr
 import com.alecstrong.sql.psi.core.psi.SqlBinaryAddExpr
 import com.alecstrong.sql.psi.core.psi.SqlBinaryExpr
@@ -137,10 +138,10 @@ internal object AnsiSqlTypeResolver : TypeResolver {
     "avg" -> IntermediateType(REAL).asNullable()
     "abs" -> encapsulatingType(exprList, INTEGER, REAL)
     "iif" -> exprList[1].type()
-    "coalesce", "ifnull" -> encapsulatingType(exprList, INTEGER, REAL, TEXT, BLOB)
+    "coalesce", "ifnull" -> encapsulatingTypePreferringKotlin(exprList, INTEGER, REAL, TEXT, BLOB)
     "nullif" -> exprList[0].type().asNullable()
-    "max" -> encapsulatingType(exprList, INTEGER, REAL, TEXT, BLOB, BOOLEAN).asNullable()
-    "min" -> encapsulatingType(exprList, BLOB, TEXT, INTEGER, REAL, BOOLEAN).asNullable()
+    "max" -> encapsulatingTypePreferringKotlin(exprList, INTEGER, REAL, TEXT, BLOB, BOOLEAN).asNullable()
+    "min" -> encapsulatingTypePreferringKotlin(exprList, BLOB, TEXT, INTEGER, REAL, BOOLEAN).asNullable()
     else -> null
   }
 

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
@@ -743,6 +743,10 @@ class QueriesTypeTest {
       |SELECT *
       |FROM soupView
       |WHERE soup_token = ?;
+      |
+      |maxSoupBroth:
+      |SELECT MAX(soup_broth)
+      |FROM soupView;
       """.trimMargin(),
       temporaryFolder,
       fileName = "MyView.sq",
@@ -798,6 +802,22 @@ class QueriesTypeTest {
       |    )
       |  }
       |
+      |  public fun <T : Any> maxSoupBroth(mapper: (MAX: ChickenSoupBase.Broth?) -> T): Query<T> =
+      |      Query(-1_892_940_684, arrayOf("soupBase", "soup"), driver, "MyView.sq", "maxSoupBroth", ""${'"'}
+      |  |SELECT MAX(soup_broth)
+      |  |FROM soupView
+      |  ""${'"'}.trimMargin()) { cursor ->
+      |    mapper(
+      |      cursor.getBytes(0)?.let { soupBaseAdapter.soup_brothAdapter.decode(it) }
+      |    )
+      |  }
+      |
+      |  public fun maxSoupBroth(): Query<MaxSoupBroth> = maxSoupBroth { MAX ->
+      |    MaxSoupBroth(
+      |      MAX
+      |    )
+      |  }
+      |
       |  private inner class ForSoupTokenQuery<out T : Any>(
       |    public val soup_token: String,
       |    mapper: (SqlCursor) -> T,