Custom dialect with extension function don't create imports

[hfhbd]

SQLDelight Version

2.0.0-alpha03

Operating System

macOS 13

Gradle Version

7.4.2

Kotlin Version

1.7.0

Dialect

Custom

AGP Version

No response

Describe the Bug

Hey, I use a custom dialect which basically just replaces the java time api with kotlinx-datetime in the mapper, so you don't need x adapter.
But converting the time types requires the usage of kotlin extension functions.

// ... removed other dialect code
    override fun cursorGetter(columnIndex: Int, cursorName: String): CodeBlock {
        return with(CodeBlock.builder()) {
            when (this@PostgreSqlType) {
                TIMESTAMP_TIMEZONE -> add(
                    "($cursorName.getObject($columnIndex) as java.time.OffsetDateTime?)?.toInstant()?.%M()",
                    MemberName("kotlinx.datetime", "toKotlinInstant", isExtension = true)
                )
                UUID -> add("($cursorName.getObject($columnIndex) as java.util.UUID?)")
                    .add("?.%M()", MemberName("kotlinx.uuid", "toKotlinUUID", isExtension = true))
            }
        }.build()
    }
}

results into this code:

mapper(
      cursor.getString(0)!!,
      (cursor.getObject(1) as
          java.time.OffsetDateTime?)?.toInstant()?.kotlinx.datetime.toKotlinInstant()!!
    )

I thought, this is a kotlinpoet issue, but I can't reproduce it with a small reproducer. So maybe sqldelight does some magic.

kotlinPoetWorkingReproducer.zip

Stacktrace

No response

Gradle Build Script

plugins {
    kotlin("jvm") version "1.7.0"
    id("app.cash.sqldelight") version "2.0.0-alpha03"
}

dependencies {
    implementation("org.jetbrains.kotlinx:kotlinx-datetime:0.4.0")
}

sqldelight {
    database("DB") {
        dialect(projects.dialect)
        packageName = "example"
    }
}

[morki]

@hfhbd could you please share the code for your custom dialect for inspiration? I thought about the same aproach to remove hundreds of adapter definitions for kotlinx.datetime and other basic classes.

[hfhbd]

Hey @morki, sure, you could use this test as a very basic dialect: https://github.com/cashapp/sqldelight/tree/master/sqldelight-gradle-plugin/src/test/custom-dialect or

Or use https://github.com/hfhbd/postgres-native-sqldelight/tree/main/postgres-native-sqldelight-dialect with postgres dialect with kotlinx types. If you want to use postgresql too, I also started adding a dialect combining/wrapping jvm and postgresql native: hfhbd/postgres-native-sqldelight#133

Feel free to ask me for more details/help. I also plan to document it in the docs (someday).

[morki]

Thank you very much for your help @hfhbd, I want to extend SqliteDialect, but not with custom types like Postgres have, but with "custom mapping" instead of always specifying ColumnAdapter, and I don't see how to do it from the examples above.

For example, I want every column, which type is INTEGER and kotlin type is kotlinx.datetime.Instant to use custom CodeBlock to map those.

Can you please make a little example with this kind of mapping? Or is it even possible in custom dialect?

[hfhbd]

For example, I want every column, which type is INTEGER and kotlin type is kotlinx.datetime.Instant to use custom CodeBlock to map those.

Do you mean this?

CREATE TABLE foo (
  bar INTEGER AS kotlinx.datetime.Instant, // replace with custom logic
  a INTEGER, // keep
  a INTEGER AS kotlin.time.Duration // keep, only for demo
);

This is not yet supported, because the AS kotlinx.datetime.Instant annotation rule is defined in the sqldelight dialect, but a (custom) dialect does only use the ANSI dialect defined in sql-psi.
Even if you implement your own type using grammar eg INSTANTINTEGER (bar INSTANTEGER NULL) you can't replace the invalid sql with a valid one, there is no method in the dialect api yet. We really should remove these hardcoded values/logic and use a better api:

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

Line 146 in f7af212

private fun PsiElement.rangesToReplace(): List<Pair<IntRange, String>> {

So the only option to use custom logic to remove column adapters are specific valid SQL types:

CREATE TABLE foo (
  bar TIMESTAMPZ // should use kotlinx.datetime.Instant and not java.time.Instant
);

with this code (for postgres):

package wipma.dialect

import app.cash.sqldelight.dialect.api.*
import app.cash.sqldelight.dialects.postgresql.*
import app.cash.sqldelight.dialects.postgresql.grammar.psi.*
import com.alecstrong.sql.psi.core.psi.*
import com.squareup.kotlinpoet.*

class PostgresKotlinxDateTimeDialect : SqlDelightDialect by PostgreSqlDialect() {
    override fun typeResolver(parentResolver: TypeResolver): TypeResolver = PostgresTypeResolver(parentResolver)
}

private class PostgresTypeResolver(parentResolver: TypeResolver) :
    TypeResolver by PostgreSqlTypeResolver(parentResolver) {
    override fun definitionType(typeName: SqlTypeName): IntermediateType = with(typeName) {
        check(this is PostgreSqlTypeName)
        val type = IntermediateType(
            when {
                smallIntDataType != null -> PostgreSqlType.SMALL_INT
                intDataType != null -> PostgreSqlType.INTEGER
                bigIntDataType != null -> PostgreSqlType.BIG_INT
                approximateNumericDataType != null -> PrimitiveType.REAL
                stringDataType != null -> PrimitiveType.TEXT
                uuidDataType != null -> PostgreSqlType.UUID
                smallSerialDataType != null -> PostgreSqlType.SMALL_INT
                serialDataType != null -> PostgreSqlType.INTEGER
                bigSerialDataType != null -> PostgreSqlType.BIG_INT
                dateDataType != null -> {
                    when (dateDataType!!.firstChild.text) {
                        "DATE" -> PostgreSqlType.DATE
                        "TIME" -> PostgreSqlType.TIME
                        "TIMESTAMP" -> if (dateDataType!!.node.getChildren(null)
                                .any { it.text == "WITH" }
                        ) PostgreSqlType.TIMESTAMP_TIMEZONE else PostgreSqlType.TIMESTAMP

                        "TIMESTAMPTZ" -> PostgreSqlType.TIMESTAMP_TIMEZONE
                        "INTERVAL" -> PostgreSqlType.INTERVAL
                        else -> throw IllegalArgumentException("Unknown date type ${dateDataType!!.text}")
                    }
                }

                jsonDataType != null -> PrimitiveType.TEXT
                booleanDataType != null -> PrimitiveType.BOOLEAN
                blobDataType != null -> PrimitiveType.BLOB
                else -> throw IllegalArgumentException("Unknown kotlin type for sql type ${this.text}")
            }
        )
        return type
    }
}

private enum class PostgreSqlType(override val javaType: TypeName) : DialectType {
    SMALL_INT(SHORT) {
        override fun decode(value: CodeBlock) = CodeBlock.of("%L.toShort()", value)

        override fun encode(value: CodeBlock) = CodeBlock.of("%L.toLong()", value)
    },
    INTEGER(INT) {
        override fun decode(value: CodeBlock) = CodeBlock.of("%L.toInt()", value)

        override fun encode(value: CodeBlock) = CodeBlock.of("%L.toLong()", value)
    },
    BIG_INT(LONG),
    DATE(ClassName("kotlinx.datetime", "LocalDate")),

    TIME(ClassName("kotlinx.datetime", "LocalTime")),
    TIMESTAMP(ClassName("kotlinx.datetime", "LocalDateTime")),
    TIMESTAMP_TIMEZONE(ClassName("kotlinx.datetime", "Instant")),
    INTERVAL(ClassName("kotlin.time", "Duration")),
    UUID(ClassName("kotlinx.uuid", "UUID"));

    override fun prepareStatementBinder(columnIndex: String, value: CodeBlock): CodeBlock {
        return CodeBlock.builder()
            .add(
                when (this) {
                    SMALL_INT, INTEGER, BIG_INT -> "bindLong"
                    DATE -> "bindObject"
                    TIME -> "bindObject"
                    TIMESTAMP -> "bindObject"
                    TIMESTAMP_TIMEZONE -> "bindObject"
                    INTERVAL -> "bindObject"
                    UUID -> "bindObject"
                }
            )
            .add(
                "($columnIndex, %L)\n",
                when (this) {
                    SMALL_INT, INTEGER, BIG_INT -> value
                    DATE -> value.toBuilder()
                        .add(".%M()", MemberName("kotlinx.datetime", "toJavaLocalDate", isExtension = true)).build()

                    TIME -> value.toBuilder()
                        .add(".%M()", MemberName("kotlinx.datetime", "toJavaLocalTime", isExtension = true)).build()

                    TIMESTAMP -> value.toBuilder()
                        .add(".%M()", MemberName("kotlinx.datetime", "toJavaLocalDateTime", isExtension = true))
                        .build()

                    TIMESTAMP_TIMEZONE -> value.toBuilder()
                        .add(".%M()", MemberName("kotlinx.datetime", "toJavaInstant", isExtension = true)).build()

                    INTERVAL -> value.toBuilder().add(".toIsoString()").build()
                    UUID -> value.toBuilder().add(".%M()", MemberName("kotlinx.uuid", "toJavaUUID", isExtension = true))
                        .build()
                }
            )
            .build()
    }

    override fun cursorGetter(columnIndex: Int, cursorName: String): CodeBlock {
        return with(CodeBlock.builder()) {
            when (this@PostgreSqlType) {
                SMALL_INT, INTEGER, BIG_INT -> add("$cursorName.getLong($columnIndex)")
                DATE -> add(
                    "($cursorName.getObject<java.time.LocalDate>($columnIndex))?.%M()",
                    MemberName("kotlinx.datetime", "toKotlinLocalDate", isExtension = true)
                )

                TIME -> add(
                    "($cursorName.getObject<java.time.LocalTime>($columnIndex))?.%M()",
                    MemberName("kotlinx.datetime", "toKotlinLocalTime", isExtension = true)
                )

                TIMESTAMP -> add(
                    "($cursorName.getObject<java.time.LocalDateTime>($columnIndex))?.%M()",
                    MemberName("kotlinx.datetime", "toKotlinLocalDateTime", isExtension = true)
                )

                TIMESTAMP_TIMEZONE -> add(
                    "($cursorName.getObject<java.time.OffsetDateTime>($columnIndex))?.toInstant()?.%M()",
                    MemberName("kotlinx.datetime", "toKotlinInstant", isExtension = true)
                )

                INTERVAL -> add(
                    "$cursorName.getObject<%T>($columnIndex)", 
                    ClassName("org.postgresql.util", "PGInterval")
                )
                UUID -> add(
                    "($cursorName.getObject($columnIndex) as java.util.UUID?)?.%M()",
                    MemberName("kotlinx.uuid", "toKotlinUUID", isExtension = true)
                )
            }
        }.build()
    }
}

[morki]

Thank you very much for your long and complete response @hfhbd. So for now it is not possible as Sqlite dialect extension in this form if I understand this correctly. Never mind, column adapters are good workaround for now, only very verbose.

Thanks again for your time :)

[hfhbd]

Yeah, it would be possible if we refactored the hardcoded replacement.