Attempt 2 to fix web console check

apache · Jul 26, 2024 · 6a1fc4e · 6a1fc4e
1 parent afa0c96
commit 6a1fc4e
Showing 1 changed file with 16 additions and 11 deletions.
diff --git a/...ti-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java b/...ti-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java
@@ -113,6 +113,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 
@@ -587,23 +588,27 @@ private Optional<SqlStatementResult> getStatementStatus(
     MSQControllerTask msqControllerTask = getMSQControllerTaskAndCheckPermission(queryId, authenticationResult, forAction);
     SqlStatementState sqlStatementState = SqlStatementResourceHelper.getSqlStatementState(statusPlus);
 
-    Optional<MSQTaskReportPayload> msqTaskReportPayload = Optional.empty();
-    try {
-      msqTaskReportPayload = Optional.ofNullable(SqlStatementResourceHelper.getPayload(contactOverlord(overlordClient.taskReportAsMap(queryId), queryId)));
-    }
-    catch (DruidException e) {
-      if (!e.getErrorCode().equals("notFound")) {
+    Supplier<Optional<MSQTaskReportPayload>> msqTaskReportPayloadSupplier = () -> {
+      try {
+        return Optional.ofNullable(SqlStatementResourceHelper.getPayload(
+            contactOverlord(overlordClient.taskReportAsMap(queryId), queryId)
+        ));
+      }
+      catch (DruidException e) {
+        if (e.getErrorCode().equals("notFound")) {
+          return Optional.empty();
+        }
         throw e;
       }
-    }
+    };
 
     if (SqlStatementState.FAILED == sqlStatementState) {
       return SqlStatementResourceHelper.getExceptionPayload(
           queryId,
           taskResponse,
           statusPlus,
           sqlStatementState,
-          msqTaskReportPayload.orElse(null),
+          msqTaskReportPayloadSupplier.get().orElse(null),
           jsonMapper,
           detail
       );
@@ -622,9 +627,9 @@ private Optional<SqlStatementResult> getStatementStatus(
               msqControllerTask.getQuerySpec().getDestination()
           ).orElse(null) : null,
           null,
-          detail ? SqlStatementResourceHelper.getQueryStagesReport(msqTaskReportPayload.orElse(null)) : null,
-          detail ? SqlStatementResourceHelper.getQueryCounters(msqTaskReportPayload.orElse(null)) : null,
-          detail ? SqlStatementResourceHelper.getQueryWarningDetails(msqTaskReportPayload.orElse(null)) : null
+          detail ? SqlStatementResourceHelper.getQueryStagesReport(msqTaskReportPayloadSupplier.get().orElse(null)) : null,
+          detail ? SqlStatementResourceHelper.getQueryCounters(msqTaskReportPayloadSupplier.get().orElse(null)) : null,
+          detail ? SqlStatementResourceHelper.getQueryWarningDetails(msqTaskReportPayloadSupplier.get().orElse(null)) : null
       ));
     }
   }