// Package store persists settings, users, group memberships and render // statistics in the shared MongoDB cluster. package store import ( "context" "fmt" "log" "strings" "time" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" "go.mongodb.org/mongo-driver/mongo/readpref" ) // Settings are the per-chat overrides. Zero values mean "use the default". type Settings struct { Preamble string `json:"preamble,omitempty" bson:"preamble,omitempty"` DPI int `json:"dpi,omitempty" bson:"dpi,omitempty"` } // User is the caller information recorded on every interaction. type User struct { ID int64 Username string FirstName string LastName string } // Chat is the group information recorded on membership changes. type Chat struct { ID int64 Type string Title string Username string } // RenderStat is one render event, kept for statistics. type RenderStat struct { Source string // chat | inline Format string // rich | png | pdf | inline ChatID int64 UserID int64 DPI int ExpressionLen int OK bool Error string DurationMs int64 } // Store provides all persistence for the bot. type Store struct { client *mongo.Client settings *mongo.Collection users *mongo.Collection groups *mongo.Collection stats *mongo.Collection files *mongo.Collection } // New connects a Store to the given database. func New(client *mongo.Client, database string) *Store { db := client.Database(database) return &Store{ client: client, settings: db.Collection("settings"), users: db.Collection("users"), groups: db.Collection("groups"), stats: db.Collection("stats"), files: db.Collection("files"), } } // Ping reports whether the cluster is reachable. func (s *Store) Ping(ctx context.Context) error { return s.client.Ping(ctx, readpref.PrimaryPreferred()) } // EnsureIndexes creates the indexes the bot relies on. func (s *Store) EnsureIndexes(ctx context.Context) error { _, err := s.stats.Indexes().CreateMany(ctx, []mongo.IndexModel{ {Keys: bson.D{{Key: "ts", Value: -1}}, Options: options.Index().SetName("stats_ts")}, {Keys: bson.D{{Key: "user_id", Value: 1}, {Key: "ts", Value: -1}}, Options: options.Index().SetName("stats_user_ts")}, }) return err } // Get returns the settings for chatID, or the zero value. func (s *Store) Get(chatID int64) Settings { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() var settings Settings err := s.settings.FindOne(ctx, bson.M{"_id": chatID}).Decode(&settings) if err != nil { if err != mongo.ErrNoDocuments { log.Printf("store: read settings for %d: %v", chatID, err) } return Settings{} } return settings } // Set stores settings for chatID. func (s *Store) Set(chatID int64, settings Settings) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() _, err := s.settings.UpdateOne(ctx, bson.M{"_id": chatID}, bson.M{"$set": settings, "$setOnInsert": bson.M{"created_at": time.Now().UTC()}}, options.Update().SetUpsert(true), ) if err != nil { return fmt.Errorf("store: write settings for %d: %w", chatID, err) } return nil } // RecordUser upserts the interacting user and refreshes last_seen_at. func (s *Store) RecordUser(user User) { if user.ID == 0 { return } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() now := time.Now().UTC() set := bson.M{"last_seen_at": now} for key, value := range map[string]string{ "username": user.Username, "first_name": user.FirstName, "last_name": user.LastName, } { if strings.TrimSpace(value) != "" { set[key] = value } } _, err := s.users.UpdateOne(ctx, bson.M{"_id": user.ID}, bson.M{"$set": set, "$setOnInsert": bson.M{"first_seen_at": now, "renders": int64(0)}}, options.Update().SetUpsert(true), ) if err != nil { log.Printf("store: record user %d: %v", user.ID, err) } } // RecordGroup upserts a chat the bot is a member of (or left). func (s *Store) RecordGroup(chat Chat, member bool) { if chat.ID == 0 { return } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() now := time.Now().UTC() status := "left" if member { status = "member" } set := bson.M{"status": status, "updated_at": now} for key, value := range map[string]string{ "type": chat.Type, "title": chat.Title, "username": chat.Username, } { if strings.TrimSpace(value) != "" { set[key] = value } } _, err := s.groups.UpdateOne(ctx, bson.M{"_id": chat.ID}, bson.M{"$set": set, "$setOnInsert": bson.M{"added_at": now}}, options.Update().SetUpsert(true), ) if err != nil { log.Printf("store: record group %d: %v", chat.ID, err) } } // Stats are the aggregate numbers reported by the /stats command. type Stats struct { Renders int64 ChatRenders int64 InlineRenders int64 FailedRenders int64 Renders24h int64 Users int64 Groups int64 } // Stats computes the aggregate counts shown by /stats. Counting is cheap // enough for a command that is used rarely. func (s *Store) Stats() Stats { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() var stats Stats if count, err := s.stats.CountDocuments(ctx, bson.M{}); err == nil { stats.Renders = count } else { log.Printf("store: count renders: %v", err) } if count, err := s.stats.CountDocuments(ctx, bson.M{"source": "chat"}); err == nil { stats.ChatRenders = count } if count, err := s.stats.CountDocuments(ctx, bson.M{"source": "inline"}); err == nil { stats.InlineRenders = count } if count, err := s.stats.CountDocuments(ctx, bson.M{"ok": false}); err == nil { stats.FailedRenders = count } if count, err := s.stats.CountDocuments(ctx, bson.M{"ts": bson.M{"$gte": time.Now().UTC().Add(-24 * time.Hour)}}); err == nil { stats.Renders24h = count } if count, err := s.users.CountDocuments(ctx, bson.M{}); err == nil { stats.Users = count } if count, err := s.groups.CountDocuments(ctx, bson.M{"status": "member"}); err == nil { stats.Groups = count } return stats } // RecordRender appends a render event to the statistics collection. func (s *Store) RecordRender(stat RenderStat) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() _, err := s.stats.InsertOne(ctx, bson.M{ "ts": time.Now().UTC(), "source": stat.Source, "format": stat.Format, "chat_id": stat.ChatID, "user_id": stat.UserID, "dpi": stat.DPI, "expression_len": stat.ExpressionLen, "ok": stat.OK, "error": stat.Error, "duration_ms": stat.DurationMs, }) if err != nil { log.Printf("store: record render stat: %v", err) } } // FileRefs are correctly typed Telegram file_ids captured from normal chat // sends. They can be reused in inline results without another upload. type FileRefs struct { PhotoID string DocID string } // FileRefs returns cached Telegram references for a rendered content hash. func (s *Store) FileRefs(hash string) FileRefs { if hash == "" { return FileRefs{} } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() var doc struct { PhotoFileID string `bson:"photo_file_id"` DocFileID string `bson:"document_file_id"` } err := s.files.FindOne(ctx, bson.M{"_id": hash}).Decode(&doc) if err != nil { if err != mongo.ErrNoDocuments { log.Printf("store: read file refs for %s: %v", hash, err) } return FileRefs{} } return FileRefs{PhotoID: doc.PhotoFileID, DocID: doc.DocFileID} } // SavePhotoID caches a Photo-type file_id for a content hash. func (s *Store) SavePhotoID(hash, fileID string) { s.saveFileID(hash, "photo_file_id", fileID) } // SaveDocID caches a Document-type file_id for a content hash. func (s *Store) SaveDocID(hash, fileID string) { s.saveFileID(hash, "document_file_id", fileID) } func (s *Store) saveFileID(hash, field, fileID string) { if hash == "" || fileID == "" { return } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() now := time.Now().UTC() _, err := s.files.UpdateOne(ctx, bson.M{"_id": hash}, bson.M{ "$set": bson.M{field: fileID, "updated_at": now}, "$setOnInsert": bson.M{"created_at": now}, }, options.Update().SetUpsert(true), ) if err != nil { log.Printf("store: save %s for %s: %v", field, hash, err) } }