Optionally include Content-Disposition header in statement results API response

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}`. The filename must not be longer than 255 characters and must not contain the characters `/`, `\`, `:`, `*`, `?`, `"`, `<`, `>`, `|`, `\0`, `\n`, or `\r`.
 
 #### Responses
 

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/resources/SqlStatementResource.java

@@ -95,6 +95,7 @@
 import org.jboss.netty.handler.codec.http.HttpResponseStatus;
 
 import javax.servlet.http.HttpServletRequest;
+import javax.validation.constraints.NotNull;
 import javax.ws.rs.Consumes;
 import javax.ws.rs.DELETE;
 import javax.ws.rs.GET;
@@ -114,6 +115,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
 
@@ -122,6 +124,8 @@ public class SqlStatementResource
 {
 
   public static final String RESULT_FORMAT = "__resultFormat";
+  public static final String CONTENT_DISPOSITION_RESPONSE_HEADER = "Content-Disposition";
+  private static final Pattern FILENAME_PATTERN = Pattern.compile("^[^/:*?><\\\\\"|\0\n\r]*$");
   private static final Logger log = new Logger(SqlStatementResource.class);
   private final SqlStatementFactory msqSqlStatementFactory;
   private final ObjectMapper jsonMapper;
@@ -277,6 +281,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 +314,12 @@ public Response doGetResults(
       );
       throwIfQueryIsNotSuccessful(queryId, statusPlus);
 
+      final String contentDispositionHeaderValue = filename != null ? StringUtils.format("attachment; filename=\"%s\"", validateFilename(filename)) : null;
+
       Optional<List<ColumnNameAndTypes>> 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();
+        return addContentDisposition(Response.ok(), contentDispositionHeaderValue).build();
       }
 
       // returning results
@@ -321,18 +328,20 @@ public Response doGetResults(
       results = getResultYielder(queryId, page, msqControllerTask, closer);
       if (!results.isPresent()) {
         // no results, return empty
-        return Response.ok().build();
+        return addContentDisposition(Response.ok(), contentDispositionHeaderValue).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();
+      ));
+
+      return addContentDisposition(responseBuilder, contentDispositionHeaderValue).build();
     }
 
 
@@ -976,6 +985,42 @@ private void checkForDurableStorageConnectorImpl()
     }
   }
 
+  private static Response.ResponseBuilder addContentDisposition(
+      Response.ResponseBuilder responseBuilder,
+      String contentDisposition
+  )
+  {
+    if (contentDisposition != null) {
+      responseBuilder.header(CONTENT_DISPOSITION_RESPONSE_HEADER, contentDisposition);
+    }
+    return responseBuilder;
+  }
+
+  /**
+   * Validates that a filename is valid. Filenames are considered to be valid if it is:
+   * <ul>
+   *   <li>Not empty.</li>
+   *   <li>Not longer than 255 characters.</li>
+   *   <li>Does not contain the characters `/`, `\`, `:`, `*`, `?`, `"`, `<`, `>`, `|`, `\0`, `\n`, or `\r`.</li>
+   * </ul>
+   */
+  @VisibleForTesting
+  static String validateFilename(@NotNull String filename)
+  {
+    if (filename.isEmpty()) {
+      throw InvalidInput.exception("Filename cannot be empty.");
+    }
+
+    if (filename.length() > 255) {
+      throw InvalidInput.exception("Filename cannot be longer than 255 characters.");
+    }
+
+    if (!FILENAME_PATTERN.matcher(filename).matches()) {
+      throw InvalidInput.exception("Filename contains invalid characters. (/, \\, :, *, ?, \", <, >, |, \0, \n, or \r)");
+    }
+    return filename;
+  }
+
   private <T> T contactOverlord(final ListenableFuture<T> future, String queryId)
   {
     try {

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,41 +488,47 @@ public void testMultipleWorkersWithPageSizeLimiting() throws IOException
         sqlStatementResult.getQueryId(),
         null,
         ResultFormat.ARRAY.name(),
+        null,
         SqlStatementResourceTest.makeOkRequest()
     )));
 
     Assert.assertEquals(rows.subList(0, 2), SqlStatementResourceTest.getResultRowsFromResponse(resource.doGetResults(
         sqlStatementResult.getQueryId(),
         0L,
         ResultFormat.ARRAY.name(),
+        null,
         SqlStatementResourceTest.makeOkRequest()
     )));
 
     Assert.assertEquals(rows.subList(2, 4), SqlStatementResourceTest.getResultRowsFromResponse(resource.doGetResults(
         sqlStatementResult.getQueryId(),
         1L,
         ResultFormat.ARRAY.name(),
+        null,
         SqlStatementResourceTest.makeOkRequest()
     )));
 
     Assert.assertEquals(rows.subList(4, 6), SqlStatementResourceTest.getResultRowsFromResponse(resource.doGetResults(
         sqlStatementResult.getQueryId(),
         2L,
         ResultFormat.ARRAY.name(),
+        null,
         SqlStatementResourceTest.makeOkRequest()
     )));
 
     Assert.assertEquals(rows.subList(6, 8), SqlStatementResourceTest.getResultRowsFromResponse(resource.doGetResults(
         sqlStatementResult.getQueryId(),
         3L,
         ResultFormat.ARRAY.name(),
+        null,
         SqlStatementResourceTest.makeOkRequest()
     )));
 
     Assert.assertEquals(rows.subList(8, 10), SqlStatementResourceTest.getResultRowsFromResponse(resource.doGetResults(
         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,13 +627,15 @@ public void testResultFormatWithParamInSelect() throws IOException
         sqlStatementResult.getQueryId(),
         null,
         ResultFormat.ARRAY.name(),
+        null,
         SqlStatementResourceTest.makeOkRequest()
     )));
 
     Assert.assertEquals(rows, SqlStatementResourceTest.getResultRowsFromResponse(resource.doGetResults(
         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());