Add support for database-generated integer entity IDs

docs/usage.md

@@ -6,7 +6,7 @@ column type from Postgres to store data. So you first have to create your table
 ```sql
 CREATE TABLE example
 (
-  -- Can have type `text` if using `StringEntityId`
+  -- Can have type `text` if using `StringEntityId`, or `bigint` if using `IntegerEntityId`
   id          uuid        NOT NULL PRIMARY KEY,
   created_at  timestamptz NOT NULL,
   modified_at timestamptz NOT NULL,

src/main/kotlin/no/liflig/documentstore/DocumentStorePlugin.kt

@@ -3,6 +3,7 @@
 package no.liflig.documentstore
 
 import java.sql.Types
+import no.liflig.documentstore.entity.IntegerEntityId
 import no.liflig.documentstore.entity.StringEntityId
 import no.liflig.documentstore.entity.UuidEntityId
 import no.liflig.documentstore.entity.Version
@@ -22,23 +23,27 @@ import org.jdbi.v3.core.spi.JdbiPlugin
  * Once this plugin is installed, the following types can be used as JDBI bind arguments:
  * - [UuidEntityId]
  * - [StringEntityId]
+ * - [IntegerEntityId]
  * - [Version]
  *
- * In addition, [UuidEntityId] and [StringEntityId] may be used as arguments to
+ * In addition, [UuidEntityId], [StringEntityId] and [IntegerEntityId] may be used as arguments to
  * [bindArray][org.jdbi.v3.core.statement.Query.bindArray], like this:
  * ```
  * bindArray("ids", UuidEntityId::class.java, ids) // where ids: List<T extends UuidEntityId>
  * bindArray("ids", StringEntityId::class.java, ids) // where ids: List<T extends StringEntityId>
+ * bindArray("ids", IntegerEntityId::class.java, ids) // where ids: List<T extends IntegerEntityId>
  * ```
  */
 class DocumentStorePlugin : JdbiPlugin.Singleton() {
   override fun customizeJdbi(jdbi: Jdbi) {
     jdbi
+        .registerArgument(VersionArgumentFactory())
         .registerArgument(UuidEntityIdArgumentFactory())
         .registerArgument(StringEntityIdArgumentFactory())
-        .registerArgument(VersionArgumentFactory())
+        .registerArgument(IntegerEntityIdArgumentFactory())
         .registerArrayType(UuidEntityId::class.java, "uuid", { id -> id.value })
         .registerArrayType(StringEntityId::class.java, "text", { id -> id.value })
+        .registerArrayType(IntegerEntityId::class.java, "bigint", { id -> id.value })
   }
 }
 
@@ -58,6 +63,15 @@ private class StringEntityIdArgumentFactory : AbstractArgumentFactory<StringEnti
       }
 }
 
+/** JDBI argument factory that lets us use [IntegerEntityId] as a bind argument. */
+private class IntegerEntityIdArgumentFactory :
+    AbstractArgumentFactory<IntegerEntityId>(Types.BIGINT) {
+  override fun build(value: IntegerEntityId, config: ConfigRegistry?): Argument =
+      Argument { position, statement, _ ->
+        statement.setLong(position, value.value)
+      }
+}
+
 /** JDBI argument factory that lets us use [Version] as a bind argument. */
 private class VersionArgumentFactory : AbstractArgumentFactory<Version>(Types.OTHER) {
   override fun build(value: Version, config: ConfigRegistry?): Argument =

src/main/kotlin/no/liflig/documentstore/entity/EntityId.kt

@@ -1,20 +1,41 @@
 package no.liflig.documentstore.entity
 
-import java.util.*
+import java.util.UUID
+import no.liflig.documentstore.repository.RepositoryJdbi
 
 /** EntityId represents the ID pointing to a specific [Entity]. */
 sealed interface EntityId
 
-/** UUID version of an [EntityId], for `id UUID` columns. */
+/** UUID version of an [EntityId], for `uuid` columns. */
 interface UuidEntityId : EntityId {
   val value: UUID
 }
 
-/** String version of an [EntityId], for `id TEXT` columns. */
+/** String version of an [EntityId], for `text` columns. */
 interface StringEntityId : EntityId {
   val value: String
 }
 
+/**
+ * Integer version of an [EntityId], for `bigint` columns.
+ *
+ * If using [RepositoryJdbi.idsGeneratedByDatabase], the `id` column can be:
+ * ```sql
+ * id bigint NOT NULL PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY
+ * ```
+ */
+interface IntegerEntityId : EntityId {
+  val value: Long
+
+  companion object {
+    /**
+     * Temporary value to use when creating entities that have their IDs generated by the database
+     * (see [RepositoryJdbi.idsGeneratedByDatabase]).
+     */
+    const val GENERATED: Long = -1
+  }
+}
+
 /**
  * In order to use JDBI's `bindArray` method to bind list arguments, we have to supply the list's
  * element type at runtime. We use this in `RepositoryJdbi.listByIds`. However, we don't have the
@@ -25,5 +46,6 @@ internal fun getEntityIdType(entityId: EntityId): Class<out EntityId> {
   return when (entityId) {
     is UuidEntityId -> UuidEntityId::class.java
     is StringEntityId -> StringEntityId::class.java
+    is IntegerEntityId -> IntegerEntityId::class.java
   }
 }

src/main/kotlin/no/liflig/documentstore/repository/RepositoryJdbi.kt

@@ -3,12 +3,15 @@
 package no.liflig.documentstore.repository
 
 import java.time.Instant
+import no.liflig.documentstore.DocumentStorePlugin
 import no.liflig.documentstore.entity.Entity
 import no.liflig.documentstore.entity.EntityId
+import no.liflig.documentstore.entity.IntegerEntityId
 import no.liflig.documentstore.entity.Version
 import no.liflig.documentstore.entity.Versioned
 import no.liflig.documentstore.entity.getEntityIdType
 import no.liflig.documentstore.utils.executeBatchOperation
+import org.jdbi.v3.core.Handle
 import org.jdbi.v3.core.Jdbi
 import org.jdbi.v3.core.mapper.RowMapper
 import org.jdbi.v3.core.statement.Query
@@ -26,7 +29,7 @@ import org.jdbi.v3.core.statement.Query
  * ```sql
  * CREATE TABLE example
  * (
- *   -- Can have type `text` if using `StringEntityId`
+ *   -- Can have type `text` if using `StringEntityId`, or `bigint` if using `IntegerEntityId`
  *   id          uuid        NOT NULL PRIMARY KEY,
  *   created_at  timestamptz NOT NULL,
  *   modified_at timestamptz NOT NULL,
@@ -35,27 +38,46 @@ import org.jdbi.v3.core.statement.Query
  * );
  * ```
  *
- * Also, the given [Jdbi] instance must have the
- * [DocumentStorePlugin][no.liflig.documentstore.DocumentStorePlugin] installed for the queries in
- * this class to work.
+ * @param jdbi Must have the [DocumentStorePlugin] installed for the queries in this class to work.
+ * @param serializationAdapter See [SerializationAdapter] for an example of how to implement this.
+ * @param idsGeneratedByDatabase When using [IntegerEntityId], one often wants the entity IDs to be
+ *   generated by the database. This affects how we perform [create], so if your `id` column is
+ *   `PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY`, you must set this flag to true.
+ *
+ *   The `create` method still takes an entity with an `id`, but when this flag is set, the `id`
+ *   given to `create` is ignored (you can use [IntegerEntityId.GENERATED] as a dummy ID value). The
+ *   `Versioned<Entity>` returned by `create` will then have the database-generated ID set on it. If
+ *   you use this entity for anything after creating it (such as returning it to the user), make
+ *   sure to use the entity returned by `create`, so that you get the actual generated ID and not
+ *   the dummy value.
+ *
+ *   Note that the `id` column cannot be `GENERATED ALWAYS` - see [createEntityWithGeneratedId] for
+ *   the reason.
  */
 open class RepositoryJdbi<EntityIdT : EntityId, EntityT : Entity<EntityIdT>>(
     protected val jdbi: Jdbi,
     protected val tableName: String,
     protected val serializationAdapter: SerializationAdapter<EntityT>,
+    protected val idsGeneratedByDatabase: Boolean = false,
 ) : Repository<EntityIdT, EntityT> {
   protected val rowMapper: RowMapper<Versioned<EntityT>> = EntityRowMapper(serializationAdapter)
 
   private val rowMapperWithTotalCount: RowMapper<MappedEntityWithTotalCount<EntityT>> =
       EntityRowMapperWithTotalCount(serializationAdapter)
 
   private val updateResultMapper: RowMapper<UpdateResult> = UpdateResultMapper()
+  private val entityDataMapper: RowMapper<EntityT> = EntityDataRowMapper(serializationAdapter)
 
   override fun create(entity: EntityT): Versioned<EntityT> {
     try {
       useHandle(jdbi) { handle ->
-        val now = Instant.now()
+        val createdAt = Instant.now()
         val version = Version.initial()
+
+        if (idsGeneratedByDatabase) {
+          return createEntityWithGeneratedId(handle, entity, createdAt, version)
+        }
+
         handle
             .createUpdate(
                 """
@@ -67,10 +89,10 @@ open class RepositoryJdbi<EntityIdT : EntityId, EntityT : Entity<EntityIdT>>(
             .bind("id", entity.id)
             .bind("data", serializationAdapter.toJson(entity))
             .bind("version", version)
-            .bind("createdAt", now)
-            .bind("modifiedAt", now)
+            .bind("createdAt", createdAt)
+            .bind("modifiedAt", createdAt)
             .execute()
-        return Versioned(entity, version, createdAt = now, modifiedAt = now)
+        return Versioned(entity, version, createdAt = createdAt, modifiedAt = createdAt)
       }
     } catch (e: Exception) {
       // Call mapDatabaseException first to handle connection-related exceptions, before calling
@@ -79,6 +101,72 @@ open class RepositoryJdbi<EntityIdT : EntityId, EntityT : Entity<EntityIdT>>(
     }
   }
 
+  /**
+   * We need a separate query to implement `create` for entities with database-generated IDs (see
+   * [idsGeneratedByDatabase]), because these entities don't have their IDs already set when given
+   * to `create`, like we assume for UUID and String IDs.
+   *
+   * We have an interesting challenge to solve here, because we store IDs in two places in our
+   * tables: both in the `id` column, and in the `data` JSONB column (since `id` is part of the
+   * [Entity] interface, so it becomes part of the JSON). When IDs are generated by the database, we
+   * need to make sure that the `id` field on `data` is set correctly to the generated `id` column.
+   * To achieve this, we:
+   * - Use the fact that `GENERATED` columns in Postgres are backed by a _sequence_ (essentially a
+   *   number generator)
+   *     - We can use the
+   *       [`nextval` function](https://www.postgresql.org/docs/17/functions-sequence.html) to get
+   *       and reserve the next value in a sequence
+   *     - And we can use the
+   *       [`pg_get_serial_sequence` function](https://www.postgresql.org/docs/17/functions-info.html)
+   *       to get the name of the sequence associated with a table's `GENERATED` column
+   * - Now that we have gotten and reserved the next generated ID, we can use that to set both the
+   *   `id` column and the `data.id` JSONB field at once in our INSERT
+   *
+   * Finally, we use `RETURNING data` to get the entity with the generated ID.
+   *
+   * This works well, though there is one drawback: since we first get the generated ID with
+   * `nextval`, and then set the `id` column explicitly, we can't use `GENERATED ALWAYS` on our `id`
+   * column, since that errors when we set `id` explictly. So we must use `GENERATED BY DEFAULT`,
+   * although the behavior we have here is more akin to `GENERATED ALWAYS`.
+   */
+  private fun createEntityWithGeneratedId(
+      handle: Handle,
+      entity: EntityT,
+      createdAt: Instant,
+      version: Version
+  ): Versioned<EntityT> {
+    val createdEntity =
+        handle
+            .createQuery(
+                """
+                  WITH generated_id AS (
+                      SELECT nextval(pg_get_serial_sequence('${tableName}', 'id')) AS value
+                  )
+                  INSERT INTO "${tableName}" (id, data, version, created_at, modified_at)
+                  SELECT
+                      generated_id.value,
+                      jsonb_set(:data::jsonb, '{id}', to_jsonb(generated_id.value)),
+                      :version,
+                      :createdAt,
+                      :modifiedAt
+                  FROM generated_id
+                  RETURNING data
+                """
+                    .trimIndent(),
+            )
+            .bind("data", serializationAdapter.toJson(entity))
+            .bind("version", version)
+            .bind("createdAt", createdAt)
+            .bind("modifiedAt", createdAt)
+            .map(entityDataMapper)
+            .firstOrNull()
+            ?: throw IllegalStateException(
+                "INSERT query for entity with generated ID did not return entity data",
+            )
+
+    return Versioned(createdEntity, version, createdAt = createdAt, modifiedAt = createdAt)
+  }
+
   override fun get(id: EntityIdT, forUpdate: Boolean): Versioned<EntityT>? {
     useHandle(jdbi) { handle ->
       return handle
@@ -194,9 +282,14 @@ open class RepositoryJdbi<EntityIdT : EntityId, EntityT : Entity<EntityIdT>>(
   override fun batchCreate(entities: Iterable<EntityT>) {
     transactional {
       useHandle(jdbi) { handle ->
-        val now = Instant.now()
+        val createdAt = Instant.now()
         val version = Version.initial()
 
+        if (idsGeneratedByDatabase) {
+          batchCreateEntitiesWithGeneratedId(handle, entities, createdAt, version)
+          return@transactional
+        }
+
         executeBatchOperation(
             handle,
             entities,
@@ -211,14 +304,48 @@ open class RepositoryJdbi<EntityIdT : EntityId, EntityT : Entity<EntityIdT>>(
                   .bind("id", entity.id)
                   .bind("data", serializationAdapter.toJson(entity))
                   .bind("version", version)
-                  .bind("createdAt", now)
-                  .bind("modifiedAt", now)
+                  .bind("createdAt", createdAt)
+                  .bind("modifiedAt", createdAt)
             },
         )
       }
     }
   }
 
+  private fun batchCreateEntitiesWithGeneratedId(
+      handle: Handle,
+      entities: Iterable<EntityT>,
+      createdAt: Instant,
+      version: Version
+  ) {
+    executeBatchOperation(
+        handle,
+        entities,
+        statement =
+            """
+              WITH generated_id AS (
+                  SELECT nextval(pg_get_serial_sequence('${tableName}', 'id')) AS value
+              )
+              INSERT INTO "${tableName}" (id, data, version, created_at, modified_at)
+              SELECT
+                  generated_id.value,
+                  jsonb_set(:data::jsonb, '{id}', to_jsonb(generated_id.value)),
+                  :version,
+                  :createdAt,
+                  :modifiedAt
+              FROM generated_id
+            """
+                .trimIndent(),
+        bindParameters = { batch, entity ->
+          batch
+              .bind("data", serializationAdapter.toJson(entity))
+              .bind("version", version)
+              .bind("createdAt", createdAt)
+              .bind("modifiedAt", createdAt)
+        },
+    )
+  }
+
   override fun batchUpdate(entities: Iterable<Versioned<EntityT>>) {
     transactional {
       useHandle(jdbi) { handle ->

src/main/kotlin/no/liflig/documentstore/repository/RowMapper.kt

@@ -38,6 +38,21 @@ internal class EntityRowMapper<EntityT : Entity<*>>(
   }
 }
 
+/**
+ * If [RepositoryJdbi.idsGeneratedByDatabase] is set to true, we need to return the entity from our
+ * INSERT query in order to get the generated ID. In this case, we don't need to return the
+ * version/createdAt/modifiedAt fields, so we use this row mapper that only maps the entity data.
+ */
+internal class EntityDataRowMapper<EntityT : Entity<*>>(
+    private val serializationAdapter: SerializationAdapter<EntityT>,
+) : RowMapper<EntityT> {
+  override fun map(resultSet: ResultSet, ctx: StatementContext): EntityT {
+    val data = getStringFromRowOrThrow(resultSet, Columns.DATA)
+
+    return serializationAdapter.fromJson(data)
+  }
+}
+
 /**
  * In [RepositoryJdbi.update], we receive an entity and its previous [Version], and want to return a
  * [Versioned] wrapper around that entity. We know the [Versioned.modifiedAt] field, since that is