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)