From 469b40d7a8bb1a69264d9e6bf4b2d1f656a9a574 Mon Sep 17 00:00:00 2001
From: Gareth Coles <gdude2002@gmail.com>
Date: Sat, 17 Aug 2024 09:59:42 +0100
Subject: [PATCH] Key data collection state by bot ID

---
 .../main/kotlin/dev/kordex/core/ExtensibleBot.kt |  8 +++++---
 .../kordex/core/datacollection/DataCollector.kt  | 16 ++++++++++------
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/kord-extensions/src/main/kotlin/dev/kordex/core/ExtensibleBot.kt b/kord-extensions/src/main/kotlin/dev/kordex/core/ExtensibleBot.kt
index 7369a3f21e..305e8245eb 100644
--- a/kord-extensions/src/main/kotlin/dev/kordex/core/ExtensibleBot.kt
+++ b/kord-extensions/src/main/kotlin/dev/kordex/core/ExtensibleBot.kt
@@ -160,8 +160,6 @@ public open class ExtensibleBot(
 			logger.warn(e) { "Unable to add shutdown hook." }
 		}
 
-		dataCollector.start()
-
 		getKoin().get<Kord>().login {
 			this.presence(settings.presenceBuilder)
 			this.intents = Intents(settings.intentsBuilder!!)
@@ -196,7 +194,7 @@ public open class ExtensibleBot(
 		@Suppress("TooGenericExceptionCaught")
 		try {
 			Runtime.getRuntime().removeShutdownHook(shutdownHook)
-		} catch (e: IllegalStateException) {
+		} catch (_: IllegalStateException) {
 			logger.debug { "Shutdown in progress, unable to remove shutdown hook." }
 		} catch (e: Exception) {
 			logger.warn(e) { "Failed to remove shutdown hook." }
@@ -224,6 +222,10 @@ public open class ExtensibleBot(
 			ignoreUnknownKeys = true
 		}
 
+		on<ReadyEvent> {
+			dataCollector.start()
+		}
+
 		on<GuildCreateEvent> {
 			withLock {  // If configured, this won't be concurrent, saving larger bots from spammy rate limits
 				if (
diff --git a/kord-extensions/src/main/kotlin/dev/kordex/core/datacollection/DataCollector.kt b/kord-extensions/src/main/kotlin/dev/kordex/core/datacollection/DataCollector.kt
index b57f7c092b..fe3f574fe1 100644
--- a/kord-extensions/src/main/kotlin/dev/kordex/core/datacollection/DataCollector.kt
+++ b/kord-extensions/src/main/kotlin/dev/kordex/core/datacollection/DataCollector.kt
@@ -26,7 +26,6 @@ import dev.kordex.data.api.types.impl.MinimalDataEntity
 import dev.kordex.data.api.types.impl.StandardDataEntity
 import io.github.oshai.kotlinlogging.KotlinLogging
 import io.ktor.client.plugins.*
-import io.ktor.utils.io.reader
 import kotlinx.coroutines.flow.count
 import kotlinx.coroutines.launch
 import org.koin.core.component.inject
@@ -69,8 +68,7 @@ public class DataCollector(public val level: DataCollection) : KordExKoinCompone
 			current.lastLevel = props.getProperty("lastLevel")?.let { DataCollection.fromDB(it) }
 			current.uuid = props.getProperty("uuid")?.let { UUID.fromString(it) }
 
-			storageUnit.save()
-
+			saveState()
 			deleteOldState()
 
 			logger.info { "Migration complete!" }
@@ -309,19 +307,25 @@ public class DataCollector(public val level: DataCollection) : KordExKoinCompone
 	}
 
 	private suspend fun getState(): State {
-		var current = storageUnit.get()
+		var current = storageUnit
+			.withUser(bot.kordRef.selfId)
+			.get()
 
 		if (current == null) {
 			current = State()
 
-			storageUnit.save(current)
+			storageUnit
+				.withUser(bot.kordRef.selfId)
+				.save(current)
 		}
 
 		return current
 	}
 
 	private suspend fun saveState() {
-		storageUnit.save()
+		storageUnit
+			.withUser(bot.kordRef.selfId)
+			.save()
 	}
 
 	/** Get the stored "last" data collection level. **/