Bind arguments with SqlBinaryExpr

sqldelight-compiler/src/main/kotlin/app/cash/sqldelight/core/lang/util/Arguments.kt

@@ -25,12 +25,16 @@ import app.cash.sqldelight.dialect.api.PrimitiveType.NULL
 import app.cash.sqldelight.dialect.api.PrimitiveType.TEXT
 import app.cash.sqldelight.dialect.api.SelectQueryable
 import com.alecstrong.sql.psi.core.psi.SqlBetweenExpr
+import com.alecstrong.sql.psi.core.psi.SqlBinaryBooleanExpr
+import com.alecstrong.sql.psi.core.psi.SqlBinaryEqualityExpr
 import com.alecstrong.sql.psi.core.psi.SqlBinaryExpr
 import com.alecstrong.sql.psi.core.psi.SqlBinaryLikeExpr
+import com.alecstrong.sql.psi.core.psi.SqlBinaryPipeExpr
 import com.alecstrong.sql.psi.core.psi.SqlBindExpr
 import com.alecstrong.sql.psi.core.psi.SqlCaseExpr
 import com.alecstrong.sql.psi.core.psi.SqlCastExpr
 import com.alecstrong.sql.psi.core.psi.SqlCollateExpr
+import com.alecstrong.sql.psi.core.psi.SqlColumnExpr
 import com.alecstrong.sql.psi.core.psi.SqlCompoundSelectStmt
 import com.alecstrong.sql.psi.core.psi.SqlExpr
 import com.alecstrong.sql.psi.core.psi.SqlFunctionExpr
@@ -114,11 +118,24 @@ internal fun SqlExpr.argumentType(argument: SqlExpr): IntermediateType {
         IntermediateType(PrimitiveType.BOOLEAN)
       }
     }
-    is SqlBetweenExpr, is SqlIsExpr, is SqlBinaryExpr -> {
+
+    is SqlBinaryPipeExpr, is SqlBinaryEqualityExpr, is SqlIsExpr, is SqlBinaryBooleanExpr, is SqlBetweenExpr -> {
       val validArg = children.lastOrNull { it is SqlExpr && it !== argument && it !is SqlBindExpr }
       validArg?.type() ?: children.last { it is SqlExpr && it !== argument }.type()
     }
 
+    is SqlBinaryExpr -> {
+      val validArg = children.lastOrNull {
+        it is SqlCastExpr && it == argument
+      } ?: children.lastOrNull {
+        it is SqlColumnExpr
+      } ?: parent.children.lastOrNull {
+        it is SqlExpr && it !== argument && it !is SqlBinaryExpr
+      }
+
+      validArg?.type() ?: children.last { it is SqlExpr && it !== argument }.type()
+    }
+
     is SqlNullExpr -> IntermediateType(NULL).asNullable()
     is SqlBinaryLikeExpr -> {
       val other = children.last { it is SqlExpr && it !== argument }.type()

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

@@ -1,16 +1,19 @@
 package app.cash.sqldelight.core
 
+import app.cash.sqldelight.core.lang.argumentType
 import app.cash.sqldelight.core.lang.types.typeResolver
 import app.cash.sqldelight.core.lang.util.argumentType
 import app.cash.sqldelight.core.lang.util.findChildrenOfType
 import app.cash.sqldelight.core.lang.util.isArrayParameter
 import app.cash.sqldelight.dialect.api.PrimitiveType
+import app.cash.sqldelight.dialects.postgresql.PostgreSqlDialect
 import app.cash.sqldelight.dialects.sqlite_3_24.SqliteDialect
 import app.cash.sqldelight.test.util.FixtureCompiler
 import com.alecstrong.sql.psi.core.psi.SqlBindExpr
 import com.alecstrong.sql.psi.core.psi.SqlColumnDef
 import com.google.common.truth.Truth.assertThat
 import com.squareup.kotlinpoet.asClassName
+import java.time.Instant
 import kotlin.test.assertFailsWith
 import org.junit.Rule
 import org.junit.Test
@@ -386,5 +389,121 @@ class BindArgsTest {
     }
   }
 
+  @Test fun `bind arg type in binary expression can be inferred from column`() {
+    val file = FixtureCompiler.parseSql(
+      """
+        |CREATE TABLE data (
+        |  datum INTEGER NOT NULL
+        |);
+        |
+        |selectData:
+        |SELECT *
+        |FROM data
+        |WHERE datum > :datum1 - 2.5 AND datum < :datum2 + 2.5;
+      """.trimMargin(),
+      tempFolder,
+    )
+    val column = file.namedQueries.first()
+    column.parameters.let { args ->
+      assertThat(args[0].dialectType).isEqualTo(PrimitiveType.INTEGER)
+      assertThat(args[0].javaType).isEqualTo(Long::class.asClassName())
+      assertThat(args[0].name).isEqualTo("datum1")
+
+      assertThat(args[1].dialectType).isEqualTo(PrimitiveType.INTEGER)
+      assertThat(args[1].javaType).isEqualTo(Long::class.asClassName())
+      assertThat(args[1].name).isEqualTo("datum2")
+    }
+  }
+
+  @Test fun `bind arg in arithmetic binary expression can be cast as type`() {
+    val file = FixtureCompiler.parseSql(
+      """
+        |CREATE TABLE data (
+        | datum INTEGER NOT NULL,
+        | point INTEGER NOT NULL
+        |);
+        |
+        |selectData:
+        |SELECT *, (datum + CAST(:datum1 AS REAL) * point) AS expected_datum
+        |FROM data;
+      """.trimMargin(),
+      tempFolder,
+    )
+
+    val column = file.namedQueries.first()
+    column.parameters.let { args ->
+      assertThat(args[0].dialectType).isEqualTo(PrimitiveType.REAL)
+      assertThat(args[0].javaType).isEqualTo(Double::class.asClassName().copy(nullable = true))
+      assertThat(args[0].name).isEqualTo("datum1")
+    }
+  }
+
+  @Test fun `bind arg in binary expression can be cast as type`() {
+    val file = FixtureCompiler.parseSql(
+      """
+        |CREATE TABLE data (
+        |  datum INTEGER NOT NULL
+        |);
+        |
+        |selectData:
+        |SELECT CAST(:datum1 AS REAL) + CAST(:datum2 AS INTEGER) - 10.5
+        |FROM data;
+      """.trimMargin(),
+      tempFolder,
+    )
+
+    val column = file.namedQueries.first()
+    column.parameters.let { args ->
+      assertThat(args[0].dialectType).isEqualTo(PrimitiveType.REAL)
+      assertThat(args[0].javaType).isEqualTo(Double::class.asClassName().copy(nullable = true))
+      assertThat(args[0].name).isEqualTo("datum1")
+
+      assertThat(args[1].dialectType).isEqualTo(PrimitiveType.INTEGER)
+      assertThat(args[1].javaType).isEqualTo(Long::class.asClassName().copy(nullable = true))
+      assertThat(args[1].name).isEqualTo("datum2")
+    }
+  }
+
+  @Test fun `bind arg in binary expression can be cast as custom type`() {
+    val file = FixtureCompiler.parseSql(
+      """
+        |import java.time.Instant;
+        |
+        |CREATE TABLE session (
+        |  id UUID PRIMARY KEY,
+        |  created_at TIMESTAMP AS Instant NOT NULL,
+        |  updated_at TIMESTAMP AS Instant NOT NULL
+        |);
+        |
+        |selectSession1:
+        |SELECT *
+        |FROM session
+        |WHERE created_at = :createdAt - INTERVAL '2 days' OR updated_at = :updatedAt + INTERVAL '2 days';
+        |
+        |selectSession2:
+        |SELECT *
+        |FROM session
+        |WHERE created_at BETWEEN :createdAt - INTERVAL '2 days' AND :createdAt + INTERVAL '2 days';
+      """.trimMargin(),
+      tempFolder,
+      dialect = PostgreSqlDialect(),
+    )
+
+    val selectSession1 = file.namedQueries[0]
+    selectSession1.parameters.let { args ->
+      assertThat(args[0].javaType).isEqualTo(Instant::class.asClassName())
+      assertThat(args[0].name).isEqualTo("createdAt")
+
+      assertThat(args[1].javaType).isEqualTo(Instant::class.asClassName())
+      assertThat(args[1].name).isEqualTo("updatedAt")
+    }
+
+    val selectSession2 = file.namedQueries[1]
+    selectSession2.parameters.let { args ->
+      assertThat(args[0].javaType).isEqualTo(Instant::class.asClassName())
+      assertThat(args[0].name).isEqualTo("createdAt")
+    }
+  }
+
   private fun SqlBindExpr.argumentType() = typeResolver.argumentType(this)
 }