diff --git a/db/db.go b/db/db.go index c57ce7a..9311814 100644 --- a/db/db.go +++ b/db/db.go @@ -11,32 +11,39 @@ import ( ) type DB struct { - cfg *config.Config + chatsColection *mongo.Collection + workoutsColection *mongo.Collection + caloriesColection *mongo.Collection } func NewDB( mongoURL string, dbName string, -) *DB { - return &DB{ - cfg: config.NewConfig(mongoURL, dbName), +) (*DB, error) { + ctx := context.Background() + opt := options.Client(). + ApplyURI(mongoURL). + SetMinPoolSize(10) + cfg := config.NewConfig(mongoURL, dbName) + + client, err := mongo.Connect(ctx, opt) + if err != nil { + return nil, err } + + return &DB{ + chatsColection: client.Database(dbName).Collection(cfg.ChatsCollectionName), + workoutsColection: client.Database(dbName).Collection(cfg.WorkoutsCollectionName), + caloriesColection: client.Database(dbName).Collection(cfg.CaloriesCollectionName), + }, nil } func (db *DB) AddChat(chatID int64) error { - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() + ctx := context.Background() - client, err := mongo.Connect(ctx, options.Client().ApplyURI(db.cfg.MongoURL)) - if err != nil { - return err - } - - collection := client.Database(db.cfg.DBName).Collection(db.cfg.ChatsCollectionName) - - result := collection.FindOne(ctx, bson.M{"chat_id": chatID}) + result := db.chatsColection.FindOne(ctx, bson.M{"chat_id": chatID}) if result.Err() == mongo.ErrNoDocuments { - if _, err := collection.InsertOne(ctx, &Chat{ChatID: chatID}); err != nil { + if _, err := db.chatsColection.InsertOne(ctx, &Chat{ChatID: chatID}); err != nil { return err } return nil @@ -45,6 +52,8 @@ func (db *DB) AddChat(chatID int64) error { } func (db *DB) AddWorkout(chatID int64, workout *Workout) error { + ctx := context.Background() + if workout.Count <= 0 { return nil } @@ -53,16 +62,7 @@ func (db *DB) AddWorkout(chatID int64, workout *Workout) error { } workout.ChatID = chatID - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - - client, err := mongo.Connect(ctx, options.Client().ApplyURI(db.cfg.MongoURL)) - if err != nil { - return err - } - - collection := client.Database(db.cfg.DBName).Collection(db.cfg.WorkoutsCollectionName) - _, err = collection.InsertOne( + _, err := db.workoutsColection.InsertOne( ctx, workout, ) @@ -70,6 +70,8 @@ func (db *DB) AddWorkout(chatID int64, workout *Workout) error { } func (db *DB) AddCalories(chatID int64, calories *Calories) error { + ctx := context.Background() + if calories.Count <= 0 { return nil } @@ -78,16 +80,7 @@ func (db *DB) AddCalories(chatID int64, calories *Calories) error { } calories.ChatID = chatID - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - - client, err := mongo.Connect(ctx, options.Client().ApplyURI(db.cfg.MongoURL)) - if err != nil { - return err - } - - collection := client.Database(db.cfg.DBName).Collection(db.cfg.CaloriesCollectionName) - _, err = collection.InsertOne( + _, err := db.caloriesColection.InsertOne( ctx, calories, ) @@ -95,36 +88,23 @@ func (db *DB) AddCalories(chatID int64, calories *Calories) error { } func (db *DB) GetChatInfo(chatID int64) (*ChatInfo, error) { + ctx := context.Background() + if err := db.AddChat(chatID); err != nil { return nil, err } - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - client, err := mongo.Connect(ctx, options.Client().ApplyURI(db.cfg.MongoURL)) - if err != nil { - return nil, err - } - - collection := client.Database(db.cfg.DBName).Collection(db.cfg.ChatsCollectionName) chatInfoDTO := &ChatInfo{} - if err := collection.FindOne(ctx, bson.M{"chat_id": chatID}).Decode(chatInfoDTO); err != nil { + if err := db.chatsColection.FindOne(ctx, bson.M{"chat_id": chatID}).Decode(chatInfoDTO); err != nil { return nil, err } - return chatInfoDTO, err + return chatInfoDTO, nil } func (db *DB) SetStatusToChat(chatID int64, status UserState) error { - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() + ctx := context.Background() - client, err := mongo.Connect(ctx, options.Client().ApplyURI(db.cfg.MongoURL)) - if err != nil { - return err - } - - collection := client.Database(db.cfg.DBName).Collection(db.cfg.ChatsCollectionName) - _, err = collection.UpdateOne( + _, err := db.chatsColection.UpdateOne( ctx, bson.M{"chat_id": chatID}, bson.M{"$set": bson.M{"status": status}}, @@ -133,14 +113,9 @@ func (db *DB) SetStatusToChat(chatID int64, status UserState) error { } func (db *DB) GetStat(chatID int64, cron bool) (map[string]int, error) { - if err := db.AddChat(chatID); err != nil { - return nil, err - } - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() + ctx := context.Background() - client, err := mongo.Connect(ctx, options.Client().ApplyURI(db.cfg.MongoURL)) - if err != nil { + if err := db.AddChat(chatID); err != nil { return nil, err } @@ -155,8 +130,7 @@ func (db *DB) GetStat(chatID int64, cron bool) (map[string]int, error) { res := map[string]int{} - collection := client.Database(db.cfg.DBName).Collection(db.cfg.WorkoutsCollectionName) - cursor, err := collection.Find( + cursor, err := db.workoutsColection.Find( ctx, bson.M{ "chat_id": chatID, @@ -177,8 +151,7 @@ func (db *DB) GetStat(chatID int64, cron bool) (map[string]int, error) { log.Fatal(err) } - caloriesCollection := client.Database(db.cfg.DBName).Collection(db.cfg.CaloriesCollectionName) - caloriesCursor, err := caloriesCollection.Find( + caloriesCursor, err := db.caloriesColection.Find( ctx, bson.M{ "chat_id": chatID, @@ -203,16 +176,9 @@ func (db *DB) GetStat(chatID int64, cron bool) (map[string]int, error) { } func (db *DB) GetAllChats() ([]int64, error) { - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() + ctx := context.Background() - client, err := mongo.Connect(ctx, options.Client().ApplyURI(db.cfg.MongoURL)) - if err != nil { - return nil, err - } - - collection := client.Database(db.cfg.DBName).Collection(db.cfg.ChatsCollectionName) - cursor, err := collection.Find(ctx, bson.M{}) + cursor, err := db.chatsColection.Find(ctx, bson.M{}) if err != nil { return nil, err } diff --git a/main.go b/main.go index d72bde7..9e952e7 100644 --- a/main.go +++ b/main.go @@ -19,6 +19,7 @@ var ( workoutTypes = []string{ "Отжимания", "Пресс", + "Подтягивания", } caloriesMap = map[string]int{ "чай": 79, @@ -27,7 +28,7 @@ var ( ) const ( - version = "v1.4.0" + version = "v1.5.0" ) type Opts struct { @@ -88,10 +89,13 @@ func run() { dbName = strings.ReplaceAll(dbName, "\n", "") fmt.Println(dbName) - dataBase := db.NewDB( + dataBase, err := db.NewDB( mongoURL, dbName, ) + if err != nil { + panic(err) + } u := tgbot.NewUpdate(0) u.Timeout = 60 @@ -230,7 +234,7 @@ func run() { log.Println(err) } case commands.Help: - _, _ = bot.Send(tgbot.NewMessage(chatID, "Вот что я умею:\n\n1) Предлагать размяться\n2) Показывать статистику\n3) Считать калории")) + _, _ = bot.Send(tgbot.NewMessage(chatID, "Вот что я умею:\n\n1) Предлагать размяться\n - несколько видов упражнений 6 раз в день\n2) Показывать статистику за день\n3) Считать калории\n - введи одно число, я запишу калории\n - введи 2 числа через пробел одно количество грамм другое количество калорий в 100 граммах (порядок не важен), я посчитаю\n - введи название того что ты съел, я посчитаю (если знаю)")) case commands.Ping: _, _ = bot.Send(tgbot.NewMessage(chatID, "pong")) case commands.Go: @@ -276,11 +280,11 @@ func calcCalories(text string) (int, error) { func sendGoToChat(bot *tgbot.BotAPI, dataBase *db.DB, chatID int64) { msg := tgbot.NewMessage(chatID, "Давай немного разомнемся, выбирай:") - row := tgbot.NewKeyboardButtonRow() + rows := make([][]tgbot.KeyboardButton, 0, len(workoutTypes)) for _, workoutType := range workoutTypes { - row = append(row, tgbot.NewKeyboardButton(workoutType)) + rows = append(rows, tgbot.NewKeyboardButtonRow(tgbot.NewKeyboardButton(workoutType))) } - msg.ReplyMarkup = tgbot.NewReplyKeyboard(row) + msg.ReplyMarkup = tgbot.NewReplyKeyboard(rows...) if _, err := bot.Send(msg); err == nil { if err := dataBase.SetStatusToChat(chatID, db.UserStateGo); err != nil { log.Println(err)