package mongo import ( "context" "fwh-toolbox-server/database" "fwh-toolbox-server/domain" "fwh-toolbox-server/dto" "time" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/bson/primitive" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" ) const ( CollectionName_BatchDownloadConfig = "BatchDownloadConfig" ) // 从数据库获取批量下载数据列表 func GetBatchDownloadManageList(page, size uint, sn, progress *string) (*dto.ResPage[domain.BatchDownloadInfo], error) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) filter := bson.M{} // 条件 if sn != nil { filter["sn"] = *sn } if progress != nil { filter["progress"] = *progress } findOptions := options.Find(). SetProjection(bson.M{ "_id": 1, "sn": 1, "description": 1, "timePeriod": 1, "options": 1, "progress": 1, "failReasons": 1, "items": 1, "totalCount": 1, "doneCount": 1, }). SetSkip(int64((page - 1) * size)). SetLimit(int64(size)). SetSort(bson.D{ {Key: "_id", Value: -1}, // 倒序 }) cursor, err := collection.Find(ctx, filter, findOptions) if err != nil { return nil, err } defer cursor.Close(ctx) var configList []domain.BatchDownloadInfo for cursor.Next(ctx) { var cfg domain.BatchDownloadInfo if err := cursor.Decode(&cfg); err != nil { return nil, err } configList = append(configList, cfg) } // 查询总数 total, err := collection.CountDocuments(ctx, filter) if err != nil { return nil, err } res := dto.ResPage[domain.BatchDownloadInfo]{ Data: configList, PageSize: size, PageNum: page, Total: uint(total), } return &res, nil } // 从数据库获取一条批量下载数据配置 func GetBatchDownloadConfig(id primitive.ObjectID) (*domain.BatchDownloadConfig, error) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) var cfg domain.BatchDownloadConfig err := collection.FindOne(ctx, bson.M{"_id": id}).Decode(&cfg) if err != nil { return nil, err } return &cfg, nil } // 插入一条批量下载配置到数据库 func InsertBatchDownloadConfig(writeConfig domain.BatchDownloadConfigWrite) (primitive.ObjectID, error) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() config := domain.BatchDownloadConfig{ Sn: writeConfig.Sn, Description: writeConfig.Description, TimePeriod: writeConfig.TimePeriod, Options: writeConfig.Options, Progress: domain.ProgressPending, Data: map[string]domain.BatchDownloadData{}, FailReasons: []string{}, } collection := database.GetCollection(CollectionName_BatchDownloadConfig) res, err := collection.InsertOne(ctx, config) if err != nil { return primitive.ObjectID{}, err } return res.InsertedID.(primitive.ObjectID), nil } // 更新一条批量下载描述描述到数据库 func UpdateBatchDownloadDescription(id primitive.ObjectID, desc string) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) // _, err := collection.InsertOne(ctx, config) filter := bson.M{"_id": id} // 按ID匹配 update := bson.M{ // 如果存在就更新全部字段 "$set": bson.M{ "description": desc, }, } opts := options.Update().SetUpsert(true) _, err := collection.UpdateOne(ctx, filter, update, opts) if err != nil { return err } return nil } // 更新一条批量下载数据owner到数据库 func UpdateBatchDownloadOwner(id primitive.ObjectID, owner []string) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) // _, err := collection.InsertOne(ctx, config) filter := bson.M{"_id": id} // 按ID匹配 update := bson.M{ // 如果存在就更新全部字段 "$set": bson.M{ "owner": owner, }, } opts := options.Update().SetUpsert(true) _, err := collection.UpdateOne(ctx, filter, update, opts) if err != nil { return err } return nil } // 更新一条批量下载结果到数据库 func UpdateBatchDownloadResult(id primitive.ObjectID, result domain.Progress, reasons []string) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) // _, err := collection.InsertOne(ctx, config) filter := bson.M{"_id": id} // 按ID匹配 update := bson.M{ // 如果存在就更新非progress字段 "$set": bson.M{ "progress": result, "failReasons": reasons, }, } opts := options.Update().SetUpsert(true) _, err := collection.UpdateOne(ctx, filter, update, opts) if err != nil { return err } return nil } // 删除一条批量下载配置 func DeleteBatchDownloadConfig(id primitive.ObjectID) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() filter := bson.M{"_id": id} // 删除条件 collection := database.GetCollection(CollectionName_BatchDownloadConfig) _, err := collection.DeleteOne(ctx, filter) if err != nil { return err } return nil } // 插入批量下载结果数据 func InsertBatchDownloadData(id primitive.ObjectID, data domain.BatchDownloadData) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) // _, err := collection.InsertOne(ctx, config) filter := bson.M{"_id": id} // 按ID匹配 update := bson.M{ // 如果存在就更新data字段 "$set": bson.M{ "data": data, }, } opts := options.Update().SetUpsert(true) _, err := collection.UpdateOne(ctx, filter, update, opts) if err != nil { return err } return nil } // 获取下载数据包路径 func GetBatchDownloadData(id primitive.ObjectID) (map[string]domain.BatchDownloadData, error) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) var cfg domain.BatchDownloadConfig err := collection.FindOne(ctx, bson.M{"_id": id}).Decode(&cfg) if err != nil { return nil, err } return cfg.Data, nil } // 插入批量下载结果数据 - PowerFlow func InsertBatchDownloadPathPowerFlow(id primitive.ObjectID, sn string, date string, path string) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) // _, err := collection.InsertOne(ctx, config) filter := bson.M{"_id": id} // 按ID匹配 update := bson.M{ // 如果存在就更新data字段 "$set": bson.M{ "data." + sn + ".powerFlow." + date: path, }, } opts := options.Update().SetUpsert(true) _, err := collection.UpdateOne(ctx, filter, update, opts) if err != nil { return err } return nil } // 插入批量下载结果数据 - MqttLogs func InsertBatchDownloadPathMqttLogs(id primitive.ObjectID, sn string, date string, logType string, path string) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) // _, err := collection.InsertOne(ctx, config) filter := bson.M{"_id": id} // 按ID匹配 update := bson.M{ // 如果存在就更新data字段 "$set": bson.M{ "data." + sn + ".mqttLogs." + date + "." + logType: path, }, } opts := options.Update().SetUpsert(true) _, err := collection.UpdateOne(ctx, filter, update, opts) if err != nil { return err } return nil } // 插入批量下载结果数据 - EventAlarm func InsertBatchDownloadPathEventAlarm(id primitive.ObjectID, sn string, date string, path string) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) // _, err := collection.InsertOne(ctx, config) filter := bson.M{"_id": id} // 按ID匹配 update := bson.M{ // 如果存在就更新data字段 "$set": bson.M{ "data." + sn + ".eventAlarm." + date: path, }, } opts := options.Update().SetUpsert(true) _, err := collection.UpdateOne(ctx, filter, update, opts) if err != nil { return err } return nil } // 插入批量下载结果数据 - AlarmSnapshot func InsertBatchDownloadPathAlarmSnapshot(id primitive.ObjectID, sn string, date string, path string) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) // _, err := collection.InsertOne(ctx, config) filter := bson.M{"_id": id} // 按ID匹配 update := bson.M{ // 如果存在就更新data字段 "$set": bson.M{ "data." + sn + ".alarmSnapshot." + date: path, }, } opts := options.Update().SetUpsert(true) _, err := collection.UpdateOne(ctx, filter, update, opts) if err != nil { return err } return nil } // 插入批量下载结果数据 - EmsLogs func InsertBatchDownloadPathEmsLogs(id primitive.ObjectID, sn string, date string, logType string, path string) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) // _, err := collection.InsertOne(ctx, config) filter := bson.M{"_id": id} // 按ID匹配 update := bson.M{ // 如果存在就更新data字段 "$set": bson.M{ "data." + sn + ".emsLogs." + date + "." + logType: path, }, } opts := options.Update().SetUpsert(true) _, err := collection.UpdateOne(ctx, filter, update, opts) if err != nil { return err } return nil } // 插入批量下载结果数据 - FaultRecord func InsertBatchDownloadPathFaultRecord(id primitive.ObjectID, sn string, fileName string, path string) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) // _, err := collection.InsertOne(ctx, config) filter := bson.M{"_id": id} // 按ID匹配 update := bson.M{ // 如果存在就更新data字段 "$set": bson.M{ "data." + sn + ".faultRecord." + fileName: path, }, } opts := options.Update().SetUpsert(true) _, err := collection.UpdateOne(ctx, filter, update, opts) if err != nil { return err } return nil } // 插入批量下载结果数据 - UpdateLog func InsertBatchDownloadPathUpdateLog(id primitive.ObjectID, sn string, path string) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) // _, err := collection.InsertOne(ctx, config) filter := bson.M{"_id": id} // 按ID匹配 update := bson.M{ // 如果存在就更新data字段 "$set": bson.M{ "data." + sn + ".updateLog": path, }, } opts := options.Update().SetUpsert(true) _, err := collection.UpdateOne(ctx, filter, update, opts) if err != nil { return err } return nil } // 插入批量下载结果数据 - Settings func InsertBatchDownloadPathSettings(id primitive.ObjectID, sn string, path string) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) // _, err := collection.InsertOne(ctx, config) filter := bson.M{"_id": id} // 按ID匹配 update := bson.M{ // 如果存在就更新data字段 "$set": bson.M{ "data." + sn + ".settings": path, }, } opts := options.Update().SetUpsert(true) _, err := collection.UpdateOne(ctx, filter, update, opts) if err != nil { return err } return nil } // ==================== 结构化 item 结果(B2 完成判定) ==================== // InitBatchDownloadItems 初始化 items 列表,并设置 totalCount / doneCount。 // doneCount 初始化为已经处于终态的 item 数量(构建期就 failed/skipped 的), // 使得后续所有 item 转终态后 doneCount 恰好等于 totalCount。 func InitBatchDownloadItems(id primitive.ObjectID, items []domain.BatchDownloadItem) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) doneInit := 0 now := time.Now() for i := range items { if items[i].UpdatedAt.IsZero() { items[i].UpdatedAt = now } if items[i].Status != domain.ItemPending { doneInit++ } } filter := bson.M{"_id": id} update := bson.M{ "$set": bson.M{ "items": items, "totalCount": len(items), "doneCount": doneInit, "progress": domain.ProgressRunning, }, } opts := options.Update().SetUpsert(true) _, err := collection.UpdateOne(ctx, filter, update, opts) if err != nil { return err } return nil } // FinishBatchDownloadItem 原子地把某个 pending item 置为终态并递增 doneCount。 // 幂等:filter 要求该 item 当前为 pending,重复调用(Asynq 重试/重复回调)不会重复计数。 // 返回 (doneCount, totalCount, error);若该 item 已非 pending(幂等跳过),返回 (-1, -1, nil)。 func FinishBatchDownloadItem(id primitive.ObjectID, key string, status domain.ItemStatus, stage, reason string) (int, int, error) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) filter := bson.M{ "_id": id, "items": bson.M{ "$elemMatch": bson.M{"key": key, "status": domain.ItemPending}, }, } update := bson.M{ "$set": bson.M{ "items.$.status": status, "items.$.stage": stage, "items.$.reason": reason, "items.$.updatedAt": time.Now(), }, "$inc": bson.M{"doneCount": 1}, } opts := options.FindOneAndUpdate().SetReturnDocument(options.After) var doc domain.BatchDownloadConfig err := collection.FindOneAndUpdate(ctx, filter, update, opts).Decode(&doc) if err != nil { if err == mongo.ErrNoDocuments { // item 已非 pending(重复回调/重试),幂等跳过 return -1, -1, nil } return 0, 0, err } return doc.DoneCount, doc.TotalCount, nil } // GetBatchDownloadItems 读取某任务的全部 item 结果,用于聚合派生终态。 func GetBatchDownloadItems(id primitive.ObjectID) ([]domain.BatchDownloadItem, error) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() collection := database.GetCollection(CollectionName_BatchDownloadConfig) var cfg domain.BatchDownloadConfig err := collection.FindOne(ctx, bson.M{"_id": id}, options.FindOne().SetProjection(bson.M{"items": 1})).Decode(&cfg) if err != nil { return nil, err } return cfg.Items, nil }