From 77304c236e28950dc2db78af04b5f0e266617aef Mon Sep 17 00:00:00 2001
From: Sergey Zhuravlev <sergey.zhuravlev@jetbrains.com>
Date: Thu, 11 Jul 2024 10:01:38 +0200
Subject: [PATCH] feat: added support for using custom types as tables ids

---
 .../org/jetbrains/exposed/sql/ColumnType.kt   |  3 +
 .../entities/CustomIdTableEntityTest.kt       | 94 +++++++++++++++++++
 2 files changed, 97 insertions(+)
 create mode 100644 exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/entities/CustomIdTableEntityTest.kt

diff --git a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/ColumnType.kt b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/ColumnType.kt
index 3144706cea..d73e1bfa40 100644
--- a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/ColumnType.kt
+++ b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/ColumnType.kt
@@ -234,6 +234,9 @@ class EntityIDColumnType<T : Comparable<T>>(
 
     override fun readObject(rs: ResultSet, index: Int): Any? = idColumn.columnType.readObject(rs, index)
 
+    override fun setParameter(stmt: PreparedStatementApi, index: Int, value: Any?) =
+        idColumn.columnType.setParameter(stmt, index, value)
+
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
 
diff --git a/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/entities/CustomIdTableEntityTest.kt b/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/entities/CustomIdTableEntityTest.kt
new file mode 100644
index 0000000000..6d59069ea1
--- /dev/null
+++ b/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/entities/CustomIdTableEntityTest.kt
@@ -0,0 +1,94 @@
+package org.jetbrains.exposed.sql.tests.shared.entities
+
+import org.jetbrains.exposed.dao.Entity
+import org.jetbrains.exposed.dao.EntityClass
+import org.jetbrains.exposed.dao.UUIDEntity
+import org.jetbrains.exposed.dao.UUIDEntityClass
+import org.jetbrains.exposed.dao.id.EntityID
+import org.jetbrains.exposed.dao.id.IdTable
+import org.jetbrains.exposed.dao.with
+import org.jetbrains.exposed.sql.*
+import org.jetbrains.exposed.sql.statements.api.PreparedStatementApi
+import org.jetbrains.exposed.sql.tests.DatabaseTestsBase
+import org.jetbrains.exposed.sql.tests.shared.assertEquals
+import org.jetbrains.exposed.sql.tests.shared.entities.UUIDTables.Cities
+import org.jetbrains.exposed.sql.tests.shared.entities.UUIDTables.City.Companion.referrersOn
+import org.jetbrains.exposed.sql.tests.shared.entities.UUIDTables.Town
+import org.jetbrains.exposed.sql.tests.shared.entities.UUIDTables.Towns
+import org.junit.Test
+import java.sql.ResultSet
+import java.util.*
+import kotlin.test.assertNotNull
+
+data class CityId(
+    val value: String
+) : Comparable<CityId> {
+    override fun toString() = value
+    override fun compareTo(other: CityId) = value.compareTo(other.value)
+}
+
+fun Table.cityId(name: String): Column<CityId> = registerColumn(name, CityIdColumnType())
+
+class CityIdColumnType : ColumnType<CityId>() {
+    private val varcharColumnType = VarCharColumnType(10)
+
+    override fun sqlType() = varcharColumnType.sqlType()
+
+    override fun valueFromDB(value: Any): CityId = varcharColumnType.valueFromDB(value).let { CityId(it) }
+
+    override fun nonNullValueToString(value: CityId) = varcharColumnType.nonNullValueToString(value.value)
+
+    override fun setParameter(stmt: PreparedStatementApi, index: Int, value: Any?): Unit =
+        varcharColumnType.setParameter(stmt, index, value?.toString())
+
+    override fun readObject(rs: ResultSet, index: Int) = varcharColumnType.readObject(rs, index)
+}
+
+
+@Suppress("MemberNameEqualsClassName")
+object CustomIdTables {
+    object Cities : IdTable<CityId>() {
+        override val id = cityId("id").entityId()
+        override val primaryKey = PrimaryKey(id)
+
+        val name = varchar("name", 50)
+    }
+
+    class City(id: EntityID<CityId>) : Entity<CityId>(id) {
+        companion object : EntityClass<CityId, City>(Cities)
+
+        var name by Cities.name
+    }
+}
+
+class CustomIdTableEntityTest : DatabaseTestsBase() {
+    @Test
+    fun `create tables`() {
+        withTables(CustomIdTables.Cities) {
+            assertEquals(true, CustomIdTables.Cities.exists())
+        }
+    }
+
+    @Test
+    fun `create records`() {
+        withTables(CustomIdTables.Cities) {
+            CustomIdTables.City.new(CityId("M")) { name = "Mumbai" }
+
+            val allCities = CustomIdTables.City.all().map { it.name }
+            assertEquals(true, allCities.contains<String>("Mumbai"))
+        }
+    }
+
+    @Test
+    fun `find by id`() {
+        withTables(CustomIdTables.Cities) {
+            CustomIdTables.City.new(CityId("M")) { name = "Mumbai" }
+
+            val city = CustomIdTables.City.findById(CityId("M"))
+
+            assertNotNull(city)
+            assertEquals("Mumbai", city.name)
+        }
+    }
+
+}