thread-master/apps/backend/internal/module/scout/repository/mongo.go

package repository

import (
	"context"

	libmongo "apps/backend/internal/lib/mongo"
	"apps/backend/internal/module/scout/domain"

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

type MonStore struct {
	brands   *mon.Model
	products *mon.Model
	active   *mon.Model
	posts    *mon.Model
	hw       *mon.Model
	crawler  *mon.Model
}

func NewMonStore(uri, database string) *MonStore {
	uri = libmongo.MustMongoURI(uri)
	return &MonStore{
		brands: mon.MustNewModel(uri, database, "scout_brands"),
		products: mon.MustNewModel(uri, database, "scout_products"),
		active: mon.MustNewModel(uri, database, "scout_active_brand"),
		posts: mon.MustNewModel(uri, database, "scout_posts"),
		hw: mon.MustNewModel(uri, database, "scout_homework"),
		crawler: mon.MustNewModel(uri, database, "scout_crawler_session"),
	}
}

func (s *MonStore) SaveBrand(ctx context.Context, b *domain.Brand) error {
	_, err := s.brands.ReplaceOne(ctx, bson.M{"_id": b.ID}, b, options.Replace().SetUpsert(true))
	return err
}
func (s *MonStore) GetBrand(ctx context.Context, id string) (*domain.Brand, error) {
	var b domain.Brand
	err := s.brands.FindOne(ctx, &b, bson.M{"_id": id})
	if err != nil {
		if err == mon.ErrNotFound {
			return nil, domain.ErrNotFound
		}
		return nil, err
	}
	return &b, nil
}
func (s *MonStore) ListBrands(ctx context.Context, ownerUID int64) ([]*domain.Brand, error) {
	var list []*domain.Brand
	err := s.brands.Find(ctx, &list, bson.M{"owner_uid": ownerUID})
	return list, err
}
func (s *MonStore) DeleteBrand(ctx context.Context, id string) error {
	res, err := s.brands.DeleteOne(ctx, bson.M{"_id": id})
	if err != nil {
		return err
	}
	if res == 0 {
		return domain.ErrNotFound
	}
	return nil
}

func (s *MonStore) SaveProduct(ctx context.Context, p *domain.Product) error {
	_, err := s.products.ReplaceOne(ctx, bson.M{"_id": p.ID}, p, options.Replace().SetUpsert(true))
	return err
}
func (s *MonStore) GetProduct(ctx context.Context, id string) (*domain.Product, error) {
	var p domain.Product
	err := s.products.FindOne(ctx, &p, bson.M{"_id": id})
	if err != nil {
		if err == mon.ErrNotFound {
			return nil, domain.ErrNotFound
		}
		return nil, err
	}
	return &p, nil
}
func (s *MonStore) ListProducts(ctx context.Context, ownerUID int64, brandID string) ([]*domain.Product, error) {
	filter := bson.M{"owner_uid": ownerUID}
	if brandID != "" {
		filter["brand_id"] = brandID
	}
	var list []*domain.Product
	err := s.products.Find(ctx, &list, filter)
	return list, err
}
func (s *MonStore) DeleteProduct(ctx context.Context, id string) error {
	res, err := s.products.DeleteOne(ctx, bson.M{"_id": id})
	if err != nil {
		return err
	}
	if res == 0 {
		return domain.ErrNotFound
	}
	return nil
}

func (s *MonStore) GetActiveBrandID(ctx context.Context, ownerUID int64) (string, error) {
	var a domain.ActiveBrand
	err := s.active.FindOne(ctx, &a, bson.M{"_id": ownerUID})
	if err != nil {
		if err == mon.ErrNotFound {
			return "", nil
		}
		return "", err
	}
	return a.BrandID, nil
}
func (s *MonStore) SetActiveBrandID(ctx context.Context, ownerUID int64, brandID string) error {
	doc := &domain.ActiveBrand{OwnerUID: ownerUID, BrandID: brandID, UpdatedAt: domain.NowNano()}
	_, err := s.active.ReplaceOne(ctx, bson.M{"_id": ownerUID}, doc, options.Replace().SetUpsert(true))
	return err
}

func (s *MonStore) SavePost(ctx context.Context, p *domain.Post) error {
	_, err := s.posts.ReplaceOne(ctx, bson.M{"_id": p.ID}, p, options.Replace().SetUpsert(true))
	return err
}
func (s *MonStore) GetPost(ctx context.Context, id string) (*domain.Post, error) {
	var p domain.Post
	err := s.posts.FindOne(ctx, &p, bson.M{"_id": id})
	if err != nil {
		if err == mon.ErrNotFound {
			return nil, domain.ErrNotFound
		}
		return nil, err
	}
	return &p, nil
}
func (s *MonStore) ListPosts(ctx context.Context, ownerUID int64, brandID string) ([]*domain.Post, error) {
	filter := bson.M{"owner_uid": ownerUID}
	if brandID != "" {
		filter["brand_id"] = brandID
	}
	var list []*domain.Post
	err := s.posts.Find(ctx, &list, filter, options.Find().SetSort(bson.D{{Key: "created_at", Value: -1}}))
	return list, err
}
func (s *MonStore) DeletePost(ctx context.Context, id string) error {
	res, err := s.posts.DeleteOne(ctx, bson.M{"_id": id})
	if err != nil {
		return err
	}
	if res == 0 {
		return domain.ErrNotFound
	}
	return nil
}
func (s *MonStore) DeletePostsByTheme(ctx context.Context, ownerUID int64, themeKey string) error {
	_, err := s.posts.DeleteMany(ctx, bson.M{"owner_uid": ownerUID, "theme_key": themeKey})
	return err
}

func (s *MonStore) SaveHomework(ctx context.Context, h *domain.Homework) error {
	// composite: store theme as id, filter by owner
	docID := formatKey(h.OwnerUID, h.ThemeKey)
	type hwDoc struct {
		ID string `bson:"_id"`
		*domain.Homework
	}
	_, err := s.hw.ReplaceOne(ctx, bson.M{"_id": docID}, &hwDoc{ID: docID, Homework: h}, options.Replace().SetUpsert(true))
	return err
}
func (s *MonStore) GetHomework(ctx context.Context, ownerUID int64, themeKey string) (*domain.Homework, error) {
	type hwDoc struct {
		ID string `bson:"_id"`
		domain.Homework `bson:",inline"`
	}
	var doc hwDoc
	err := s.hw.FindOne(ctx, &doc, bson.M{"_id": formatKey(ownerUID, themeKey)})
	if err != nil {
		if err == mon.ErrNotFound {
			return nil, domain.ErrNotFound
		}
		return nil, err
	}
	h := doc.Homework
	return &h, nil
}
func (s *MonStore) ListHomework(ctx context.Context, ownerUID int64) ([]*domain.Homework, error) {
	var list []*domain.Homework
	err := s.hw.Find(ctx, &list, bson.M{"owner_uid": ownerUID})
	return list, err
}
func (s *MonStore) DeleteHomework(ctx context.Context, ownerUID int64, themeKey string) error {
	res, err := s.hw.DeleteOne(ctx, bson.M{"_id": formatKey(ownerUID, themeKey)})
	if err != nil {
		return err
	}
	if res == 0 {
		return domain.ErrNotFound
	}
	return nil
}

func (s *MonStore) GetCrawlerSession(ctx context.Context, ownerUID int64) (*domain.CrawlerSession, error) {
	var c domain.CrawlerSession
	err := s.crawler.FindOne(ctx, &c, bson.M{"_id": ownerUID})
	if err != nil {
		if err == mon.ErrNotFound {
			return nil, domain.ErrNotFound
		}
		return nil, err
	}
	if c.Token == "" {
		return nil, domain.ErrNotFound
	}
	return &c, nil
}
func (s *MonStore) SetCrawlerSession(ctx context.Context, sess *domain.CrawlerSession) error {
	_, err := s.crawler.ReplaceOne(ctx, bson.M{"_id": sess.OwnerUID}, sess, options.Replace().SetUpsert(true))
	return err
}
func (s *MonStore) ClearCrawlerSession(ctx context.Context, ownerUID int64) error {
	_, err := s.crawler.DeleteOne(ctx, bson.M{"_id": ownerUID})
	return err
}

func formatKey(uid int64, theme string) string {
	return formatUID(uid) + "|" + theme
}

var _ domain.Repository = (*MonStore)(nil)