153 lines
3.9 KiB
Go
153 lines
3.9 KiB
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 {
|
||
j.UpdatedAt = domain.NowNano()
|
||
res, err := s.jobs.ReplaceOne(ctx, bson.M{"_id": j.ID}, j)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if res.MatchedCount == 0 {
|
||
return domain.ErrNotFound
|
||
}
|
||
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: "created_at", Value: -1}}))
|
||
return list, err
|
||
}
|
||
|
||
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 + "）",
|
||
}}
|
||
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)
|