You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
110 lines
3.2 KiB
110 lines
3.2 KiB
"""SQLite 消息持久化"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import sqlite3
|
|
import time
|
|
from pathlib import Path
|
|
|
|
from ..models.message import Message
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_CREATE_TABLE = """
|
|
CREATE TABLE IF NOT EXISTS messages (
|
|
msg_id TEXT PRIMARY KEY,
|
|
conversation TEXT NOT NULL,
|
|
sender TEXT NOT NULL,
|
|
content TEXT NOT NULL,
|
|
msg_type TEXT NOT NULL,
|
|
timestamp REAL NOT NULL,
|
|
is_self INTEGER NOT NULL DEFAULT 0,
|
|
created_at REAL NOT NULL
|
|
)
|
|
"""
|
|
|
|
_CREATE_INDEXES = [
|
|
"CREATE INDEX IF NOT EXISTS idx_conv ON messages(conversation)",
|
|
"CREATE INDEX IF NOT EXISTS idx_ts ON messages(timestamp)",
|
|
]
|
|
|
|
|
|
class SQLiteStorage:
|
|
"""基于 SQLite 的消息存储"""
|
|
|
|
def __init__(self, db_path: str = "wechat_messages.db") -> None:
|
|
self._db_path = db_path
|
|
self._conn: sqlite3.Connection | None = None
|
|
|
|
def connect(self) -> None:
|
|
self._conn = sqlite3.connect(self._db_path, check_same_thread=False)
|
|
self._conn.execute(_CREATE_TABLE)
|
|
for idx_sql in _CREATE_INDEXES:
|
|
self._conn.execute(idx_sql)
|
|
self._conn.commit()
|
|
logger.info("SQLite 存储已初始化: %s", self._db_path)
|
|
|
|
def close(self) -> None:
|
|
if self._conn:
|
|
self._conn.close()
|
|
self._conn = None
|
|
|
|
def save_message(self, msg: Message) -> bool:
|
|
"""保存单条消息（幂等）"""
|
|
if self._conn is None:
|
|
logger.error("数据库未连接")
|
|
return False
|
|
try:
|
|
self._conn.execute(
|
|
"INSERT OR IGNORE INTO messages "
|
|
"(msg_id, conversation, sender, content, msg_type, timestamp, is_self, created_at) "
|
|
"VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
|
|
(
|
|
msg.msg_id,
|
|
msg.conversation,
|
|
msg.sender,
|
|
msg.content,
|
|
msg.msg_type.value,
|
|
msg.timestamp,
|
|
int(msg.is_self),
|
|
time.time(),
|
|
),
|
|
)
|
|
self._conn.commit()
|
|
return True
|
|
except Exception:
|
|
logger.exception("保存消息失败")
|
|
return False
|
|
|
|
def save_messages(self, messages: list[Message]) -> int:
|
|
"""批量保存，返回成功数量"""
|
|
count = 0
|
|
for msg in messages:
|
|
if self.save_message(msg):
|
|
count += 1
|
|
return count
|
|
|
|
def query_messages(
|
|
self,
|
|
conversation: str = "",
|
|
limit: int = 100,
|
|
since: float = 0,
|
|
) -> list[dict]:
|
|
"""查询消息"""
|
|
if self._conn is None:
|
|
return []
|
|
sql = "SELECT * FROM messages WHERE 1=1"
|
|
params: list = []
|
|
if conversation:
|
|
sql += " AND conversation = ?"
|
|
params.append(conversation)
|
|
if since:
|
|
sql += " AND timestamp > ?"
|
|
params.append(since)
|
|
sql += " ORDER BY timestamp DESC LIMIT ?"
|
|
params.append(limit)
|
|
|
|
cursor = self._conn.execute(sql, params)
|
|
cols = [d[0] for d in cursor.description]
|
|
return [dict(zip(cols, row)) for row in cursor.fetchall()]
|
|
|