[BI-1831] Add DatasetID to ObservationUnits

src/main/java/org/breedinginsight/brapi/v2/services/BrAPITrialService.java

@@ -296,18 +296,14 @@ private StreamedFile zipFiles(List<DownloadFile> files) throws IOException {
     }
 
     public Dataset getDatasetData(Program program, UUID experimentId, UUID datsetId, Boolean stats) throws ApiException, DoesNotExistException {
-        BrAPITrial experiment = this.getExperiment(program, experimentId);
-
-        // TODO: Once BI-1831 is complete and OUs in a dataset can be identified using the datasetId stored as a xref
-        // the expOUs needs to be replaced with datasetOUs, as was done with datasetObsVars
-        List<BrAPIObservationUnit> expOUs = ouDAO.getObservationUnitsForTrialDbId(program.getId(), experiment.getTrialDbId(), true);
+        List<BrAPIObservationUnit> datasetOUs = ouDAO.getObservationUnitsForDataset(datsetId.toString(), program);
         List<BrAPIObservationVariable> datasetObsVars = getDatasetObsVars(datsetId.toString(), program);
-        List<String> ouDbIds = expOUs.stream().map(BrAPIObservationUnit::getObservationUnitDbId).collect(Collectors.toList());
+        List<String> ouDbIds = datasetOUs.stream().map(BrAPIObservationUnit::getObservationUnitDbId).collect(Collectors.toList());
         List<String> obsVarDbIds = datasetObsVars.stream().map(BrAPIObservationVariable::getObservationVariableDbId).collect(Collectors.toList());
         List<BrAPIObservation> data = observationDAO.getObservationsByObservationUnitsAndVariables(ouDbIds, obsVarDbIds, program);
-        Dataset dataset = new Dataset(experimentId.toString(), data, expOUs, datasetObsVars);
+        Dataset dataset = new Dataset(experimentId.toString(), data, datasetOUs, datasetObsVars);
         if (stats) {
-            Integer ouCount = expOUs.size();
+            Integer ouCount = datasetOUs.size();
             Integer obsVarCount = datasetObsVars.size();
             Integer obsCount = ouCount * obsVarCount;
             Integer obsDataCount = data.size();

src/main/java/org/breedinginsight/brapps/importer/daos/BrAPIObservationUnitDAO.java

@@ -23,8 +23,10 @@
 import com.google.gson.reflect.TypeToken;
 import io.micronaut.context.annotation.Property;
 import org.brapi.client.v2.JSON;
+import io.micronaut.http.server.exceptions.InternalServerException;
 import org.brapi.client.v2.model.exceptions.ApiException;
 import org.brapi.client.v2.modules.phenotype.ObservationUnitsApi;
+import org.brapi.v2.model.BrAPIExternalReference;
 import org.brapi.v2.model.germ.BrAPIGermplasm;
 import org.brapi.v2.model.pheno.BrAPIObservationUnit;
 import org.brapi.v2.model.pheno.request.BrAPIObservationUnitSearchRequest;
@@ -42,6 +44,7 @@
 import org.breedinginsight.services.brapi.BrAPIEndpointProvider;
 import org.breedinginsight.services.exceptions.DoesNotExistException;
 import org.breedinginsight.utilities.BrAPIDAOUtil;
+import org.breedinginsight.utilities.Utilities;
 
 import javax.inject.Inject;
 import javax.inject.Singleton;
@@ -64,7 +67,13 @@ public class BrAPIObservationUnitDAO {
     private final Type treatmentlistType = new TypeToken<ArrayList<BrAPIObservationTreatment>>(){}.getType();
 
     @Inject
-    public BrAPIObservationUnitDAO(ProgramDAO programDAO, ImportDAO importDAO, BrAPIDAOUtil brAPIDAOUtil, BrAPIEndpointProvider brAPIEndpointProvider, BrAPIGermplasmService germplasmService, ProgramService programService, @Property(name = "brapi.server.reference-source") String referenceSource) {
+    public BrAPIObservationUnitDAO(ProgramDAO programDAO,
+                                   ImportDAO importDAO,
+                                   BrAPIDAOUtil brAPIDAOUtil,
+                                   BrAPIEndpointProvider brAPIEndpointProvider,
+                                   BrAPIGermplasmService germplasmService,
+                                   ProgramService programService,
+                                   @Property(name = "brapi.server.reference-source") String referenceSource) {
         this.programDAO = programDAO;
         this.importDAO = importDAO;
         this.brAPIDAOUtil = brAPIDAOUtil;
@@ -97,6 +106,19 @@ public List<BrAPIObservationUnit> createBrAPIObservationUnits(List<BrAPIObservat
         return ous;
     }
 
+    public BrAPIObservationUnit updateBrAPIObservationUnit(String ouDbId, BrAPIObservationUnit ou, UUID programId) {
+        ObservationUnitsApi api = brAPIEndpointProvider.get(programDAO.getCoreClient(programId), ObservationUnitsApi.class);
+        BrAPIObservationUnit updatedOu = null;
+        try {
+            if (ou != null && !ouDbId.isBlank()) {
+                updatedOu = brAPIDAOUtil.put(ouDbId, ou, api::observationunitsObservationUnitDbIdPut);
+            }
+            return updatedOu;
+        } catch (Exception e) {
+            throw new InternalServerException("Unknown error has occurred: " + e.getMessage(), e);
+        }
+    }
+
     public List<BrAPIObservationUnit> getObservationUnitsById(Collection<String> observationUnitExternalIds, Program program) throws ApiException {
         if(observationUnitExternalIds.isEmpty()) {
             return Collections.emptyList();
@@ -124,6 +146,18 @@ public List<BrAPIObservationUnit> getObservationUnitsForTrialDbId(@NotNull UUID
         return getObservationUnitsForTrialDbId(programId, trialDbId, false);
     }
 
+    public List<BrAPIObservationUnit> getObservationUnitsForDataset(@NotNull String datasetId, @NotNull Program program) throws ApiException {
+        String datasetReferenceSource = Utilities.generateReferenceSource(referenceSource, ExternalReferenceSource.DATASET);
+        BrAPIObservationUnitSearchRequest ouSearchRequest = new BrAPIObservationUnitSearchRequest();
+        ouSearchRequest.programDbIds(List.of(program.getBrapiProgram().getProgramDbId()));
+        ouSearchRequest.externalReferenceSources(List.of(datasetReferenceSource));
+        ouSearchRequest.externalReferenceIDs(List.of(datasetId));
+        ObservationUnitsApi api = brAPIEndpointProvider.get(programDAO.getCoreClient(program.getId()), ObservationUnitsApi.class);
+        return brAPIDAOUtil.search(api::searchObservationunitsPost,
+                api::searchObservationunitsSearchResultsDbIdGet,
+                ouSearchRequest);
+    }
+
     public List<BrAPIObservationUnit> getObservationUnitsForTrialDbId(@NotNull UUID programId, @NotNull String trialDbId, boolean withGID) throws ApiException, DoesNotExistException {
         Program program = programService.getById(programId).orElseThrow(() -> new DoesNotExistException("Program id does not exist"));
 

src/main/java/org/breedinginsight/brapps/importer/model/imports/experimentObservation/ExperimentObservation.java

@@ -223,6 +223,7 @@ public BrAPIObservationUnit constructBrAPIObservationUnit(
             String germplasmName,
             String referenceSource,
             UUID trialID,
+            UUID datasetId,
             UUID studyID,
             UUID id
     ) {
@@ -232,7 +233,7 @@ public BrAPIObservationUnit constructBrAPIObservationUnit(
             observationUnit.setObservationUnitName(Utilities.appendProgramKey(getExpUnitId(), program.getKey(), seqVal));
 
             // Set external reference
-            observationUnit.setExternalReferences(getObsUnitExternalReferences(program, referenceSource, trialID, studyID, id));
+            observationUnit.setExternalReferences(getObsUnitExternalReferences(program, referenceSource, trialID, datasetId, studyID, id));
         } else {
             observationUnit.setObservationUnitName(getExpUnitId());
         }
@@ -345,13 +346,16 @@ public BrAPIObservation constructBrAPIObservation(
     }
 
     private List<BrAPIExternalReference> getBrAPIExternalReferences(
-            Program program, String referenceSourceBaseName, UUID trialId, UUID studyId, UUID obsUnitId, UUID observationId) {
+            Program program, String referenceSourceBaseName, UUID trialId, UUID datasetId, UUID studyId, UUID obsUnitId, UUID observationId) {
         List<BrAPIExternalReference> refs = new ArrayList<>();
 
         addReference(refs, program.getId(), referenceSourceBaseName, ExternalReferenceSource.PROGRAMS);
         if (trialId != null) {
             addReference(refs, trialId, referenceSourceBaseName, ExternalReferenceSource.TRIALS);
         }
+        if (datasetId != null) {
+            addReference(refs, datasetId, referenceSourceBaseName, ExternalReferenceSource.DATASET);
+        }
         if (studyId != null) {
             addReference(refs, studyId, referenceSourceBaseName, ExternalReferenceSource.STUDIES);
         }
@@ -367,22 +371,22 @@ private List<BrAPIExternalReference> getBrAPIExternalReferences(
 
     private List<BrAPIExternalReference> getTrialExternalReferences(
             Program program, String referenceSourceBaseName, UUID trialId) {
-        return getBrAPIExternalReferences(program, referenceSourceBaseName, trialId, null, null, null);
+        return getBrAPIExternalReferences(program, referenceSourceBaseName, trialId, null, null, null, null);
     }
 
     private List<BrAPIExternalReference> getStudyExternalReferences(
             Program program, String referenceSourceBaseName, UUID trialId, UUID studyId) {
-        return getBrAPIExternalReferences(program, referenceSourceBaseName, trialId, studyId, null, null);
+        return getBrAPIExternalReferences(program, referenceSourceBaseName, trialId, null, studyId, null, null);
     }
 
     private List<BrAPIExternalReference> getObsUnitExternalReferences(
-            Program program, String referenceSourceBaseName, UUID trialId, UUID studyId, UUID obsUnitId) {
-        return getBrAPIExternalReferences(program, referenceSourceBaseName, trialId, studyId, obsUnitId, null);
+            Program program, String referenceSourceBaseName, UUID trialId, UUID datasetId, UUID studyId, UUID obsUnitId) {
+        return getBrAPIExternalReferences(program, referenceSourceBaseName, trialId, datasetId, studyId, obsUnitId, null);
     }
 
     private List<BrAPIExternalReference> getObservationExternalReferences(
             Program program, String referenceSourceBaseName, UUID trialId, UUID studyId, UUID obsUnitId, UUID observationId) {
-        return getBrAPIExternalReferences(program, referenceSourceBaseName, trialId, studyId, obsUnitId, observationId);
+        return getBrAPIExternalReferences(program, referenceSourceBaseName, trialId, null, studyId, obsUnitId, observationId);
     }
 
 

src/main/java/org/breedinginsight/brapps/importer/services/processors/ExperimentProcessor.java

@@ -41,6 +41,7 @@
 import org.breedinginsight.brapi.v2.dao.BrAPIGermplasmDAO;
 import org.breedinginsight.brapps.importer.daos.*;
 import org.breedinginsight.brapps.importer.model.ImportUpload;
+import org.breedinginsight.brapps.importer.model.base.AdditionalInfo;
 import org.breedinginsight.brapps.importer.model.imports.BrAPIImport;
 import org.breedinginsight.brapps.importer.model.imports.PendingImport;
 import org.breedinginsight.brapps.importer.model.imports.experimentObservation.ExperimentObservation;
@@ -805,10 +806,16 @@ private PendingImportObject<BrAPIObservationUnit> fetchOrCreateObsUnitPIO(Progra
             }
             PendingImportObject<BrAPITrial> trialPIO = this.trialByNameNoScope.get(importRow.getExpTitle());
             UUID trialID = trialPIO.getId();
+            UUID datasetId = null;
+            if (commit) {
+                datasetId = UUID.fromString(trialPIO.getBrAPIObject()
+                        .getAdditionalInfo().getAsJsonObject()
+                        .get(BrAPIAdditionalInfoFields.OBSERVATION_DATASET_ID).getAsString());
+            }
             PendingImportObject<BrAPIStudy> studyPIO = this.studyByNameNoScope.get(importRow.getEnv());
             UUID studyID = studyPIO.getId();
             UUID id = UUID.randomUUID();
-            BrAPIObservationUnit newObservationUnit = importRow.constructBrAPIObservationUnit(program, envSeqValue, commit, germplasmName, BRAPI_REFERENCE_SOURCE, trialID, studyID, id);
+            BrAPIObservationUnit newObservationUnit = importRow.constructBrAPIObservationUnit(program, envSeqValue, commit, germplasmName, BRAPI_REFERENCE_SOURCE, trialID, datasetId, studyID, id);
             pio = new PendingImportObject<>(ImportObjectState.NEW, newObservationUnit, id);
             this.observationUnitByNameNoScope.put(key, pio);
         }

src/main/java/org/breedinginsight/db/migration/V1_0_13__Update_BrAPI_Locations_XRefs.java

@@ -49,7 +49,7 @@ public void migrate(Context context) throws Exception {
         String referenceSource = placeholders.get(BRAPI_REFERENCE_SOURCE_KEY);
 
         // Get all the programs
-        List<Program> programs = getAllPrograms(context, defaultUrl);
+        List<Program> programs = Utilities.getAllProgramsFlyway(context, defaultUrl);
         Map<UUID, LocationsApi> locationsApiForProgram = new HashMap<>();
         for (Program program : programs) {
             BrAPIClient client = new BrAPIClient(program.getBrapiUrl(), 240000);
@@ -103,22 +103,4 @@ private List<ProgramLocation> getAllLocations(Context context) throws SQLExcepti
         }
         return locations;
     }
-
-    private List<Program> getAllPrograms(Context context, String defaultUrl) throws Exception {
-        List<Program> programs = new ArrayList<>();
-        try (Statement select = context.getConnection().createStatement()) {
-            try (ResultSet rows = select.executeQuery("SELECT id, brapi_url, key FROM program where active = true ORDER BY id")) {
-                while (rows.next()) {
-                    Program program = new Program();
-                    program.setId(UUID.fromString(rows.getString(1)));
-                    String brapi_url = rows.getString(2);
-                    if (brapi_url == null) brapi_url = defaultUrl;
-                    program.setBrapiUrl(brapi_url);
-                    program.setKey(rows.getString(3));
-                    programs.add(program);
-                }
-            }
-        }
-        return programs;
-    }
 }