feat: Add ability to pass custom sequence to auto-increment column

JetBrains · Aug 13, 2024 · 2eb546d · 2eb546d
1 parent e63f35b
commit 2eb546d
Show file tree

Hide file tree

Showing 7 changed files with 146 additions and 70 deletions.
diff --git a/exposed-core/api/exposed-core.api b/exposed-core/api/exposed-core.api
@@ -221,11 +221,13 @@ public final class org/jetbrains/exposed/sql/ArrayColumnType : org/jetbrains/exp
 
 public final class org/jetbrains/exposed/sql/AutoIncColumnType : org/jetbrains/exposed/sql/IColumnType {
 	public fun <init> (Lorg/jetbrains/exposed/sql/ColumnType;Ljava/lang/String;Ljava/lang/String;)V
+	public fun <init> (Lorg/jetbrains/exposed/sql/ColumnType;Lorg/jetbrains/exposed/sql/Sequence;)V
 	public fun equals (Ljava/lang/Object;)Z
 	public final fun getAutoincSeq ()Ljava/lang/String;
 	public final fun getDelegate ()Lorg/jetbrains/exposed/sql/ColumnType;
 	public final fun getNextValExpression ()Lorg/jetbrains/exposed/sql/NextVal;
 	public fun getNullable ()Z
+	public final fun getSequence ()Lorg/jetbrains/exposed/sql/Sequence;
 	public fun hashCode ()I
 	public fun nonNullValueAsDefaultString (Ljava/lang/Object;)Ljava/lang/String;
 	public fun nonNullValueToString (Ljava/lang/Object;)Ljava/lang/String;
@@ -2132,6 +2134,7 @@ public final class org/jetbrains/exposed/sql/Sequence {
 	public final fun getIncrementBy ()Ljava/lang/Long;
 	public final fun getMaxValue ()Ljava/lang/Long;
 	public final fun getMinValue ()Ljava/lang/Long;
+	public final fun getName ()Ljava/lang/String;
 	public final fun getStartWith ()Ljava/lang/Long;
 }
 
@@ -2424,6 +2427,7 @@ public class org/jetbrains/exposed/sql/Table : org/jetbrains/exposed/sql/ColumnS
 	public static synthetic fun array$default (Lorg/jetbrains/exposed/sql/Table;Ljava/lang/String;Lorg/jetbrains/exposed/sql/ColumnType;Ljava/lang/Integer;ILjava/lang/Object;)Lorg/jetbrains/exposed/sql/Column;
 	public final fun autoGenerate (Lorg/jetbrains/exposed/sql/Column;)Lorg/jetbrains/exposed/sql/Column;
 	public final fun autoIncrement (Lorg/jetbrains/exposed/sql/Column;Ljava/lang/String;)Lorg/jetbrains/exposed/sql/Column;
+	public final fun autoIncrement (Lorg/jetbrains/exposed/sql/Column;Lorg/jetbrains/exposed/sql/Sequence;)Lorg/jetbrains/exposed/sql/Column;
 	public static synthetic fun autoIncrement$default (Lorg/jetbrains/exposed/sql/Table;Lorg/jetbrains/exposed/sql/Column;Ljava/lang/String;ILjava/lang/Object;)Lorg/jetbrains/exposed/sql/Column;
 	public final fun autoinc (Lorg/jetbrains/exposed/sql/Column;Ljava/lang/String;)Lorg/jetbrains/exposed/sql/Column;
 	public static synthetic fun autoinc$default (Lorg/jetbrains/exposed/sql/Table;Lorg/jetbrains/exposed/sql/Column;Ljava/lang/String;ILjava/lang/Object;)Lorg/jetbrains/exposed/sql/Column;

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
@@ -96,7 +96,8 @@ interface IColumnType<T> {
      * [value] can be of any type (including [Expression])
      * */
     @Throws(IllegalArgumentException::class)
-    fun validateValueBeforeUpdate(value: T?) {}
+    fun validateValueBeforeUpdate(value: T?) {
+    }
 
     /**
      * Defines the appearance of parameter markers in prepared SQL statements.
@@ -136,9 +137,21 @@ class AutoIncColumnType<T>(
     private val fallbackSeqName: String
 ) : IColumnType<T> by delegate {
 
-    private val nextValValue = run {
-        val sequence = Sequence(_autoincSeq ?: fallbackSeqName)
-        if (delegate is IntegerColumnType) sequence.nextIntVal() else sequence.nextLongVal()
+    private var _sequence: Sequence? = null
+
+    /** The sequence used to generate new values for this auto-increment column. */
+    val sequence: Sequence?
+        get() = _sequence ?: autoincSeq?.let {
+            Sequence(
+                it,
+                startWith = 1,
+                minValue = 1,
+                maxValue = Long.MAX_VALUE
+            )
+        }
+
+    constructor(delegate: ColumnType<T>, sequence: Sequence) : this(delegate, sequence.name, sequence.name) {
+        _sequence = sequence
     }
 
     /** The name of the sequence used to generate new values for this auto-increment column. */
@@ -148,7 +161,9 @@ class AutoIncColumnType<T>(
 
     /** The SQL expression that advances the sequence of this auto-increment column. */
     val nextValExpression: NextVal<*>?
-        get() = nextValValue.takeIf { autoincSeq != null }
+        get() = autoincSeq?.let {
+            if (delegate is IntegerColumnType) sequence?.nextIntVal() else sequence?.nextLongVal()
+        }
 
     private fun resolveAutoIncType(columnType: IColumnType<*>): String = when {
         columnType is EntityIDColumnType<*> -> resolveAutoIncType(columnType.idColumn.columnType)
@@ -178,6 +193,7 @@ class AutoIncColumnType<T>(
             delegate != other.delegate -> false
             _autoincSeq != other._autoincSeq -> false
             fallbackSeqName != other.fallbackSeqName -> false
+            sequence != other.sequence -> false
             else -> true
         }
     }
@@ -186,6 +202,7 @@ class AutoIncColumnType<T>(
         var result = delegate.hashCode()
         result = 31 * result + (_autoincSeq?.hashCode() ?: 0)
         result = 31 * result + fallbackSeqName.hashCode()
+        result = 31 * result + (sequence?.hashCode() ?: 0)
         return result
     }
 }
@@ -265,7 +282,7 @@ interface ColumnTransformer<Unwrapped, Wrapped> {
     fun wrap(value: Unwrapped): Wrapped
 }
 
-fun <Unwrapped, Wrapped>columnTransformer(unwrap: (value: Wrapped) -> Unwrapped, wrap: (value: Unwrapped) -> Wrapped): ColumnTransformer<Unwrapped, Wrapped> {
+fun <Unwrapped, Wrapped> columnTransformer(unwrap: (value: Wrapped) -> Unwrapped, wrap: (value: Unwrapped) -> Wrapped): ColumnTransformer<Unwrapped, Wrapped> {
     return object : ColumnTransformer<Unwrapped, Wrapped> {
         override fun unwrap(value: Wrapped): Unwrapped = unwrap(value)
         override fun wrap(value: Unwrapped): Wrapped = wrap(value)

diff --git a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/SchemaUtils.kt b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/SchemaUtils.kt
@@ -109,12 +109,12 @@ object SchemaUtils {
     }
 
     private fun tableDdlWithoutExistingSequence(table: Table): List<String> {
-        val existingAutoIncSeq = table.autoIncColumn?.autoIncColumnType?.autoincSeq
-            ?.takeIf { currentDialect.sequenceExists(Sequence(it)) }
+        val existingAutoIncSeq = table.autoIncColumn?.autoIncColumnType?.sequence
+            ?.takeIf { currentDialect.sequenceExists(it) }
 
         return table.ddl.filter { statement ->
             if (existingAutoIncSeq != null) {
-                !statement.lowercase().startsWith("create sequence") || !statement.contains(existingAutoIncSeq)
+                !statement.lowercase().startsWith("create sequence") || !statement.contains(existingAutoIncSeq.name)
             } else {
                 true
             }

diff --git a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/Sequence.kt b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/Sequence.kt
@@ -17,7 +17,7 @@ import org.jetbrains.exposed.sql.vendors.currentDialect
  * @param cache Specifies how many sequence numbers are to be pre-allocated and stored in memory for faster access.
  */
 class Sequence(
-    private val name: String,
+    val name: String,
     val startWith: Long? = null,
     val incrementBy: Long? = null,
     val minValue: Long? = null,

diff --git a/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/Table.kt b/exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/Table.kt
@@ -920,6 +920,17 @@ open class Table(name: String = "") : ColumnSet(), DdlAware {
     fun <N : Any> Column<N>.autoIncrement(idSeqName: String? = null): Column<N> =
         cloneWithAutoInc(idSeqName).also { replaceColumn(this, it) }
 
+    /**
+     * Make @receiver column an auto-increment column to generate its values in a database.
+     * **Note:** Only integer and long columns are supported (signed and unsigned types).
+     * Some databases, like PostgreSQL, support auto-increment via sequences.
+     * In this case, a sequence should be provided using the [sequence] param.
+     *
+     * @param sequence a parameter to provide a sequence
+     */
+    fun <N : Any> Column<N>.autoIncrement(sequence: Sequence): Column<N> =
+        cloneWithAutoInc(sequence).also { replaceColumn(this, it) }
+
     /**
      * Make @receiver column an auto-increment column to generate its values in a database.
      * **Note:** Only integer and long columns are supported (signed and unsigned types).
@@ -1571,6 +1582,12 @@ open class Table(name: String = "") : ColumnSet(), DdlAware {
         else -> error("Unsupported column type for auto-increment $columnType")
     }
 
+    private fun <T> Column<T>.cloneWithAutoInc(sequence: Sequence): Column<T> = when (columnType) {
+        is AutoIncColumnType -> this
+        is ColumnType -> this.withColumnType(AutoIncColumnType(columnType, sequence))
+        else -> error("Unsupported column type for auto-increment $columnType")
+    }
+
     // DDL statements
 
     internal fun primaryKeyConstraint(): String? {
@@ -1636,16 +1653,7 @@ open class Table(name: String = "") : ColumnSet(), DdlAware {
     }
 
     private fun createAutoIncColumnSequence(): List<String> {
-        return autoIncColumn?.autoIncColumnType?.autoincSeq?.let {
-            Sequence(
-                it,
-                startWith = 1,
-                minValue = 1,
-                maxValue = Long.MAX_VALUE
-            )
-        }
-            ?.createStatement()
-            .orEmpty()
+        return autoIncColumn?.autoIncColumnType?.sequence?.createStatement().orEmpty()
     }
 
     override fun modifyStatement(): List<String> =
@@ -1665,7 +1673,7 @@ open class Table(name: String = "") : ColumnSet(), DdlAware {
             }
         }
 
-        val dropSequence = autoIncColumn?.autoIncColumnType?.autoincSeq?.let { Sequence(it).dropStatement() }.orEmpty()
+        val dropSequence = autoIncColumn?.autoIncColumnType?.sequence?.dropStatement().orEmpty()
 
         return listOf(dropTable) + dropSequence
     }

diff --git a/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/ddl/CreateTableTests.kt b/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/ddl/CreateTableTests.kt
@@ -308,14 +308,7 @@ class CreateTableTests : DatabaseTestsBase() {
         withDb {
             val t = TransactionManager.current()
             val expected = listOfNotNull(
-                child.autoIncColumn?.autoIncColumnType?.autoincSeq?.let {
-                    Sequence(
-                        it,
-                        startWith = 1,
-                        minValue = 1,
-                        maxValue = Long.MAX_VALUE
-                    ).createStatement().single()
-                },
+                child.autoIncColumn?.autoIncColumnType?.sequence?.createStatement()?.single(),
                 "CREATE TABLE " + addIfNotExistsIfSupported() + "${t.identity(child)} (" +
                     "${child.columns.joinToString { it.descriptionDdl(false) }}," +
                     " CONSTRAINT ${t.db.identifierManager.cutIfNecessaryAndQuote(fkName).inProperCase()}" +
@@ -389,14 +382,7 @@ class CreateTableTests : DatabaseTestsBase() {
         withDb {
             val t = TransactionManager.current()
             val expected = listOfNotNull(
-                child.autoIncColumn?.autoIncColumnType?.autoincSeq?.let {
-                    Sequence(
-                        it,
-                        startWith = 1,
-                        minValue = 1,
-                        maxValue = Long.MAX_VALUE
-                    ).createStatement().single()
-                },
+                child.autoIncColumn?.autoIncColumnType?.sequence?.createStatement()?.single(),
                 "CREATE TABLE " + addIfNotExistsIfSupported() + "${t.identity(child)} (" +
                     "${child.columns.joinToString { it.descriptionDdl(false) }}," +
                     " CONSTRAINT ${t.db.identifierManager.cutIfNecessaryAndQuote(fkName).inProperCase()}" +
@@ -424,14 +410,7 @@ class CreateTableTests : DatabaseTestsBase() {
         withDb {
             val t = TransactionManager.current()
             val expected = listOfNotNull(
-                child.autoIncColumn?.autoIncColumnType?.autoincSeq?.let {
-                    Sequence(
-                        it,
-                        startWith = 1,
-                        minValue = 1,
-                        maxValue = Long.MAX_VALUE
-                    ).createStatement().single()
-                },
+                child.autoIncColumn?.autoIncColumnType?.sequence?.createStatement()?.single(),
                 "CREATE TABLE " + addIfNotExistsIfSupported() + "${t.identity(child)} (" +
                     "${child.columns.joinToString { it.descriptionDdl(false) }}," +
                     " CONSTRAINT ${t.db.identifierManager.cutIfNecessaryAndQuote(fkName).inProperCase()}" +
@@ -462,14 +441,7 @@ class CreateTableTests : DatabaseTestsBase() {
         withDb {
             val t = TransactionManager.current()
             val expected = listOfNotNull(
-                child.autoIncColumn?.autoIncColumnType?.autoincSeq?.let {
-                    Sequence(
-                        it,
-                        startWith = 1,
-                        minValue = 1,
-                        maxValue = Long.MAX_VALUE
-                    ).createStatement().single()
-                },
+                child.autoIncColumn?.autoIncColumnType?.sequence?.createStatement()?.single(),
                 "CREATE TABLE " + addIfNotExistsIfSupported() + "${t.identity(child)} (" +
                     "${child.columns.joinToString { it.descriptionDdl(false) }}," +
                     " CONSTRAINT ${t.db.identifierManager.cutIfNecessaryAndQuote(fkName).inProperCase()}" +
@@ -507,14 +479,7 @@ class CreateTableTests : DatabaseTestsBase() {
             val t = TransactionManager.current()
             val updateCascadePart = if (testDb != TestDB.ORACLE) " ON UPDATE CASCADE" else ""
             val expected = listOfNotNull(
-                child.autoIncColumn?.autoIncColumnType?.autoincSeq?.let {
-                    Sequence(
-                        it,
-                        startWith = 1,
-                        minValue = 1,
-                        maxValue = Long.MAX_VALUE
-                    ).createStatement().single()
-                },
+                child.autoIncColumn?.autoIncColumnType?.sequence?.createStatement()?.single(),
                 "CREATE TABLE " + addIfNotExistsIfSupported() + "${t.identity(child)} (" +
                     "${child.columns.joinToString { it.descriptionDdl(false) }}," +
                     " CONSTRAINT ${t.db.identifierManager.cutIfNecessaryAndQuote(fkName).inProperCase()}" +
@@ -553,14 +518,7 @@ class CreateTableTests : DatabaseTestsBase() {
         withDb {
             val t = TransactionManager.current()
             val expected = listOfNotNull(
-                child.autoIncColumn?.autoIncColumnType?.autoincSeq?.let {
-                    Sequence(
-                        it,
-                        startWith = 1,
-                        minValue = 1,
-                        maxValue = Long.MAX_VALUE
-                    ).createStatement().single()
-                },
+                child.autoIncColumn?.autoIncColumnType?.sequence?.createStatement()?.single(),
                 "CREATE TABLE " + addIfNotExistsIfSupported() + "${t.identity(child)} (" +
                     "${child.columns.joinToString { it.descriptionDdl(false) }}," +
                     " CONSTRAINT ${t.db.identifierManager.cutIfNecessaryAndQuote(fkName).inProperCase()}" +

diff --git a/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/ddl/SequencesTests.kt b/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/ddl/SequencesTests.kt
@@ -9,6 +9,7 @@ import org.jetbrains.exposed.sql.tests.TestDB
 import org.jetbrains.exposed.sql.tests.currentDialectTest
 import org.jetbrains.exposed.sql.tests.inProperCase
 import org.jetbrains.exposed.sql.tests.shared.assertEquals
+import org.jetbrains.exposed.sql.tests.shared.assertFalse
 import org.jetbrains.exposed.sql.tests.shared.assertTrue
 import org.jetbrains.exposed.sql.vendors.currentDialect
 import org.junit.Test
@@ -61,6 +62,39 @@ class SequencesTests : DatabaseTestsBase() {
         }
     }
 
+    @Test
+    fun `testInsertWithCustomSequence`() {
+        val tester = object : Table("tester") {
+            val id = integer("id").autoIncrement(myseq)
+            var name = varchar("name", 25)
+
+            override val primaryKey = PrimaryKey(id, name)
+        }
+        withDb {
+            if (currentDialectTest.supportsSequenceAsGeneratedKeys) {
+                try {
+                    SchemaUtils.create(tester)
+                    assertTrue(myseq.exists())
+
+                    var testerId = tester.insert {
+                        it[name] = "Hichem"
+                    } get tester.id
+
+                    assertEquals(myseq.startWith, testerId.toLong())
+
+                    testerId = tester.insert {
+                        it[name] = "Andrey"
+                    } get tester.id
+
+                    assertEquals(myseq.startWith!! + myseq.incrementBy!!, testerId.toLong())
+                } finally {
+                    SchemaUtils.drop(tester)
+                    assertFalse(myseq.exists())
+                }
+            }
+        }
+    }
+
     @Test
     fun `test insert int IdTable with sequences`() {
         withTables(DeveloperWithLongId) {
@@ -88,7 +122,40 @@ class SequencesTests : DatabaseTestsBase() {
     }
 
     @Test
-    fun `test insert LongIdTable with auth-increment with sequence`() {
+    fun `testInsertInIdTableWithCustomSequence`() {
+        val tester = object : IdTable<Long>("tester") {
+            override val id = long("id").autoIncrement(myseq).entityId()
+            var name = varchar("name", 25)
+
+            override val primaryKey = PrimaryKey(id, name)
+        }
+        withDb {
+            if (currentDialectTest.supportsSequenceAsGeneratedKeys) {
+                try {
+                    SchemaUtils.create(tester)
+                    assertTrue(myseq.exists())
+
+                    var testerId = tester.insert {
+                        it[name] = "Hichem"
+                    } get tester.id
+
+                    assertEquals(myseq.startWith, testerId.value)
+
+                    testerId = tester.insert {
+                        it[name] = "Andrey"
+                    } get tester.id
+
+                    assertEquals(myseq.startWith!! + myseq.incrementBy!!, testerId.value)
+                } finally {
+                    SchemaUtils.drop(tester)
+                    assertFalse(myseq.exists())
+                }
+            }
+        }
+    }
+
+    @Test
+    fun `test insert LongIdTable with auto-increment with sequence`() {
         withDb {
             if (currentDialectTest.supportsSequenceAsGeneratedKeys) {
                 try {
@@ -152,6 +219,28 @@ class SequencesTests : DatabaseTestsBase() {
         }
     }
 
+    @Test
+    fun testExistingSequencesForAutoIncrementWithCustomSequence() {
+        val tableWithExplicitSequenceName = object : IdTable<Long>() {
+            override val id: Column<EntityID<Long>> = long("id").autoIncrement(myseq).entityId()
+        }
+
+        withDb {
+            if (currentDialectTest.supportsSequenceAsGeneratedKeys) {
+                try {
+                    SchemaUtils.create(tableWithExplicitSequenceName)
+
+                    val sequences = currentDialectTest.sequences()
+
+                    assertTrue(sequences.isNotEmpty())
+                    assertTrue(sequences.any { it == myseq.name.inProperCase() })
+                } finally {
+                    SchemaUtils.drop(tableWithExplicitSequenceName)
+                }
+            }
+        }
+    }
+
     @Test
     fun testExistingSequencesForAutoIncrementWithExplicitSequenceName() {
         val sequenceName = "id_seq"