matrix-veles/internal/db/db.go

package db

import (
	"context"
	"github.com/Unkn0wnCat/matrix-veles/internal/db/model"
	"github.com/spf13/viper"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/event"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"log"
	"regexp"
	"time"
)

var DbClient *mongo.Client
var Db *mongo.Database

func Connect() {
	if viper.GetString("bot.mongo.uri") == "" {
		log.Println("Skipping database login...")
		return
	}

	cmdMonitor := &event.CommandMonitor{
		Started: func(_ context.Context, evt *event.CommandStartedEvent) {
			log.Print(evt.Command)
		},
	}

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	newClient, err := mongo.Connect(ctx, options.Client().ApplyURI(viper.GetString("bot.mongo.uri")).SetMonitor(cmdMonitor))
	if err != nil {
		log.Println("Could not connect to DB")
		log.Panicln(err)
	}

	DbClient = newClient
	Db = DbClient.Database(viper.GetString("bot.mongo.database"))
}

func SaveEntry(entry *model.DBEntry) error {
	db := DbClient.Database(viper.GetString("bot.mongo.database"))

	opts := options.Replace().SetUpsert(true)

	filter := bson.D{{"_id", entry.ID}}

	_, err := db.Collection(viper.GetString("bot.mongo.collection.entries")).ReplaceOne(context.TODO(), filter, entry, opts)

	return err
}

func GetEntryByID(id primitive.ObjectID) (*model.DBEntry, error) {
	db := DbClient.Database(viper.GetString("bot.mongo.database"))

	res := db.Collection(viper.GetString("bot.mongo.collection.entries")).FindOne(context.TODO(), bson.D{{"_id", id}})
	if res.Err() != nil {
		return nil, res.Err()
	}

	object := model.DBEntry{}

	err := res.Decode(&object)
	if err != nil {
		return nil, err
	}

	return &object, nil
}

func GetEntries(first int64, cursor *primitive.ObjectID) ([]*model.DBEntry, error) {
	db := DbClient.Database(viper.GetString("bot.mongo.database"))

	opts := options.FindOptions{
		Limit: &first,
	}

	filter := bson.M{}

	if cursor != nil {
		filter = bson.M{
			"_id": bson.M{
				"$gt": *cursor,
			},
		}
		log.Println(filter)
	}

	res, err := db.Collection(viper.GetString("bot.mongo.collection.entries")).Find(context.TODO(), filter, &opts)
	if err != nil {
		return nil, res.Err()
	}

	var object []*model.DBEntry

	err = res.All(context.TODO(), &object)
	if err != nil {
		return nil, err
	}

	return object, nil
}

func GetEntryByHash(hash string) (*model.DBEntry, error) {
	db := DbClient.Database(viper.GetString("bot.mongo.database"))

	res := db.Collection(viper.GetString("bot.mongo.collection.entries")).FindOne(context.TODO(), bson.D{{"hash_value", hash}})
	if res.Err() != nil {
		return nil, res.Err()
	}

	object := model.DBEntry{}

	err := res.Decode(&object)
	if err != nil {
		return nil, err
	}

	return &object, nil
}

func SaveList(list *model.DBHashList) error {
	db := DbClient.Database(viper.GetString("bot.mongo.database"))

	opts := options.Replace().SetUpsert(true)

	filter := bson.D{{"_id", list.ID}}

	_, err := db.Collection(viper.GetString("bot.mongo.collection.lists")).ReplaceOne(context.TODO(), filter, list, opts)

	return err
}

func GetListByID(id primitive.ObjectID) (*model.DBHashList, error) {
	db := DbClient.Database(viper.GetString("bot.mongo.database"))

	res := db.Collection(viper.GetString("bot.mongo.collection.lists")).FindOne(context.TODO(), bson.D{{"_id", id}})
	if res.Err() != nil {
		return nil, res.Err()
	}

	object := model.DBHashList{}

	err := res.Decode(&object)
	if err != nil {
		return nil, err
	}

	return &object, nil
}

func GetListByName(name string) (*model.DBHashList, error) {
	db := DbClient.Database(viper.GetString("bot.mongo.database"))

	res := db.Collection(viper.GetString("bot.mongo.collection.lists")).FindOne(context.TODO(), bson.D{{"name", name}})
	if res.Err() != nil {
		return nil, res.Err()
	}

	object := model.DBHashList{}

	err := res.Decode(&object)
	if err != nil {
		return nil, err
	}

	return &object, nil
}

func GetLists(first int64, cursor *primitive.ObjectID) ([]*model.DBHashList, error) {
	db := DbClient.Database(viper.GetString("bot.mongo.database"))

	opts := options.FindOptions{
		Limit: &first,
	}

	filter := bson.M{}

	if cursor != nil {
		filter = bson.M{
			"_id": bson.M{
				"$gt": *cursor,
			},
		}
		log.Println(filter)
	}

	res, err := db.Collection(viper.GetString("bot.mongo.collection.lists")).Find(context.TODO(), filter, &opts)
	if err != nil {
		return nil, res.Err()
	}

	var object []*model.DBHashList

	err = res.All(context.TODO(), &object)
	if err != nil {
		return nil, err
	}

	return object, nil
}

func SaveUser(user *model.DBUser) error {
	db := DbClient.Database(viper.GetString("bot.mongo.database"))

	opts := options.Replace().SetUpsert(true)

	filter := bson.D{{"_id", user.ID}}

	_, err := db.Collection(viper.GetString("bot.mongo.collection.users")).ReplaceOne(context.TODO(), filter, user, opts)

	return err
}

func GetUserByID(id primitive.ObjectID) (*model.DBUser, error) {
	db := DbClient.Database(viper.GetString("bot.mongo.database"))

	res := db.Collection(viper.GetString("bot.mongo.collection.users")).FindOne(context.TODO(), bson.D{{"_id", id}})
	if res.Err() != nil {
		return nil, res.Err()
	}

	object := model.DBUser{}

	err := res.Decode(&object)
	if err != nil {
		return nil, err
	}

	return &object, nil
}

func GetUserByUsername(username string) (*model.DBUser, error) {
	db := DbClient.Database(viper.GetString("bot.mongo.database"))

	regexPattern := "^" + regexp.QuoteMeta(username) + "$"

	res := db.Collection(viper.GetString("bot.mongo.collection.users")).FindOne(context.TODO(), bson.D{{"username", primitive.Regex{
		Pattern: regexPattern,
		Options: "i",
	}}})
	if res.Err() != nil {
		return nil, res.Err()
	}

	object := model.DBUser{}

	err := res.Decode(&object)
	if err != nil {
		return nil, err
	}

	return &object, nil
}

func GetUsers(first int64, cursor *primitive.ObjectID) ([]*model.DBUser, error) {
	db := DbClient.Database(viper.GetString("bot.mongo.database"))

	opts := options.FindOptions{
		Limit: &first,
	}

	filter := bson.M{}

	if cursor != nil {
		filter = bson.M{
			"_id": bson.M{
				"$gt": *cursor,
			},
		}
		log.Println(filter)
	}

	res, err := db.Collection(viper.GetString("bot.mongo.collection.users")).Find(context.TODO(), filter, &opts)
	if err != nil {
		return nil, res.Err()
	}

	var object []*model.DBUser

	err = res.All(context.TODO(), &object)
	if err != nil {
		return nil, err
	}

	return object, nil
}