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

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 {
	return s.update(ctx, j, nil)
}

func (s *MonStore) UpdateOwned(ctx context.Context, j *domain.Job, leaseOwner string) error {
	filter := bson.M{"status": domain.StatusRunning, "lease_owner": leaseOwner}
	return s.update(ctx, j, filter)
}

func (s *MonStore) update(ctx context.Context, j *domain.Job, guard bson.M) error {
	if j == nil {
		return domain.ErrNotFound
	}
	expectedVersion := j.Version
	j.UpdatedAt = domain.NowNano()
	candidate := *j
	candidate.Version = expectedVersion + 1
	filter := bson.M{"_id": j.ID, "version": expectedVersion}
	for key, value := range guard {
		filter[key] = value
	}
	res, err := s.jobs.ReplaceOne(ctx, filter, &candidate)
	if err != nil {
		return err
	}
	if res.MatchedCount == 0 {
		return domain.ErrIllegalStatus
	}
	j.Version = candidate.Version
	return nil
}

func (s *MonStore) RenewLease(ctx context.Context, id, leaseOwner string, leaseExpiresAt int64) error {
	res, err := s.jobs.UpdateOne(ctx, bson.M{
		"_id": id, "status": domain.StatusRunning, "lease_owner": leaseOwner,
	}, bson.M{
		"$set": bson.M{"lease_expires_at": leaseExpiresAt, "updated_at": domain.NowNano()},
	})
	if err != nil {
		return err
	}
	if res.MatchedCount == 0 {
		return domain.ErrIllegalStatus
	}
	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()
	filter := bson.M{
		"$or": []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}},
				},
			},
			{
				"status": domain.StatusRunning,
				"$or": []bson.M{
					{"lease_expires_at": bson.M{"$exists": false}},
					{"lease_expires_at": 0},
					{"lease_expires_at": bson.M{"$lte": now}},
				},
			},
		},
	}
	update := bson.M{"$set": bson.M{
		"status":           domain.StatusRunning,
		"worker_id":        workerID,
		"lease_owner":      workerID,
		"lease_expires_at": now + domain.DefaultLeaseDurationNs,
		"updated_at":       now,
		"progress_summary": "已由 worker 領取（" + workerID + "）",
	}, "$inc": bson.M{"version": 1, "attempt": 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)