From 724afa9acd8141c7794c0e4df33d2e2fc17a8b7f Mon Sep 17 00:00:00 2001
From: pfurio <pedrofurio@gmail.com>
Date: Thu, 26 Sep 2024 16:43:11 +0200
Subject: [PATCH] app: fix cohort-sample reference migration, #TASK-6998

---
 .../SyncCohortsAndSamplesMigration.java       | 27 +++++++++++++++----
 1 file changed, 22 insertions(+), 5 deletions(-)

diff --git a/opencga-app/src/main/java/org/opencb/opencga/app/migrations/v2_12_6/SyncCohortsAndSamplesMigration.java b/opencga-app/src/main/java/org/opencb/opencga/app/migrations/v2_12_6/SyncCohortsAndSamplesMigration.java
index edfdbb7c82..d0d0cc931c 100644
--- a/opencga-app/src/main/java/org/opencb/opencga/app/migrations/v2_12_6/SyncCohortsAndSamplesMigration.java
+++ b/opencga-app/src/main/java/org/opencb/opencga/app/migrations/v2_12_6/SyncCohortsAndSamplesMigration.java
@@ -9,20 +9,24 @@
 import org.bson.conversions.Bson;
 import org.opencb.opencga.catalog.db.api.CohortDBAdaptor;
 import org.opencb.opencga.catalog.db.api.SampleDBAdaptor;
+import org.opencb.opencga.catalog.db.api.StudyDBAdaptor;
 import org.opencb.opencga.catalog.db.mongodb.MongoDBAdaptor;
 import org.opencb.opencga.catalog.db.mongodb.MongoDBAdaptorFactory;
 import org.opencb.opencga.catalog.migration.Migration;
 import org.opencb.opencga.catalog.migration.MigrationTool;
 
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.stream.Collectors;
 
 @Migration(id = "syncCohortsAndSamplesMigration" ,
-        description = "Sync array of samples from cohort with array of cohortIds from Sample",
+        description = "Sync array of samples from cohort with array of cohortIds from Sample.",
         version = "2.12.6",
         domain = Migration.MigrationDomain.CATALOG,
         language = Migration.MigrationLanguage.JAVA,
-        date = 20240621
+        date = 20240621,
+        patch = 2 // TASK-6998
 )
 public class SyncCohortsAndSamplesMigration extends MigrationTool {
 
@@ -31,8 +35,18 @@ protected void run() throws Exception {
         MongoCollection<Document> sampleCollection = getMongoCollection(MongoDBAdaptorFactory.SAMPLE_COLLECTION);
         MongoCollection<Document> sampleArchiveCollection = getMongoCollection(MongoDBAdaptorFactory.SAMPLE_ARCHIVE_COLLECTION);
 
+        // Fill map study uid - fqn
+        Map<Long, String> uidFqnMap = new HashMap<>();
+        Bson studyProjection = Projections.include(StudyDBAdaptor.QueryParams.UID.key(), StudyDBAdaptor.QueryParams.FQN.key());
+        queryMongo(MongoDBAdaptorFactory.STUDY_COLLECTION, new Document(), studyProjection, study -> {
+            long studyUid = study.get(StudyDBAdaptor.QueryParams.UID.key(), Number.class).longValue();
+            String studyFqn = study.getString(StudyDBAdaptor.QueryParams.FQN.key());
+            uidFqnMap.put(studyUid, studyFqn);
+        });
+
         queryMongo(MongoDBAdaptorFactory.COHORT_COLLECTION, new Document(),
-                Projections.include(CohortDBAdaptor.QueryParams.ID.key(), CohortDBAdaptor.QueryParams.SAMPLES.key()),
+                Projections.include(CohortDBAdaptor.QueryParams.ID.key(), CohortDBAdaptor.QueryParams.SAMPLES.key(),
+                        CohortDBAdaptor.QueryParams.STUDY_UID.key()),
                 cohortDoc -> {
                     String cohortId = cohortDoc.getString(CohortDBAdaptor.QueryParams.ID.key());
                     List<Document> samples = cohortDoc.getList(CohortDBAdaptor.QueryParams.SAMPLES.key(), Document.class);
@@ -50,8 +64,11 @@ protected void run() throws Exception {
                         long addedMissingCohort = sampleCollection.updateMany(query, update).getModifiedCount();
                         sampleArchiveCollection.updateMany(query, update);
 
+                        long studyUid = cohortDoc.get(CohortDBAdaptor.QueryParams.STUDY_UID.key(), Number.class).longValue();
+
                         // Ensure there aren't any samples pointing to this cohort that are not in the samples array
                         query = Filters.and(
+                                Filters.eq(SampleDBAdaptor.QueryParams.STUDY_UID.key(), studyUid),
                                 Filters.nin(SampleDBAdaptor.QueryParams.UID.key(), sampleUids),
                                 Filters.eq(SampleDBAdaptor.QueryParams.COHORT_IDS.key(), cohortId),
                                 Filters.eq(MongoDBAdaptor.LAST_OF_VERSION, true)
@@ -61,10 +78,10 @@ protected void run() throws Exception {
                         sampleArchiveCollection.updateMany(query, update);
 
                         if (addedMissingCohort > 0 || removedNonAssociatedCohort > 0) {
-                            logger.info("Fixed cohort '{}' references. "
+                            logger.info("Fixed cohort '{}' references from study '{}'. "
                                             + "Added missing reference to {} samples. "
                                             + "Removed non-associated reference from {} samples.",
-                                    cohortId, addedMissingCohort, removedNonAssociatedCohort);
+                                    cohortId, uidFqnMap.get(studyUid), addedMissingCohort, removedNonAssociatedCohort);
                         }
                     }
                 });