dataDownloadRepo.go 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658
  1. package mongo
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "time"
  7. "fwh-toolbox-server/database"
  8. "fwh-toolbox-server/domain"
  9. "go.mongodb.org/mongo-driver/bson"
  10. "go.mongodb.org/mongo-driver/bson/primitive"
  11. "go.mongodb.org/mongo-driver/mongo"
  12. "go.mongodb.org/mongo-driver/mongo/options"
  13. )
  14. const (
  15. ddTasks = "data_download_tasks"
  16. ddItems = "data_download_items"
  17. ddDocuments = "data_download_documents"
  18. ddTaskDocs = "data_download_task_documents"
  19. ddRenderSlots = "data_download_render_slots"
  20. ddArtifacts = "data_download_artifacts"
  21. ddPackages = "data_download_package_jobs"
  22. ddFaultEntries = "data_download_fault_record_entries"
  23. ddFileAttempts = "data_download_file_attempts"
  24. ddFaultCalls = "data_download_fault_calls"
  25. )
  26. func InitDataDownloadIndexes(ctx context.Context) error {
  27. indexes := map[string][]mongo.IndexModel{
  28. ddTasks: {
  29. {Keys: bson.D{{Key: "ownerId", Value: 1}, {Key: "createdAt", Value: -1}}},
  30. {Keys: bson.D{{Key: "status", Value: 1}, {Key: "updatedAt", Value: 1}}},
  31. },
  32. ddItems: {
  33. {Keys: bson.D{{Key: "taskId", Value: 1}, {Key: "itemKey", Value: 1}}, Options: options.Index().SetUnique(true)},
  34. {Keys: bson.D{{Key: "taskId", Value: 1}, {Key: "status", Value: 1}}},
  35. },
  36. ddFileAttempts: {
  37. {Keys: bson.D{{Key: "taskId", Value: 1}, {Key: "itemId", Value: 1}, {Key: "sourceFileId", Value: 1}}, Options: options.Index().SetUnique(true)},
  38. {Keys: bson.D{{Key: "taskId", Value: 1}, {Key: "status", Value: 1}}},
  39. },
  40. ddFaultCalls: {
  41. {Keys: bson.D{{Key: "expiresAt", Value: 1}}, Options: options.Index().SetExpireAfterSeconds(0)},
  42. },
  43. ddDocuments: {
  44. {Keys: bson.D{{Key: "requestKey", Value: 1}}, Options: options.Index().SetUnique(true)},
  45. {Keys: bson.D{{Key: "kind", Value: 1}, {Key: "siteSn", Value: 1}, {Key: "date", Value: 1}}},
  46. },
  47. ddFaultEntries: {
  48. {Keys: bson.D{{Key: "documentId", Value: 1}, {Key: "sourceFileId", Value: 1}}, Options: options.Index().SetUnique(true)},
  49. {Keys: bson.D{{Key: "documentId", Value: 1}, {Key: "fileName", Value: 1}}},
  50. },
  51. ddTaskDocs: {
  52. {Keys: bson.D{{Key: "taskId", Value: 1}, {Key: "itemKey", Value: 1}, {Key: "documentId", Value: 1}}, Options: options.Index().SetUnique(true)},
  53. {Keys: bson.D{{Key: "documentId", Value: 1}}},
  54. },
  55. ddRenderSlots: {
  56. {Keys: bson.D{{Key: "documentId", Value: 1}, {Key: "format", Value: 1}}, Options: options.Index().SetUnique(true)},
  57. },
  58. ddArtifacts: {
  59. {Keys: bson.D{{Key: "documentId", Value: 1}, {Key: "format", Value: 1}, {Key: "contentGeneration", Value: 1}}},
  60. {Keys: bson.D{{Key: "expiresAt", Value: 1}}},
  61. },
  62. ddPackages: {
  63. {Keys: bson.D{{Key: "ownerId", Value: 1}, {Key: "createdAt", Value: -1}}},
  64. {Keys: bson.D{{Key: "taskIds", Value: 1}}},
  65. {Keys: bson.D{{Key: "purgeAt", Value: 1}}},
  66. },
  67. }
  68. for collection, models := range indexes {
  69. if _, err := database.GetCollection(collection).Indexes().CreateMany(ctx, models); err != nil {
  70. return fmt.Errorf("create %s indexes: %w", collection, err)
  71. }
  72. }
  73. return nil
  74. }
  75. func InsertDataDownloadTask(ctx context.Context, task *domain.DataDownloadTask) error {
  76. result, err := database.GetCollection(ddTasks).InsertOne(ctx, task)
  77. if err == nil {
  78. task.ID = result.InsertedID.(primitive.ObjectID)
  79. }
  80. return err
  81. }
  82. func GetDataDownloadTask(ctx context.Context, id primitive.ObjectID) (*domain.DataDownloadTask, error) {
  83. var task domain.DataDownloadTask
  84. err := database.GetCollection(ddTasks).FindOne(ctx, bson.M{"_id": id}).Decode(&task)
  85. return &task, err
  86. }
  87. func ListDataDownloadTasks(ctx context.Context, ownerID string, mineOnly bool, description, sn, status string, createdFrom, createdTo *time.Time, page, pageSize int64) ([]domain.DataDownloadTask, int64, error) {
  88. filter := bson.M{}
  89. if mineOnly {
  90. filter["ownerId"] = ownerID
  91. }
  92. if description != "" {
  93. filter["description"] = primitive.Regex{Pattern: regexpQuote(description), Options: "i"}
  94. }
  95. if sn != "" {
  96. filter["sns"] = sn
  97. }
  98. if status != "" {
  99. filter["status"] = status
  100. }
  101. if createdFrom != nil || createdTo != nil {
  102. rangeFilter := bson.M{}
  103. if createdFrom != nil {
  104. rangeFilter["$gte"] = *createdFrom
  105. }
  106. if createdTo != nil {
  107. rangeFilter["$lte"] = *createdTo
  108. }
  109. filter["createdAt"] = rangeFilter
  110. }
  111. coll := database.GetCollection(ddTasks)
  112. total, err := coll.CountDocuments(ctx, filter)
  113. if err != nil {
  114. return nil, 0, err
  115. }
  116. cursor, err := coll.Find(ctx, filter, options.Find().SetSort(bson.D{{Key: "createdAt", Value: -1}}).SetSkip((page-1)*pageSize).SetLimit(pageSize))
  117. if err != nil {
  118. return nil, 0, err
  119. }
  120. defer cursor.Close(ctx)
  121. tasks := make([]domain.DataDownloadTask, 0)
  122. err = cursor.All(ctx, &tasks)
  123. return tasks, total, err
  124. }
  125. func regexpQuote(value string) string {
  126. quoted := ""
  127. for _, r := range value {
  128. if r == '\\' || r == '.' || r == '+' || r == '*' || r == '?' || r == '(' || r == ')' || r == '[' || r == ']' || r == '{' || r == '}' || r == '^' || r == '$' || r == '|' {
  129. quoted += "\\"
  130. }
  131. quoted += string(r)
  132. }
  133. return quoted
  134. }
  135. func UpdateDataDownloadTaskStatus(ctx context.Context, id primitive.ObjectID, status domain.DataDownloadStatus) error {
  136. update := bson.M{"$set": bson.M{"status": status, "updatedAt": time.Now().UTC()}}
  137. if status == domain.DataDownloadSucceeded || status == domain.DataDownloadPartial || status == domain.DataDownloadFailed || status == domain.DataDownloadError || status == domain.DataDownloadCancelled {
  138. now := time.Now().UTC()
  139. update["$set"].(bson.M)["finishedAt"] = now
  140. }
  141. result, err := database.GetCollection(ddTasks).UpdateByID(ctx, id, update)
  142. if err == nil && result.MatchedCount != 1 {
  143. return mongo.ErrNoDocuments
  144. }
  145. return err
  146. }
  147. func InsertDataDownloadItems(ctx context.Context, items []domain.DataDownloadItem) error {
  148. if len(items) == 0 {
  149. return nil
  150. }
  151. models := make([]mongo.WriteModel, len(items))
  152. keys := make([]string, len(items))
  153. for i := range items {
  154. keys[i] = items[i].ItemKey
  155. models[i] = mongo.NewUpdateOneModel().
  156. SetFilter(bson.M{"taskId": items[i].TaskID, "itemKey": items[i].ItemKey}).
  157. SetUpdate(bson.M{"$setOnInsert": items[i]}).
  158. SetUpsert(true)
  159. }
  160. coll := database.GetCollection(ddItems)
  161. if _, err := coll.BulkWrite(ctx, models, options.BulkWrite().SetOrdered(false)); err != nil {
  162. return err
  163. }
  164. cursor, err := coll.Find(ctx, bson.M{"taskId": items[0].TaskID, "itemKey": bson.M{"$in": keys}})
  165. if err != nil {
  166. return err
  167. }
  168. defer cursor.Close(ctx)
  169. var persisted []domain.DataDownloadItem
  170. if err := cursor.All(ctx, &persisted); err != nil {
  171. return err
  172. }
  173. byKey := make(map[string]primitive.ObjectID, len(persisted))
  174. for i := range persisted {
  175. byKey[persisted[i].ItemKey] = persisted[i].ID
  176. }
  177. for i := range items {
  178. id, ok := byKey[items[i].ItemKey]
  179. if !ok {
  180. return fmt.Errorf("data-download item %q was not persisted", items[i].ItemKey)
  181. }
  182. items[i].ID = id
  183. }
  184. return nil
  185. }
  186. func GetDataDownloadItem(ctx context.Context, id primitive.ObjectID) (*domain.DataDownloadItem, error) {
  187. var item domain.DataDownloadItem
  188. err := database.GetCollection(ddItems).FindOne(ctx, bson.M{"_id": id}).Decode(&item)
  189. return &item, err
  190. }
  191. func ListDataDownloadItems(ctx context.Context, taskID primitive.ObjectID) ([]domain.DataDownloadItem, error) {
  192. cursor, err := database.GetCollection(ddItems).Find(ctx, bson.M{"taskId": taskID}, options.Find().SetSort(bson.D{{Key: "itemKey", Value: 1}}))
  193. if err != nil {
  194. return nil, err
  195. }
  196. defer cursor.Close(ctx)
  197. items := make([]domain.DataDownloadItem, 0)
  198. err = cursor.All(ctx, &items)
  199. return items, err
  200. }
  201. func UpdateDataDownloadItem(ctx context.Context, id primitive.ObjectID, status domain.DataDownloadStatus, stage, reason string) error {
  202. update := bson.M{"$set": bson.M{"status": status, "stage": stage, "failureReason": reason, "updatedAt": time.Now().UTC()}}
  203. if status == domain.DataDownloadRunning {
  204. update["$inc"] = bson.M{"attempt": 1}
  205. }
  206. result, err := database.GetCollection(ddItems).UpdateByID(ctx, id, update)
  207. if err == nil && result.MatchedCount != 1 {
  208. return mongo.ErrNoDocuments
  209. }
  210. return err
  211. }
  212. func UpdateDataDownloadFaultCall(ctx context.Context, id primitive.ObjectID, requestedAt *time.Time, requestedCount int, recordID int64, status *int) error {
  213. set := bson.M{"updatedAt": time.Now().UTC()}
  214. if requestedAt != nil {
  215. set["callRequestedAt"] = *requestedAt
  216. }
  217. if requestedCount > 0 {
  218. set["callRequestedCount"] = requestedCount
  219. }
  220. if recordID != 0 {
  221. set["callRecordId"] = recordID
  222. }
  223. if status != nil {
  224. set["callRecordStatus"] = *status
  225. }
  226. result, err := database.GetCollection(ddItems).UpdateByID(ctx, id, bson.M{"$set": set})
  227. if err == nil && result.MatchedCount != 1 {
  228. return mongo.ErrNoDocuments
  229. }
  230. return err
  231. }
  232. func ResetDataDownloadFaultCall(ctx context.Context, id primitive.ObjectID) error {
  233. result, err := database.GetCollection(ddItems).UpdateByID(ctx, id, bson.M{
  234. "$unset": bson.M{"callRequestedAt": "", "callRequestedCount": "", "callRecordId": "", "callRecordStatus": ""},
  235. "$set": bson.M{"updatedAt": time.Now().UTC()},
  236. })
  237. if err == nil && result.MatchedCount != 1 {
  238. return mongo.ErrNoDocuments
  239. }
  240. return err
  241. }
  242. // TryOwnDataDownloadFaultCall keeps the same logical owner across delayed
  243. // queue executions. A different item cannot start another call for the same
  244. // document until the owner releases it or the original 2-hour window expires.
  245. func TryOwnDataDownloadFaultCall(ctx context.Context, documentID string, itemID primitive.ObjectID, requestedAt time.Time) (bool, error) {
  246. now := time.Now().UTC()
  247. expiresAt := requestedAt.Add(2*time.Hour + 10*time.Minute)
  248. filter := bson.M{"_id": documentID, "$or": bson.A{bson.M{"itemId": itemID}, bson.M{"expiresAt": bson.M{"$lte": now}}}}
  249. update := bson.M{"$set": bson.M{"itemId": itemID, "requestedAt": requestedAt, "expiresAt": expiresAt, "updatedAt": now}}
  250. result, err := database.GetCollection(ddFaultCalls).UpdateOne(ctx, filter, update, options.Update().SetUpsert(true))
  251. if mongo.IsDuplicateKeyError(err) {
  252. return false, nil
  253. }
  254. if err != nil {
  255. return false, err
  256. }
  257. return result.MatchedCount == 1 || result.UpsertedCount == 1, nil
  258. }
  259. func OtherDataDownloadFaultCallActive(ctx context.Context, documentID string, itemID primitive.ObjectID) (bool, error) {
  260. var owner struct {
  261. ItemID primitive.ObjectID `bson:"itemId"`
  262. ExpiresAt time.Time `bson:"expiresAt"`
  263. }
  264. err := database.GetCollection(ddFaultCalls).FindOne(ctx, bson.M{"_id": documentID}).Decode(&owner)
  265. if errors.Is(err, mongo.ErrNoDocuments) {
  266. return false, nil
  267. }
  268. if err != nil {
  269. return false, err
  270. }
  271. return owner.ItemID != itemID && owner.ExpiresAt.After(time.Now().UTC()), nil
  272. }
  273. func ReleaseDataDownloadFaultCall(ctx context.Context, documentID string, itemID primitive.ObjectID) error {
  274. _, err := database.GetCollection(ddFaultCalls).DeleteOne(ctx, bson.M{"_id": documentID, "itemId": itemID})
  275. return err
  276. }
  277. func DeferDataDownloadItem(ctx context.Context, id primitive.ObjectID) error {
  278. result, err := database.GetCollection(ddItems).UpdateByID(ctx, id, bson.M{
  279. "$set": bson.M{"status": domain.DataDownloadPending, "stage": "waiting", "failureReason": "", "updatedAt": time.Now().UTC()},
  280. "$inc": bson.M{"attempt": -1},
  281. })
  282. if err == nil && result.MatchedCount != 1 {
  283. return mongo.ErrNoDocuments
  284. }
  285. return err
  286. }
  287. func RecountDataDownloadTask(ctx context.Context, taskID primitive.ObjectID) error {
  288. pipeline := mongo.Pipeline{
  289. {{Key: "$match", Value: bson.M{"taskId": taskID}}},
  290. {{Key: "$group", Value: bson.M{"_id": "$status", "count": bson.M{"$sum": 1}}}},
  291. }
  292. cursor, err := database.GetCollection(ddItems).Aggregate(ctx, pipeline)
  293. if err != nil {
  294. return err
  295. }
  296. defer cursor.Close(ctx)
  297. type count struct {
  298. ID domain.DataDownloadStatus `bson:"_id"`
  299. Count int64 `bson:"count"`
  300. }
  301. var counts []count
  302. if err := cursor.All(ctx, &counts); err != nil {
  303. return err
  304. }
  305. set := bson.M{"updatedAt": time.Now().UTC(), "succeeded": int64(0), "failed": int64(0), "empty": int64(0)}
  306. var total int64
  307. for _, c := range counts {
  308. total += c.Count
  309. switch c.ID {
  310. case domain.DataDownloadItemSucceeded, domain.DataDownloadSkipped, domain.DataDownloadPartial:
  311. set["succeeded"] = set["succeeded"].(int64) + c.Count
  312. case domain.DataDownloadFailed:
  313. set["failed"] = c.Count
  314. }
  315. }
  316. empty, err := database.GetCollection(ddTaskDocs).CountDocuments(ctx, bson.M{"taskId": taskID, "empty": true})
  317. if err != nil {
  318. return err
  319. }
  320. set["empty"] = empty
  321. set["total"] = total
  322. _, err = database.GetCollection(ddTasks).UpdateByID(ctx, taskID, bson.M{"$set": set})
  323. return err
  324. }
  325. func UpsertDataDownloadDocument(ctx context.Context, document *domain.DataDownloadDocument) error {
  326. _, err := database.GetCollection(ddDocuments).ReplaceOne(ctx, bson.M{"_id": document.ID}, document, options.Replace().SetUpsert(true))
  327. return err
  328. }
  329. func GetDataDownloadDocument(ctx context.Context, id string) (*domain.DataDownloadDocument, error) {
  330. var document domain.DataDownloadDocument
  331. err := database.GetCollection(ddDocuments).FindOne(ctx, bson.M{"_id": id}).Decode(&document)
  332. document.HasBinary = document.BinaryObjectKey != "" || document.FileCount > 0
  333. return &document, err
  334. }
  335. func ListDataDownloadDocuments(ctx context.Context, page, pageSize int64) ([]domain.DataDownloadDocument, int64, error) {
  336. coll := database.GetCollection(ddDocuments)
  337. total, err := coll.CountDocuments(ctx, bson.M{})
  338. if err != nil {
  339. return nil, 0, err
  340. }
  341. cursor, err := coll.Find(ctx, bson.M{}, options.Find().SetSort(bson.D{{Key: "updatedAt", Value: -1}}).SetSkip((page-1)*pageSize).SetLimit(pageSize))
  342. if err != nil {
  343. return nil, 0, err
  344. }
  345. defer cursor.Close(ctx)
  346. documents := make([]domain.DataDownloadDocument, 0)
  347. err = cursor.All(ctx, &documents)
  348. for i := range documents {
  349. documents[i].HasBinary = documents[i].BinaryObjectKey != "" || documents[i].FileCount > 0
  350. }
  351. return documents, total, err
  352. }
  353. func CountDataDownloadDocumentReferences(ctx context.Context, documentID string) (int64, error) {
  354. return database.GetCollection(ddTaskDocs).CountDocuments(ctx, bson.M{"documentId": documentID})
  355. }
  356. func ListDataDownloadArtifacts(ctx context.Context, documentID string) ([]domain.DataDownloadArtifact, error) {
  357. cursor, err := database.GetCollection(ddArtifacts).Find(ctx, bson.M{"documentId": documentID})
  358. if err != nil {
  359. return nil, err
  360. }
  361. defer cursor.Close(ctx)
  362. artifacts := make([]domain.DataDownloadArtifact, 0)
  363. err = cursor.All(ctx, &artifacts)
  364. return artifacts, err
  365. }
  366. func DeleteDataDownloadDocumentMetadata(ctx context.Context, documentID string) error {
  367. if _, err := database.GetCollection(ddRenderSlots).DeleteMany(ctx, bson.M{"documentId": documentID}); err != nil {
  368. return err
  369. }
  370. if _, err := database.GetCollection(ddArtifacts).DeleteMany(ctx, bson.M{"documentId": documentID}); err != nil {
  371. return err
  372. }
  373. if _, err := database.GetCollection(ddFaultEntries).DeleteMany(ctx, bson.M{"documentId": documentID}); err != nil {
  374. return err
  375. }
  376. result, err := database.GetCollection(ddDocuments).DeleteOne(ctx, bson.M{"_id": documentID})
  377. if err == nil && result.DeletedCount != 1 {
  378. return mongo.ErrNoDocuments
  379. }
  380. return err
  381. }
  382. func LinkDataDownloadDocument(ctx context.Context, link *domain.DataDownloadTaskDocument) error {
  383. _, err := database.GetCollection(ddTaskDocs).InsertOne(ctx, link)
  384. if mongo.IsDuplicateKeyError(err) {
  385. return nil
  386. }
  387. return err
  388. }
  389. func ListDataDownloadTaskDocuments(ctx context.Context, taskID primitive.ObjectID) ([]domain.DataDownloadTaskDocument, error) {
  390. cursor, err := database.GetCollection(ddTaskDocs).Find(ctx, bson.M{"taskId": taskID}, options.Find().SetSort(bson.D{{Key: "logicalPath", Value: 1}}))
  391. if err != nil {
  392. return nil, err
  393. }
  394. defer cursor.Close(ctx)
  395. links := make([]domain.DataDownloadTaskDocument, 0)
  396. err = cursor.All(ctx, &links)
  397. return links, err
  398. }
  399. func DetachDataDownloadRenderSlots(ctx context.Context, documentID string) error {
  400. _, err := database.GetCollection(ddRenderSlots).UpdateMany(ctx, bson.M{"documentId": documentID}, bson.M{"$unset": bson.M{"artifactId": ""}, "$set": bson.M{"updatedAt": time.Now().UTC()}})
  401. return err
  402. }
  403. func GetDataDownloadRenderSlot(ctx context.Context, documentID, format string) (*domain.DataDownloadRenderSlot, error) {
  404. var slot domain.DataDownloadRenderSlot
  405. err := database.GetCollection(ddRenderSlots).FindOne(ctx, bson.M{"documentId": documentID, "format": format}).Decode(&slot)
  406. return &slot, err
  407. }
  408. func GetDataDownloadArtifact(ctx context.Context, id primitive.ObjectID) (*domain.DataDownloadArtifact, error) {
  409. var artifact domain.DataDownloadArtifact
  410. err := database.GetCollection(ddArtifacts).FindOne(ctx, bson.M{"_id": id}).Decode(&artifact)
  411. return &artifact, err
  412. }
  413. func PublishDataDownloadArtifact(ctx context.Context, artifact *domain.DataDownloadArtifact) error {
  414. result, err := database.GetCollection(ddArtifacts).InsertOne(ctx, artifact)
  415. if err != nil {
  416. return err
  417. }
  418. artifact.ID = result.InsertedID.(primitive.ObjectID)
  419. _, err = database.GetCollection(ddRenderSlots).UpdateOne(ctx,
  420. bson.M{"documentId": artifact.DocumentID, "format": artifact.Format},
  421. bson.M{"$set": bson.M{"contentGeneration": artifact.ContentGeneration, "artifactId": artifact.ID, "updatedAt": time.Now().UTC()}},
  422. options.Update().SetUpsert(true))
  423. return err
  424. }
  425. func UpdateDataDownloadDescription(ctx context.Context, id primitive.ObjectID, description string) error {
  426. result, err := database.GetCollection(ddTasks).UpdateByID(ctx, id, bson.M{"$set": bson.M{"description": description, "updatedAt": time.Now().UTC()}})
  427. if err == nil && result.MatchedCount != 1 {
  428. return mongo.ErrNoDocuments
  429. }
  430. return err
  431. }
  432. func DeleteDataDownloadTask(ctx context.Context, id primitive.ObjectID) error {
  433. if _, err := database.GetCollection(ddTaskDocs).DeleteMany(ctx, bson.M{"taskId": id}); err != nil {
  434. return err
  435. }
  436. if _, err := database.GetCollection(ddItems).DeleteMany(ctx, bson.M{"taskId": id}); err != nil {
  437. return err
  438. }
  439. if _, err := database.GetCollection(ddFileAttempts).DeleteMany(ctx, bson.M{"taskId": id}); err != nil {
  440. return err
  441. }
  442. result, err := database.GetCollection(ddTasks).DeleteOne(ctx, bson.M{"_id": id})
  443. if err == nil && result.DeletedCount != 1 {
  444. return mongo.ErrNoDocuments
  445. }
  446. return err
  447. }
  448. func GetDataDownloadFileAttempt(ctx context.Context, taskID, itemID primitive.ObjectID, sourceFileID string) (*domain.DataDownloadFileAttempt, error) {
  449. var attempt domain.DataDownloadFileAttempt
  450. err := database.GetCollection(ddFileAttempts).FindOne(ctx, bson.M{"taskId": taskID, "itemId": itemID, "sourceFileId": sourceFileID}).Decode(&attempt)
  451. return &attempt, err
  452. }
  453. func RecordDataDownloadFileAttempt(ctx context.Context, attempt *domain.DataDownloadFileAttempt) error {
  454. now := time.Now().UTC()
  455. set := bson.M{
  456. "documentId": attempt.DocumentID, "fileName": attempt.FileName, "status": attempt.Status,
  457. "entryId": attempt.EntryID, "failureStage": attempt.FailureStage,
  458. "failureReason": attempt.FailureReason, "updatedAt": now,
  459. }
  460. _, err := database.GetCollection(ddFileAttempts).UpdateOne(ctx,
  461. bson.M{"taskId": attempt.TaskID, "itemId": attempt.ItemID, "sourceFileId": attempt.SourceFileID},
  462. bson.M{"$set": set, "$setOnInsert": bson.M{"createdAt": now}, "$inc": bson.M{"attemptCount": 1}},
  463. options.Update().SetUpsert(true))
  464. return err
  465. }
  466. func ListDataDownloadFileAttempts(ctx context.Context, taskID primitive.ObjectID) ([]domain.DataDownloadFileAttempt, error) {
  467. cursor, err := database.GetCollection(ddFileAttempts).Find(ctx, bson.M{"taskId": taskID}, options.Find().SetSort(bson.D{{Key: "itemId", Value: 1}, {Key: "sourceFileId", Value: 1}}))
  468. if err != nil {
  469. return nil, err
  470. }
  471. defer cursor.Close(ctx)
  472. attempts := make([]domain.DataDownloadFileAttempt, 0)
  473. err = cursor.All(ctx, &attempts)
  474. return attempts, err
  475. }
  476. func DeleteDataDownloadFileAttempt(ctx context.Context, taskID, itemID primitive.ObjectID, sourceFileID string) error {
  477. _, err := database.GetCollection(ddFileAttempts).DeleteOne(ctx, bson.M{
  478. "taskId": taskID, "itemId": itemID, "sourceFileId": sourceFileID,
  479. })
  480. return err
  481. }
  482. func HasDataDownloadTaskDocument(ctx context.Context, taskID primitive.ObjectID, documentID string) (bool, error) {
  483. count, err := database.GetCollection(ddTaskDocs).CountDocuments(ctx, bson.M{"taskId": taskID, "documentId": documentID})
  484. return count > 0, err
  485. }
  486. func UpsertDataDownloadFaultRecordEntry(ctx context.Context, entry *domain.DataDownloadFaultRecordEntry) (string, error) {
  487. var existing domain.DataDownloadFaultRecordEntry
  488. err := database.GetCollection(ddFaultEntries).FindOne(ctx, bson.M{"documentId": entry.DocumentID, "sourceFileId": entry.SourceFileID}).Decode(&existing)
  489. oldObjectKey := ""
  490. if err == nil {
  491. entry.ID = existing.ID
  492. entry.CreatedAt = existing.CreatedAt
  493. oldObjectKey = existing.ObjectKey
  494. } else if !errors.Is(err, mongo.ErrNoDocuments) {
  495. return "", err
  496. }
  497. if entry.ID.IsZero() {
  498. entry.ID = primitive.NewObjectID()
  499. }
  500. if entry.CreatedAt.IsZero() {
  501. entry.CreatedAt = entry.UpdatedAt
  502. }
  503. _, err = database.GetCollection(ddFaultEntries).ReplaceOne(ctx,
  504. bson.M{"documentId": entry.DocumentID, "sourceFileId": entry.SourceFileID}, entry, options.Replace().SetUpsert(true))
  505. return oldObjectKey, err
  506. }
  507. func GetDataDownloadFaultRecordEntry(ctx context.Context, documentID, sourceFileID string) (*domain.DataDownloadFaultRecordEntry, error) {
  508. var entry domain.DataDownloadFaultRecordEntry
  509. err := database.GetCollection(ddFaultEntries).FindOne(ctx, bson.M{"documentId": documentID, "sourceFileId": sourceFileID}).Decode(&entry)
  510. return &entry, err
  511. }
  512. func GetDataDownloadFaultRecordEntryByID(ctx context.Context, id primitive.ObjectID) (*domain.DataDownloadFaultRecordEntry, error) {
  513. var entry domain.DataDownloadFaultRecordEntry
  514. err := database.GetCollection(ddFaultEntries).FindOne(ctx, bson.M{"_id": id}).Decode(&entry)
  515. return &entry, err
  516. }
  517. func ListDataDownloadFaultRecordEntries(ctx context.Context, documentID string) ([]domain.DataDownloadFaultRecordEntry, error) {
  518. cursor, err := database.GetCollection(ddFaultEntries).Find(ctx, bson.M{"documentId": documentID}, options.Find().SetSort(bson.D{{Key: "fileName", Value: 1}, {Key: "sourceFileId", Value: 1}}))
  519. if err != nil {
  520. return nil, err
  521. }
  522. defer cursor.Close(ctx)
  523. entries := make([]domain.DataDownloadFaultRecordEntry, 0)
  524. err = cursor.All(ctx, &entries)
  525. return entries, err
  526. }
  527. func InsertDataDownloadPackageJob(ctx context.Context, job *domain.DataDownloadPackageJob) error {
  528. result, err := database.GetCollection(ddPackages).InsertOne(ctx, job)
  529. if err == nil {
  530. job.ID = result.InsertedID.(primitive.ObjectID)
  531. }
  532. return err
  533. }
  534. func GetDataDownloadPackageJob(ctx context.Context, id primitive.ObjectID) (*domain.DataDownloadPackageJob, error) {
  535. var job domain.DataDownloadPackageJob
  536. err := database.GetCollection(ddPackages).FindOne(ctx, bson.M{"_id": id}).Decode(&job)
  537. return &job, err
  538. }
  539. func ListDataDownloadPackageJobs(ctx context.Context, ownerID string, mineOnly bool, page, pageSize int64) ([]domain.DataDownloadPackageJob, int64, error) {
  540. filter := bson.M{}
  541. if mineOnly {
  542. filter["ownerId"] = ownerID
  543. }
  544. coll := database.GetCollection(ddPackages)
  545. total, err := coll.CountDocuments(ctx, filter)
  546. if err != nil {
  547. return nil, 0, err
  548. }
  549. cursor, err := coll.Find(ctx, filter, options.Find().SetSort(bson.D{{Key: "createdAt", Value: -1}}).SetSkip((page-1)*pageSize).SetLimit(pageSize))
  550. if err != nil {
  551. return nil, 0, err
  552. }
  553. defer cursor.Close(ctx)
  554. jobs := make([]domain.DataDownloadPackageJob, 0)
  555. err = cursor.All(ctx, &jobs)
  556. return jobs, total, err
  557. }
  558. func UpdateDataDownloadPackageJob(ctx context.Context, id primitive.ObjectID, set bson.M) error {
  559. set["updatedAt"] = time.Now().UTC()
  560. result, err := database.GetCollection(ddPackages).UpdateByID(ctx, id, bson.M{"$set": set})
  561. if err == nil && result.MatchedCount != 1 {
  562. return mongo.ErrNoDocuments
  563. }
  564. return err
  565. }
  566. func DeleteDataDownloadPackageJob(ctx context.Context, id primitive.ObjectID) error {
  567. result, err := database.GetCollection(ddPackages).DeleteOne(ctx, bson.M{"_id": id})
  568. if err == nil && result.DeletedCount != 1 {
  569. return mongo.ErrNoDocuments
  570. }
  571. return err
  572. }
  573. func FindExpiredDataDownloadPackageJobs(ctx context.Context, now time.Time) ([]domain.DataDownloadPackageJob, error) {
  574. cursor, err := database.GetCollection(ddPackages).Find(ctx, bson.M{"status": domain.DataDownloadSucceeded, "expiresAt": bson.M{"$lte": now}})
  575. if err != nil {
  576. return nil, err
  577. }
  578. defer cursor.Close(ctx)
  579. jobs := make([]domain.DataDownloadPackageJob, 0)
  580. err = cursor.All(ctx, &jobs)
  581. return jobs, err
  582. }
  583. func PurgeExpiredDataDownloadPackageJobs(ctx context.Context, now time.Time) error {
  584. _, err := database.GetCollection(ddPackages).DeleteMany(ctx, bson.M{"status": "expired", "purgeAt": bson.M{"$lte": now}})
  585. return err
  586. }
  587. func FindExpiredDataDownloadArtifacts(ctx context.Context, now time.Time) ([]domain.DataDownloadArtifact, error) {
  588. cursor, err := database.GetCollection(ddArtifacts).Find(ctx, bson.M{"expiresAt": bson.M{"$lte": now}})
  589. if err != nil {
  590. return nil, err
  591. }
  592. defer cursor.Close(ctx)
  593. artifacts := make([]domain.DataDownloadArtifact, 0)
  594. err = cursor.All(ctx, &artifacts)
  595. return artifacts, err
  596. }
  597. func DetachDataDownloadArtifact(ctx context.Context, artifact domain.DataDownloadArtifact) error {
  598. _, err := database.GetCollection(ddRenderSlots).UpdateMany(ctx, bson.M{"artifactId": artifact.ID}, bson.M{"$unset": bson.M{"artifactId": ""}, "$set": bson.M{"updatedAt": time.Now().UTC()}})
  599. return err
  600. }
  601. func DeleteDataDownloadArtifact(ctx context.Context, artifact domain.DataDownloadArtifact) error {
  602. _, err := database.GetCollection(ddArtifacts).DeleteOne(ctx, bson.M{"_id": artifact.ID})
  603. return err
  604. }