javax.sql.DataSource Syntax Mix-in for Kotlin

DataSourceSyntax.kt

package com.github.nomisrev

import arrow.fx.coroutines.Resource
import arrow.fx.coroutines.fromAutoCloseable
import java.sql.Connection
import java.sql.PreparedStatement
import java.sql.ResultSet
import java.sql.Types
import javax.sql.DataSource

fun DataSource.connection(): Resource<Connection> =
  Resource.fromAutoClosable { connection }

fun DataSource.prepareStatement(sql: String, binders: (SqlPreparedStatement.() -> Unit)? = null): Resource<PreparedStatement> =
  connection().flatMap { connection ->
    Resource.fromAutoClosable {
      connection.prepareStatement(sql)
        .apply { if (binders != null) SqlPreparedStatement(this).binders() }
    }
  }

suspend fun DataSource.query(sql: String): Unit =
  prepareStatement(sql).flatMap { preparedStatement ->
    Resource { preparedStatement.executeUpdate() }
  }.use { }

suspend fun <A> DataSource.queryOneOrNull(
  sql: String,
  binders: (SqlPreparedStatement.() -> Unit)? = null,
  mapper: SqlCursor.() -> A
): A? =
  prepareStatement(sql).flatMap { preparedStatement ->
    Resource.fromAutoClosable {
      preparedStatement
        .apply { if (binders != null) SqlPreparedStatement(this).binders() }
        .executeQuery()
    }
  }.use { rs ->
    if (rs.next()) mapper(SqlCursor(rs)) else null
  }

suspend fun <A> DataSource.queryAsList(
  sql: String,
  binders: (SqlPreparedStatement.() -> Unit)? = null,
  mapper: SqlCursor.() -> A?
): List<A> =
  prepareStatement(sql).flatMap { preparedStatement ->
    Resource.fromAutoClosable {
      preparedStatement
        .apply { if (binders != null) SqlPreparedStatement(this).binders() }
        .executeQuery()
    }
  }.use { rs ->
    val buffer = mutableListOf<A>()
    while (rs.next()) {
      mapper(SqlCursor(rs))?.let(buffer::add)
    }
    buffer
  }

class SqlPreparedStatement(private val preparedStatement: PreparedStatement) {
  private var index: Int = 1

  fun bind(short: Short?): Unit = bind(short?.toLong())
  fun bind(byte: Byte?): Unit = bind(byte?.toLong())
  fun bind(int: Int?): Unit = bind(int?.toLong())
  fun bind(char: Char?): Unit = bind(char?.toString())

  fun bind(bytes: ByteArray?): Unit =
    if (bytes == null) preparedStatement.setNull(index++, Types.BLOB)
    else preparedStatement.setBytes(index++, bytes)

  fun bind(long: Long?): Unit =
    if (long == null) preparedStatement.setNull(index++, Types.INTEGER)
    else preparedStatement.setLong(index++, long)

  fun bind(double: Double?): Unit =
    if (double == null) preparedStatement.setNull(index++, Types.REAL)
    else preparedStatement.setDouble(index++, double)

  fun bind(string: String?): Unit =
    if (string == null) preparedStatement.setNull(index++, Types.VARCHAR)
    else preparedStatement.setString(index++, string)
}

class SqlCursor(private val resultSet: ResultSet) {
  private var index: Int = 1
  fun int(): Int? = long()?.toInt()
  fun string(): String? = resultSet.getString(index++)
  fun bytes(): ByteArray? = resultSet.getBytes(index++)
  fun long(): Long? = resultSet.getLong(index++).takeUnless { resultSet.wasNull() }
  fun double(): Double? = resultSet.getDouble(index++).takeUnless { resultSet.wasNull() }
  fun nextRow(): Boolean = resultSet.next()
}

Usage.kt

private val TABLE_NAME = "testTable"

suspend fun DataSource.createTable() =
  query(
    """
      CREATE TABLE $TABLE_NAME(
        id SERIAL PRIMARY KEY,
        name VARCHAR(20),
        address TEXT,
        age INT
      );
    """.trimIndent()
  )

suspend fun DataSource.insertValue(name: String, address: String, age: Int): Long? =
  queryOneOrNull("""
      INSERT INTO $TABLE_NAME(name, address, age)
      VALUES(?, ?, ?)
      RETURNING id;
    """.trimIndent(), {
    bind(name)
    bind(address)
    bind(age)
  }) {
    long()
  }

suspend fun DataSource.selectWithId(id: Long): Tuple4<Long, String, String, Int>? =
  queryOneOrNull(
    """
      SELECT id, name, address, age
      FROM $TABLE_NAME
      WHERE id = ?;
    """.trimIndent(),
    { bind(id) }
  ) {
    Nullable.zip(long(), string(), string(), int(), ::Tuple4)
  }


suspend fun DataSource.selectAll(): List<Tuple4<Long, String, String, Int>> =
  queryAsList(
    """
      SELECT id, name, address, age
      FROM $TABLE_NAME;
    """.trimIndent()
  ) {
    Nullable.zip(long(), string(), string(), int(), ::Tuple4)
  }

/** Utility function to clear all rows for a certain table */
suspend fun DataSource.truncateTable(name: String): Unit =
  query("TRUNCATE $name RESTART IDENTITY;")

/** Utility function to clear all rows for a set of tables */
suspend fun DataSource.truncateTables(vararg names: String): Unit =
  query("TRUNCATE ${names.joinToString()} RESTART IDENTITY;")