go-cqhttp/db/mongodb.go

package db

import (
	"context"

	"github.com/pkg/errors"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type MongoDBImpl struct {
	uri   string
	db    string
	mongo *mongo.Database
}

const (
	MongoGroupMessageCollection   = "group-messages"
	MongoPrivateMessageCollection = "private-messages"
)

func UseMongoDB(uri, db string) *MongoDBImpl {
	return &MongoDBImpl{uri: uri, db: db}
}

func (db *MongoDBImpl) Open() error {
	cli, err := mongo.Connect(context.Background(), options.Client().ApplyURI(db.uri))
	if err != nil {
		return errors.Wrap(err, "open mongo connection error")
	}
	db.mongo = cli.Database(db.db)
	return nil
}

func (db *MongoDBImpl) GetMessageByGlobalID(id int32) (IStoredMessage, error) {
	if r, err := db.GetGroupMessageByGlobalID(id); err == nil {
		return r, nil
	}
	return db.GetPrivateMessageByGlobalID(id)
}

func (db *MongoDBImpl) GetGroupMessageByGlobalID(id int32) (*StoredGroupMessage, error) {
	coll := db.mongo.Collection(MongoGroupMessageCollection)
	var ret StoredGroupMessage
	if err := coll.FindOne(context.Background(), bson.D{{"globalId", id}}).Decode(&ret); err != nil {
		return nil, errors.Wrap(err, "query error")
	}
	return &ret, nil
}

func (db *MongoDBImpl) GetPrivateMessageByGlobalID(id int32) (*StoredPrivateMessage, error) {
	coll := db.mongo.Collection(MongoPrivateMessageCollection)
	var ret StoredPrivateMessage
	if err := coll.FindOne(context.Background(), bson.D{{"globalId", id}}).Decode(&ret); err != nil {
		return nil, errors.Wrap(err, "query error")
	}
	return &ret, nil
}

func (db *MongoDBImpl) InsertGroupMessage(msg *StoredGroupMessage) error {
	coll := db.mongo.Collection(MongoGroupMessageCollection)
	_, err := coll.InsertOne(context.Background(), msg)
	return errors.Wrap(err, "insert error")
}

func (db *MongoDBImpl) InsertPrivateMessage(msg *StoredPrivateMessage) error {
	coll := db.mongo.Collection(MongoPrivateMessageCollection)
	_, err := coll.InsertOne(context.Background(), msg)
	return errors.Wrap(err, "insert error")
}