From 12be5341d26c11f8d9030540054589d576a1895c Mon Sep 17 00:00:00 2001
From: Alex Arvanitidis <alex.arvanitidis.ntua.gr@gmail.com>
Date: Fri, 29 Nov 2024 10:45:22 +0200
Subject: [PATCH] fix: dataset result query only when needed (#126)

We should query the s3 dataset bucket for results only when the dataset status is success
---
 src/main/kotlin/org/jaqpot/api/entity/Dataset.kt   |  8 +++++++-
 .../jaqpot/api/service/dataset/DatasetService.kt   | 14 ++++++++++++--
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/src/main/kotlin/org/jaqpot/api/entity/Dataset.kt b/src/main/kotlin/org/jaqpot/api/entity/Dataset.kt
index 9def881a..9cdf53e6 100644
--- a/src/main/kotlin/org/jaqpot/api/entity/Dataset.kt
+++ b/src/main/kotlin/org/jaqpot/api/entity/Dataset.kt
@@ -47,4 +47,10 @@ class Dataset(
     var executedAt: OffsetDateTime? = null,
 
     var executionFinishedAt: OffsetDateTime? = null
-) : BaseEntity()
+) : BaseEntity() {
+    fun shouldHaveResult(): Boolean {
+        return this.type == DatasetType.PREDICTION && this.status == DatasetStatus.SUCCESS
+    }
+}
+
+
diff --git a/src/main/kotlin/org/jaqpot/api/service/dataset/DatasetService.kt b/src/main/kotlin/org/jaqpot/api/service/dataset/DatasetService.kt
index 4f5d7d43..0082af0b 100644
--- a/src/main/kotlin/org/jaqpot/api/service/dataset/DatasetService.kt
+++ b/src/main/kotlin/org/jaqpot/api/service/dataset/DatasetService.kt
@@ -36,7 +36,11 @@ class DatasetService(
 
         return dataset.map {
             val input = storageService.readRawDatasetInput(it)
-            val result = storageService.readRawDatasetResult(it)
+            val result = if (it.shouldHaveResult()) {
+                storageService.readRawDatasetResult(it)
+            } else {
+                null
+            }
             ResponseEntity.ok(it.toDto(input, result))
         }
             .orElse(ResponseEntity.notFound().build())
@@ -48,7 +52,13 @@ class DatasetService(
         val datasets = datasetRepository.findAllByUserId(userId, pageable)
 
         val inputsMap = storageService.readRawDatasetInputs(datasets.content)
-        val resultsMap = storageService.readRawDatasetResults(datasets.content)
+        val resultsMap: MutableMap<String, List<Any>?> =
+            datasets.content.associateBy { it.id.toString() }.mapValues { null }.toMutableMap()
+
+        val datasetsWithResults = datasets.content.filter { it.shouldHaveResult() }
+        storageService.readRawDatasetResults(datasetsWithResults).forEach {
+            resultsMap[it.key] = it.value
+        }
 
         return ResponseEntity.ok().body(datasets.toGetDatasets200ResponseDto(inputsMap, resultsMap))
     }