Skip to content

Commit

Permalink
Add set operations JetBrains#402 / UNION and UNION ALL
Browse files Browse the repository at this point in the history
  • Loading branch information
Tapac authored and SchweinchenFuntik committed Oct 23, 2021
1 parent ad00553 commit 2be392a
Show file tree
Hide file tree
Showing 6 changed files with 409 additions and 108 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
package org.jetbrains.exposed.sql

import org.jetbrains.exposed.sql.statements.Statement
import org.jetbrains.exposed.sql.statements.StatementType
import org.jetbrains.exposed.sql.transactions.TransactionManager
import java.sql.ResultSet

abstract class AbstractQuery<T : AbstractQuery<T>>(targets: List<Table>) : SizedIterable<ResultRow>, Statement<ResultSet>(StatementType.SELECT, targets) {
protected val transaction get() = TransactionManager.current()

var orderByExpressions: List<Pair<Expression<*>, SortOrder>> = mutableListOf()
private set
var distinct: Boolean = false
protected set
var limit: Int? = null
protected set
var offset: Long = 0
private set
var fetchSize: Int? = null
private set

abstract val set: FieldSet

protected fun copyTo(other: AbstractQuery<T>) {
other.orderByExpressions = orderByExpressions.toMutableList()
other.distinct = distinct
other.limit = limit
other.offset = offset
other.fetchSize = fetchSize
}

override fun prepareSQL(transaction: Transaction) = prepareSQL(QueryBuilder(true))

internal abstract fun prepareSQL(builder: QueryBuilder): String

override fun arguments() = QueryBuilder(true).let {
prepareSQL(it)
if (it.args.isNotEmpty()) listOf(it.args) else emptyList()
}

fun withDistinct(value: Boolean = true): T = apply {
distinct = value
} as T

override fun limit(n: Int, offset: Long): T = apply {
limit = n
this.offset = offset
} as T

fun orderBy(column: Expression<*>, order: SortOrder = SortOrder.ASC): T = orderBy(column to order)

override fun orderBy(vararg order: Pair<Expression<*>, SortOrder>): T = apply {
(orderByExpressions as MutableList).addAll(order)
} as T

fun fetchSize(n: Int): T = apply {
fetchSize = n
} as T

protected var count: Boolean = false

protected abstract val queryToExecute: Statement<ResultSet>

override fun iterator(): Iterator<ResultRow> {
val resultIterator = ResultIterator(transaction.exec(queryToExecute)!!)
return if (transaction.db.supportsMultipleResultSets) resultIterator
else {
Iterable { resultIterator }.toList().iterator()
}
}

protected inner class ResultIterator(val rs: ResultSet) : Iterator<ResultRow> {
private var hasNext: Boolean? = null

private val fields = set.realFields

override operator fun next(): ResultRow {
if (hasNext == null) hasNext()
if (hasNext == false) throw NoSuchElementException()
hasNext = null
return ResultRow.create(rs, fields)
}

override fun hasNext(): Boolean {
if (hasNext == null) hasNext = rs.next()
if (hasNext == false) rs.close()
return hasNext!!
}
}
}
142 changes: 34 additions & 108 deletions exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/Query.kt
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
package org.jetbrains.exposed.sql

import org.jetbrains.exposed.sql.statements.Statement
import org.jetbrains.exposed.sql.statements.StatementType
import org.jetbrains.exposed.sql.statements.api.PreparedStatementApi
import org.jetbrains.exposed.sql.transactions.TransactionManager
import org.jetbrains.exposed.sql.vendors.currentDialect
Expand All @@ -12,37 +11,43 @@ enum class SortOrder {
ASC, DESC
}

open class Query(set: FieldSet, where: Op<Boolean>?): SizedIterable<ResultRow>, Statement<ResultSet>(StatementType.SELECT, set.source.targetTables()) {
private val transaction get() = TransactionManager.current()
open class Query(override var set: FieldSet, where: Op<Boolean>?) : AbstractQuery<Query>(set.source.targetTables()) {

var groupedByColumns: List<Expression<*>> = mutableListOf()
private set
var orderByExpressions: List<Pair<Expression<*>, SortOrder>> = mutableListOf()
private set

var having: Op<Boolean>? = null
private set
var distinct: Boolean = false
private set

private var forUpdate: Boolean? = null
var set: FieldSet = set
private set

//private set
var where: Op<Boolean>? = where
private set
var limit: Int? = null
private set
var offset: Long = 0
private set
var fetchSize: Int? = null
private set

override val queryToExecute: Statement<ResultSet> get() {
val distinctExpressions = set.fields.distinct()
return if (distinctExpressions.size < set.fields.size) {
copy().adjustSlice { slice(distinctExpressions) }
} else
this
}

override fun copy(): Query = Query(set, where).also { copy ->
copyTo(copy)
copy.groupedByColumns = groupedByColumns.toMutableList()
copy.orderByExpressions = orderByExpressions.toMutableList()
copy.having = having
copy.distinct = distinct
copy.forUpdate = forUpdate
copy.limit = limit
copy.offset = offset
copy.fetchSize = fetchSize
}

override fun forUpdate(): Query {
this.forUpdate = true
return this
}

override fun notForUpdate(): Query {
forUpdate = false
return this
}

/**
Expand Down Expand Up @@ -79,21 +84,13 @@ open class Query(set: FieldSet, where: Op<Boolean>?): SizedIterable<ResultRow>,
return executeQuery()
}

override fun arguments() = QueryBuilder(true).let {
prepareSQL(it)
if (it.args.isNotEmpty()) listOf(it.args) else emptyList()
}

override fun prepareSQL(transaction: Transaction): String = prepareSQL(QueryBuilder(true))

fun prepareSQL(builder: QueryBuilder): String {
override fun prepareSQL(builder: QueryBuilder): String {
builder {
append("SELECT ")

if (count) {
append("COUNT(*)")
}
else {
} else {
if (distinct) {
append("DISTINCT ")
}
Expand Down Expand Up @@ -140,97 +137,26 @@ open class Query(set: FieldSet, where: Op<Boolean>?): SizedIterable<ResultRow>,
return builder.toString()
}

override fun forUpdate() : Query {
this.forUpdate = true
return this
}

override fun notForUpdate(): Query {
forUpdate = false
return this
}

fun withDistinct(value: Boolean = true) : Query {
distinct = value
return this
}

fun groupBy(vararg columns: Expression<*>): Query {
for (column in columns) {
(groupedByColumns as MutableList).add(column)
}
return this
}

fun having(op: SqlExpressionBuilder.() -> Op<Boolean>) : Query {
val oop = Op.build { op() }
fun having(op: SqlExpressionBuilder.() -> Op<Boolean>): Query {
val oop = SqlExpressionBuilder.op()
if (having != null) {
error ("HAVING clause is specified twice. Old value = '$having', new value = '$oop'")
error("HAVING clause is specified twice. Old value = '$having', new value = '$oop'")
}
having = oop
return this
}

fun orderBy(column: Expression<*>, order: SortOrder = SortOrder.ASC) = orderBy(column to order)

override fun orderBy(vararg columns: Pair<Expression<*>, SortOrder>) : Query {
(orderByExpressions as MutableList).addAll(columns)
return this
}

override fun limit(n: Int, offset: Long): Query {
this.limit = n
this.offset = offset
return this
}

fun fetchSize(n: Int): Query {
this.fetchSize = n
return this
}

private inner class ResultIterator(val rs: ResultSet): Iterator<ResultRow> {
private var hasNext: Boolean? = null

private val fields = set.realFields

override operator fun next(): ResultRow {
if (hasNext == null) hasNext()
if (hasNext == false) throw NoSuchElementException()
hasNext = null
return ResultRow.create(rs, fields)
}

override fun hasNext(): Boolean {
if (hasNext == null) hasNext = rs.next()
if (hasNext == false) rs.close()
return hasNext!!
}
}


override operator fun iterator(): Iterator<ResultRow> {
val distinctExpressions = this.set.fields.distinct()
val queryToExecute = if (distinctExpressions.size < set.fields.size) {
copy().adjustSlice { slice(distinctExpressions) }
} else
this
val resultIterator = ResultIterator(transaction.exec(queryToExecute)!!)
return if (transaction.db.supportsMultipleResultSets)
resultIterator
else {
arrayListOf<ResultRow>().apply {
resultIterator.forEach {
this += it
}
}.iterator()
}
}

private var count: Boolean = false
override fun count(): Long {
return if (distinct || groupedByColumns.isNotEmpty() || limit != null) {
fun Column<*>.makeAlias() = alias(transaction.db.identifierManager.quoteIfNecessary("${table.tableName}_$name"))
fun Column<*>.makeAlias() =
alias(transaction.db.identifierManager.quoteIfNecessary("${table.tableName}_$name"))

val originalSet = set
try {
Expand Down Expand Up @@ -277,7 +203,7 @@ open class Query(set: FieldSet, where: Op<Boolean>?): SizedIterable<ResultRow>,
*/
fun Query.andWhere(andPart: SqlExpressionBuilder.() -> Op<Boolean>) = adjustWhere {
val expr = Op.build { andPart() }
if(this == null) expr
if (this == null) expr
else this and expr
}

Expand All @@ -287,6 +213,6 @@ fun Query.andWhere(andPart: SqlExpressionBuilder.() -> Op<Boolean>) = adjustWher
*/
fun Query.orWhere(andPart: SqlExpressionBuilder.() -> Op<Boolean>) = adjustWhere {
val expr = Op.build { andPart() }
if(this == null) expr
if (this == null) expr
else this or expr
}
Loading

0 comments on commit 2be392a

Please sign in to comment.