add struct
continuous-integration/drone/push Build is passing

This commit is contained in:
2023-04-06 17:36:57 +07:00
parent b771745c47
commit 64345c9350
13 changed files with 7 additions and 34 deletions
+22
View File
@@ -0,0 +1,22 @@
package db
import "time"
type Calories struct {
ChatID int64 `bson:"chat_id"`
Count int `bson:"count"`
CreatedAt time.Time `bson:"created_at"`
Username string `bson:"username"`
}
func NewCalories(
count int,
username string,
) *Calories {
loc, _ := time.LoadLocation("Asia/Novosibirsk")
return &Calories{
Count: count,
Username: username,
CreatedAt: time.Now().In(loc),
}
}
+23
View File
@@ -0,0 +1,23 @@
package db
const (
UserStateNone = UserState("")
UserStateGo = UserState("Go")
UserStateEat = UserState("Eat")
UserStatePause = UserState("Pause")
)
type UserState string
type Chat struct {
ChatID int64 `bson:"chat_id"`
}
type ChatInfo struct {
ChatID int64 `bson:"chat_id"`
Status string `bson:"status"`
}
func (c *ChatInfo) GetStatus() UserState {
return UserState(c.Status)
}
+225
View File
@@ -0,0 +1,225 @@
package db
import (
"context"
"log"
"time"
"valera/internal/config"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
)
type DB struct {
chatsColection *mongo.Collection
workoutsColection *mongo.Collection
caloriesColection *mongo.Collection
}
func NewDB(
mongoURL string,
dbName string,
) (*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 := context.Background()
result := db.chatsColection.FindOne(ctx, bson.M{"chat_id": chatID})
if result.Err() == mongo.ErrNoDocuments {
if _, err := db.chatsColection.InsertOne(ctx, &Chat{ChatID: chatID}); err != nil {
return err
}
return nil
}
return result.Err()
}
func (db *DB) AddWorkout(chatID int64, workout *Workout) error {
ctx := context.Background()
if workout.Count <= 0 {
return nil
}
if err := db.AddChat(chatID); err != nil {
return err
}
workout.ChatID = chatID
_, err := db.workoutsColection.InsertOne(
ctx,
workout,
)
return err
}
func (db *DB) AddCalories(chatID int64, calories *Calories) error {
ctx := context.Background()
if calories.Count <= 0 {
return nil
}
if err := db.AddChat(chatID); err != nil {
return err
}
calories.ChatID = chatID
_, err := db.caloriesColection.InsertOne(
ctx,
calories,
)
return err
}
func (db *DB) GetChatInfo(chatID int64) (*ChatInfo, error) {
ctx := context.Background()
if err := db.AddChat(chatID); err != nil {
return nil, err
}
chatInfoDTO := &ChatInfo{}
if err := db.chatsColection.FindOne(ctx, bson.M{"chat_id": chatID}).Decode(chatInfoDTO); err != nil {
return nil, err
}
return chatInfoDTO, nil
}
func (db *DB) SetStatusToChat(chatID int64, status UserState) error {
ctx := context.Background()
_, err := db.chatsColection.UpdateOne(
ctx,
bson.M{"chat_id": chatID},
bson.M{"$set": bson.M{"status": status}},
)
return err
}
func (db *DB) SetPause(chatID int64, duration time.Duration) error {
ctx := context.Background()
loc, _ := time.LoadLocation("Asia/Novosibirsk")
t := time.Now().In(loc).Add(+duration)
_, err := db.chatsColection.UpdateOne(
ctx,
bson.M{"chat_id": chatID},
bson.M{"$set": bson.M{"pause_until": t}},
)
return err
}
func (db *DB) GetStat(chatID int64, cron bool) (map[string]int, error) {
ctx := context.Background()
if err := db.AddChat(chatID); err != nil {
return nil, err
}
loc, _ := time.LoadLocation("Asia/Novosibirsk")
t := time.Now().In(loc)
if cron {
t = t.Add(-24 * time.Hour)
}
if !cron {
t = time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, loc)
}
res := map[string]int{}
cursor, err := db.workoutsColection.Find(
ctx,
bson.M{
"chat_id": chatID,
"created_at": bson.M{"$gt": t},
},
)
if err != nil {
return nil, err
}
for cursor.Next(context.Background()) {
var result Workout
if err := cursor.Decode(&result); err != nil {
log.Fatal(err)
}
res[result.Name] += result.Count
}
if err := cursor.Err(); err != nil {
log.Fatal(err)
}
caloriesCursor, err := db.caloriesColection.Find(
ctx,
bson.M{
"chat_id": chatID,
"created_at": bson.M{"$gt": t},
},
)
if err != nil {
return nil, err
}
for caloriesCursor.Next(context.Background()) {
var result Calories
if err := caloriesCursor.Decode(&result); err != nil {
log.Fatal(err)
}
res["Калории"] += result.Count
}
if err := caloriesCursor.Err(); err != nil {
log.Fatal(err)
}
return res, err
}
func (db *DB) GetAllChats() ([]int64, error) {
ctx := context.Background()
loc, _ := time.LoadLocation("Asia/Novosibirsk")
t := time.Now().In(loc)
cursor, err := db.chatsColection.Find(
ctx,
bson.M{
"$or": bson.A{
bson.M{
"pause_until": bson.M{"$lt": t},
},
bson.M{
"pause_until": bson.M{"$exists": false},
},
},
},
)
if err != nil {
return nil, err
}
var res []int64
for cursor.Next(context.Background()) {
var result ChatInfo
if err := cursor.Decode(&result); err != nil {
log.Fatal(err)
}
res = append(res, result.ChatID)
}
if err := cursor.Err(); err != nil {
log.Fatal(err)
}
return res, err
}
+25
View File
@@ -0,0 +1,25 @@
package db
import "time"
type Workout struct {
ChatID int64 `bson:"chat_id"`
Name string `bson:"name"`
Count int `bson:"count"`
CreatedAt time.Time `bson:"created_at"`
Username string `bson:"username"`
}
func NewWorkout(
name string,
count int,
username string,
) *Workout {
loc, _ := time.LoadLocation("Asia/Novosibirsk")
return &Workout{
Name: name,
Count: count,
Username: username,
CreatedAt: time.Now().In(loc),
}
}