mirror of
https://github.com/jie65535/JChatGPT.git
synced 2026-09-15 02:56:10 +08:00
history: replace external recorder with SQLite
This commit is contained in:
@@ -0,0 +1,368 @@
|
||||
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<ChatMessageRecord> {
|
||||
check(initialized) { "聊天记录数据库尚未初始化" }
|
||||
require(start <= end) { "start must not be after end" }
|
||||
|
||||
val conditions = mutableListOf<String>()
|
||||
val parameters = mutableListOf<Any>()
|
||||
|
||||
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"),
|
||||
)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user