library-go/storage/mongo/doc-db-with-cache.go

package mongo

import (
	"context"
	"fmt"
	"github.com/zeromicro/go-zero/core/stores/cache"
	"github.com/zeromicro/go-zero/core/stores/mon"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type DocumentDBWithCache struct {
	DocumentDBUseCase
	Cache cache.Cache
}

func MustDocumentDBWithCache(conf *Conf, cacheConf cache.CacheConf, dbOpts []mon.Option, cacheOpts []cache.Option) (DocumentDBWithCacheUseCase, error) {
	documentDB, err := NewDocumentDB(conf, dbOpts...)
	if err != nil {
		return nil, fmt.Errorf("failed to initialize DocumentDB: %w", err)
	}

	c := MustModelCache(cacheConf, cacheOpts...)

	return &DocumentDBWithCache{
		DocumentDBUseCase: documentDB,
		Cache:             c,
	}, nil
}

func (dc *DocumentDBWithCache) DelCache(ctx context.Context, keys ...string) error {
	return dc.Cache.DelCtx(ctx, keys...)
}

func (dc *DocumentDBWithCache) GetCache(key string, v any) error {
	return dc.Cache.Get(key, v)
}

func (dc *DocumentDBWithCache) SetCache(key string, v any) error {
	return dc.Cache.Set(key, v)
}

func (dc *DocumentDBWithCache) DeleteOne(ctx context.Context, key string, filter any, opts ...*options.DeleteOptions) (int64, error) {
	val, err := dc.GetClient().DeleteOne(ctx, filter, opts...)
	if err != nil {
		return 0, err
	}

	if err := dc.DelCache(ctx, key); err != nil {
		return 0, err
	}

	return val, nil
}

func (dc *DocumentDBWithCache) FindOne(ctx context.Context, key string, v, filter any, opts ...*options.FindOneOptions) error {
	return dc.Cache.TakeCtx(ctx, v, key, func(v any) error {
		return dc.GetClient().FindOne(ctx, v, filter, opts...)
	})
}

func (dc *DocumentDBWithCache) FindOneAndDelete(ctx context.Context, key string, v, filter any, opts ...*options.FindOneAndDeleteOptions) error {
	if err := dc.GetClient().FindOneAndDelete(ctx, v, filter, opts...); err != nil {
		return err
	}

	return dc.DelCache(ctx, key)
}

func (dc *DocumentDBWithCache) FindOneAndReplace(ctx context.Context, key string, v, filter, replacement any, opts ...*options.FindOneAndReplaceOptions) error {
	if err := dc.GetClient().FindOneAndReplace(ctx, v, filter, replacement, opts...); err != nil {
		return err
	}

	return dc.DelCache(ctx, key)
}

func (dc *DocumentDBWithCache) InsertOne(ctx context.Context, key string, document any, opts ...*options.InsertOneOptions) (*mongo.InsertOneResult, error) {
	res, err := dc.GetClient().InsertOne(ctx, document, opts...)
	if err != nil {
		return nil, err
	}

	if err = dc.DelCache(ctx, key); err != nil {
		return nil, err
	}

	return res, nil
}

func (dc *DocumentDBWithCache) UpdateByID(ctx context.Context, key string, id, update any, opts ...*options.UpdateOptions) (*mongo.UpdateResult, error) {
	res, err := dc.GetClient().UpdateByID(ctx, id, update, opts...)
	if err != nil {
		return nil, err
	}

	if err = dc.DelCache(ctx, key); err != nil {
		return nil, err
	}

	return res, nil
}

func (dc *DocumentDBWithCache) UpdateMany(ctx context.Context, keys []string, filter, update any, opts ...*options.UpdateOptions) (*mongo.UpdateResult, error) {
	res, err := dc.GetClient().UpdateMany(ctx, filter, update, opts...)
	if err != nil {
		return nil, err
	}

	if err = dc.DelCache(ctx, keys...); err != nil {
		return nil, err
	}

	return res, nil
}

func (dc *DocumentDBWithCache) UpdateOne(ctx context.Context, key string, filter, update any, opts ...*options.UpdateOptions) (*mongo.UpdateResult, error) {
	res, err := dc.GetClient().UpdateOne(ctx, filter, update, opts...)
	if err != nil {
		return nil, err
	}

	if err = dc.DelCache(ctx, key); err != nil {
		return nil, err
	}

	return res, nil
}

// ========================

// MustModelCache returns a cache cluster.
func MustModelCache(conf cache.CacheConf, opts ...cache.Option) cache.Cache {
	return cache.New(conf, singleFlight, stats, mongo.ErrNoDocuments, opts...)
}