CTE.kt

/*
Author: Salomon BRYS
License: MIT (do whatever you want with that code!)

Allows using Common Table Expressions with Jetbrains Exposed JDBC.
Provides primitives that implement CTEs, and a DSL that allows simpler usage.
Inspired by https://github.com/bog-walk/bootrack/blob/main/backend-jdbc/src/main/kotlin/dev/bogwalk/bootrack/backend/repository/custom/WithCTE.kt

/!\ Because the result rows of a CTE are referencing the With temporary table(s) fields and not the original query
table(s) fields, the results in a ResultRow are not easily accessible: you need to use either CTETable.get or wrap a
ResultRow into a RowValues that will automatically use CTETable.get.

Following examples use the following table:
    object Employees : UUIDTable("employees") {
        val name = varchar("name", 150)
        val yearOfBirth = integer("year_of_birth")
        val gender = varchar("gender", 10)
        val salary = integer("salary")
        val boss = reference("boss", id).nullable()
    }

Simple anonymous CTE (lowest salary of gen Y employees by gender):
    val results = CTE {
        val genY = With(
            name = "gen_y_employees",
            query = Employees.selectAll()
                .where { (Employees.yearOfBirth greaterEq 1981) and (Employees.yearOfBirth lessEq 1996) }
        )
        // Note that we must use the CTE fields and not the Employees fields.
        // We must therefore find the corresponding columns with genY[Employees.field].
        genY.select(genY[Employees.gender], genY[Employees.salary].min().alias("min_salary"))
            .groupBy(genY[Employees.gender])
    }
    // The CTE function returns a list of RowValues, not ResultRow, which automatically makes the correspondance
    // between the CTE table fields and the CTE query table fields.
    results.forEach { println("${it[Employees.gender]}: ${it["min_salary"]}") }

Simple object CTE (lowest salary of gen Y employees by gender):
    val GenYEmployees = object : CTETable(
        name = "gen_y_employees",
        query = Employees.selectAll()
            .where { (Employees.yearOfBirth greaterEq 1981) and (Employees.yearOfBirth lessEq 1996) }
    ) {
        val gender = get(Employees.gender)
        val salary = get(Employees.salary)
    }
    val minSalary = GenYEmployees.salary.min()
    val results = With(GenYEmployees)
        .select(GenYEmployees.gender, minSalary)
        .groupBy(GenYEmployees.gender)
    results.forEach { println("${it[GenYEmployees.gender]}: ${it[minSalary]}") }

Recursive anonymous CTE (list of employees that work under rootBossId):
    val results = CTE {
        val columns = listOf(Employees.id, Employees.name, Employees.boss)
        val hierarchy = With(
            name = "employee_hierarchy",
            query = Employees.select(columns).where { Employees.boss eq rootBossId },
            recursive = { table ->
                Employees.innerJoin(
                    otherTable = table,
                    onColumn = { boss },
                    otherColumn = { get(Employees.id) }
                ).select(columns)
            }
        )
        // Note that we must use the CTE fields and not the Employees fields.
        // We must therefore find the corresponding columns with hierarchy[Employees.field].
        hierarchy.selectAll()
            .orderBy(hierarchy[Employees.name])
    }
    // The CTE function returns a list of RowValues, not ResultRow, which automatically makes the correspondance
    // between the CTE table fields and the CTE query table fields.
    results.forEach { println("${it[Employees.name]}: ${it[Employees.id]}") }

Recursive object CTE (list of employees that work under rootBossId):
    val employeeColumns = listOf(Employees.id, Employees.name, Employees.boss)
    val hierarchy = object : CTETable(
        name = "employee_hierarchy",
        query = Employees.select(employeeColumns).where { Employees.boss eq rootBossId },
        recursive = { table ->
            Employees.innerJoin(
                otherTable = table,
                onColumn = { boss },
                // Find the CTE table field corresponding its query field Employees.id
                otherColumn = { get(Employees.id) }
            ).select(employeeColumns)
        }
    ) {
        val id = get(Employees.id)
        val name = get(Employees.name)
    }
    val results = With(hierarchy)
        .selectAll()
        .orderBy(hierarchy.name)
    results.forEach { println("${it[hierarchy.name]}: ${it[hierarchy.id]}") }
*/

import org.jetbrains.exposed.v1.core.*
import org.jetbrains.exposed.v1.jdbc.Query
import org.jetbrains.exposed.v1.jdbc.transactions.TransactionManager


/**
 * A temporary table, built by a query, to be used in a CTE query.
 *
 * @param name The name (alias) of the temporary table.
 * @param query The initial query that builds the temporary table.
 * @param recursive An optional lambda that uses this CTETable to build a recursive query.
 *                  Defining this parameter means that this is a recursive CTE.
 * @param unionAll Only used if [recursive] is not null. Whether to recurse with `UNION` or `UNION ALL`.
 */
open class CTETable(
    name: String,
    val query: AbstractQuery<*>,
    val recursive: ((CTETable) -> AbstractQuery<*>)? = null,
    val unionAll: Boolean = false,
) : Table(name) {

    private val fieldMap = query.set.fields.associateWith { field ->
        when (field) {
            is Column<*> -> Column(this, field.name, field.columnType)
            is IExpressionAlias<*> -> field.aliasOnlyExpression()
            else -> field
        }
    }

    /**
     * Whether this [queryField] is mapped to a field in this table.
     */
    operator fun contains(queryField: Expression<*>): Boolean = queryField in fieldMap

    final override val fields: List<Expression<*>> get() = fieldMap.values.toList()

    final override val columns: List<Column<*>> get() = fields.filterIsInstance<Column<*>>()

    final override fun createStatement(): List<String> = error("CREATE statements are not supported by CTEs")

    final override fun modifyStatement(): List<String> = error("ALTER statements are not supported by CTEs")

    final override fun dropStatement(): List<String> = error("DROP statements are not supported by CTEs")

    internal fun describeWith(builder: QueryBuilder) {
        builder {
            [email protected](TransactionManager.current(), builder)
            append(" AS (")
            query.prepareSQL(builder)
            if (recursive != null) {
                append(" UNION ")
                if (unionAll) append("ALL ")
                recursive(this@CTETable).prepareSQL(builder)
            }
            append(") ")
        }
    }

    /**
     * Get a field on this table that corresponds to a field that was used in [query].
     */
    @Suppress("UNCHECKED_CAST")
    operator fun <T> get(queryField: Expression<T>): Expression<T> =
        fieldMap[queryField] as? Expression<T>
            ?: error("$queryField is not in query set")

    /**
     * Get a field on this table that corresponds to a field that was used in [query].
     */
    @Suppress("UNCHECKED_CAST")
    operator fun <T> get(queryField: Column<T>): Column<T> =
        fieldMap[queryField] as? Column<T>
            ?: error("$queryField is not in query set")

    /**
     * Get a field on this table that corresponds to a field that was used in [query].
     */
    @Suppress("UNCHECKED_CAST")
    operator fun <T> get(queryField: IExpressionAlias<T>): Expression<T> =
        fieldMap[queryField as Expression<*>] as? Expression<T>
            ?: error("$queryField is not in query set")
}


/**
 * Starts a CTE query.
 *
 * @param tables The list of [CTETable]s that are necessary to build the final query.
 * @param columnSet The set on which the final query will be constructed.
 */
class With(
    val tables: List<CTETable>,
    val columnSet: ColumnSet,
) {
    constructor(table: CTETable) : this(listOf(table), table)
}

/**
 * Creates a `SELECT` [Query] by selecting all columns from the provided [With.tables]s.
 */
fun With.selectAll(): Query = CTEQuery(
    tables = tables,
    fieldSet = columnSet.source,
)

/**
 * Creates a `SELECT` [Query] by selecting either a single [column], or a subset of [columns],
 * from the provided [With.tables]s.
 */
fun With.select(column: Expression<*>, vararg columns: Expression<*>): Query = CTEQuery(
    tables = tables,
    fieldSet = Select(columnSet, listOf(column) + columns.toList()),
)

/**
 * Creates a `SELECT` [Query] using a list of [columns] or expressions from the provided [With.tables]s.
 */
fun With.select(columns: List<Expression<*>>): Query = CTEQuery(
    tables = tables,
    fieldSet = Select(columnSet, columns),
)

private class CTEQuery(
    val tables: List<CTETable>,
    val fieldSet: FieldSet,
) : Query(fieldSet, null) {

    override fun prepareSQL(builder: QueryBuilder): String {
        builder {
            append("WITH ")
            if (tables.any { it.recursive != null }) {
                append("RECURSIVE ")
            }
            tables.forEachIndexed { index, with ->
                if (index > 0) append(", ")
                with.describeWith(builder)
            }
        }
        return super.prepareSQL(builder)
    }

    override fun copy(): CTEQuery = CTEQuery(tables, fieldSet).also { copy ->
        copyTo(copy)
    }
}


/**
 * An accessor for values of a result row.
 * This "enhances" ResultRow by allowing:
 * - Access to a value by its string name.
 * - Implementation to intercept row access and change the column if needed.
 *
 * For example, with a CTE [Query], you probably want to access using the field declaration of the table you are
 * querying, not using the fields of the temporary table. However, the result row contains values that are referenced
 * by the fields of the [CTETable] itself. This provides a way to "translate" columns on the fly.
 */
interface RowValues {
    operator fun <T> get(column: Expression<T>): T
    operator fun contains(column: Expression<*>): Boolean
    operator fun get(name: String): Any
    operator fun contains(name: String): Boolean
}

private open class ResultRowValues(
    val row: ResultRow,
) : RowValues {

    open fun <T> delegate(column: Expression<T>): Expression<T> = column

    override fun <T> get(column: Expression<T>): T = row[delegate(column)]

    override fun contains(column: Expression<*>): Boolean = row.hasValue(delegate(column))

    private val nameCache = HashMap<String, Expression<*>>()

    private fun getExpression(name: String): Expression<*>? {
        nameCache[name]?.let { return it }
        val expression = row.fieldIndex.keys.firstOrNull {
            (it is Column<*> && it.name == name) || (it is IExpressionAlias<*> && it.alias == name)
        } ?: return null
        nameCache[name] = expression
        return expression
    }

    override fun get(name: String): Any {
        val expression = getExpression(name)
            ?: error("$name is not in record set")
        return row[expression] as Any
    }

    override fun contains(name: String): Boolean = getExpression(name) != null

    override fun toString() = row.toString()
}

/**
 * Wraps this [ResultRow] into a [RowValues] that will access the values directly on the provided row.
 */
fun ResultRow.values(): RowValues = ResultRowValues(this)

private open class ResultRowCTEValues(
    row: ResultRow,
    val tables: List<CTETable>,
) : ResultRowValues(row) {

    override fun <T> delegate(column: Expression<T>): Expression<T> {
        if (row.hasValue(column)) return column
        return tables.firstOrNull { column in it }?.get(column)
            ?: error("$column is not in record set")
    }

}

/**
 * Wraps this [ResultRow] into a [RowValues] that will provide values for both the [tables] fields AND the fields
 * of their [CTETable.query].
 *
 * This allows to call [RowValues.get] with a field that was used in a [CTETable.query], and get the corresponding
 * value, without manually translating the field to a [CTETable] field.
 */
fun ResultRow.values(tables: List<CTETable>): RowValues = ResultRowCTEValues(this, tables)

/**
 * Wraps this [ResultRow] into a [RowValues] that will provide values for both the [table] fields AND the fields
 * of their [CTETable.query].
 *
 * This allows to call [RowValues.get] with a field that was used in a [CTETable.query], and get the corresponding
 * value, without manually translating the field to a [CTETable] field.
 */
fun ResultRow.values(table: CTETable, vararg tables: CTETable): RowValues = values(listOf(table) + tables)


/**
 *
 */
@Suppress("FunctionName")
class CTEBuilder internal constructor(
    private val tables: MutableList<CTETable>
) {

    /**
     * Creates a [CTETable] and registers it to be part of this [CTE] query.
     *
     * @param name The name (alias) of the temporary table.
     * @param query The initial query that builds the temporary table.
     * @param recursive An optional lambda that uses this CTETable to build a recursive query.
     *                  Defining this parameter means that this is a recursive CTE.
     * @param unionAll Only used if [recursive] is not null. Whether to recurse with `UNION` or `UNION ALL`.
     */
    fun With(
        name: String,
        query: AbstractQuery<*>,
        recursive: ((CTETable) -> AbstractQuery<*>)? = null,
        unionAll: Boolean = false,
    ) = CTETable(name, query, recursive, unionAll).also { tables.add(it) }

    /**
     * Registers this [table] to be part of this [CTE] query.
     */
    fun <T : CTETable> With(table: T): T = table.also { tables.add(it) }
}

/**
 * Builds an anonymous CTE query.
 */
fun CTE(build: CTEBuilder.() -> Query): List<RowValues> {
    val tables = ArrayList<CTETable>()
    var query = build(CTEBuilder(tables))
    if (query !is CTEQuery) {
        val cte = CTEQuery(tables, query.set)
        query.copyTo(cte)
        query = cte
    }
    return query.map { it.values(tables) }
}