diff --git a/docs/api-reference/sql-api.md b/docs/api-reference/sql-api.md index 2b25bdd02f0b..cead5f6040d1 100644 --- a/docs/api-reference/sql-api.md +++ b/docs/api-reference/sql-api.md @@ -1286,6 +1286,9 @@ Getting the query results for an ingestion query returns an empty response. * `resultFormat` (optional) * Type: String * Defines the format in which the results are presented. The following options are supported `arrayLines`,`objectLines`,`array`,`object`, and `csv`. The default is `object`. +* `filename` (optional) + * Type: String + * If set, attaches a `Content-Disposition` header to the response with the value of `attachment; filename={filename}`. #### Responses 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 dcaa447539b1..eac08c953a54 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 @@ -122,6 +122,7 @@ public class SqlStatementResource { public static final String RESULT_FORMAT = "__resultFormat"; + public static final String CONTENT_DISPOSITION_RESPONSE_HEADER = "Content-Disposition"; private static final Logger log = new Logger(SqlStatementResource.class); private final SqlStatementFactory msqSqlStatementFactory; private final ObjectMapper jsonMapper; @@ -277,6 +278,7 @@ public Response doGetResults( @PathParam("id") final String queryId, @QueryParam("page") Long page, @QueryParam("resultFormat") String resultFormat, + @QueryParam("filename") String filename, @Context final HttpServletRequest req ) { @@ -309,10 +311,18 @@ public Response doGetResults( ); throwIfQueryIsNotSuccessful(queryId, statusPlus); + final String contentDispositionHeaderValue = filename != null ? StringUtils.format("attachment; filename=%s", filename) : null; + Optional> signature = SqlStatementResourceHelper.getSignature(msqControllerTask); if (!signature.isPresent() || MSQControllerTask.isIngestion(msqControllerTask.getQuerySpec())) { // Since it's not a select query, nothing to return. - return Response.ok().build(); + final Response.ResponseBuilder responseBuilder = Response.ok(); + + if (contentDispositionHeaderValue != null) { + responseBuilder.header(CONTENT_DISPOSITION_RESPONSE_HEADER, contentDispositionHeaderValue); + } + + return responseBuilder.build(); } // returning results @@ -321,18 +331,30 @@ public Response doGetResults( results = getResultYielder(queryId, page, msqControllerTask, closer); if (!results.isPresent()) { // no results, return empty - return Response.ok().build(); + final Response.ResponseBuilder responseBuilder = Response.ok(); + + if (contentDispositionHeaderValue != null) { + responseBuilder.header(CONTENT_DISPOSITION_RESPONSE_HEADER, contentDispositionHeaderValue); + } + + return responseBuilder.build(); } ResultFormat preferredFormat = getPreferredResultFormat(resultFormat, msqControllerTask.getQuerySpec()); - return Response.ok((StreamingOutput) outputStream -> resultPusher( + final Response.ResponseBuilder responseBuilder = Response.ok((StreamingOutput) outputStream -> resultPusher( queryId, signature, closer, results, new CountingOutputStream(outputStream), preferredFormat - )).build(); + )); + + if (contentDispositionHeaderValue != null) { + responseBuilder.header(CONTENT_DISPOSITION_RESPONSE_HEADER, contentDispositionHeaderValue); + } + + return responseBuilder.build(); } diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlMSQStatementResourcePostTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlMSQStatementResourcePostTest.java index e810de56f4e1..0437fbccda25 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlMSQStatementResourcePostTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlMSQStatementResourcePostTest.java @@ -393,6 +393,7 @@ public void testWithDurableStorage() throws IOException sqlStatementResult.getQueryId(), null, ResultFormat.OBJECTLINES.name(), + null, SqlStatementResourceTest.makeOkRequest() ), objectMapper @@ -406,6 +407,7 @@ public void testWithDurableStorage() throws IOException sqlStatementResult.getQueryId(), 0L, ResultFormat.OBJECTLINES.name(), + null, SqlStatementResourceTest.makeOkRequest() ), objectMapper @@ -419,6 +421,7 @@ public void testWithDurableStorage() throws IOException sqlStatementResult.getQueryId(), 2L, ResultFormat.OBJECTLINES.name(), + null, SqlStatementResourceTest.makeOkRequest() ), objectMapper @@ -485,6 +488,7 @@ public void testMultipleWorkersWithPageSizeLimiting() throws IOException sqlStatementResult.getQueryId(), null, ResultFormat.ARRAY.name(), + null, SqlStatementResourceTest.makeOkRequest() ))); @@ -492,6 +496,7 @@ public void testMultipleWorkersWithPageSizeLimiting() throws IOException sqlStatementResult.getQueryId(), 0L, ResultFormat.ARRAY.name(), + null, SqlStatementResourceTest.makeOkRequest() ))); @@ -499,6 +504,7 @@ public void testMultipleWorkersWithPageSizeLimiting() throws IOException sqlStatementResult.getQueryId(), 1L, ResultFormat.ARRAY.name(), + null, SqlStatementResourceTest.makeOkRequest() ))); @@ -506,6 +512,7 @@ public void testMultipleWorkersWithPageSizeLimiting() throws IOException sqlStatementResult.getQueryId(), 2L, ResultFormat.ARRAY.name(), + null, SqlStatementResourceTest.makeOkRequest() ))); @@ -513,6 +520,7 @@ public void testMultipleWorkersWithPageSizeLimiting() throws IOException sqlStatementResult.getQueryId(), 3L, ResultFormat.ARRAY.name(), + null, SqlStatementResourceTest.makeOkRequest() ))); @@ -520,6 +528,7 @@ public void testMultipleWorkersWithPageSizeLimiting() throws IOException sqlStatementResult.getQueryId(), 4L, ResultFormat.ARRAY.name(), + null, SqlStatementResourceTest.makeOkRequest() ))); } @@ -565,6 +574,7 @@ public void testResultFormat() throws Exception sqlStatementResult.getQueryId(), null, resultFormat.name(), + null, SqlStatementResourceTest.makeOkRequest() ), objectMapper ) @@ -577,6 +587,7 @@ public void testResultFormat() throws Exception sqlStatementResult.getQueryId(), 0L, resultFormat.name(), + null, SqlStatementResourceTest.makeOkRequest() ), objectMapper ) @@ -616,6 +627,7 @@ public void testResultFormatWithParamInSelect() throws IOException sqlStatementResult.getQueryId(), null, ResultFormat.ARRAY.name(), + null, SqlStatementResourceTest.makeOkRequest() ))); @@ -623,6 +635,7 @@ public void testResultFormatWithParamInSelect() throws IOException sqlStatementResult.getQueryId(), 0L, ResultFormat.ARRAY.name(), + null, SqlStatementResourceTest.makeOkRequest() ))); } @@ -695,6 +708,7 @@ public void testInsert() actual.getQueryId(), 0L, null, + null, SqlStatementResourceTest.makeOkRequest() ); Assert.assertEquals(Response.Status.OK.getStatusCode(), resultsResponse.getStatus()); @@ -738,6 +752,7 @@ public void testReplaceAll() actual.getQueryId(), 0L, null, + null, SqlStatementResourceTest.makeOkRequest() ); Assert.assertEquals(Response.Status.OK.getStatusCode(), resultsResponse.getStatus()); diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlStatementResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlStatementResourceTest.java index 40dcb303b1dc..64ba45a979ba 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlStatementResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/resources/SqlStatementResourceTest.java @@ -22,6 +22,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Iterables; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.SettableFuture; import org.apache.calcite.sql.type.SqlTypeName; @@ -92,6 +93,7 @@ import org.mockito.Mock; import org.mockito.Mockito; +import javax.annotation.Nullable; import javax.ws.rs.core.Response; import java.io.IOException; import java.nio.charset.StandardCharsets; @@ -683,6 +685,16 @@ private static AuthenticationResult makeAuthResultForUser(String user) ); } + @Nullable + private Object getHeader(Response resp, String header) + { + final List objects = resp.getMetadata().get(header); + if (objects == null) { + return null; + } + return Iterables.getOnlyElement(objects); + } + @BeforeEach public void init() { @@ -716,7 +728,7 @@ public void testMSQSelectAcceptedQuery() ); assertExceptionMessage( - resource.doGetResults(ACCEPTED_SELECT_MSQ_QUERY, 0L, null, makeOkRequest()), + resource.doGetResults(ACCEPTED_SELECT_MSQ_QUERY, 0L, null, null, makeOkRequest()), StringUtils.format( "Query[%s] is currently in [%s] state. Please wait for it to complete.", ACCEPTED_SELECT_MSQ_QUERY, @@ -750,7 +762,7 @@ public void testMSQSelectRunningQuery() ); assertExceptionMessage( - resource.doGetResults(RUNNING_SELECT_MSQ_QUERY, 0L, null, makeOkRequest()), + resource.doGetResults(RUNNING_SELECT_MSQ_QUERY, 0L, null, null, makeOkRequest()), StringUtils.format( "Query[%s] is currently in [%s] state. Please wait for it to complete.", RUNNING_SELECT_MSQ_QUERY, @@ -816,7 +828,7 @@ public void testFinishedSelectMSQQuery() throws Exception null )), objectMapper.writeValueAsString(response.getEntity())); - Response resultsResponse = resource.doGetResults(FINISHED_SELECT_MSQ_QUERY, 0L, ResultFormat.OBJECTLINES.name(), makeOkRequest()); + Response resultsResponse = resource.doGetResults(FINISHED_SELECT_MSQ_QUERY, 0L, ResultFormat.OBJECTLINES.name(), null, makeOkRequest()); Assert.assertEquals(Response.Status.OK.getStatusCode(), resultsResponse.getStatus()); String expectedResult = "{\"_time\":123,\"alias\":\"foo\",\"market\":\"bar\"}\n" @@ -824,6 +836,8 @@ public void testFinishedSelectMSQQuery() throws Exception assertExpectedResults(expectedResult, resultsResponse); + Assert.assertNull(getHeader(resultsResponse, SqlStatementResource.CONTENT_DISPOSITION_RESPONSE_HEADER)); + Assert.assertEquals( Response.Status.OK.getStatusCode(), resource.deleteQuery(FINISHED_SELECT_MSQ_QUERY, makeOkRequest()).getStatus() @@ -835,6 +849,7 @@ public void testFinishedSelectMSQQuery() throws Exception FINISHED_SELECT_MSQ_QUERY, 0L, ResultFormat.OBJECTLINES.name(), + null, makeOkRequest() ) ); @@ -845,13 +860,22 @@ public void testFinishedSelectMSQQuery() throws Exception FINISHED_SELECT_MSQ_QUERY, null, ResultFormat.OBJECTLINES.name(), + null, makeOkRequest() ) ); Assert.assertEquals( Response.Status.BAD_REQUEST.getStatusCode(), - resource.doGetResults(FINISHED_SELECT_MSQ_QUERY, -1L, null, makeOkRequest()).getStatus() + resource.doGetResults(FINISHED_SELECT_MSQ_QUERY, -1L, null, null, makeOkRequest()).getStatus() + ); + + Assert.assertEquals( + "attachment; filename=my-file.ndjson", + getHeader( + resource.doGetResults(FINISHED_SELECT_MSQ_QUERY, 0L, ResultFormat.OBJECTLINES.name(), "my-file.ndjson", makeOkRequest()), + SqlStatementResource.CONTENT_DISPOSITION_RESPONSE_HEADER + ) ); } @@ -867,7 +891,7 @@ public void testFailedMSQQuery() for (String queryID : ImmutableList.of(ERRORED_SELECT_MSQ_QUERY, ERRORED_INSERT_MSQ_QUERY)) { assertExceptionMessage(resource.doGetStatus(queryID, false, makeOkRequest()), FAILURE_MSG, Response.Status.OK); assertExceptionMessage( - resource.doGetResults(queryID, 0L, null, makeOkRequest()), + resource.doGetResults(queryID, 0L, null, null, makeOkRequest()), StringUtils.format( "Query[%s] failed. Check the status api for more details.", queryID @@ -897,18 +921,29 @@ public void testFinishedInsertMSQQuery() null ), (SqlStatementResult) response.getEntity()); + final Response resultResponse = resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, 0L, null, null, makeOkRequest()); + Assert.assertEquals( Response.Status.OK.getStatusCode(), - resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, 0L, null, makeOkRequest()).getStatus() + resultResponse.getStatus() ); + Assert.assertNull(getHeader(resultResponse, SqlStatementResource.CONTENT_DISPOSITION_RESPONSE_HEADER)); Assert.assertEquals( Response.Status.OK.getStatusCode(), - resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, null, null, makeOkRequest()).getStatus() + resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, null, null, null, makeOkRequest()).getStatus() ); Assert.assertEquals( Response.Status.BAD_REQUEST.getStatusCode(), - resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, -1L, null, makeOkRequest()).getStatus() + resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, -1L, null, null, makeOkRequest()).getStatus() + ); + + Assert.assertEquals( + "attachment; filename=my-file.ndjson", + getHeader( + resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, 0L, null, "my-file.ndjson", makeOkRequest()), + SqlStatementResource.CONTENT_DISPOSITION_RESPONSE_HEADER + ) ); } @@ -917,7 +952,7 @@ public void testNonMSQTasks() { for (String queryID : ImmutableList.of(RUNNING_NON_MSQ_TASK, FAILED_NON_MSQ_TASK, FINISHED_NON_MSQ_TASK)) { assertNotFound(resource.doGetStatus(queryID, false, makeOkRequest()), queryID); - assertNotFound(resource.doGetResults(queryID, 0L, null, makeOkRequest()), queryID); + assertNotFound(resource.doGetResults(queryID, 0L, null, null, makeOkRequest()), queryID); assertNotFound(resource.deleteQuery(queryID, makeOkRequest()), queryID); } } @@ -941,7 +976,7 @@ public void testMSQInsertAcceptedQuery() ); assertExceptionMessage( - resource.doGetResults(ACCEPTED_INSERT_MSQ_TASK, 0L, null, makeOkRequest()), + resource.doGetResults(ACCEPTED_INSERT_MSQ_TASK, 0L, null, null, makeOkRequest()), StringUtils.format( "Query[%s] is currently in [%s] state. Please wait for it to complete.", ACCEPTED_INSERT_MSQ_TASK, @@ -974,7 +1009,7 @@ public void testMSQInsertRunningQuery() ); assertExceptionMessage( - resource.doGetResults(RUNNING_INSERT_MSQ_QUERY, 0L, null, makeOkRequest()), + resource.doGetResults(RUNNING_INSERT_MSQ_QUERY, 0L, null, null, makeOkRequest()), StringUtils.format( "Query[%s] is currently in [%s] state. Please wait for it to complete.", RUNNING_INSERT_MSQ_QUERY, @@ -1005,6 +1040,7 @@ public void testAPIBehaviourWithSuperUsers() RUNNING_SELECT_MSQ_QUERY, 1L, null, + null, makeExpectedReq(makeAuthResultForUser(SUPERUSER)) ).getStatus() ); @@ -1035,6 +1071,7 @@ public void testAPIBehaviourWithDifferentUserAndNoStatePermission() RUNNING_SELECT_MSQ_QUERY, 1L, null, + null, makeExpectedReq(differentUserAuthResult) ).getStatus() ); @@ -1065,6 +1102,7 @@ public void testAPIBehaviourWithDifferentUserAndStateRPermission() RUNNING_SELECT_MSQ_QUERY, 1L, null, + null, makeExpectedReq(differentUserAuthResult) ).getStatus() ); @@ -1095,6 +1133,7 @@ public void testAPIBehaviourWithDifferentUserAndStateWPermission() RUNNING_SELECT_MSQ_QUERY, 1L, null, + null, makeExpectedReq(differentUserAuthResult) ).getStatus() ); @@ -1125,6 +1164,7 @@ public void testAPIBehaviourWithDifferentUserAndStateRWPermission() RUNNING_SELECT_MSQ_QUERY, 1L, null, + null, makeExpectedReq(differentUserAuthResult) ).getStatus() ); @@ -1156,7 +1196,7 @@ public void testTaskIdNotFound() ); Assert.assertEquals( Response.Status.NOT_FOUND.getStatusCode(), - resource.doGetResults(taskIdNotFound, null, null, makeOkRequest()).getStatus() + resource.doGetResults(taskIdNotFound, null, null, null, makeOkRequest()).getStatus() ); Assert.assertEquals( Response.Status.NOT_FOUND.getStatusCode(),