diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java index 98a6e207b432f..1ee5f1030a403 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java +++ b/extensions-core/multi-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,15 +588,19 @@ private Optional getStatementStatus( MSQControllerTask msqControllerTask = getMSQControllerTaskAndCheckPermission(queryId, authenticationResult, forAction); SqlStatementState sqlStatementState = SqlStatementResourceHelper.getSqlStatementState(statusPlus); - Optional msqTaskReportPayload = Optional.empty(); - try { - msqTaskReportPayload = Optional.ofNullable(SqlStatementResourceHelper.getPayload(contactOverlord(overlordClient.taskReportAsMap(queryId), queryId))); - } - catch (DruidException e) { - if (!e.getErrorCode().equals("notFound")) { + Supplier> 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( @@ -603,7 +608,7 @@ private Optional getStatementStatus( taskResponse, statusPlus, sqlStatementState, - msqTaskReportPayload.orElse(null), + msqTaskReportPayloadSupplier.get().orElse(null), jsonMapper, detail ); @@ -622,9 +627,9 @@ private Optional 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 )); } }