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

package repository

import (
	"context"

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

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

type MonStore struct {
	personas *mon.Model
	active   *mon.Model
	plays    *mon.Model
	outbox   *mon.Model
	ownPosts *mon.Model
	mentions *mon.Model
	syncMeta *mon.Model
}

func NewMonStore(uri, database string) *MonStore {
	uri = libmongo.MustMongoURI(uri)
	return &MonStore{
		personas: mon.MustNewModel(uri, database, "studio_personas"),
		active:   mon.MustNewModel(uri, database, "studio_active_persona"),
		plays:    mon.MustNewModel(uri, database, "studio_plays"),
		outbox:   mon.MustNewModel(uri, database, "studio_outbox"),
		ownPosts: mon.MustNewModel(uri, database, "studio_own_posts"),
		mentions: mon.MustNewModel(uri, database, "studio_mentions"),
		syncMeta: mon.MustNewModel(uri, database, "studio_sync_meta"),
	}
}

func (s *MonStore) SavePersona(ctx context.Context, p *domain.Persona) error {
	_, err := s.personas.ReplaceOne(ctx, bson.M{"_id": p.ID}, p, options.Replace().SetUpsert(true))
	return err
}

func (s *MonStore) GetPersona(ctx context.Context, id string) (*domain.Persona, error) {
	var p domain.Persona
	err := s.personas.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) ListPersonas(ctx context.Context, ownerUID int64) ([]*domain.Persona, error) {
	var list []*domain.Persona
	err := s.personas.Find(ctx, &list, bson.M{"owner_uid": ownerUID},
		options.Find().SetSort(bson.D{{Key: "updated_at", Value: -1}}))
	return list, err
}

func (s *MonStore) DeletePersona(ctx context.Context, id string) error {
	res, err := s.personas.DeleteOne(ctx, bson.M{"_id": id})
	if err != nil {
		return err
	}
	if res == 0 {
		return domain.ErrNotFound
	}
	return nil
}

func (s *MonStore) GetActivePersonaID(ctx context.Context, ownerUID int64) (string, error) {
	var a domain.ActivePersona
	err := s.active.FindOne(ctx, &a, bson.M{"_id": ownerUID})
	if err != nil {
		if err == mon.ErrNotFound {
			return "", nil
		}
		return "", err
	}
	return a.PersonaID, nil
}

func (s *MonStore) SetActivePersonaID(ctx context.Context, ownerUID int64, personaID string) error {
	doc := &domain.ActivePersona{OwnerUID: ownerUID, PersonaID: personaID, UpdatedAt: domain.NowNano()}
	_, err := s.active.ReplaceOne(ctx, bson.M{"_id": ownerUID}, doc, options.Replace().SetUpsert(true))
	return err
}

func (s *MonStore) SavePlay(ctx context.Context, p *domain.Play) error {
	_, err := s.plays.ReplaceOne(ctx, bson.M{"_id": p.ID}, p, options.Replace().SetUpsert(true))
	return err
}

func (s *MonStore) GetPlay(ctx context.Context, id string) (*domain.Play, error) {
	var p domain.Play
	err := s.plays.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) ListPlays(ctx context.Context, ownerUID int64) ([]*domain.Play, error) {
	var list []*domain.Play
	err := s.plays.Find(ctx, &list, bson.M{"owner_uid": ownerUID},
		options.Find().SetSort(bson.D{{Key: "updated_at", Value: -1}}))
	return list, err
}

func (s *MonStore) DeletePlay(ctx context.Context, id string) error {
	res, err := s.plays.DeleteOne(ctx, bson.M{"_id": id})
	if err != nil {
		return err
	}
	if res == 0 {
		return domain.ErrNotFound
	}
	return nil
}

func (s *MonStore) SaveOutbox(ctx context.Context, b *domain.OutboxBundle) error {
	_, err := s.outbox.ReplaceOne(ctx, bson.M{"_id": b.ID}, b, options.Replace().SetUpsert(true))
	return err
}

func (s *MonStore) GetOutbox(ctx context.Context, id string) (*domain.OutboxBundle, error) {
	var b domain.OutboxBundle
	err := s.outbox.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) ListOutbox(ctx context.Context, ownerUID int64) ([]*domain.OutboxBundle, error) {
	var list []*domain.OutboxBundle
	err := s.outbox.Find(ctx, &list, bson.M{"owner_uid": ownerUID},
		options.Find().SetSort(bson.D{{Key: "updated_at", Value: -1}}))
	return list, err
}

func (s *MonStore) DeleteOutbox(ctx context.Context, id string) error {
	res, err := s.outbox.DeleteOne(ctx, bson.M{
		"_id": id, "steps": bson.M{"$not": bson.M{"$elemMatch": bson.M{"status": domain.StepPublishing}}},
	})
	if err != nil {
		return err
	}
	if res == 0 {
		if _, getErr := s.GetOutbox(ctx, id); getErr == nil {
			return domain.ErrIllegalStatus
		} else {
			return getErr
		}
	}
	return nil
}

func (s *MonStore) ClaimOutboxStep(ctx context.Context, bundleID, stepID, leaseOwner string, now, leaseExpiresAt int64) (*domain.OutboxBundle, error) {
	claimable := bson.M{"id": stepID, "$or": bson.A{
		bson.M{"status": domain.StepScheduled, "scheduled_at": bson.M{"$lte": now}},
		bson.M{"status": domain.StepPublishing, "$or": bson.A{
			bson.M{"lease_expires_at": bson.M{"$lte": now}}, bson.M{"lease_expires_at": bson.M{"$exists": false}},
		}},
	}}
	filter := bson.M{"_id": bundleID, "steps": bson.M{"$elemMatch": claimable}}
	update := bson.M{"$set": bson.M{
		"steps.$.status": domain.StepPublishing, "steps.$.error": "",
		"steps.$.lease_owner": leaseOwner, "steps.$.lease_expires_at": leaseExpiresAt,
		"status": domain.OBActive, "updated_at": now,
	}}
	var b domain.OutboxBundle
	err := s.outbox.FindOneAndUpdate(ctx, &b, filter, update, options.FindOneAndUpdate().SetReturnDocument(options.After))
	if err == mon.ErrNotFound {
		return nil, domain.ErrIllegalStatus
	}
	return &b, err
}

func (s *MonStore) RenewOutboxStepLease(ctx context.Context, bundleID, stepID, leaseOwner string, now, leaseExpiresAt int64) error {
	filter := ownedPublishingStep(bundleID, stepID, leaseOwner, now)
	res, err := s.outbox.UpdateOne(ctx, filter, bson.M{"$set": bson.M{
		"steps.$.lease_expires_at": leaseExpiresAt, "updated_at": now,
	}})
	if err != nil {
		return err
	}
	if res.MatchedCount == 0 {
		return domain.ErrLeaseLost
	}
	return nil
}

func (s *MonStore) FinishOutboxStep(ctx context.Context, bundleID, stepID, leaseOwner string, step *domain.OutboxStep, now int64) error {
	filter := ownedPublishingStep(bundleID, stepID, leaseOwner, now)
	set := bson.M{
		"steps.$[claimed].status": step.Status, "steps.$[claimed].error": step.Error,
		"steps.$[claimed].published_at": step.PublishedAt, "steps.$[claimed].media_id": step.MediaID,
		"steps.$[claimed].image_urls": step.ImageURLs, "updated_at": now,
	}
	unset := bson.M{"steps.$[claimed].lease_owner": "", "steps.$[claimed].lease_expires_at": ""}
	filters := options.ArrayFilters{Filters: []interface{}{bson.M{
		"claimed.id": stepID, "claimed.status": domain.StepPublishing, "claimed.lease_owner": leaseOwner,
	}}}
	if step.Kind == domain.StepRoot && (step.Status == domain.StepFailed || step.Status == domain.StepPublished) {
		replyFilter := bson.M{"reply.kind": domain.StepReply, "reply.status": bson.M{"$nin": bson.A{domain.StepPublished, domain.StepPublishing}}}
		if step.Status == domain.StepFailed {
			set["steps.$[reply].status"], set["steps.$[reply].error"] = domain.StepBlocked, "waiting for root success"
		} else {
			replyFilter["reply.status"] = domain.StepBlocked
			set["steps.$[reply].status"], set["steps.$[reply].error"] = domain.StepScheduled, ""
		}
		filters.Filters = append(filters.Filters, replyFilter)
	}
	res, err := s.outbox.UpdateOne(ctx, filter, bson.M{"$set": set, "$unset": unset}, options.Update().SetArrayFilters(filters))
	if err != nil {
		return err
	}
	if res.MatchedCount == 0 {
		return domain.ErrLeaseLost
	}
	return nil
}

func ownedPublishingStep(bundleID, stepID, leaseOwner string, now int64) bson.M {
	return bson.M{"_id": bundleID, "steps": bson.M{"$elemMatch": bson.M{
		"id": stepID, "status": domain.StepPublishing, "lease_owner": leaseOwner,
		"lease_expires_at": bson.M{"$gt": now},
	}}}
}

func (s *MonStore) RecomputeOutboxStatus(ctx context.Context, bundleID string, now int64) error {
	statuses := bson.M{"$map": bson.M{"input": "$steps", "as": "step", "in": "$$step.status"}}
	pipeline := mongo.Pipeline{bson.D{{Key: "$set", Value: bson.M{
		"status": bson.M{"$switch": bson.M{"branches": bson.A{
			bson.M{"case": bson.M{"$eq": bson.A{statuses, bson.A{}}}, "then": domain.OBCancelled},
			bson.M{"case": bson.M{"$allElementsTrue": bson.A{bson.M{"$map": bson.M{"input": "$steps", "as": "step", "in": bson.M{"$eq": bson.A{"$$step.status", domain.StepPublished}}}}}}, "then": domain.OBCompleted},
			bson.M{"case": bson.M{"$in": bson.A{domain.StepFailed, statuses}}, "then": domain.OBPartial},
			bson.M{"case": bson.M{"$in": bson.A{domain.StepPublishing, statuses}}, "then": domain.OBActive},
		}, "default": domain.OBScheduling}}, "updated_at": now,
	}}}}
	res, err := s.outbox.UpdateOne(ctx, bson.M{"_id": bundleID}, pipeline)
	if err != nil {
		return err
	}
	if res.MatchedCount == 0 {
		return domain.ErrNotFound
	}
	return nil
}

func (s *MonStore) RetryOutboxStep(ctx context.Context, bundleID, stepID string, now int64) (*domain.OutboxBundle, error) {
	filter := bson.M{"_id": bundleID, "steps": bson.M{"$elemMatch": bson.M{"id": stepID, "status": domain.StepFailed}}}
	update := bson.M{
		"$set": bson.M{
			"steps.$.status": domain.StepScheduled, "steps.$.error": "", "steps.$.scheduled_at": now,
			"status": domain.OBScheduling, "updated_at": now,
		},
		"$unset": bson.M{"steps.$.lease_owner": "", "steps.$.lease_expires_at": ""},
	}
	var b domain.OutboxBundle
	err := s.outbox.FindOneAndUpdate(ctx, &b, filter, update, options.FindOneAndUpdate().SetReturnDocument(options.After))
	if err == mon.ErrNotFound {
		return nil, domain.ErrIllegalStatus
	}
	return &b, err
}

func (s *MonStore) ListAllOutbox(ctx context.Context) ([]*domain.OutboxBundle, error) {
	var list []*domain.OutboxBundle
	err := s.outbox.Find(ctx, &list, bson.M{}, options.Find().SetLimit(500))
	return list, err
}

func (s *MonStore) SaveOwnPost(ctx context.Context, p *domain.OwnPost) error {
	_, err := s.ownPosts.ReplaceOne(ctx, bson.M{"_id": p.ID}, p, options.Replace().SetUpsert(true))
	return err
}

func (s *MonStore) GetOwnPost(ctx context.Context, id string) (*domain.OwnPost, error) {
	var p domain.OwnPost
	err := s.ownPosts.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) ListOwnPosts(ctx context.Context, ownerUID int64, accountID string) ([]*domain.OwnPost, error) {
	filter := bson.M{"owner_uid": ownerUID}
	if accountID != "" {
		filter["account_id"] = accountID
	}
	var list []*domain.OwnPost
	err := s.ownPosts.Find(ctx, &list, filter,
		options.Find().SetSort(bson.D{{Key: "published_at", Value: -1}}))
	return list, err
}

func (s *MonStore) DeleteOwnPost(ctx context.Context, id string) error {
	_, err := s.ownPosts.DeleteOne(ctx, bson.M{"_id": id})
	return err
}

func (s *MonStore) GetSyncMeta(ctx context.Context, ownerUID int64) (*domain.SyncMeta, error) {
	var m domain.SyncMeta
	err := s.syncMeta.FindOne(ctx, &m, bson.M{"_id": ownerUID})
	if err != nil {
		if err == mon.ErrNotFound {
			return &domain.SyncMeta{OwnerUID: ownerUID}, nil
		}
		return nil, err
	}
	return &m, nil
}

func (s *MonStore) SetSyncMeta(ctx context.Context, m *domain.SyncMeta) error {
	_, err := s.syncMeta.ReplaceOne(ctx, bson.M{"_id": m.OwnerUID}, m, options.Replace().SetUpsert(true))
	return err
}

func (s *MonStore) SaveMention(ctx context.Context, m *domain.Mention) error {
	_, err := s.mentions.ReplaceOne(ctx, bson.M{"_id": m.ID}, m, options.Replace().SetUpsert(true))
	return err
}

func (s *MonStore) GetMention(ctx context.Context, id string) (*domain.Mention, error) {
	var m domain.Mention
	err := s.mentions.FindOne(ctx, &m, bson.M{"_id": id})
	if err != nil {
		if err == mon.ErrNotFound {
			return nil, domain.ErrNotFound
		}
		return nil, err
	}
	return &m, nil
}

func (s *MonStore) ListMentions(ctx context.Context, ownerUID int64, accountID string) ([]*domain.Mention, error) {
	filter := bson.M{"owner_uid": ownerUID}
	if accountID != "" {
		filter["account_id"] = accountID
	}
	var list []*domain.Mention
	err := s.mentions.Find(ctx, &list, filter,
		options.Find().SetSort(bson.D{{Key: "created_at", Value: -1}}))
	return list, err
}

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