com.lightningkite.lightningdb.MongoFieldCollection.kt Maven / Gradle / Ivy
package com.lightningkite.lightningdb
import com.lightningkite.lightningserver.exceptions.BadRequestException
import com.lightningkite.lightningserver.serialization.Serialization
import com.mongodb.*
import com.mongodb.client.model.*
import com.mongodb.client.model.Aggregates.group
import com.mongodb.client.model.Aggregates.match
import com.mongodb.kotlin.client.coroutine.MongoCollection
import kotlinx.coroutines.*
import kotlinx.coroutines.flow.*
import kotlinx.serialization.ExperimentalSerializationApi
import kotlinx.serialization.KSerializer
import kotlinx.serialization.Serializable
import org.bson.BsonDocument
import org.bson.conversions.Bson
import java.nio.channels.ClosedChannelException
import java.util.concurrent.TimeUnit
import kotlin.reflect.KProperty1
/**
* MongoFieldCollection implements FieldCollection specifically for a MongoDB server.
*/
class MongoFieldCollection(
val serializer: KSerializer,
private val getMongo: () -> MongoCollection,
private val onConnectionError: ()->Unit,
) : AbstractSignalFieldCollection() {
val mongo: MongoCollection get() = getMongo()
private fun sort(orderBy: List>, lastly: Bson? = null): Bson = Sorts.orderBy(orderBy.map {
if (it.ascending)
Sorts.ascending(it.field.mongo)
else
Sorts.descending(it.field.mongo)
} + listOfNotNull(lastly))
private inline fun exceptionWrap(operation: () -> T): T {
try {
return operation()
} catch (e: Throwable) {
handleException(e)
}
}
private fun handleException(e: Throwable): Nothing {
when {
e is MongoBulkWriteException && e.writeErrors.all { ErrorCategory.fromErrorCode(it.code) == ErrorCategory.DUPLICATE_KEY } -> {
throw UniqueViolationException(cause = e, collection = mongo.namespace.collectionName)
}
e is MongoException && ErrorCategory.fromErrorCode(e.code) == ErrorCategory.DUPLICATE_KEY -> {
throw UniqueViolationException(cause = e, collection = mongo.namespace.collectionName)
}
e is MongoSocketWriteException && e.cause is ClosedChannelException -> {
onConnectionError()
throw e
}
else -> throw e
}
}
override suspend fun find(
condition: Condition,
orderBy: List>,
skip: Int,
limit: Int,
maxQueryMs: Long,
): Flow {
val cs = condition.simplify()
if (cs is Condition.Never) return emptyFlow()
prepare.await()
return mongo.find(cs.bson(serializer))
.let {
if (skip != 0) it.skip(skip)
else it
}
.let {
if (limit != Int.MAX_VALUE) it.limit(limit)
else it
}
.maxTime(maxQueryMs, TimeUnit.MILLISECONDS)
.let {
var anyFts = false
condition.walk { if (it is Condition.FullTextSearch) anyFts = true }
val mts = if (anyFts) {
it.projection(Projections.metaTextScore("text_search_score"))
Sorts.metaTextScore("text_search_score")
} else null
it.sort(sort(orderBy, mts))
}
.let {
if (orderBy.any { it.ignoreCase }) {
it.collation(Collation.builder().locale("en").build())
} else it
}
.map { Serialization.Internal.bson.load(serializer, it) }
.catch { handleException(it) }
}
@Serializable
data class KeyHolder(val _id: Key)
override suspend fun count(condition: Condition): Int {
val cs = condition.simplify()
if (cs is Condition.Never) return 0
prepare.await()
return exceptionWrap { mongo.countDocuments(cs.bson(serializer)).toInt() }
}
override suspend fun groupCount(
condition: Condition,
groupBy: DataClassPath,
): Map {
val cs = condition.simplify()
if (cs is Condition.Never) return mapOf()
prepare.await()
return exceptionWrap {
mongo.aggregate(
listOf(
match(cs.bson(serializer)),
group("\$" + groupBy.mongo, Accumulators.sum("count", 1))
)
)
.toList()
.associate {
Serialization.Internal.bson.load(
KeyHolder.serializer(serializer.fieldSerializer(groupBy)!!),
it
)._id to it.getNumber("count").intValue()
}
}
}
private fun Aggregate.asValueBson(propertyName: String) = when (this) {
Aggregate.Sum -> Accumulators.sum("value", "\$" + propertyName)
Aggregate.Average -> Accumulators.avg("value", "\$" + propertyName)
Aggregate.StandardDeviationPopulation -> Accumulators.stdDevPop("value", "\$" + propertyName)
Aggregate.StandardDeviationSample -> Accumulators.stdDevSamp("value", "\$" + propertyName)
}
override suspend fun aggregate(
aggregate: Aggregate,
condition: Condition,
property: DataClassPath,
): Double? {
val cs = condition.simplify()
if (cs is Condition.Never) return null
prepare.await()
return exceptionWrap {
mongo.aggregate(listOf(match(cs.bson(serializer)), group(null, aggregate.asValueBson(property.mongo))))
.toList()
.map {
if (it.isNull("value")) null
else it.getNumber("value").doubleValue()
}
.firstOrNull()
}
}
override suspend fun groupAggregate(
aggregate: Aggregate,
condition: Condition,
groupBy: DataClassPath,
property: DataClassPath,
): Map {
val cs = condition.simplify()
if (cs is Condition.Never) return mapOf()
prepare.await()
return exceptionWrap {
mongo.aggregate(
listOf(
match(cs.bson(serializer)),
group("\$" + groupBy.mongo, aggregate.asValueBson(property.mongo))
)
)
.toList()
.associate {
Serialization.Internal.bson.load(
KeyHolder.serializer(serializer.fieldSerializer(groupBy)!!),
it
)._id to (if (it.isNull("value")) null else it.getNumber("value").doubleValue())
}
}
}
override suspend fun insertImpl(
models: Iterable,
): List {
prepare.await()
if (models.none()) return emptyList()
val asList = models.toList()
exceptionWrap { mongo.insertMany(asList.map { Serialization.Internal.bson.stringify(serializer, it) }) }
return asList
}
override suspend fun replaceOneImpl(
condition: Condition,
model: Model,
orderBy: List>,
): EntryChange {
val cs = condition.simplify()
if (cs is Condition.Never) return EntryChange(null, null)
prepare.await()
return updateOne(cs, Modification.Assign(model), orderBy)
}
override suspend fun replaceOneIgnoringResultImpl(
condition: Condition,
model: Model,
orderBy: List>,
): Boolean {
val cs = condition.simplify()
if (cs is Condition.Never) return false
prepare.await()
if (orderBy.isNotEmpty()) return updateOneIgnoringResultImpl(cs, Modification.Assign(model), orderBy)
return exceptionWrap {
mongo.replaceOne(
cs.bson(serializer),
Serialization.Internal.bson.stringify(serializer, model)
).matchedCount != 0L
}
}
override suspend fun upsertOneImpl(
condition: Condition,
modification: Modification,
model: Model,
): EntryChange {
val cs = condition.simplify()
if (cs is Condition.Never) return EntryChange(null, null)
if (modification is Modification.Chain && modification.modifications.isEmpty()) return EntryChange(null, null)
prepare.await()
val m = modification.simplify().bson(serializer)
return exceptionWrap {
// TODO: Ugly hack for handling weird upserts
if (m.upsert(model, serializer)) {
mongo.findOneAndUpdate(
cs.bson(serializer),
m.document,
FindOneAndUpdateOptions()
.returnDocument(ReturnDocument.BEFORE)
.upsert(m.options.isUpsert)
.bypassDocumentValidation(m.options.bypassDocumentValidation)
.collation(m.options.collation)
.arrayFilters(m.options.arrayFilters)
.hint(m.options.hint)
.hintString(m.options.hintString)
)?.let { Serialization.Internal.bson.load(serializer, it) }?.let { EntryChange(it, modification(it)) }
?: EntryChange(null, model)
} else {
mongo.findOneAndUpdate(
cs.bson(serializer),
m.document,
FindOneAndUpdateOptions()
.returnDocument(ReturnDocument.BEFORE)
.upsert(m.options.isUpsert)
.bypassDocumentValidation(m.options.bypassDocumentValidation)
.collation(m.options.collation)
.arrayFilters(m.options.arrayFilters)
.hint(m.options.hint)
.hintString(m.options.hintString)
)?.let { Serialization.Internal.bson.load(serializer, it) }?.let { EntryChange(it, modification(it)) }
?: run { mongo.insertOne(Serialization.Internal.bson.stringify(serializer, model)); EntryChange(null, model) }
}
}
}
override suspend fun upsertOneIgnoringResultImpl(
condition: Condition,
modification: Modification,
model: Model,
): Boolean {
val cs = condition.simplify()
if (cs is Condition.Never) return false
if (modification is Modification.Chain && modification.modifications.isEmpty()) return false
prepare.await()
return exceptionWrap {
val m = modification.simplify().bson(serializer)
// TODO: Ugly hack for handling weird upserts
if (m.upsert(model, serializer)) {
mongo.updateOne(cs.bson(serializer), m.document, m.options).matchedCount > 0
} else {
if (mongo.updateOne(cs.bson(serializer), m.document, m.options).matchedCount != 0L) {
true
} else {
mongo.insertOne(Serialization.Internal.bson.stringify(serializer, model))
false
}
}
}
}
override suspend fun updateOneImpl(
condition: Condition,
modification: Modification,
orderBy: List>,
): EntryChange {
val cs = condition.simplify()
if (cs is Condition.Never) return EntryChange(null, null)
if (modification is Modification.Chain && modification.modifications.isEmpty()) return EntryChange(null, null)
prepare.await()
val m = modification.simplify().bson(serializer)
val before = exceptionWrap {
mongo.findOneAndUpdate(
cs.bson(serializer),
m.document,
FindOneAndUpdateOptions()
.returnDocument(ReturnDocument.BEFORE)
.let { if (orderBy.isEmpty()) it else it.sort(sort(orderBy)) }
.upsert(m.options.isUpsert)
.bypassDocumentValidation(m.options.bypassDocumentValidation)
.collation(m.options.collation)
.arrayFilters(m.options.arrayFilters)
.hint(m.options.hint)
.hintString(m.options.hintString)
) ?: return EntryChange(null, null)
}.let { Serialization.Internal.bson.load(serializer,it) }
val after = modification(before)
return EntryChange(before, after)
}
override suspend fun updateOneIgnoringResultImpl(
condition: Condition,
modification: Modification,
orderBy: List>,
): Boolean {
val cs = condition.simplify()
if (cs is Condition.Never) return false
if (modification is Modification.Chain && modification.modifications.isEmpty()) return false
prepare.await()
val m = modification.simplify().bson(serializer)
return exceptionWrap { mongo.updateOne(cs.bson(serializer), m.document, m.options).matchedCount != 0L }
}
override suspend fun updateManyImpl(
condition: Condition,
modification: Modification,
): CollectionChanges {
val cs = condition.simplify()
if (cs is Condition.Never) return CollectionChanges()
if (modification is Modification.Chain && modification.modifications.isEmpty()) return CollectionChanges()
prepare.await()
val m = modification.simplify().bson(serializer)
val changes = ArrayList>()
// TODO: Don't love that we have to do this in chunks, but I guess we'll live. Could this be done with pipelines?
exceptionWrap {
mongo.find(cs.bson(serializer)).collectChunked(1000) { list ->
mongo.updateMany(Filters.`in`("_id", list.map { it["_id"] }), m.document, m.options)
list.asSequence().map { Serialization.Internal.bson.load(serializer, it) }
.forEach {
changes.add(EntryChange(it, modification(it)))
}
}
}
return CollectionChanges(changes = changes)
}
override suspend fun updateManyIgnoringResultImpl(
condition: Condition,
modification: Modification,
): Int {
val cs = condition.simplify()
if (cs is Condition.Never) return 0
if (modification is Modification.Chain && modification.modifications.isEmpty()) return 0
prepare.await()
val m = modification.simplify().bson(serializer)
return exceptionWrap {
mongo.updateMany(
cs.bson(serializer),
m.document,
m.options
).matchedCount.toInt()
}
}
override suspend fun deleteOneImpl(condition: Condition, orderBy: List>): Model? {
val cs = condition.simplify()
if (cs is Condition.Never) return null
prepare.await()
return exceptionWrap {
// TODO: Hack, needs some retry logic at a minimum
mongo.withDocumentClass().find(cs.bson(serializer))
.let { if (orderBy.isEmpty()) it else it.sort(sort(orderBy)) }
.limit(1).firstOrNull()?.let {
val id = it["_id"]
mongo.deleteOne(Filters.eq("_id", id))
Serialization.Internal.bson.load(serializer, it)
}
}
}
override suspend fun deleteOneIgnoringOldImpl(
condition: Condition,
orderBy: List>,
): Boolean {
val cs = condition.simplify()
if (cs is Condition.Never) return false
if (orderBy.isNotEmpty()) return deleteOneImpl(condition, orderBy) != null
prepare.await()
return exceptionWrap { mongo.deleteOne(cs.bson(serializer)).deletedCount > 0 }
}
override suspend fun deleteManyImpl(condition: Condition): List {
val cs = condition.simplify()
if (cs is Condition.Never) return listOf()
prepare.await()
val remove = ArrayList()
exceptionWrap {
// TODO: Don't love that we have to do this in chunks, but I guess we'll live. Could this be done with pipelines?
mongo.withDocumentClass().find(cs.bson(serializer)).collectChunked(1000) { list ->
mongo.deleteMany(Filters.`in`("_id", list.map { it["_id"] }))
list.asSequence().map { Serialization.Internal.bson.load(serializer, it) }
.forEach {
remove.add(it)
}
}
}
return remove
}
override suspend fun deleteManyIgnoringOldImpl(
condition: Condition,
): Int {
val cs = condition.simplify()
if (cs is Condition.Never) return 0
prepare.await()
return exceptionWrap { mongo.deleteMany(cs.bson(serializer)).deletedCount.toInt() }
}
@OptIn(DelicateCoroutinesApi::class, ExperimentalSerializationApi::class)
val prepare = GlobalScope.async(Dispatchers.Unconfined, start = CoroutineStart.LAZY) {
val requireCompletion = ArrayList()
serializer.descriptor.annotations.filterIsInstance().firstOrNull()?.let {
requireCompletion += launch {
val name = "${mongo.namespace.fullName}TextIndex"
val keys = documentOf(*it.fields.map { it to "text" }.toTypedArray())
val options = IndexOptions().name(name)
try {
mongo.createIndex(keys, options)
} catch (e: MongoCommandException) {
//there is an exception if the parameters of an existing index are changed.
//then drop the index and create a new one
mongo.dropIndex(name)
mongo.createIndex(
keys,
options
)
}
}
}
serializer.descriptor.indexes().forEach {
if (it.unique) {
requireCompletion += launch {
val keys = Sorts.ascending(it.fields)
val options = IndexOptions().unique(true).name(it.name)
try {
mongo.createIndex(keys, options)
} catch (e: MongoCommandException) {
// Reform index if it already exists but with some difference in options
if (e.errorCode == 85) {
mongo.dropIndex(keys)
mongo.createIndex(keys, options)
}
}
}
} else {
launch {
val keys = Sorts.ascending(it.fields)
val options = IndexOptions().unique(false).background(true).name(it.name)
try {
mongo.createIndex(keys, options)
} catch (e: MongoCommandException) {
// Reform index if it already exists but with some difference in options
if (e.errorCode == 85) {
mongo.dropIndex(keys)
mongo.createIndex(keys, options)
}
}
}
}
}
requireCompletion.forEach { it.join() }
}
}
© 2015 - 2024 Weber Informatics LLC | Privacy Policy