package top.jie65535.mirai import net.mamoe.mirai.Bot import net.mamoe.mirai.contact.Contact import net.mamoe.mirai.contact.Friend import net.mamoe.mirai.contact.Group import net.mamoe.mirai.contact.Member import net.mamoe.mirai.contact.Stranger import net.mamoe.mirai.event.events.MessageEvent import net.mamoe.mirai.event.events.MessagePostSendEvent import net.mamoe.mirai.event.events.MessageRecallEvent import net.mamoe.mirai.message.data.MessageSource import net.mamoe.mirai.message.data.MessageSourceKind import net.mamoe.mirai.message.data.source import net.mamoe.mirai.message.data.toMessageChain import java.io.File import java.sql.Connection import java.sql.DriverManager import java.sql.ResultSet /** * 基于 SQLite 的聊天记录存储。 * * 写连接由 [writeLock] 串行保护;SQLite 使用 WAL + NORMAL synchronous,允许并发读取。 */ object ChatHistoryStore { private const val SCHEMA_VERSION = 2 private const val BUSY_TIMEOUT_MS = 30_000 private val lifecycleLock = Any() private val writeLock = Any() @Volatile private var initialized = false private lateinit var databaseFile: File private var writeConnection: Connection? = null val isAvailable: Boolean get() = initialized fun init(dataFolder: File) { synchronized(lifecycleLock) { if (initialized) return Class.forName("org.sqlite.JDBC") dataFolder.mkdirs() databaseFile = dataFolder.resolve("chat-history.sqlite") val connection = openConnection() try { configureWriteConnection(connection) createSchema(connection) writeConnection = connection initialized = true } catch (cause: Throwable) { connection.close() throw cause } } } fun close() { synchronized(lifecycleLock) { if (!initialized) return synchronized(writeLock) { writeConnection?.let { connection -> runCatching { connection.createStatement().use { statement -> statement.execute("PRAGMA wal_checkpoint(TRUNCATE)") } }.onFailure { JChatGPT.logger.warning("SQLite WAL checkpoint 失败", it) } connection.close() } writeConnection = null initialized = false } } } fun record(event: MessageEvent) { if (!initialized) return val message = event.message.asSequence() .filterNot { it is MessageSource } .toMessageChain() insert(ChatMessageRecord.fromSuccess(event.message.source, message)) } fun record(event: MessagePostSendEvent<*>) { if (!initialized) return val source = event.receipt?.source ?: return val message = event.message.asSequence() .filterNot { it is MessageSource } .toMessageChain() insert(ChatMessageRecord.fromSuccess(source, message)) } fun markRecalled(event: MessageRecallEvent) { if (!initialized) return val (kind, targetId, recalled) = when (event) { is MessageRecallEvent.FriendRecall -> Triple( MessageSourceKind.FRIEND, event.bot.id, 2, ) is MessageRecallEvent.GroupRecall -> Triple( MessageSourceKind.GROUP, event.group.id, if ((event.operator?.id ?: event.bot.id) == event.authorId) 2 else 3, ) } val messageIds = event.messageIds.joinToString(",") val messageInternalIds = event.messageInternalIds.joinToString(",") withWriteConnection { connection -> connection.prepareStatement( """ UPDATE message_record SET recalled = ? WHERE id = ( SELECT id FROM message_record WHERE bot_id = ? AND kind = ? AND from_id = ? AND target_id = ? AND (ids = ? OR internal_ids = ?) ORDER BY ABS(time - ?) ASC, id DESC LIMIT 1 ) """.trimIndent() ).use { statement -> statement.setInt(1, recalled) statement.setLong(2, event.bot.id) statement.setInt(3, kind.ordinal) statement.setLong(4, event.authorId) statement.setLong(5, targetId) statement.setString(6, messageIds) statement.setString(7, messageInternalIds) statement.setInt(8, event.messageTime) val updated = statement.executeUpdate() if (updated == 0) { JChatGPT.logger.warning( "未在 SQLite 中找到撤回消息: bot=${event.bot.id}, " + "author=${event.authorId}, target=$targetId, ids=$messageIds, " + "internalIds=$messageInternalIds, " + "time=${event.messageTime}" ) } } } } fun query( contact: Contact, start: Int, end: Int, limit: Int, fromId: Long? = null, ): List { check(initialized) { "聊天记录数据库尚未初始化" } require(start <= end) { "start must not be after end" } val conditions = mutableListOf() val parameters = mutableListOf() conditions += "bot_id = ?" parameters += contact.bot.id conditions += "time BETWEEN ? AND ?" parameters += start parameters += end when (contact) { is Group -> { conditions += "kind = ?" parameters += MessageSourceKind.GROUP.ordinal conditions += "target_id = ?" parameters += contact.id } is Member -> { conditions += "kind = ?" parameters += MessageSourceKind.GROUP.ordinal conditions += "target_id = ?" parameters += contact.group.id conditions += "from_id = ?" parameters += contact.id } is Friend -> { conditions += "kind = ?" parameters += MessageSourceKind.FRIEND.ordinal conditions += "(from_id = ? OR target_id = ?)" parameters += contact.id parameters += contact.id } is Stranger -> { conditions += "kind = ?" parameters += MessageSourceKind.STRANGER.ordinal conditions += "(from_id = ? OR target_id = ?)" parameters += contact.id parameters += contact.id } is Bot -> Unit else -> error("不支持查询的联系人 $contact") } if (fromId != null && contact !is Member) { conditions += "from_id = ?" parameters += fromId } val sql = buildString { append( """ SELECT id, bot_id, from_id, target_id, ids, internal_ids, time, kind, code, recalled FROM message_record WHERE """.trimIndent() ) append(' ') append(conditions.joinToString(" AND ")) append(" ORDER BY time DESC, id DESC LIMIT ?") } return openReadConnection().use { connection -> connection.prepareStatement(sql).use { statement -> parameters.forEachIndexed { index, value -> when (value) { is Int -> statement.setInt(index + 1, value) is Long -> statement.setLong(index + 1, value) else -> error("不支持的查询参数类型 ${value::class}") } } statement.setInt(parameters.size + 1, limit.coerceAtLeast(1)) statement.executeQuery().use { results -> buildList { while (results.next()) { add(results.toRecord()) } } } } } } private fun insert(record: ChatMessageRecord) { withWriteConnection { connection -> connection.prepareStatement( """ INSERT INTO message_record( bot_id, from_id, target_id, ids, internal_ids, time, kind, code, recalled ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) """.trimIndent() ).use { statement -> statement.setLong(1, record.botId) statement.setLong(2, record.fromId) statement.setLong(3, record.targetId) statement.setString(4, record.ids) statement.setString(5, record.internalIds) statement.setInt(6, record.time) statement.setInt(7, record.kind.ordinal) statement.setString(8, record.code) statement.setInt(9, record.recalled) statement.executeUpdate() } } } private fun withWriteConnection(block: (Connection) -> Unit) { synchronized(writeLock) { check(initialized) { "聊天记录数据库尚未初始化" } val connection = writeConnection?.takeUnless(Connection::isClosed) ?: openConnection().also { configureWriteConnection(it) writeConnection = it } block(connection) } } private fun openConnection(): Connection { return DriverManager.getConnection("jdbc:sqlite:${databaseFile.absolutePath}") } private fun openReadConnection(): Connection { return openConnection().also { connection -> connection.createStatement().use { statement -> statement.execute("PRAGMA busy_timeout=$BUSY_TIMEOUT_MS") statement.execute("PRAGMA query_only=ON") } } } private fun configureWriteConnection(connection: Connection) { connection.createStatement().use { statement -> statement.execute("PRAGMA journal_mode=WAL") statement.execute("PRAGMA synchronous=NORMAL") statement.execute("PRAGMA busy_timeout=$BUSY_TIMEOUT_MS") statement.execute("PRAGMA wal_autocheckpoint=1000") } } private fun createSchema(connection: Connection) { connection.createStatement().use { statement -> statement.executeUpdate( """ CREATE TABLE IF NOT EXISTS message_record( id INTEGER PRIMARY KEY AUTOINCREMENT, bot_id INTEGER NOT NULL, from_id INTEGER NOT NULL, target_id INTEGER NOT NULL, ids TEXT, internal_ids TEXT, time INTEGER NOT NULL, kind INTEGER NOT NULL, code TEXT NOT NULL, recalled INTEGER NOT NULL DEFAULT 0 ) """.trimIndent() ) statement.executeUpdate( "CREATE INDEX IF NOT EXISTS idx_message_subject_time " + "ON message_record(bot_id, kind, target_id, time DESC, id DESC)" ) statement.executeUpdate( "CREATE INDEX IF NOT EXISTS idx_message_sender_subject_time " + "ON message_record(bot_id, kind, target_id, from_id, time DESC, id DESC)" ) statement.executeUpdate( "CREATE INDEX IF NOT EXISTS idx_message_recall_identity " + "ON message_record(bot_id, kind, target_id, from_id, time, ids)" ) statement.executeUpdate( """ CREATE TABLE IF NOT EXISTS chat_history_meta( key TEXT PRIMARY KEY, value TEXT NOT NULL ) """.trimIndent() ) } connection.prepareStatement( "INSERT INTO chat_history_meta(key, value) VALUES ('schema_version', ?) " + "ON CONFLICT(key) DO UPDATE SET value = excluded.value" ).use { statement -> statement.setString(1, SCHEMA_VERSION.toString()) statement.executeUpdate() } } private fun ResultSet.toRecord(): ChatMessageRecord { val kindOrdinal = getInt("kind") val kind = MessageSourceKind.values().getOrNull(kindOrdinal) ?: error("未知的消息类型序号 $kindOrdinal") return ChatMessageRecord( id = getLong("id"), botId = getLong("bot_id"), fromId = getLong("from_id"), targetId = getLong("target_id"), ids = getString("ids"), internalIds = getString("internal_ids"), time = getInt("time"), kind = kind, code = getString("code"), recalled = getInt("recalled"), ) } }