jvmMain.MongoKeyValueStore.kt Maven / Gradle / Ivy

Go to download
package io.fluidsonic.raptor.keyvaluestore.mongo

import com.mongodb.client.model.*
import com.mongodb.client.model.Filters.*
import com.mongodb.client.model.Updates.*
import io.fluidsonic.mongo.*
import io.fluidsonic.raptor.*
import io.fluidsonic.raptor.bson.*
import io.fluidsonic.raptor.keyvaluestore.*
import io.fluidsonic.raptor.keyvaluestore.RaptorKeyValueStore.*
import io.fluidsonic.raptor.mongo.*
import kotlin.reflect.*
import kotlinx.coroutines.flow.*
import org.bson.*
import org.bson.codecs.*
import org.bson.codecs.configuration.*


internal class MongoKeyValueStore(
	collection: MongoCollection,
	private val keyClass: KClass,
	private val valueClass: KClass,
) : RaptorKeyValueStore {

	@Suppress("UNCHECKED_CAST")
	private val collection: MongoCollection> = collection
		.withCodecRegistry(CodecRegistries.fromRegistries(
			CodecRegistries.fromCodecs(EntryCodec(
				keyCodec = collection.codecRegistry.get(keyClass.java),
				valueCodec = collection.codecRegistry.get(valueClass.java),
			)),
			collection.codecRegistry,
		))
		.withDocumentClass(Entry::class) as MongoCollection>


	override suspend fun clear() {
		collection.deleteMany(Document())
	}


	override fun entries(): Flow> =
		collection.find().map { it.toPair() }


	override fun keys(): Flow =
		collection.findOneField(Fields.key, keyClass)


	override fun values(): Flow =
		collection.findOneField(Fields.value, valueClass)


	override suspend fun set(key: Key, value: Value) {
		collection.replaceOneById(key, Entry(key, value), ReplaceOptions().upsert(true))
	}


	override suspend fun setIfAbsent(key: Key, value: Value): Boolean =
		collection.updateOne(
			filter = eq(Fields.key, key),
			update = setOnInsert(Fields.value, value),
			options = UpdateOptions().upsert(true),
		).upsertedId != null


	override suspend fun remove(key: Key): Boolean =
		collection.deleteOneById(key).deletedCount > 0


	override suspend fun get(key: Key): Value? =
		collection.findOneById(key)?.value


	private data class Entry(
		val key: Key,
		val value: Value,
	) {

		fun toPair(): Pair =
			key to value
	}


	private class EntryCodec(
		private val keyCodec: Codec,
		private val valueCodec: Codec,
	) : Codec> {

		override fun decode(reader: BsonReader, decoderContext: DecoderContext): Entry {
			reader.readStartDocument()

			reader.readName(Fields.key)
			val key = keyCodec.decode(reader, decoderContext)

			reader.readName(Fields.value)
			val value = valueCodec.decode(reader, decoderContext)

			reader.readEndDocument()

			return Entry(key = key, value = value)
		}


		override fun encode(writer: BsonWriter, value: Entry, encoderContext: EncoderContext) {
			writer.writeStartDocument()

			writer.writeName(Fields.key)
			keyCodec.encode(writer, value.key, encoderContext)

			writer.writeName(Fields.value)
			valueCodec.encode(writer, value.value, encoderContext)

			writer.writeEndDocument()
		}


		@Suppress("UNCHECKED_CAST")
		override fun getEncoderClass(): Class> =
			Entry::class.java as Class>
	}


	private object Fields {

		const val key = "_id"
		const val value = "value"
	}
}


@Suppress("FunctionName")
internal inline fun  MongoKeyValueStore(
	collection: MongoCollection,
): RaptorKeyValueStore =
	MongoKeyValueStore(collection = collection, keyClass = Key::class, valueClass = Value::class)