template-monorepo/internal/library/mongo/doc-db-with-cache.go

package mongo

import (
	"context"
	"fmt"

	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/core/stores/cache"
	"github.com/zeromicro/go-zero/core/stores/mon"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

type DocumentDBWithCache struct {
	DocumentDBUseCase
	Cache cache.Cache
}

func MustDocumentDBWithCache(
	conf *Conf,
	collection string,
	cacheConf cache.CacheConf,
	dbOpts []mon.Option,
	cacheOpts []cache.Option,
) (DocumentDBWithCacheUseCase, error) {
	documentDB, err := NewDocumentDB(conf, collection, dbOpts...)
	if err != nil {
		return nil, fmt.Errorf("mongo: document db: %w", err)
	}
	return &DocumentDBWithCache{
		DocumentDBUseCase: documentDB,
		Cache:             MustModelCache(cacheConf, cacheOpts...),
	}, nil
}

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

func (dc *DocumentDBWithCache) GetCache(ctx context.Context, key string, v any) error {
	return dc.Cache.GetCtx(ctx, key, v)
}

func (dc *DocumentDBWithCache) SetCache(ctx context.Context, key string, v any) error {
	return dc.Cache.SetCtx(ctx, key, v)
}

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

func (dc *DocumentDBWithCache) FindOne(ctx context.Context, key string, v, filter any, opts ...options.Lister[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.Lister[options.FindOneAndDeleteOptions]) error {
	if err := dc.GetClient().FindOneAndDelete(ctx, v, filter, opts...); err != nil {
		return err
	}
	dc.delCacheBestEffort(ctx, key)
	return nil
}

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

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

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

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

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

func (dc *DocumentDBWithCache) delCacheBestEffort(ctx context.Context, keys ...string) {
	if len(keys) == 0 {
		return
	}
	if err := dc.DelCache(ctx, keys...); err != nil {
		logx.WithContext(ctx).Errorf("[DocumentDBWithCache] del cache keys=%v: %v", keys, err)
	}
}

func MustModelCache(conf cache.CacheConf, opts ...cache.Option) cache.Cache {
	return cache.New(conf, singleFlight, stats, ErrNotFound, opts...)
}