mirror of
https://github.com/tiennm99/miti99bot.git
synced 2026-07-25 16:21:55 +00:00
466 lines
12 KiB
Go
466 lines
12 KiB
Go
package stats
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"strconv"
|
|
"sync"
|
|
|
|
"go.mongodb.org/mongo-driver/v2/bson"
|
|
"go.mongodb.org/mongo-driver/v2/mongo"
|
|
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
|
|
|
"github.com/tiennm99/miti99bot/internal/storage"
|
|
)
|
|
|
|
func bsonField(key string, value any) bson.E {
|
|
return bson.E{Key: key, Value: value}
|
|
}
|
|
|
|
// usageEntry is the only stats payload shape. One document represents either
|
|
// a command-level anonymous bucket or a concrete (command, Telegram user) pair.
|
|
type usageEntry struct {
|
|
Cmd string `json:"cmd" bson:"cmd"`
|
|
UserID int64 `json:"uid,omitempty" bson:"uid,omitempty"`
|
|
Username string `json:"user,omitempty" bson:"user,omitempty"`
|
|
N int64 `json:"n" bson:"n"`
|
|
Deleted bool `json:"deleted,omitempty" bson:"deleted,omitempty"`
|
|
}
|
|
|
|
type usageUser struct {
|
|
ID int64
|
|
Username string
|
|
}
|
|
|
|
type usageStore interface {
|
|
Increment(ctx context.Context, cmd string, user usageUser, hasUser bool) error
|
|
TopCommands(ctx context.Context, limit int) ([]row, error)
|
|
TopUsers(ctx context.Context, limit int) ([]row, error)
|
|
CommandsByUser(ctx context.Context, username string, limit int) ([]row, bool, error)
|
|
UsersByCommand(ctx context.Context, cmd string, limit int) ([]row, error)
|
|
}
|
|
|
|
func newUsageStore(coll storage.Collection) usageStore {
|
|
if mongoColl, ok := storage.MongoCollection(coll); ok {
|
|
return &mongoUsageStore{coll: mongoColl}
|
|
}
|
|
return &docUsageStore{docs: storage.Typed[usageEntry](coll)}
|
|
}
|
|
|
|
func usageKey(cmd string, userID int64) string {
|
|
if userID == 0 {
|
|
return cmd
|
|
}
|
|
return cmd + ":" + strconv.FormatInt(userID, 10)
|
|
}
|
|
|
|
type docUsageStore struct {
|
|
mu sync.Mutex
|
|
docs storage.DocStore[usageEntry]
|
|
}
|
|
|
|
func (s *docUsageStore) Increment(ctx context.Context, cmd string, user usageUser, hasUser bool) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
userID := int64(0)
|
|
if hasUser {
|
|
userID = user.ID
|
|
}
|
|
key := usageKey(cmd, userID)
|
|
entry, _, err := s.docs.Get(ctx, key)
|
|
switch {
|
|
case errors.Is(err, storage.ErrNotFound):
|
|
entry = usageEntry{}
|
|
case err != nil:
|
|
return fmt.Errorf("stats get %s: %w", key, err)
|
|
}
|
|
|
|
entry.Cmd = cmd
|
|
entry.N++
|
|
entry.Deleted = false
|
|
if hasUser {
|
|
entry.UserID = user.ID
|
|
entry.Username = user.Username
|
|
} else {
|
|
entry.UserID = 0
|
|
entry.Username = ""
|
|
}
|
|
if err := s.docs.Put(ctx, key, entry); err != nil {
|
|
return fmt.Errorf("stats put %s: %w", key, err)
|
|
}
|
|
if hasUser {
|
|
return s.refreshUsernameLocked(ctx, user)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *docUsageStore) refreshUsernameLocked(ctx context.Context, user usageUser) error {
|
|
entries, err := s.loadEntriesLocked(ctx)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, e := range entries {
|
|
if e.UserID != user.ID || e.Username == user.Username {
|
|
continue
|
|
}
|
|
e.Username = user.Username
|
|
if err := s.docs.Put(ctx, usageKey(e.Cmd, e.UserID), e); err != nil {
|
|
return fmt.Errorf("stats refresh username %d: %w", user.ID, err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *docUsageStore) TopCommands(ctx context.Context, limit int) ([]row, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
entries, err := s.loadEntriesLocked(ctx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
totals := make(map[string]int64)
|
|
for _, e := range entries {
|
|
if e.Deleted {
|
|
continue
|
|
}
|
|
totals[e.Cmd] += e.N
|
|
}
|
|
rows := make([]row, 0, len(totals))
|
|
for cmd, n := range totals {
|
|
rows = append(rows, row{display: "/" + cmd, n: n})
|
|
}
|
|
sortRows(rows)
|
|
return limitRows(rows, limit), nil
|
|
}
|
|
|
|
func (s *docUsageStore) TopUsers(ctx context.Context, limit int) ([]row, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
entries, err := s.loadEntriesLocked(ctx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
type total struct {
|
|
username string
|
|
n int64
|
|
}
|
|
totals := make(map[int64]total)
|
|
for _, e := range entries {
|
|
if e.Deleted || e.UserID == 0 || e.Username == "" {
|
|
continue
|
|
}
|
|
t := totals[e.UserID]
|
|
t.username = e.Username
|
|
t.n += e.N
|
|
totals[e.UserID] = t
|
|
}
|
|
rows := make([]row, 0, len(totals))
|
|
for _, t := range totals {
|
|
rows = append(rows, row{display: "@" + t.username, n: t.n})
|
|
}
|
|
sortRows(rows)
|
|
return limitRows(rows, limit), nil
|
|
}
|
|
|
|
func (s *docUsageStore) CommandsByUser(ctx context.Context, username string, limit int) ([]row, bool, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
entries, err := s.loadEntriesLocked(ctx)
|
|
if err != nil {
|
|
return nil, false, err
|
|
}
|
|
var (
|
|
userID int64
|
|
found bool
|
|
)
|
|
for _, e := range entries {
|
|
if !e.Deleted && e.UserID != 0 && e.Username == username {
|
|
userID = e.UserID
|
|
found = true
|
|
break
|
|
}
|
|
}
|
|
if !found {
|
|
return nil, false, nil
|
|
}
|
|
|
|
rows := make([]row, 0)
|
|
for _, e := range entries {
|
|
if !e.Deleted && e.UserID == userID {
|
|
rows = append(rows, row{display: "/" + e.Cmd, n: e.N})
|
|
}
|
|
}
|
|
sortRows(rows)
|
|
return limitRows(rows, limit), true, nil
|
|
}
|
|
|
|
func (s *docUsageStore) UsersByCommand(ctx context.Context, cmd string, limit int) ([]row, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
entries, err := s.loadEntriesLocked(ctx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
rows := make([]row, 0)
|
|
for _, e := range entries {
|
|
if !e.Deleted && e.Cmd == cmd && e.UserID != 0 && e.Username != "" {
|
|
rows = append(rows, row{display: "@" + e.Username, n: e.N})
|
|
}
|
|
}
|
|
sortRows(rows)
|
|
return limitRows(rows, limit), nil
|
|
}
|
|
|
|
func (s *docUsageStore) loadEntriesLocked(ctx context.Context) ([]usageEntry, error) {
|
|
keys, err := s.docs.List(ctx, "")
|
|
if err != nil {
|
|
return nil, fmt.Errorf("stats list: %w", err)
|
|
}
|
|
entries := make([]usageEntry, 0, len(keys))
|
|
for _, key := range keys {
|
|
entry, _, err := s.docs.Get(ctx, key)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("stats get %s: %w", key, err)
|
|
}
|
|
if entry.Cmd == "" {
|
|
continue
|
|
}
|
|
entries = append(entries, entry)
|
|
}
|
|
return entries, nil
|
|
}
|
|
|
|
type mongoUsageStore struct {
|
|
coll *mongo.Collection
|
|
}
|
|
|
|
func (s *mongoUsageStore) Increment(ctx context.Context, cmd string, user usageUser, hasUser bool) error {
|
|
userID := int64(0)
|
|
if hasUser {
|
|
userID = user.ID
|
|
}
|
|
|
|
set := bson.M{"cmd": cmd}
|
|
unset := bson.M{"deleted": ""}
|
|
update := bson.M{
|
|
"$set": set,
|
|
"$inc": bson.M{"n": int64(1), "version": int64(1)},
|
|
"$unset": unset,
|
|
"$currentDate": bson.M{"updatedAt": true},
|
|
}
|
|
if hasUser {
|
|
set["uid"] = user.ID
|
|
set["user"] = user.Username
|
|
} else {
|
|
unset["uid"] = ""
|
|
unset["user"] = ""
|
|
}
|
|
|
|
key := usageKey(cmd, userID)
|
|
if _, err := s.coll.UpdateOne(ctx,
|
|
bson.M{"_id": key},
|
|
update,
|
|
options.UpdateOne().SetUpsert(true),
|
|
); err != nil {
|
|
return fmt.Errorf("mongo stats increment %s: %w", key, err)
|
|
}
|
|
if !hasUser {
|
|
return nil
|
|
}
|
|
if _, err := s.coll.UpdateMany(ctx,
|
|
bson.M{"uid": user.ID, "user": bson.M{"$ne": user.Username}},
|
|
bson.M{
|
|
"$set": bson.M{"user": user.Username},
|
|
"$inc": bson.M{"version": int64(1)},
|
|
"$currentDate": bson.M{"updatedAt": true},
|
|
},
|
|
); err != nil {
|
|
return fmt.Errorf("mongo stats refresh username %d: %w", user.ID, err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *mongoUsageStore) TopCommands(ctx context.Context, limit int) ([]row, error) {
|
|
pipeline := mongo.Pipeline{
|
|
bson.D{bsonField("$match", bson.D{
|
|
bsonField("cmd", bson.D{bsonField("$type", "string")}),
|
|
bsonField("deleted", bson.D{bsonField("$ne", true)}),
|
|
})},
|
|
bson.D{bsonField("$group", bson.D{bsonField("_id", "$cmd"), bsonField("n", bson.D{bsonField("$sum", "$n")})})},
|
|
bson.D{bsonField("$sort", bson.D{bsonField("n", -1), bsonField("_id", 1)})},
|
|
}
|
|
pipeline = withLimit(pipeline, limit)
|
|
|
|
cur, err := s.coll.Aggregate(ctx, pipeline)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("mongo stats top commands: %w", err)
|
|
}
|
|
defer func() { _ = cur.Close(ctx) }()
|
|
|
|
var rows []row
|
|
for cur.Next(ctx) {
|
|
var doc struct {
|
|
Cmd string `bson:"_id"`
|
|
N int64 `bson:"n"`
|
|
}
|
|
if err := cur.Decode(&doc); err != nil {
|
|
return nil, fmt.Errorf("mongo stats top commands decode: %w", err)
|
|
}
|
|
rows = append(rows, row{display: "/" + doc.Cmd, n: doc.N})
|
|
}
|
|
if err := cur.Err(); err != nil {
|
|
return nil, fmt.Errorf("mongo stats top commands cursor: %w", err)
|
|
}
|
|
return rows, nil
|
|
}
|
|
|
|
func (s *mongoUsageStore) TopUsers(ctx context.Context, limit int) ([]row, error) {
|
|
pipeline := mongo.Pipeline{
|
|
bson.D{bsonField("$match", bson.D{
|
|
bsonField("uid", bson.D{bsonField("$gt", int64(0))}),
|
|
bsonField("user", bson.D{bsonField("$type", "string"), bsonField("$ne", "")}),
|
|
bsonField("deleted", bson.D{bsonField("$ne", true)}),
|
|
})},
|
|
bson.D{bsonField("$sort", bson.D{bsonField("updatedAt", -1)})},
|
|
bson.D{bsonField("$group", bson.D{
|
|
bsonField("_id", "$uid"),
|
|
bsonField("user", bson.D{bsonField("$first", "$user")}),
|
|
bsonField("n", bson.D{bsonField("$sum", "$n")}),
|
|
})},
|
|
bson.D{bsonField("$sort", bson.D{bsonField("n", -1), bsonField("user", 1)})},
|
|
}
|
|
pipeline = withLimit(pipeline, limit)
|
|
|
|
cur, err := s.coll.Aggregate(ctx, pipeline)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("mongo stats top users: %w", err)
|
|
}
|
|
defer func() { _ = cur.Close(ctx) }()
|
|
|
|
var rows []row
|
|
for cur.Next(ctx) {
|
|
var doc struct {
|
|
User string `bson:"user"`
|
|
N int64 `bson:"n"`
|
|
}
|
|
if err := cur.Decode(&doc); err != nil {
|
|
return nil, fmt.Errorf("mongo stats top users decode: %w", err)
|
|
}
|
|
rows = append(rows, row{display: "@" + doc.User, n: doc.N})
|
|
}
|
|
if err := cur.Err(); err != nil {
|
|
return nil, fmt.Errorf("mongo stats top users cursor: %w", err)
|
|
}
|
|
return rows, nil
|
|
}
|
|
|
|
func (s *mongoUsageStore) CommandsByUser(ctx context.Context, username string, limit int) ([]row, bool, error) {
|
|
userID, found, err := s.userIDByUsername(ctx, username)
|
|
if err != nil || !found {
|
|
return nil, found, err
|
|
}
|
|
|
|
opts := options.Find().
|
|
SetProjection(bson.M{"cmd": 1, "n": 1}).
|
|
SetSort(bson.D{bsonField("n", -1), bsonField("cmd", 1)})
|
|
if limit > 0 {
|
|
opts.SetLimit(int64(limit))
|
|
}
|
|
cur, err := s.coll.Find(ctx, bson.M{
|
|
"uid": userID,
|
|
"cmd": bson.M{"$type": "string"},
|
|
"deleted": bson.M{"$ne": true},
|
|
}, opts)
|
|
if err != nil {
|
|
return nil, true, fmt.Errorf("mongo stats commands by user %s: %w", username, err)
|
|
}
|
|
defer func() { _ = cur.Close(ctx) }()
|
|
|
|
var rows []row
|
|
for cur.Next(ctx) {
|
|
var doc usageEntry
|
|
if err := cur.Decode(&doc); err != nil {
|
|
return nil, true, fmt.Errorf("mongo stats commands by user decode: %w", err)
|
|
}
|
|
rows = append(rows, row{display: "/" + doc.Cmd, n: doc.N})
|
|
}
|
|
if err := cur.Err(); err != nil {
|
|
return nil, true, fmt.Errorf("mongo stats commands by user cursor: %w", err)
|
|
}
|
|
return rows, true, nil
|
|
}
|
|
|
|
func (s *mongoUsageStore) UsersByCommand(ctx context.Context, cmd string, limit int) ([]row, error) {
|
|
opts := options.Find().
|
|
SetProjection(bson.M{"user": 1, "n": 1}).
|
|
SetSort(bson.D{bsonField("n", -1), bsonField("user", 1)})
|
|
if limit > 0 {
|
|
opts.SetLimit(int64(limit))
|
|
}
|
|
cur, err := s.coll.Find(ctx, bson.M{
|
|
"cmd": cmd,
|
|
"uid": bson.M{"$gt": int64(0)},
|
|
"user": bson.M{"$type": "string", "$ne": ""},
|
|
"deleted": bson.M{"$ne": true},
|
|
}, opts)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("mongo stats users by command %s: %w", cmd, err)
|
|
}
|
|
defer func() { _ = cur.Close(ctx) }()
|
|
|
|
var rows []row
|
|
for cur.Next(ctx) {
|
|
var doc usageEntry
|
|
if err := cur.Decode(&doc); err != nil {
|
|
return nil, fmt.Errorf("mongo stats users by command decode: %w", err)
|
|
}
|
|
rows = append(rows, row{display: "@" + doc.Username, n: doc.N})
|
|
}
|
|
if err := cur.Err(); err != nil {
|
|
return nil, fmt.Errorf("mongo stats users by command cursor: %w", err)
|
|
}
|
|
return rows, nil
|
|
}
|
|
|
|
func (s *mongoUsageStore) userIDByUsername(ctx context.Context, username string) (int64, bool, error) {
|
|
var doc struct {
|
|
UserID int64 `bson:"uid"`
|
|
}
|
|
err := s.coll.FindOne(ctx,
|
|
bson.M{
|
|
"uid": bson.M{"$gt": int64(0)},
|
|
"user": username,
|
|
"deleted": bson.M{"$ne": true},
|
|
},
|
|
options.FindOne().
|
|
SetProjection(bson.M{"uid": 1}),
|
|
).Decode(&doc)
|
|
if err != nil {
|
|
if errors.Is(err, mongo.ErrNoDocuments) {
|
|
return 0, false, nil
|
|
}
|
|
return 0, false, fmt.Errorf("mongo stats find user %s: %w", username, err)
|
|
}
|
|
return doc.UserID, true, nil
|
|
}
|
|
|
|
func withLimit(pipeline mongo.Pipeline, limit int) mongo.Pipeline {
|
|
if limit <= 0 {
|
|
return pipeline
|
|
}
|
|
return append(pipeline, bson.D{bsonField("$limit", int64(limit))})
|
|
}
|
|
|
|
func limitRows(rows []row, limit int) []row {
|
|
if limit <= 0 || len(rows) <= limit {
|
|
return rows
|
|
}
|
|
return rows[:limit]
|
|
}
|