fix: dataset result query only when needed (#126)

We should query the s3 dataset bucket for results only when the dataset status is success
ntua-unit-of-control-and-informatics · Nov 29, 2024 · 12be534 · 12be534
1 parent 49949ec
commit 12be534
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git 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
@@ -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))
     }