142 lines
3.3 KiB
Go
142 lines
3.3 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
|
|
"apps/backend/internal/module/radar/domain"
|
|
|
|
"github.com/zeromicro/go-zero/core/stores/mon"
|
|
"go.mongodb.org/mongo-driver/bson"
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
|
)
|
|
|
|
func (s *MonStore) CreateSweep(ctx context.Context, sw *domain.RadarSweep) error {
|
|
if sw == nil {
|
|
return domain.ErrValidation
|
|
}
|
|
if err := sw.Normalize(); err != nil {
|
|
return err
|
|
}
|
|
if sw.ID == "" {
|
|
sw.ID = domain.NewID()
|
|
}
|
|
if sw.StartedAt == 0 {
|
|
sw.StartedAt = domain.NowNano()
|
|
}
|
|
_, err := s.sweeps.InsertOne(ctx, sw)
|
|
return err
|
|
}
|
|
|
|
func (s *MonStore) UpdateSweep(ctx context.Context, id string, delta domain.SweepDelta) (*domain.RadarSweep, error) {
|
|
inc := bson.M{}
|
|
if delta.HitCount != 0 {
|
|
inc["hit_count"] = delta.HitCount
|
|
}
|
|
if delta.JudgedCount != 0 {
|
|
inc["judged_count"] = delta.JudgedCount
|
|
}
|
|
if delta.CreatedCount != 0 {
|
|
inc["created_count"] = delta.CreatedCount
|
|
}
|
|
if delta.TruncatedCount != 0 {
|
|
inc["truncated_count"] = delta.TruncatedCount
|
|
}
|
|
if delta.CreditsUsed != 0 {
|
|
inc["credits_used"] = delta.CreditsUsed
|
|
}
|
|
|
|
update := bson.M{}
|
|
if len(inc) > 0 {
|
|
update["$inc"] = inc
|
|
}
|
|
set := bson.M{}
|
|
if delta.FailedReason != nil {
|
|
set["failed_reason"] = strings.TrimSpace(*delta.FailedReason)
|
|
}
|
|
if delta.EndedAt > 0 {
|
|
set["ended_at"] = delta.EndedAt
|
|
}
|
|
if len(set) > 0 {
|
|
update["$set"] = set
|
|
}
|
|
ids := domain.MergeJudgedExternalIDs(nil, delta.JudgedExternalIDs)
|
|
if len(ids) > 0 {
|
|
update["$addToSet"] = bson.M{"judged_external_ids": bson.M{"$each": ids}}
|
|
}
|
|
if len(update) == 0 {
|
|
return s.GetSweep(ctx, id)
|
|
}
|
|
|
|
res, err := s.sweeps.UpdateOne(ctx, bson.M{"_id": id}, update)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if res.MatchedCount == 0 {
|
|
return nil, domain.ErrNotFound
|
|
}
|
|
return s.GetSweep(ctx, id)
|
|
}
|
|
|
|
func (s *MonStore) GetSweep(ctx context.Context, id string) (*domain.RadarSweep, error) {
|
|
var sw domain.RadarSweep
|
|
err := s.sweeps.FindOne(ctx, &sw, bson.M{"_id": id})
|
|
if err == mon.ErrNotFound {
|
|
return nil, domain.ErrNotFound
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &sw, nil
|
|
}
|
|
|
|
func (s *MonStore) GetSweepByJobID(ctx context.Context, jobID string) (*domain.RadarSweep, error) {
|
|
if strings.TrimSpace(jobID) == "" {
|
|
return nil, domain.ErrNotFound
|
|
}
|
|
var sw domain.RadarSweep
|
|
err := s.sweeps.FindOne(ctx, &sw, bson.M{"job_id": jobID})
|
|
if err == mon.ErrNotFound {
|
|
return nil, domain.ErrNotFound
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &sw, nil
|
|
}
|
|
|
|
func (s *MonStore) SetSweepPath(ctx context.Context, id, path string) error {
|
|
res, err := s.sweeps.UpdateOne(ctx, bson.M{"_id": id}, bson.M{"$set": bson.M{"path": path}})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if res.MatchedCount == 0 {
|
|
return domain.ErrNotFound
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *MonStore) ListSweeps(ctx context.Context, ownerUID int64, f domain.SweepListFilter) ([]*domain.RadarSweep, int64, error) {
|
|
q := bson.M{"owner_uid": ownerUID}
|
|
if f.WatchID != "" {
|
|
q["watch_id"] = f.WatchID
|
|
}
|
|
total, err := s.sweeps.CountDocuments(ctx, q)
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
page, ps := f.Page, f.PageSize
|
|
if page < 1 {
|
|
page = 1
|
|
}
|
|
if ps < 1 {
|
|
ps = 20
|
|
}
|
|
var list []*domain.RadarSweep
|
|
err = s.sweeps.Find(ctx, &list, q, options.Find().
|
|
SetSort(bson.D{{Key: "started_at", Value: -1}}).
|
|
SetSkip(int64((page-1)*ps)).
|
|
SetLimit(int64(ps)))
|
|
return list, total, err
|
|
}
|