backend/pkg/library/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/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("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)
}

// DeleteOne deletes a single document and invalidates cache
func (dc *DocumentDBWithCache) DeleteOne(ctx context.Context, key string, filter any, opts ...*options.DeleteOneOptions) (int64, error) {
	// Convert options to Builder format
	var listerOpts []options.Lister[options.DeleteOneOptions]
	for _, opt := range opts {
		if opt != nil {
			builder := options.DeleteOne()
			if opt.Collation != nil {
				builder.SetCollation(opt.Collation)
			}
			if opt.Comment != nil {
				builder.SetComment(opt.Comment)
			}
			if opt.Hint != nil {
				builder.SetHint(opt.Hint)
			}
			listerOpts = append(listerOpts, builder)
		}
	}

	val, err := dc.GetClient().DeleteOne(ctx, filter, listerOpts...)
	if err != nil {
		return 0, err
	}

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

	return val, nil
}

// FindOne finds a single document with cache support
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 {
		// Convert options to Builder format
		var listerOpts []options.Lister[options.FindOneOptions]
		for _, opt := range opts {
			if opt != nil {
				builder := options.FindOne()
				if opt.Collation != nil {
					builder.SetCollation(opt.Collation)
				}
				if opt.Comment != nil {
					builder.SetComment(opt.Comment)
				}
				if opt.Hint != nil {
					builder.SetHint(opt.Hint)
				}
				if opt.Projection != nil {
					builder.SetProjection(opt.Projection)
				}
				if opt.Skip != nil {
					builder.SetSkip(*opt.Skip)
				}
				if opt.Sort != nil {
					builder.SetSort(opt.Sort)
				}
				listerOpts = append(listerOpts, builder)
			}
		}

		return dc.GetClient().FindOne(ctx, v, filter, listerOpts...)
	})
}

// FindOneAndDelete finds and deletes a single document with cache invalidation
func (dc *DocumentDBWithCache) FindOneAndDelete(ctx context.Context, key string, v, filter any, opts ...*options.FindOneAndDeleteOptions) error {
	// Convert options to Builder format
	var listerOpts []options.Lister[options.FindOneAndDeleteOptions]
	for _, opt := range opts {
		if opt != nil {
			builder := options.FindOneAndDelete()
			if opt.Collation != nil {
				builder.SetCollation(opt.Collation)
			}
			if opt.Comment != nil {
				builder.SetComment(opt.Comment)
			}
			if opt.Hint != nil {
				builder.SetHint(opt.Hint)
			}
			if opt.Projection != nil {
				builder.SetProjection(opt.Projection)
			}
			if opt.Sort != nil {
				builder.SetSort(opt.Sort)
			}
			listerOpts = append(listerOpts, builder)
		}
	}

	if err := dc.GetClient().FindOneAndDelete(ctx, v, filter, listerOpts...); err != nil {
		return err
	}

	return dc.DelCache(ctx, key)
}

// FindOneAndReplace finds and replaces a single document with cache invalidation
func (dc *DocumentDBWithCache) FindOneAndReplace(ctx context.Context, key string, v, filter, replacement any, opts ...*options.FindOneAndReplaceOptions) error {
	// Convert options to Builder format
	var listerOpts []options.Lister[options.FindOneAndReplaceOptions]
	for _, opt := range opts {
		if opt != nil {
			builder := options.FindOneAndReplace()
			if opt.Collation != nil {
				builder.SetCollation(opt.Collation)
			}
			if opt.Comment != nil {
				builder.SetComment(opt.Comment)
			}
			if opt.Hint != nil {
				builder.SetHint(opt.Hint)
			}
			if opt.Projection != nil {
				builder.SetProjection(opt.Projection)
			}
			if opt.ReturnDocument != nil {
				builder.SetReturnDocument(*opt.ReturnDocument)
			}
			if opt.Sort != nil {
				builder.SetSort(opt.Sort)
			}
			if opt.Upsert != nil {
				builder.SetUpsert(*opt.Upsert)
			}
			listerOpts = append(listerOpts, builder)
		}
	}

	if err := dc.GetClient().FindOneAndReplace(ctx, v, filter, replacement, listerOpts...); err != nil {
		return err
	}

	return dc.DelCache(ctx, key)
}

// InsertOne inserts a single document and invalidates cache
func (dc *DocumentDBWithCache) InsertOne(ctx context.Context, key string, document any, opts ...*options.InsertOneOptions) (*mongo.InsertOneResult, error) {
	// Convert options to Builder format
	var listerOpts []options.Lister[options.InsertOneOptions]
	for _, opt := range opts {
		if opt != nil {
			builder := options.InsertOne()
			if opt.BypassDocumentValidation != nil {
				builder.SetBypassDocumentValidation(*opt.BypassDocumentValidation)
			}
			if opt.Comment != nil {
				builder.SetComment(opt.Comment)
			}
			listerOpts = append(listerOpts, builder)
		}
	}

	res, err := dc.GetClient().Collection.InsertOne(ctx, document, listerOpts...)
	if err != nil {
		return nil, err
	}

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

	return res, nil
}

// UpdateByID updates a document by ID and invalidates cache
func (dc *DocumentDBWithCache) UpdateByID(ctx context.Context, key string, id, update any, opts ...*options.UpdateOneOptions) (*mongo.UpdateResult, error) {
	// Convert options to Builder format
	var listerOpts []options.Lister[options.UpdateOneOptions]
	for _, opt := range opts {
		if opt != nil {
			builder := options.UpdateOne()
			if opt.ArrayFilters != nil {
				builder.SetArrayFilters(opt.ArrayFilters)
			}
			if opt.BypassDocumentValidation != nil {
				builder.SetBypassDocumentValidation(*opt.BypassDocumentValidation)
			}
			if opt.Collation != nil {
				builder.SetCollation(opt.Collation)
			}
			if opt.Comment != nil {
				builder.SetComment(opt.Comment)
			}
			if opt.Hint != nil {
				builder.SetHint(opt.Hint)
			}
			if opt.Upsert != nil {
				builder.SetUpsert(*opt.Upsert)
			}
			listerOpts = append(listerOpts, builder)
		}
	}

	res, err := dc.GetClient().Collection.UpdateByID(ctx, id, update, listerOpts...)
	if err != nil {
		return nil, err
	}

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

	return res, nil
}

// UpdateMany updates multiple documents and invalidates cache
func (dc *DocumentDBWithCache) UpdateMany(ctx context.Context, keys []string, filter, update any, opts ...*options.UpdateManyOptions) (*mongo.UpdateResult, error) {
	// Convert options to Builder format
	var listerOpts []options.Lister[options.UpdateManyOptions]
	for _, opt := range opts {
		if opt != nil {
			builder := options.UpdateMany()
			if opt.ArrayFilters != nil {
				builder.SetArrayFilters(opt.ArrayFilters)
			}
			if opt.BypassDocumentValidation != nil {
				builder.SetBypassDocumentValidation(*opt.BypassDocumentValidation)
			}
			if opt.Collation != nil {
				builder.SetCollation(opt.Collation)
			}
			if opt.Comment != nil {
				builder.SetComment(opt.Comment)
			}
			if opt.Hint != nil {
				builder.SetHint(opt.Hint)
			}
			if opt.Upsert != nil {
				builder.SetUpsert(*opt.Upsert)
			}
			listerOpts = append(listerOpts, builder)
		}
	}

	res, err := dc.GetClient().Collection.UpdateMany(ctx, filter, update, listerOpts...)
	if err != nil {
		return nil, err
	}

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

	return res, nil
}

// UpdateOne updates a single document and invalidates cache
func (dc *DocumentDBWithCache) UpdateOne(ctx context.Context, key string, filter, update any, opts ...*options.UpdateOneOptions) (*mongo.UpdateResult, error) {
	// Convert options to Builder format
	var listerOpts []options.Lister[options.UpdateOneOptions]
	for _, opt := range opts {
		if opt != nil {
			builder := options.UpdateOne()
			if opt.ArrayFilters != nil {
				builder.SetArrayFilters(opt.ArrayFilters)
			}
			if opt.BypassDocumentValidation != nil {
				builder.SetBypassDocumentValidation(*opt.BypassDocumentValidation)
			}
			if opt.Collation != nil {
				builder.SetCollation(opt.Collation)
			}
			if opt.Comment != nil {
				builder.SetComment(opt.Comment)
			}
			if opt.Hint != nil {
				builder.SetHint(opt.Hint)
			}
			if opt.Upsert != nil {
				builder.SetUpsert(*opt.Upsert)
			}
			listerOpts = append(listerOpts, builder)
		}
	}

	res, err := dc.GetClient().Collection.UpdateOne(ctx, filter, update, listerOpts...)
	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...)
}