feat(status): mongo support (#160)

Co-authored-by: Shigma <1700011071@pku.edu.cn>
koishijs · Mar 18, 2021 · 6956c81 · 6956c81
1 parent eb6104f
commit 6956c81
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 5 deletions.
diff --git a/packages/plugin-status/server/mongo.ts b/packages/plugin-status/server/mongo.ts
@@ -1,5 +1,14 @@
-import { Database } from 'koishi-core'
-import {} from 'koishi-plugin-mongo'
+import { Database, Logger, Time } from 'koishi-core'
+import type MongoDatabase from 'koishi-plugin-mongo'
+import { RECENT_LENGTH, StatRecord, Synchronizer } from './stats'
+
+const logger = new Logger('status')
+
+declare module 'koishi-plugin-mongo' {
+  interface Tables {
+    'plugin-status': any
+  }
+}
 
 Database.extend('koishi-plugin-mongo', {
   async getProfile() {
@@ -13,4 +22,67 @@ Database.extend('koishi-plugin-mongo', {
     ])
     return { allGroups, activeGroups, allUsers, activeUsers, storageSize }
   },
+
+  async setChannels(data) {
+    await Promise.all(data.map(ch => this.setChannel(ch.type, ch.id, ch)))
+  },
+
+  Synchronizer: class {
+    groups: StatRecord
+    daily: any
+    hourly: Record<Synchronizer.HourlyField, number>
+    longterm: Record<Synchronizer.LongtermField, number>
+
+    constructor(private db: MongoDatabase) {
+      this.reset()
+    }
+
+    reset() {
+      this.hourly = Object.fromEntries(Synchronizer.hourlyFields.map(i => [i, 0])) as any
+      this.daily = {}
+      this.longterm = Object.fromEntries(Synchronizer.longtermFields.map(i => [i, 0])) as any
+      this.groups = {}
+    }
+
+    addDaily(field: Synchronizer.DailyField, key: string | number) {
+      if (!this.daily[field]) this.daily[field] = {}
+      const stat: Record<string, number> = this.daily[field]
+      stat[key] = (stat[key] || 0) + 1
+    }
+
+    async upload(date: Date): Promise<void> {
+      logger.debug(this.hourly, this.daily, this.longterm, this.groups)
+      const coll = this.db.collection('plugin-status')
+      const _date = new Date(date)
+      _date.setMinutes(0)
+      _date.setSeconds(0)
+      _date.setMilliseconds(0)
+      await coll.updateOne({ type: 'hourly', time: _date }, { $inc: this.hourly }, { upsert: true })
+      _date.setHours(0)
+      const $inc = {}
+      for (const key in this.daily) {
+        for (const subkey in this.daily[key]) {
+          $inc[`${key}.${subkey}`] = this.daily[key][subkey]
+        }
+      }
+      if (Object.keys($inc).length) await coll.updateOne({ type: 'daily', time: _date }, { $inc }, { upsert: true })
+      await coll.updateOne({ type: 'longterm', time: _date }, { $inc: this.longterm }, { upsert: true })
+      for (const id in this.groups) {
+        await this.db.channel.updateOne({ id }, { $inc: { ['activity.' + Time.getDateNumber(date)]: this.groups[id] } } as any)
+      }
+      this.reset()
+      logger.debug('stats updated')
+    }
+
+    async download(date: Date) {
+      const time = { $lt: new Date(date) }
+      const coll = this.db.collection('plugin-status')
+      const hourly = await coll.find({ type: 'hourly', time }).sort({ time: -1 }).limit(24 * RECENT_LENGTH).toArray()
+      const daily = await coll.find({ type: 'daily', time }).sort({ time: -1 }).limit(RECENT_LENGTH).toArray()
+      const longterm = await coll.find({ type: 'longterm', time }).sort({ time: -1 }).toArray()
+      const groups = await this.db.channel.find({}).project({ type: 1, pid: 1, name: 1, assignee: 1 })
+        .map(data => ({ ...data, id: `${data.type}:${data.pid}` })).toArray()
+      return { daily, hourly, longterm, groups }
+    }
+  },
 })
diff --git a/packages/plugin-status/server/stats.ts b/packages/plugin-status/server/stats.ts
@@ -120,7 +120,7 @@ async function download(ctx: Context, date: Date) {
     const groups = await bot.getGroupList()
     for (const { groupId, groupName: name } of groups) {
       const id = `${bot.platform}:${groupId}`
-      if (!messageMap[id] || groupSet.has(id)) continue
+      if (!messageMap[id] || !groupMap[id] || groupSet.has(id)) continue
       groupSet.add(id)
       const { name: oldName, assignee } = groupMap[id]
       if (name !== oldName) updateList.push({ id, name })
@@ -145,7 +145,7 @@ async function download(ctx: Context, date: Date) {
         name: name || key,
         value: messageMap[key],
         last: daily[0].group[key],
-        assignee: ctx.bots[`${platform}:${assignee}`].selfId,
+        assignee: ctx.bots[`${platform}:${assignee}`]?.selfId || '',
       })
     }
   }
@@ -159,7 +159,7 @@ async function download(ctx: Context, date: Date) {
   // dialogue
   if (ctx.database.getDialoguesById) {
     const dialogueMap = average(daily.map(data => data.dialogue))
-    const dialogues = await ctx.database.getDialoguesById(Object.keys(dialogueMap) as any, ['id', 'original'])
+    const dialogues = await ctx.database.getDialoguesById(Object.keys(dialogueMap).map(i => +i), ['id', 'original'])
     const questionMap: Record<string, QuestionData> = {}
     for (const dialogue of dialogues) {
       const { id, original: name } = dialogue