167 lines
8.2 KiB
Python
167 lines
8.2 KiB
Python
import json
|
|
from datetime import datetime, timezone
|
|
|
|
import aiosqlite
|
|
|
|
|
|
SCHEMA = """
|
|
PRAGMA journal_mode=WAL;
|
|
PRAGMA synchronous=FULL;
|
|
CREATE TABLE IF NOT EXISTS messages (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
group_id TEXT NOT NULL, message_id TEXT NOT NULL, sender_id TEXT NOT NULL,
|
|
sender_name TEXT NOT NULL, sent_at TEXT NOT NULL, text TEXT NOT NULL,
|
|
raw_json TEXT NOT NULL, media_json TEXT NOT NULL,
|
|
analysis_status TEXT NOT NULL DEFAULT 'pending', analysis_attempts INTEGER NOT NULL DEFAULT 0,
|
|
next_attempt_at TEXT, last_error TEXT, created_at TEXT NOT NULL,
|
|
UNIQUE(group_id, message_id)
|
|
);
|
|
CREATE TABLE IF NOT EXISTS analyses (
|
|
message_pk INTEGER PRIMARY KEY REFERENCES messages(id), importance INTEGER NOT NULL,
|
|
category TEXT NOT NULL, need_action INTEGER NOT NULL, deadline TEXT,
|
|
summary TEXT NOT NULL, reason TEXT NOT NULL, result_json TEXT NOT NULL,
|
|
notified INTEGER NOT NULL DEFAULT 0, notify_attempts INTEGER NOT NULL DEFAULT 0,
|
|
notify_next_at TEXT, notify_error TEXT
|
|
);
|
|
CREATE TABLE IF NOT EXISTS daily_summaries (
|
|
local_date TEXT PRIMARY KEY, summary TEXT NOT NULL, result_json TEXT NOT NULL,
|
|
notified INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS daily_group_summaries (
|
|
local_date TEXT NOT NULL, group_id TEXT NOT NULL,
|
|
summary TEXT NOT NULL, result_json TEXT NOT NULL,
|
|
notified INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL,
|
|
PRIMARY KEY(local_date, group_id)
|
|
);
|
|
"""
|
|
|
|
|
|
class Storage:
|
|
def __init__(self, path: str):
|
|
self.path = path
|
|
|
|
async def connect(self):
|
|
db = await aiosqlite.connect(self.path)
|
|
db.row_factory = aiosqlite.Row
|
|
await db.execute("PRAGMA busy_timeout=5000")
|
|
await db.execute("PRAGMA foreign_keys=ON")
|
|
return db
|
|
|
|
async def init(self):
|
|
db = await self.connect()
|
|
try:
|
|
await db.executescript(SCHEMA)
|
|
# A process can stop after claiming a row. Requeue it on startup.
|
|
await db.execute("UPDATE messages SET analysis_status='retry' WHERE analysis_status='processing'")
|
|
await db.commit()
|
|
finally:
|
|
await db.close()
|
|
|
|
async def save_message(self, item: dict, text: str,
|
|
media: list[dict]) -> tuple[int, bool]:
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
db = await self.connect()
|
|
try:
|
|
cur = await db.execute(
|
|
"INSERT OR IGNORE INTO messages(group_id,message_id,sender_id,sender_name,sent_at,text,raw_json,media_json,created_at) VALUES(?,?,?,?,?,?,?,?,?)",
|
|
(str(item["group_id"]), str(item["message_id"]), str(item.get("user_id", "")), item.get("sender", {}).get("card") or item.get("sender", {}).get("nickname") or "未知", datetime.fromtimestamp(item.get("time", 0), timezone.utc).isoformat(), text, json.dumps(item, ensure_ascii=False), json.dumps(media, ensure_ascii=False), now),
|
|
)
|
|
inserted = cur.rowcount == 1
|
|
if inserted:
|
|
pk = cur.lastrowid
|
|
else:
|
|
row = await (await db.execute("SELECT id FROM messages WHERE group_id=? AND message_id=?", (str(item["group_id"]), str(item["message_id"])))).fetchone()
|
|
pk = row["id"]
|
|
await db.commit()
|
|
return int(pk), inserted
|
|
finally:
|
|
await db.close()
|
|
|
|
async def claim_pending(self):
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
db = await self.connect()
|
|
try:
|
|
await db.execute("BEGIN IMMEDIATE")
|
|
row = await (await db.execute("SELECT * FROM messages WHERE analysis_status IN ('pending','retry') AND (next_attempt_at IS NULL OR next_attempt_at<=?) ORDER BY id LIMIT 1", (now,))).fetchone()
|
|
if row:
|
|
await db.execute("UPDATE messages SET analysis_status='processing',analysis_attempts=analysis_attempts+1 WHERE id=?", (row["id"],))
|
|
await db.commit()
|
|
return dict(row) if row else None
|
|
finally:
|
|
await db.close()
|
|
|
|
async def context(self, row: dict, count: int):
|
|
db = await self.connect()
|
|
try:
|
|
rows = await (await db.execute("SELECT m.sender_name,m.sent_at,m.text FROM messages m LEFT JOIN analyses a ON a.message_pk=m.id WHERE m.group_id=? AND m.id<? AND (a.category IS NULL OR a.category<>'回执') ORDER BY m.id DESC LIMIT ?", (row["group_id"], row["id"], count))).fetchall()
|
|
return [dict(x) for x in reversed(rows)]
|
|
finally:
|
|
await db.close()
|
|
|
|
async def save_analysis(self, pk: int, result: dict):
|
|
db = await self.connect()
|
|
try:
|
|
await db.execute("INSERT OR REPLACE INTO analyses(message_pk,importance,category,need_action,deadline,summary,reason,result_json) VALUES(?,?,?,?,?,?,?,?)", (pk, result["importance"], result["category"], int(result["need_action"]), result.get("deadline"), result["summary"], result["reason"], json.dumps(result, ensure_ascii=False)))
|
|
await db.execute("UPDATE messages SET analysis_status='done',last_error=NULL,next_attempt_at=NULL WHERE id=?", (pk,))
|
|
await db.commit()
|
|
finally:
|
|
await db.close()
|
|
|
|
async def retry_analysis(self, pk: int, error: str, next_at: str):
|
|
db = await self.connect()
|
|
try:
|
|
await db.execute("UPDATE messages SET analysis_status='retry',last_error=?,next_attempt_at=? WHERE id=?", (error[:500], next_at, pk))
|
|
await db.commit()
|
|
finally:
|
|
await db.close()
|
|
|
|
async def pending_notifications(self, min_importance: int = 1):
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
db = await self.connect()
|
|
try:
|
|
rows = await (await db.execute("SELECT m.*,a.* FROM messages m JOIN analyses a ON a.message_pk=m.id WHERE a.notified=0 AND a.importance>=? AND (a.notify_next_at IS NULL OR a.notify_next_at<=?) ORDER BY m.id LIMIT 10", (min_importance, now))).fetchall()
|
|
return [dict(x) for x in rows]
|
|
finally:
|
|
await db.close()
|
|
|
|
async def notification_result(self, pk: int, ok: bool, error: str = "", next_at: str | None = None):
|
|
db = await self.connect()
|
|
try:
|
|
if ok:
|
|
await db.execute("UPDATE analyses SET notified=1,notify_error=NULL,notify_next_at=NULL WHERE message_pk=?", (pk,))
|
|
else:
|
|
await db.execute("UPDATE analyses SET notify_attempts=notify_attempts+1,notify_error=?,notify_next_at=? WHERE message_pk=?", (error[:500], next_at, pk))
|
|
await db.commit()
|
|
finally:
|
|
await db.close()
|
|
|
|
async def messages_for_day(self, start_utc: str, end_utc: str, group_id: str):
|
|
db = await self.connect()
|
|
try:
|
|
rows = await (await db.execute("SELECT m.*,a.importance,a.category,a.need_action,a.deadline,a.summary,a.reason FROM messages m LEFT JOIN analyses a ON a.message_pk=m.id WHERE m.sent_at>=? AND m.sent_at<? AND m.group_id=? AND (a.category IS NULL OR a.category<>'回执') ORDER BY m.sent_at", (start_utc, end_utc, group_id))).fetchall()
|
|
return [dict(x) for x in rows]
|
|
finally:
|
|
await db.close()
|
|
|
|
async def has_daily(self, day: str, group_id: str) -> bool:
|
|
db = await self.connect()
|
|
try:
|
|
return bool(await (await db.execute("SELECT 1 FROM daily_group_summaries WHERE local_date=? AND group_id=? AND notified=1", (day, group_id))).fetchone())
|
|
finally:
|
|
await db.close()
|
|
|
|
async def save_daily(self, day: str, group_id: str, summary: str, result: dict, notified: bool):
|
|
db = await self.connect()
|
|
try:
|
|
await db.execute("INSERT OR REPLACE INTO daily_group_summaries(local_date,group_id,summary,result_json,notified,created_at) VALUES(?,?,?,?,?,?)", (day, group_id, summary, json.dumps(result, ensure_ascii=False), int(notified), datetime.now(timezone.utc).isoformat()))
|
|
await db.commit()
|
|
finally:
|
|
await db.close()
|
|
|
|
async def stats(self):
|
|
db = await self.connect()
|
|
try:
|
|
return dict(await (await db.execute("SELECT COUNT(*) total,COALESCE(SUM(analysis_status='done'),0) done,COALESCE(SUM(analysis_status IN ('pending','retry','processing')),0) pending,COALESCE(SUM(analysis_status='failed'),0) failed FROM messages")).fetchone())
|
|
finally:
|
|
await db.close()
|