fix sql system

RegadPoleCN · Aug 29, 2024 · 8dfc386 · 8dfc386
1 parent 486df94
commit 8dfc386
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 24 deletions.
diff --git a/common/src/main/kotlin/me/regadpole/plumbot/PlumBot.kt b/common/src/main/kotlin/me/regadpole/plumbot/PlumBot.kt
@@ -6,6 +6,9 @@ import me.regadpole.plumbot.bot.QQBot
 import me.regadpole.plumbot.config.ConfigLoader
 import me.regadpole.plumbot.config.ConfigResolver
 import me.regadpole.plumbot.config.LangConfig
+import me.regadpole.plumbot.database.Database
+import me.regadpole.plumbot.database.MySQL
+import me.regadpole.plumbot.database.SQLite
 import taboolib.common.platform.Plugin
 import taboolib.common.platform.function.disablePlugin
 import taboolib.common.platform.function.info
@@ -21,6 +24,8 @@ object PlumBot : Plugin() {
 
     private lateinit var lang: LangConfig
 
+    private lateinit var database: Database
+
     override fun onLoad() {
         ConfigResolver.loadConfig()
         config = ConfigResolver.getConfigLoader()
@@ -30,6 +35,13 @@ object PlumBot : Plugin() {
 
     // 项目使用TabooLib Start Jar 创建!
     override fun onEnable() {
+        database = when(config.getConfig().database.type?.lowercase()) {
+            "sqlite" -> SQLite()
+            "mysql" -> MySQL()
+            else -> error("Unknown database type.")
+        }
+        database.initialize()
+        info("Loaded database")
         info("Successfully running PlumBot!")
     }
 
@@ -54,15 +66,8 @@ object PlumBot : Plugin() {
     }
 
     override fun onDisable() {
-        when (config.getConfig().bot.mode) {
-            "go-cqhttp", "kook" -> {
-                bot.shutdown()
-            }
-            else -> {
-                warning("无法正常关闭服务，将在服务器关闭后强制关闭")
-                disablePlugin()
-            }
-        }
+        bot.shutdown()
+        database.close()
     }
 
     fun reloadConfig() {
@@ -81,5 +86,7 @@ object PlumBot : Plugin() {
     fun getLangConfig(): LangConfig {
         return lang
     }
-
+    fun getDatabase(): Database {
+        return database
+    }
 }
diff --git a/common/src/main/kotlin/me/regadpole/plumbot/database/Database.kt b/common/src/main/kotlin/me/regadpole/plumbot/database/Database.kt
@@ -26,10 +26,10 @@ interface Database {
 
     fun addBind(user: String, name: String)
 
-    fun removeBind(user: String, name: String)
+    fun removeBind(user: String, id: Int)
     fun removeBind(user: String)
     fun removeBindByName(name: String)
 
-    fun getBind(user: String): List<String>
+    fun getBind(user: String): MutableMap<Int, String>
     fun getBindByName(name: String): String?
 }
diff --git a/common/src/main/kotlin/me/regadpole/plumbot/database/MySQL.kt b/common/src/main/kotlin/me/regadpole/plumbot/database/MySQL.kt
@@ -2,7 +2,6 @@ package me.regadpole.plumbot.database
 
 import me.regadpole.plumbot.PlumBot
 import taboolib.common.platform.function.info
-import taboolib.common.util.asList
 import taboolib.module.database.ColumnOptionSQL
 import taboolib.module.database.ColumnTypeSQL
 import taboolib.module.database.Table
@@ -55,9 +54,9 @@ class MySQL: Database {
         }
     }
 
-    override fun removeBind(user: String, name: String) {
+    override fun removeBind(user: String, id: Int) {
         table.delete(dataSource) {
-            where { "user" eq user and ("name" eq name) }
+            where { "user" eq user and ("id" eq id) }
         }
     }
 
@@ -73,11 +72,15 @@ class MySQL: Database {
         }
     }
 
-    override fun getBind(user: String): List<String> {
-        return table.select(dataSource) {
+    override fun getBind(user: String): MutableMap<Int, String> {
+        val map: MutableMap<Int, String> = LinkedHashMap()
+        val result = table.select(dataSource) {
+            rows("id")
             rows("name")
             where("user" eq user)
-        }.asList()
+        }
+        result.forEach { map[getInt("id")] = getString("name") }
+        return map
     }
 
     override fun getBindByName(name: String): String? {

diff --git a/common/src/main/kotlin/me/regadpole/plumbot/database/SQLite.kt b/common/src/main/kotlin/me/regadpole/plumbot/database/SQLite.kt
@@ -2,7 +2,6 @@ package me.regadpole.plumbot.database
 
 import me.regadpole.plumbot.PlumBot
 import taboolib.common.platform.function.info
-import taboolib.common.util.asList
 import taboolib.module.database.ColumnTypeSQLite
 import taboolib.module.database.Table
 import javax.sql.DataSource
@@ -52,9 +51,9 @@ class SQLite: Database {
         }
     }
 
-    override fun removeBind(user: String, name: String) {
+    override fun removeBind(user: String, id: Int) {
         table.delete(dataSource) {
-            where { "user" eq user and ("name" eq name) }
+            where { "user" eq user and ("id" eq id) }
         }
     }
 
@@ -70,11 +69,15 @@ class SQLite: Database {
         }
     }
 
-    override fun getBind(user: String): List<String> {
-        return table.select(dataSource) {
+    override fun getBind(user: String): MutableMap<Int, String> {
+        val map: MutableMap<Int, String> = LinkedHashMap()
+        val result = table.select(dataSource) {
+            rows("id")
             rows("name")
             where("user" eq user)
-        }.asList()
+        }
+        result.forEach { map[getInt("id")] = getString("name") }
+        return map
     }
 
     override fun getBindByName(name: String): String? {

diff --git a/common/src/main/kotlin/me/regadpole/plumbot/internal/WhitelistHelper.kt b/common/src/main/kotlin/me/regadpole/plumbot/internal/WhitelistHelper.kt
@@ -0,0 +1,27 @@
+package me.regadpole.plumbot.internal
+
+import me.regadpole.plumbot.PlumBot
+
+
+class WhitelistHelper {
+    fun checkCount(user: String): Boolean {
+        val idList = PlumBot.getDatabase().getBind(user).values
+        val maxCount: Int = PlumBot.getConfig().getConfig().whiteList.maxCount
+        if (idList.isEmpty()) return true
+        return idList.size < maxCount
+    }
+
+    fun checkIDNotExist(name: String): Boolean {
+        return PlumBot.getDatabase().getBindByName(name).isNullOrEmpty()
+    }
+
+    fun addAndGet(user: String, name: String): MutableMap<Int, String> {
+        PlumBot.getDatabase().addBind(user, name)
+        return PlumBot.getDatabase().getBind(user)
+    }
+
+    fun removeAndGet(user: String, id: Int): MutableMap<Int, String> {
+        PlumBot.getDatabase().removeBind(user, id)
+        return PlumBot.getDatabase().getBind(user)
+    }
+}