visual-driver

1
@file:Suppress("UNCHECKED_CAST")
2

3
import com.mongodb.ClientSessionOptions
4
import com.mongodb.ConnectionString
5
import com.mongodb.MongoClientSettings
6
import com.mongodb.client.model.Filters
7
import com.mongodb.client.model.Indexes
8
import com.mongodb.reactivestreams.client.ClientSession
9
import data.FeatureUserStorage
10
import kotlinx.coroutines.runBlocking
11
import kotlinx.coroutines.withTimeout
12
import me.func.protocol.Unique
13
import org.bson.UuidRepresentation
14
import org.litote.kmongo.coroutine.CoroutineCollection
15
import org.litote.kmongo.coroutine.CoroutineDatabase
16
import org.litote.kmongo.coroutine.coroutine
17
import org.litote.kmongo.reactivestreams.KMongo
18
import org.litote.kmongo.replaceUpsert
19
import java.util.UUID
20
import kotlin.reflect.KClass
21

22
class CollectionTypeNotRegisteredException(name: String) : RuntimeException(name)
23

24
class MongoAdapter(private val url: String, private val databaseName: String, private val collectionName: String) {
25

26
    val collections: MutableMap<KClass<out Unique>, CoroutineCollection<out Unique>> = hashMapOf()
27
    private var database: CoroutineDatabase
28
    var session: ClientSession
29

30
    init {
31
        runBlocking {
32
            withTimeout(10000L) {
33
                val client = KMongo.createClient(
34
                    MongoClientSettings.builder()
35
                        .uuidRepresentation(UuidRepresentation.JAVA_LEGACY)
36
                        .applyConnectionString(ConnectionString(url))
37
                        .build()
38
                ).coroutine
39
                session = client.startSession(ClientSessionOptions.builder().causallyConsistent(true).build())
40
                database = client.getDatabase(databaseName)
41

42
                // Регистрирует типы, которые могут быть в коллекции $collectionName
43
                registerCollection<FeatureUserStorage>()
44

45
                // Создание индекса UUID для быстрого поиска
46
                collections.forEach { (_, value) ->
47
                    value.createIndex(Indexes.hashed("uuid"))
48
                    println("Created index for collection.")
49
                    println("Documents total: ${value.countDocuments()}")
50
                }
51
            }
52
        }
53
    }
54

55
    private inline fun <reified T : Unique> registerCollection() =
56
        collections.put(T::class, database.getCollection<T>(collectionName))
57

58
    inline fun <reified T : Unique> findCollection(): CoroutineCollection<T> = collections[T::class]?.run {
59
        return this as CoroutineCollection<T>
60
    } ?: throw CollectionTypeNotRegisteredException(T::class.simpleName ?: "null")
61

62
    suspend inline fun <reified T : Unique> find(uuid: UUID) =
63
        findCollection<T>().findOne(session, Filters.eq(uuid.toString()))
64

65
    suspend inline fun <reified T : Unique> save(unique: T) =
66
        findCollection<T>().replaceOne(session, Filters.eq(unique.uuid.toString()), unique, replaceUpsert())
67
}
68