mirror of
https://github.com/jie65535/JChatGPT.git
synced 2026-09-15 02:56:10 +08:00
369 lines
13 KiB
Kotlin
369 lines
13 KiB
Kotlin
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"),
|
|
)
|
|
}
|
|
}
|