Merge pull request #11027 from kalletlak/km-changes

Resolve KM plot censored event bug and add caching
cBioPortal · Sep 30, 2024 · 8206a7c · 8206a7c
2 parents 840b05a + fef066a
commit 8206a7c
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 26 deletions.
diff --git a/src/main/java/org/cbioportal/service/impl/ClinicalEventServiceImpl.java b/src/main/java/org/cbioportal/service/impl/ClinicalEventServiceImpl.java
@@ -150,12 +150,13 @@ public List<ClinicalData> getSurvivalData(List<String> studyIds,
             ).toList();
 
         List<ClinicalEvent> patientCensoredEvents = filterClinicalEvents(patientStartEvents, survivalRequest.getCensoredEventRequestIdentifier());
+        Map<String, ClinicalEvent> patientCensoredEventsById = patientCensoredEvents.stream().collect(Collectors.toMap(ClinicalEventServiceImpl::getKey, Function.identity()));
 
         return patientStartEvents.stream()
             .flatMap(event -> {
-                ClinicalData clinicalDataMonths = buildClinicalSurvivalMonths(attributeIdPrefix, event, survivalRequest, patientEndEvents, patientCensoredEvents);
+                ClinicalData clinicalDataMonths = buildClinicalSurvivalMonths(attributeIdPrefix, event, survivalRequest, patientEndEventsById, patientCensoredEventsById);
                 if (clinicalDataMonths == null) return Stream.empty();
-                ClinicalData clinicalDataStatus = buildClinicalSurvivalStatus(attributeIdPrefix, event, patientEndEvents);
+                ClinicalData clinicalDataStatus = buildClinicalSurvivalStatus(attributeIdPrefix, event, patientEndEventsById);
 
                 return Stream.of(clinicalDataMonths, clinicalDataStatus);
             }).toList();
@@ -200,22 +201,19 @@ private List<ClinicalEvent> filterClinicalEvents(List<ClinicalEvent> patientEven
 
         // only fetch end timeline events for patients that have endClinicalEventsMeta and start timeline events
         List<ClinicalEvent> queriedPatientEvents = new ArrayList<>();
-        if (CollectionUtils.isNotEmpty(clinicalEventsMeta) && CollectionUtils.isNotEmpty(filteredStudyIds)) {
+        if (CollectionUtils.isNotEmpty(filteredStudyIds)) {
             queriedPatientEvents = clinicalEventRepository.getTimelineEvents(filteredStudyIds, filteredPatientIds, clinicalEventsMeta);
         }
         return queriedPatientEvents;
     }
 
-    private ClinicalData buildClinicalSurvivalMonths(String attributeIdPrefix, ClinicalEvent event, SurvivalRequest survivalRequest, List<ClinicalEvent> patientEndEvents, List<ClinicalEvent> patientCensoredEvents) {
+    private ClinicalData buildClinicalSurvivalMonths(String attributeIdPrefix, ClinicalEvent event, SurvivalRequest survivalRequest, Map<String, ClinicalEvent> patientEndEventsById, Map<String, ClinicalEvent> patientCensoredEventsById) {
         final String SURVIVAL_MONTH_ATTRIBUTE = attributeIdPrefix + "_MONTHS";
         ClinicalData clinicalDataMonths = new ClinicalData();
         clinicalDataMonths.setStudyId(event.getStudyId());
         clinicalDataMonths.setPatientId(event.getPatientId());
         clinicalDataMonths.setAttrId(SURVIVAL_MONTH_ATTRIBUTE);
 
-        Map<String, ClinicalEvent> patientEndEventsById = patientEndEvents.stream().collect(Collectors.toMap(ClinicalEventServiceImpl::getKey, Function.identity()));
-        Map<String, ClinicalEvent> patientCensoredEventsById = patientCensoredEvents.stream().collect(Collectors.toMap(ClinicalEventServiceImpl::getKey, Function.identity()));
-
         ToIntFunction<ClinicalEvent> startPositionIdentifier = getPositionIdentifier(survivalRequest.getStartEventRequestIdentifier().getPosition());
         ToIntFunction<ClinicalEvent> endPositionIdentifier = survivalRequest.getEndEventRequestIdentifier() == null ? ClinicalEvent::getStopDate : getPositionIdentifier(survivalRequest.getEndEventRequestIdentifier().getPosition());
         ToIntFunction<ClinicalEvent> censoredPositionIdentifier = survivalRequest.getCensoredEventRequestIdentifier() == null ? ClinicalEvent::getStopDate : getPositionIdentifier(survivalRequest.getCensoredEventRequestIdentifier().getPosition());
@@ -241,8 +239,7 @@ private ClinicalData buildClinicalSurvivalMonths(String attributeIdPrefix, Clini
         return clinicalDataMonths;
     }
 
-    private ClinicalData buildClinicalSurvivalStatus(String attributeIdPrefix, ClinicalEvent event, List<ClinicalEvent> patientEndEvents) {
-        Map<String, ClinicalEvent> patientEndEventsById = patientEndEvents.stream().collect(Collectors.toMap(ClinicalEventServiceImpl::getKey, Function.identity()));
+    private ClinicalData buildClinicalSurvivalStatus(String attributeIdPrefix, ClinicalEvent event, Map<String, ClinicalEvent> patientEndEventsById) {
 
         ClinicalData clinicalDataStatus = new ClinicalData();
         clinicalDataStatus.setStudyId(event.getStudyId());

diff --git a/src/main/java/org/cbioportal/web/ClinicalEventController.java b/src/main/java/org/cbioportal/web/ClinicalEventController.java
@@ -16,14 +16,10 @@
 import org.cbioportal.service.exception.StudyNotFoundException;
 import org.cbioportal.web.config.InternalApiTags;
 import org.cbioportal.web.config.annotation.InternalApi;
-import org.cbioportal.web.parameter.ClinicalEventAttributeRequest;
-import org.cbioportal.web.parameter.Direction;
-import org.cbioportal.web.parameter.HeaderKeyConstants;
-import org.cbioportal.web.parameter.PagingConstants;
-import org.cbioportal.web.parameter.PatientIdentifier;
-import org.cbioportal.web.parameter.Projection;
+import org.cbioportal.web.parameter.*;
 import org.cbioportal.web.parameter.sort.ClinicalEventSortBy;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.cache.annotation.Cacheable;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
@@ -144,6 +140,15 @@ public ResponseEntity<List<ClinicalEvent>> fetchClinicalEventsMeta(
         @Parameter(hidden = true) // prevent reference to this attribute in the swagger-ui interface. This attribute is needed for the @PreAuthorize tag above.
         @Valid @RequestAttribute(required = false, value = "interceptedClinicalEventAttributeRequest") ClinicalEventAttributeRequest interceptedClinicalEventAttributeRequest) {
 
+        return new ResponseEntity<>(cachedClinicalEventsMeta(interceptedClinicalEventAttributeRequest), HttpStatus.OK);
+
+    }
+
+    @Cacheable(
+        cacheResolver = "generalRepositoryCacheResolver",
+        condition = "@cacheEnabledConfig.getEnabled()"
+    )
+    public List<ClinicalEvent> cachedClinicalEventsMeta(ClinicalEventAttributeRequest interceptedClinicalEventAttributeRequest) {
         List<String> studyIds = new ArrayList<>();
         List<String> patientIds = new ArrayList<>();
         for (PatientIdentifier patientIdentifier : interceptedClinicalEventAttributeRequest.getPatientIdentifiers()) {
@@ -160,9 +165,7 @@ public ResponseEntity<List<ClinicalEvent>> fetchClinicalEventsMeta(
                 return clinicalEvent;
             })
             .toList();
-
-        return new ResponseEntity<>(clinicalEventService.getClinicalEventsMeta(
-            studyIds, patientIds, clinicalEventsRequest), HttpStatus.OK);
-
+        return clinicalEventService.getClinicalEventsMeta(
+            studyIds, patientIds, clinicalEventsRequest);
     }
 }
diff --git a/src/main/java/org/cbioportal/web/SurvivalController.java b/src/main/java/org/cbioportal/web/SurvivalController.java
@@ -14,6 +14,7 @@
 import org.cbioportal.web.parameter.PatientIdentifier;
 import org.cbioportal.web.parameter.SurvivalRequest;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.cache.annotation.Cacheable;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;
@@ -36,7 +37,7 @@
 @Tag(name = "Survival", description = " ")
 public class SurvivalController {
     private final ClinicalEventService clinicalEventService;
-    
+
     @Autowired
     public SurvivalController(ClinicalEventService clinicalEventService) {
         this.clinicalEventService = clinicalEventService;
@@ -59,18 +60,25 @@ public ResponseEntity<List<ClinicalData>> fetchSurvivalData(
         // prevent reference to this attribute in the swagger-ui interface. this attribute is needed for the @PreAuthorize tag above.
         @Valid @RequestAttribute(required = false, value = "interceptedSurvivalRequest") SurvivalRequest interceptedSurvivalRequest) {
 
+        return new ResponseEntity<>(cachedSurvivalData(interceptedSurvivalRequest),
+                                    HttpStatus.OK);
+    }
+
+    @Cacheable(
+        cacheResolver = "generalRepositoryCacheResolver",
+        condition = "@cacheEnabledConfig.getEnabled()"
+    )
+    public List<ClinicalData> cachedSurvivalData(SurvivalRequest interceptedSurvivalRequest) {
         List<String> studyIds = new ArrayList<>();
         List<String> patientIds = new ArrayList<>();
         for (PatientIdentifier patientIdentifier : interceptedSurvivalRequest.getPatientIdentifiers()) {
             studyIds.add(patientIdentifier.getStudyId());
             patientIds.add(patientIdentifier.getPatientId());
         }
 
-        List<ClinicalData> result = clinicalEventService.getSurvivalData(studyIds,
-            patientIds,
-            interceptedSurvivalRequest.getAttributeIdPrefix(),
-            interceptedSurvivalRequest);
-
-        return new ResponseEntity<>(result, HttpStatus.OK);
+        return clinicalEventService.getSurvivalData(studyIds,
+                                                                         patientIds,
+                                                                         interceptedSurvivalRequest.getAttributeIdPrefix(),
+                                                                         interceptedSurvivalRequest);
     }
 }