TASK-6345 - Cohort and sample double association is missing

opencga-app/src/main/java/org/opencb/opencga/app/migrations/v2/v2_12_6/syncCohortsAndSamplesMigration.java

@@ -0,0 +1,73 @@
+package org.opencb.opencga.app.migrations.v2.v2_12_6;
+
+import com.mongodb.client.MongoCollection;
+import com.mongodb.client.model.Filters;
+import com.mongodb.client.model.Projections;
+import com.mongodb.client.model.Updates;
+import org.apache.commons.collections4.CollectionUtils;
+import org.bson.Document;
+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.mongodb.MongoDBAdaptor;
+import org.opencb.opencga.catalog.db.mongodb.OrganizationMongoDBAdaptorFactory;
+import org.opencb.opencga.catalog.migration.Migration;
+import org.opencb.opencga.catalog.migration.MigrationTool;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+@Migration(id = "syncCohortsAndSamplesMigration" ,
+        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
+)
+public class syncCohortsAndSamplesMigration extends MigrationTool {
+
+    @Override
+    protected void run() throws Exception {
+        MongoCollection<Document> sampleCollection = getMongoCollection(OrganizationMongoDBAdaptorFactory.SAMPLE_COLLECTION);
+        MongoCollection<Document> sampleArchiveCollection = getMongoCollection(OrganizationMongoDBAdaptorFactory.SAMPLE_ARCHIVE_COLLECTION);
+
+        queryMongo(OrganizationMongoDBAdaptorFactory.COHORT_COLLECTION, new Document(),
+                Projections.include(CohortDBAdaptor.QueryParams.ID.key(), CohortDBAdaptor.QueryParams.SAMPLES.key()),
+                cohortDoc -> {
+                    String cohortId = cohortDoc.getString(CohortDBAdaptor.QueryParams.ID.key());
+                    List<Document> samples = cohortDoc.getList(CohortDBAdaptor.QueryParams.SAMPLES.key(), Document.class);
+                    if (CollectionUtils.isNotEmpty(samples)) {
+                        List<Long> sampleUids = samples
+                                .stream()
+                                .map(s -> s.get(SampleDBAdaptor.QueryParams.UID.key(), Number.class).longValue())
+                                .collect(Collectors.toList());
+                        // Ensure all those samples have a reference to the cohortId
+                        Bson query = Filters.and(
+                                Filters.in(SampleDBAdaptor.QueryParams.UID.key(), sampleUids),
+                                Filters.eq(MongoDBAdaptor.LAST_OF_VERSION, true)
+                        );
+                        Bson update = Updates.addToSet(SampleDBAdaptor.QueryParams.COHORT_IDS.key(), cohortId);
+                        long addedMissingCohort = sampleCollection.updateMany(query, update).getModifiedCount();
+                        sampleArchiveCollection.updateMany(query, update);
+
+                        // Ensure there aren't any samples pointing to this cohort that are not in the samples array
+                        query = Filters.and(
+                                Filters.nin(SampleDBAdaptor.QueryParams.UID.key(), sampleUids),
+                                Filters.eq(SampleDBAdaptor.QueryParams.COHORT_IDS.key(), cohortId),
+                                Filters.eq(MongoDBAdaptor.LAST_OF_VERSION, true)
+                        );
+                        update = Updates.pull(SampleDBAdaptor.QueryParams.COHORT_IDS.key(), cohortId);
+                        long removedNonAssociatedCohort = sampleCollection.updateMany(query, update).getModifiedCount();
+                        sampleArchiveCollection.updateMany(query, update);
+
+                        if (addedMissingCohort > 0 || removedNonAssociatedCohort > 0) {
+                            logger.info("Fixed cohort '{}' references. "
+                                            + "Added missing reference to {} samples. "
+                                            + "Removed non-associated reference from {} samples.",
+                                    cohortId, addedMissingCohort, removedNonAssociatedCohort);
+                        }
+                    }
+                });
+    }
+
+}

opencga-catalog/src/main/java/org/opencb/opencga/catalog/db/mongodb/CohortMongoDBAdaptor.java

@@ -235,22 +235,21 @@ public OpenCGAResult update(long cohortId, ObjectMap parameters, QueryOptions qu
     @Override
     public OpenCGAResult update(long cohortUid, ObjectMap parameters, List<VariableSet> variableSetList, QueryOptions queryOptions)
             throws CatalogDBException, CatalogParameterException, CatalogAuthorizationException {
-        Query query = new Query(QueryParams.UID.key(), cohortUid);
-        QueryOptions options = new QueryOptions(QueryOptions.INCLUDE,
-                Arrays.asList(QueryParams.ID.key(), QueryParams.UID.key(), QueryParams.STUDY_UID.key(),
-                        QueryParams.SAMPLES.key() + "." + QueryParams.ID.key()));
-        OpenCGAResult<Cohort> documentResult = get(query, options);
-        if (documentResult.getNumResults() == 0) {
-            throw new CatalogDBException("Could not update cohort. Cohort uid '" + cohortUid + "' not found.");
-        }
-        String cohortId = documentResult.first().getId();
-
         try {
-            return runTransaction(clientSession -> transactionalUpdate(clientSession, documentResult.first(), parameters, variableSetList,
-                    queryOptions));
-        } catch (CatalogDBException e) {
-            logger.error("Could not update cohort {}: {}", cohortId, e.getMessage(), e);
-            throw new CatalogDBException("Could not update cohort " + cohortId + ": " + e.getMessage(), e.getCause());
+            return runTransaction(clientSession -> {
+                Query query = new Query(QueryParams.UID.key(), cohortUid);
+                QueryOptions options = new QueryOptions(QueryOptions.INCLUDE,
+                        Arrays.asList(QueryParams.ID.key(), QueryParams.UID.key(), QueryParams.STUDY_UID.key(),
+                                QueryParams.SAMPLES.key() + "." + QueryParams.ID.key()));
+                OpenCGAResult<Cohort> documentResult = get(clientSession, query, options);
+                if (documentResult.getNumResults() == 0) {
+                    throw new CatalogDBException("Could not update cohort. Cohort uid '" + cohortUid + "' not found.");
+                }
+                return transactionalUpdate(clientSession, documentResult.first(), parameters, variableSetList, queryOptions);
+            });
+        } catch (Exception e) {
+            logger.error("Could not update cohort {}: {}", cohortUid, e.getMessage(), e);
+            throw new CatalogDBException("Could not update cohort " + cohortUid + ": " + e.getMessage(), e);
         }
     }
 

opencga-catalog/src/test/java/org/opencb/opencga/catalog/managers/CatalogManagerTest.java

@@ -16,10 +16,12 @@
 
 package org.opencb.opencga.catalog.managers;
 
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import com.mongodb.BasicDBObject;
 import org.apache.commons.collections4.CollectionUtils;
 import org.apache.commons.lang3.RandomStringUtils;
 import org.apache.commons.lang3.StringUtils;
+import org.apache.commons.lang3.time.StopWatch;
 import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
@@ -66,6 +68,9 @@
 import javax.naming.NamingException;
 import java.io.IOException;
 import java.util.*;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
 import static org.hamcrest.CoreMatchers.allOf;
@@ -1650,6 +1655,88 @@ public void updateSampleCohortTest() throws Exception {
         }
     }
 
+    @Test
+    public void updateSampleCohortWithThreadsTest() throws Exception {
+        Sample sampleId1 = catalogManager.getSampleManager().create(studyFqn, new Sample().setId("SAMPLE_1"), INCLUDE_RESULT, ownerToken).first();
+        Sample sampleId2 = catalogManager.getSampleManager().create(studyFqn, new Sample().setId("SAMPLE_2"), INCLUDE_RESULT, ownerToken).first();
+        Sample sampleId3 = catalogManager.getSampleManager().create(studyFqn, new Sample().setId("SAMPLE_3"), INCLUDE_RESULT, ownerToken).first();
+        catalogManager.getCohortManager().create(studyFqn, new Cohort().setId("MyCohort1")
+                .setSamples(Arrays.asList(sampleId1)), null, ownerToken).first();
+        catalogManager.getCohortManager().create(studyFqn, new Cohort().setId("MyCohort2")
+                .setSamples(Arrays.asList(sampleId2, sampleId3)), null, ownerToken).first();
+
+        ExecutorService executorService = Executors.newFixedThreadPool(10,
+                new ThreadFactoryBuilder()
+                        .setNameFormat("executor-service-%d")
+                        .build());
+
+        StopWatch stopWatch = StopWatch.createStarted();
+        List<List<String>> sampleIds = new ArrayList<>(5);
+        List<String> innerArray = new ArrayList<>(50);
+        for (int i = 0; i < 250; i++) {
+            if (i % 50 == 0) {
+                System.out.println("i = " + i);
+            }
+
+            String sampleId = "SAMPLE_AUTO_" + i;
+            executorService.submit(() -> {
+                try {
+                    catalogManager.getSampleManager().create(studyFqn, new Sample().setId(sampleId), QueryOptions.empty(), ownerToken);
+                } catch (CatalogException e) {
+                    throw new RuntimeException(e);
+                }
+            });
+            if (innerArray.size() == 50) {
+                sampleIds.add(new ArrayList<>(innerArray));
+                innerArray.clear();
+            }
+            innerArray.add(sampleId);
+        }
+        sampleIds.add(new ArrayList<>(innerArray));
+        executorService.shutdown();
+        executorService.awaitTermination(1, TimeUnit.MINUTES);
+
+        System.out.println("Creating 250 samples took " + stopWatch.getTime(TimeUnit.SECONDS) + " seconds");
+
+        stopWatch.stop();
+        stopWatch.reset();
+        stopWatch.start();
+        executorService = Executors.newFixedThreadPool(3);
+        int execution = 0;
+        Map<String, Object> actionMap = new HashMap<>();
+        actionMap.put(CohortDBAdaptor.QueryParams.SAMPLES.key(), ParamUtils.BasicUpdateAction.SET);
+        QueryOptions queryOptions = new QueryOptions();
+        queryOptions.put(Constants.ACTIONS, actionMap);
+        for (List<String> innerSampleIds : sampleIds) {
+            Cohort myCohort1 = catalogManager.getCohortManager().get(studyFqn, "MyCohort1", null, ownerToken).first();
+            List<SampleReferenceParam> sampleReferenceParamList = new ArrayList<>(myCohort1.getNumSamples() + innerSampleIds.size());
+            sampleReferenceParamList.addAll(myCohort1.getSamples().stream().map(s -> new SampleReferenceParam().setId(s.getId())).collect(Collectors.toList()));
+            sampleReferenceParamList.addAll(innerSampleIds.stream().map(s -> new SampleReferenceParam().setId(s)).collect(Collectors.toList()));
+            int executionId = execution++;
+            executorService.submit(() -> {
+                try {
+                    catalogManager.getCohortManager().update(studyFqn, "MyCohort1",
+                            new CohortUpdateParams().setSamples(sampleReferenceParamList),
+                            queryOptions, ownerToken);
+                    System.out.println("Execution: " + executionId);
+                } catch (CatalogException e) {
+                    throw new RuntimeException(e);
+                }
+            });
+        }
+        executorService.shutdown();
+        executorService.awaitTermination(1, TimeUnit.MINUTES);
+        System.out.println("Attaching 250 samples took " + stopWatch.getTime(TimeUnit.SECONDS) + " seconds");
+
+        // Ensure persistence
+        Query sampleQuery = new Query(SampleDBAdaptor.QueryParams.COHORT_IDS.key(), "MyCohort1");
+        OpenCGAResult<Sample> search = catalogManager.getSampleManager().search(studyFqn, sampleQuery, SampleManager.INCLUDE_SAMPLE_IDS, ownerToken);
+        Cohort myCohort1 = catalogManager.getCohortManager().get(studyFqn, "MyCohort1", null, ownerToken).first();
+        assertEquals(search.getNumResults(), myCohort1.getNumSamples());
+        Set<String> sampleIdSet = search.getResults().stream().map(Sample::getId).collect(Collectors.toSet());
+        assertTrue(myCohort1.getSamples().stream().map(Sample::getId).collect(Collectors.toSet()).containsAll(sampleIdSet));
+    }
+
     @Test
     public void deleteSampleCohortTest() throws Exception {
         Sample sampleId1 = catalogManager.getSampleManager().create(studyFqn, new Sample().setId("SAMPLE_1"), INCLUDE_RESULT, ownerToken).first();