package repository import ( "context" libmongo "apps/backend/internal/lib/mongo" "apps/backend/internal/module/job/domain" "github.com/zeromicro/go-zero/core/stores/mon" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo/options" ) const colJobs = "jobs" type MonStore struct { jobs *mon.Model } func NewMonStore(uri, database string) *MonStore { uri = libmongo.MustMongoURI(uri) return &MonStore{jobs: mon.MustNewModel(uri, database, colJobs)} } func (s *MonStore) Insert(ctx context.Context, j *domain.Job) error { _, err := s.jobs.InsertOne(ctx, j) return err } func (s *MonStore) Update(ctx context.Context, j *domain.Job) error { if j == nil { return domain.ErrNotFound } expectedVersion := j.Version j.UpdatedAt = domain.NowNano() candidate := *j candidate.Version = expectedVersion + 1 res, err := s.jobs.ReplaceOne(ctx, bson.M{"_id": j.ID, "version": expectedVersion}, &candidate) if err != nil { return err } if res.MatchedCount == 0 { return domain.ErrIllegalStatus } j.Version = candidate.Version return nil } func (s *MonStore) FindByID(ctx context.Context, id string) (*domain.Job, error) { var j domain.Job err := s.jobs.FindOne(ctx, &j, bson.M{"_id": id}) if err != nil { if err == mon.ErrNotFound { return nil, domain.ErrNotFound } return nil, err } return &j, nil } func (s *MonStore) ListByOwner(ctx context.Context, ownerUID int64) ([]*domain.Job, error) { var list []*domain.Job err := s.jobs.Find(ctx, &list, bson.M{"owner_uid": ownerUID}, options.Find().SetSort(bson.D{{Key: "updated_at", Value: -1}})) return list, err } func jobTabFilter(ownerUID int64, tab string, nowNs int64) bson.M { base := bson.M{"owner_uid": ownerUID} switch domain.NormalizeListTab(tab) { case domain.TabHistory: base["status"] = bson.M{"$in": []string{ domain.StatusSucceeded, domain.StatusFailed, domain.StatusCancelled, }} case domain.TabRecurring: base["status"] = bson.M{"$in": []string{domain.StatusPending, domain.StatusQueued}} base["run_after"] = bson.M{"$gt": nowNs} default: // active base["$or"] = []bson.M{ {"status": domain.StatusRunning}, { "status": bson.M{"$in": []string{domain.StatusPending, domain.StatusQueued}}, "$or": []bson.M{ {"run_after": bson.M{"$exists": false}}, {"run_after": 0}, {"run_after": bson.M{"$lte": nowNs}}, }, }, } } return base } func (s *MonStore) ListByOwnerTab(ctx context.Context, ownerUID int64, tab string, page, pageSize int) ([]*domain.Job, int64, error) { page, pageSize = domain.ClampPage(page, pageSize) now := domain.NowNano() filter := jobTabFilter(ownerUID, tab, now) total, err := s.jobs.CountDocuments(ctx, filter) if err != nil { return nil, 0, err } sort := bson.D{{Key: "updated_at", Value: -1}} if domain.NormalizeListTab(tab) == domain.TabRecurring { sort = bson.D{{Key: "run_after", Value: 1}, {Key: "created_at", Value: 1}} } else if domain.NormalizeListTab(tab) == domain.TabHistory { sort = bson.D{{Key: "completed_at", Value: -1}, {Key: "updated_at", Value: -1}} } skip := int64((page - 1) * pageSize) var list []*domain.Job err = s.jobs.Find(ctx, &list, filter, options.Find().SetSort(sort).SetSkip(skip).SetLimit(int64(pageSize))) if err != nil { return nil, 0, err } return list, total, nil } func (s *MonStore) ClaimNext(ctx context.Context, workerID string) (*domain.Job, error) { now := domain.NowNano() // due: run_after missing/0 OR run_after <= now filter := bson.M{ "status": bson.M{"$in": []string{domain.StatusPending, domain.StatusQueued}}, "$or": []bson.M{ {"run_after": bson.M{"$exists": false}}, {"run_after": 0}, {"run_after": bson.M{"$lte": now}}, }, } update := bson.M{"$set": bson.M{ "status": domain.StatusRunning, "worker_id": workerID, "updated_at": now, "progress_summary": "已由 worker 領取(" + workerID + ")", }, "$inc": bson.M{"version": 1}} opts := options.FindOneAndUpdate(). SetSort(bson.D{{Key: "run_after", Value: 1}, {Key: "created_at", Value: 1}}). SetReturnDocument(options.After) var j domain.Job err := s.jobs.FindOneAndUpdate(ctx, &j, filter, update, opts) if err != nil { if err == mon.ErrNotFound { return nil, domain.ErrNotFound } return nil, err } return &j, nil } func (s *MonStore) CancelPendingByRef(ctx context.Context, ownerUID int64, template, refID string) error { if refID == "" { return nil } filter := bson.M{ "owner_uid": ownerUID, "template_type": template, "ref_id": refID, "status": bson.M{"$in": []string{domain.StatusPending, domain.StatusQueued}}, } now := domain.NowNano() // list then update each (mon may not expose UpdateMany consistently) var list []*domain.Job if err := s.jobs.Find(ctx, &list, filter); err != nil { return err } for _, j := range list { j.Status = domain.StatusCancelled j.ProgressSummary = "已由新的定期排程取代" j.CompletedAt = now j.UpdatedAt = now _ = s.Update(ctx, j) } return nil } func (s *MonStore) Delete(ctx context.Context, id string) error { res, err := s.jobs.DeleteOne(ctx, bson.M{"_id": id}) if err != nil { return err } if res == 0 { return domain.ErrNotFound } return nil } func (s *MonStore) DeleteTerminalBefore(ctx context.Context, beforeNs int64) (int64, error) { // completed_at 優先;缺則用 updated_at(舊資料) filter := bson.M{ "status": bson.M{"$in": []string{ domain.StatusSucceeded, domain.StatusFailed, domain.StatusCancelled, }}, "$or": []bson.M{ {"completed_at": bson.M{"$gt": 0, "$lt": beforeNs}}, { "$and": []bson.M{ {"$or": []bson.M{ {"completed_at": bson.M{"$exists": false}}, {"completed_at": 0}, }}, {"updated_at": bson.M{"$lt": beforeNs}}, }, }, }, } // go-zero mon.DeleteMany returns deleted count (int64) return s.jobs.DeleteMany(ctx, filter) } var _ domain.Repository = (*MonStore)(nil)