From 1a764b3f58199d2f74c8ee2b93e7aa4ad05b47ab Mon Sep 17 00:00:00 2001
From: Oleg Smirnov <oleg.smirnov@exactprosystems.com>
Date: Wed, 20 Sep 2023 11:14:56 +0400
Subject: [PATCH] [TH2-5046] Correct error reporting for group download (#63)

* [TH2-5046] Correct error handling when loading group

* [TH2-5046] Update readme
---
 README.md                                     |  4 +++
 .../handlers/SearchMessagesHandler.kt         | 11 ++++---
 .../http/FileDownloadHandler.kt               | 21 +++++++++++--
 .../http/TestFileDownloadHandler.kt           | 30 +++++++++++++++++++
 4 files changed, 57 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 6c6720c7..4e26a1c5 100644
--- a/README.md
+++ b/README.md
@@ -223,6 +223,10 @@ spec:
 + Add support for requesting message groups in reversed order
 + Add filter by stream to gRPC API for group search request
 
+### Fixed:
+
++ error reporting when executing group search (if error occurs during a call to the Cradle API stream were closed before the error was reported)
+
 ## 2.1.0
 
 + Updated bom: `4.5.0-dev`
diff --git a/src/main/kotlin/com/exactpro/th2/lwdataprovider/handlers/SearchMessagesHandler.kt b/src/main/kotlin/com/exactpro/th2/lwdataprovider/handlers/SearchMessagesHandler.kt
index 9598c790..64de8dbe 100644
--- a/src/main/kotlin/com/exactpro/th2/lwdataprovider/handlers/SearchMessagesHandler.kt
+++ b/src/main/kotlin/com/exactpro/th2/lwdataprovider/handlers/SearchMessagesHandler.kt
@@ -264,9 +264,8 @@ class SearchMessagesHandler(
                 markerAsGroup = true,
                 limit = request.limit,
             )
-            try {
-                rootSink.use { sink ->
-
+            rootSink.use { sink ->
+                try {
                     val parameters = CradleGroupRequest(
                         preFilter = createInitialPrefilter(request),
                     )
@@ -301,10 +300,10 @@ class SearchMessagesHandler(
                             }
                         } while (keepPulling)
                     }
+                } catch (ex: Exception) {
+                    logger.error("Error getting messages group", ex)
+                    rootSink.onError(ex)
                 }
-            } catch (ex: Exception) {
-                logger.error("Error getting messages group", ex)
-                rootSink.onError(ex)
             }
         }
     }
diff --git a/src/main/kotlin/com/exactpro/th2/lwdataprovider/http/FileDownloadHandler.kt b/src/main/kotlin/com/exactpro/th2/lwdataprovider/http/FileDownloadHandler.kt
index 10675fb2..d8dee96c 100644
--- a/src/main/kotlin/com/exactpro/th2/lwdataprovider/http/FileDownloadHandler.kt
+++ b/src/main/kotlin/com/exactpro/th2/lwdataprovider/http/FileDownloadHandler.kt
@@ -200,9 +200,19 @@ class FileDownloadHandler(
     ) {
         val matchedPath = ctx.matchedPath()
         var dataSent = 0
-        ctx.status(HttpStatus.OK)
-            .contentType(JSON_STREAM_CONTENT_TYPE)
-            .header(Header.TRANSFER_ENCODING, "chunked")
+
+        var writeHeader = true
+        var status: HttpStatus = HttpStatus.OK
+
+        fun writeHeader() {
+            if (writeHeader) {
+                ctx.status(status)
+                    .contentType(JSON_STREAM_CONTENT_TYPE)
+                    .header(Header.TRANSFER_ENCODING, "chunked")
+                writeHeader = false
+            }
+        }
+
         val output = ctx.res().outputStream.buffered()
         try {
             do {
@@ -210,6 +220,11 @@ class FileDownloadHandler(
                     val nextEvent = queue.take()
                     ResponseQueue.currentSize(matchedPath, queue.size)
                     val sseEvent = dataMeasurement.start("await_convert_to_json").use { nextEvent.get() }
+                    if (writeHeader && sseEvent is SseEvent.ErrorData.SimpleError) {
+                        // something happened during request
+                        status = HttpStatus.INTERNAL_SERVER_ERROR
+                    }
+                    writeHeader()
                     when (sseEvent.event) {
                         EventType.KEEP_ALIVE -> output.flush()
                         EventType.CLOSE -> {
diff --git a/src/test/kotlin/com/exactpro/th2/lwdataprovider/http/TestFileDownloadHandler.kt b/src/test/kotlin/com/exactpro/th2/lwdataprovider/http/TestFileDownloadHandler.kt
index 62d3c936..459aeb02 100644
--- a/src/test/kotlin/com/exactpro/th2/lwdataprovider/http/TestFileDownloadHandler.kt
+++ b/src/test/kotlin/com/exactpro/th2/lwdataprovider/http/TestFileDownloadHandler.kt
@@ -19,6 +19,7 @@ package com.exactpro.th2.lwdataprovider.http
 import com.exactpro.cradle.Direction
 import com.exactpro.cradle.Order
 import com.exactpro.cradle.messages.StoredMessageIdUtils
+import com.exactpro.cradle.utils.CradleStorageException
 import com.exactpro.th2.common.message.addField
 import com.exactpro.th2.common.message.message
 import com.exactpro.th2.common.message.setMetadata
@@ -27,8 +28,10 @@ import com.exactpro.th2.lwdataprovider.util.GroupBatch
 import com.exactpro.th2.lwdataprovider.util.createCradleStoredMessage
 import io.javalin.http.HttpStatus
 import org.junit.jupiter.api.Test
+import org.mockito.kotlin.any
 import org.mockito.kotlin.argThat
 import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.doThrow
 import org.mockito.kotlin.whenever
 import strikt.api.expectThat
 import strikt.assertions.isEqualTo
@@ -297,4 +300,31 @@ class TestFileDownloadHandler : AbstractHttpHandlerTest<FileDownloadHandler>() {
             }
         }
     }
+
+    @Test
+    fun `respond with error and correct status if first cradle call throws an exception`() {
+        whenever(storage.getGroupedMessageBatches(any())) doThrow CradleStorageException("ignore")
+
+        startTest { _, client ->
+            val now = Instant.now().toEpochMilli()
+            val response = client.get(
+                "/download/messages?" +
+                        "startTimestamp=${now}&endTimestamp=${now + 100}" +
+                        "&group=test-group" +
+                        "&bookId=test-book" +
+                        "&responseFormat=BASE_64"
+            )
+
+            expectThat(response) {
+                get { code } isEqualTo HttpStatus.INTERNAL_SERVER_ERROR.code
+                get { body?.bytes()?.toString(Charsets.UTF_8) }
+                    .isNotNull()
+                    .isEqualTo(
+                        """{"error":"ignore"}
+                          |
+                        """.trimMargin()
+                    )
+            }
+        }
+    }
 }
\ No newline at end of file