From f94fe1e5e95254d9f18e1bb36db1e93fdb5aa461 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Fri, 30 Jun 2023 23:27:30 +0530 Subject: [PATCH 1/5] Add page information to druid statement resource --- .../druid/msq/sql/entity/PageInformation.java | 100 ++++++++++++++++++ .../msq/sql/entity/ResultSetInformation.java | 59 ++++------- .../sql/resources/SqlStatementResource.java | 47 +++++--- .../sql/SqlMsqStatementResourcePostTest.java | 6 +- .../msq/sql/SqlStatementResourceTest.java | 82 ++++---------- .../sql/entity/ResultSetInformationTest.java | 19 ++-- 6 files changed, 185 insertions(+), 128 deletions(-) create mode 100644 extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/PageInformation.java diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/PageInformation.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/PageInformation.java new file mode 100644 index 000000000000..81abb49486b2 --- /dev/null +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/PageInformation.java @@ -0,0 +1,100 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.msq.sql.entity; + + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; + +import javax.annotation.Nullable; +import java.util.Objects; + +public class PageInformation +{ + @Nullable + private final Long numRows; + @Nullable + private final Long sizeInBytes; + private final long id; + + @JsonCreator + public PageInformation( + @JsonProperty("numRows") @Nullable Long numRows, + @JsonProperty("sizeInBytes") @Nullable Long sizeInBytes, + @JsonProperty("id") long id + ) + { + this.numRows = numRows; + this.sizeInBytes = sizeInBytes; + this.id = id; + } + + @JsonProperty + @Nullable + public Long getNumRows() + { + return numRows; + } + + @JsonProperty + @Nullable + public Long getSizeInBytes() + { + return sizeInBytes; + } + + @JsonProperty + public long getId() + { + return id; + } + + @Override + public boolean equals(Object o) + { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + PageInformation that = (PageInformation) o; + return id == that.id && Objects.equals(numRows, that.numRows) && Objects.equals( + sizeInBytes, + that.sizeInBytes + ); + } + + @Override + public int hashCode() + { + return Objects.hash(numRows, sizeInBytes, id); + } + + @Override + public String toString() + { + return "PageInformation{" + + "numRows=" + numRows + + ", sizeInBytes=" + sizeInBytes + + ", id=" + id + + '}'; + } +} diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/ResultSetInformation.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/ResultSetInformation.java index 43201fdac6e0..64228523e741 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/ResultSetInformation.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/ResultSetInformation.java @@ -31,51 +31,27 @@ public class ResultSetInformation { - @Nullable - private final Long numRows; - @Nullable - private final Long sizeInBytes; - @Nullable private final ResultFormat resultFormat; - @Nullable private final List records; - @Nullable private final String dataSource; + @Nullable + private final List pageInformationList; @JsonCreator public ResultSetInformation( @JsonProperty("resultFormat") @Nullable ResultFormat resultFormat, - @JsonProperty("numRows") @Nullable Long numRows, - @JsonProperty("sizeInBytes") @Nullable Long sizeInBytes, @JsonProperty("dataSource") @Nullable String dataSource, - @JsonProperty("sampleRecords") @Nullable - List records + @JsonProperty("sampleRecords") @Nullable List records, + @JsonProperty("pageInformationList") @Nullable List pageInformationList ) { - this.numRows = numRows; - this.sizeInBytes = sizeInBytes; this.resultFormat = resultFormat; this.dataSource = dataSource; this.records = records; - } - - @Nullable - @JsonProperty - @JsonInclude(JsonInclude.Include.NON_NULL) - public Long getNumRows() - { - return numRows; - } - - @Nullable - @JsonProperty - @JsonInclude(JsonInclude.Include.NON_NULL) - public Long getSizeInBytes() - { - return sizeInBytes; + this.pageInformationList = pageInformationList; } @JsonProperty @@ -94,14 +70,21 @@ public String getDataSource() return dataSource; } - @Nullable @JsonProperty("sampleRecords") + @Nullable @JsonInclude(JsonInclude.Include.NON_NULL) public List getRecords() { return records; } + @JsonProperty + @Nullable + @JsonInclude(JsonInclude.Include.NON_NULL) + public List getPageInformationList() + { + return pageInformationList; + } @Override public boolean equals(Object o) @@ -113,30 +96,26 @@ public boolean equals(Object o) return false; } ResultSetInformation that = (ResultSetInformation) o; - return Objects.equals(numRows, that.numRows) - && Objects.equals(sizeInBytes, that.sizeInBytes) - && resultFormat == that.resultFormat + return resultFormat == that.resultFormat && Objects.equals(records, that.records) - && Objects.equals(dataSource, that.dataSource); + && Objects.equals(dataSource, that.dataSource) + && Objects.equals(pageInformationList, that.pageInformationList); } @Override public int hashCode() { - return Objects.hash(numRows, sizeInBytes, resultFormat, records, dataSource); + return Objects.hash(resultFormat, records, dataSource, pageInformationList); } @Override public String toString() { return "ResultSetInformation{" + - "totalRows=" + numRows + - ", totalSize=" + sizeInBytes + - ", resultFormat=" + resultFormat + + "resultFormat=" + resultFormat + ", records=" + records + ", dataSource='" + dataSource + '\'' + + ", pageInformationList=" + pageInformationList + '}'; } - } - 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 6ce5c7800525..770126311eb9 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 @@ -21,6 +21,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.io.CountingOutputStream; import com.google.common.util.concurrent.ListenableFuture; @@ -46,6 +47,7 @@ import org.apache.druid.msq.sql.MSQTaskSqlEngine; import org.apache.druid.msq.sql.SqlStatementState; import org.apache.druid.msq.sql.entity.ColumnNameAndTypes; +import org.apache.druid.msq.sql.entity.PageInformation; import org.apache.druid.msq.sql.entity.ResultSetInformation; import org.apache.druid.msq.sql.entity.SqlStatementResult; import org.apache.druid.msq.util.SqlStatementResourceHelper; @@ -268,8 +270,7 @@ public Response doGetStatus( @Produces(MediaType.APPLICATION_JSON) public Response doGetResults( @PathParam("id") final String queryId, - @QueryParam("offset") Long offset, - @QueryParam("numRows") Long numberOfRows, + @QueryParam("page") Long page, @Context final HttpServletRequest req ) { @@ -284,28 +285,25 @@ public Response doGetResults( } final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(req); - if (offset != null && offset < 0) { + if (page == null) { return buildNonOkResponse( DruidException.forPersona(DruidException.Persona.USER) .ofCategory(DruidException.Category.INVALID_INPUT) .build( - "offset cannot be negative. Please pass a positive number." + "Page cannot be null." ) ); } - if (numberOfRows != null && numberOfRows < 0) { + if (page < 0) { return buildNonOkResponse( DruidException.forPersona(DruidException.Persona.USER) .ofCategory(DruidException.Category.INVALID_INPUT) .build( - "numRows cannot be negative. Please pass a positive number." + "Page cannot be negative. Please pass a positive number." ) ); } - final long start = offset == null ? 0 : offset; - final long last = SqlStatementResourceHelper.getLastIndex(numberOfRows, start); - TaskStatusResponse taskResponse = contactOverlord(overlordClient.taskStatus(queryId)); if (taskResponse == null) { return Response.status(Response.Status.NOT_FOUND).build(); @@ -343,8 +341,21 @@ public Response doGetResults( if (!signature.isPresent()) { return Response.ok().build(); } - Optional> results = SqlStatementResourceHelper.getResults(SqlStatementResourceHelper.getPayload( - contactOverlord(overlordClient.taskReportAsMap(queryId)))); + + if (page > 0) { + // Results from task report are only present as one page. + return buildNonOkResponse( + DruidException.forPersona(DruidException.Persona.USER) + .ofCategory(DruidException.Category.INVALID_INPUT) + .build("Page number is out of range of the results.") + ); + } + + Optional> results = SqlStatementResourceHelper.getResults( + SqlStatementResourceHelper.getPayload( + contactOverlord(overlordClient.taskReportAsMap(queryId)) + ) + ); return Response.ok((StreamingOutput) outputStream -> { CountingOutputStream os = new CountingOutputStream(outputStream); @@ -353,7 +364,7 @@ public Response doGetResults( List rowSignature = signature.get(); writer.writeResponseStart(); - for (long k = start; k < Math.min(last, results.get().size()); k++) { + for (long k = 0; k < results.get().size(); k++) { writer.writeRowStart(); for (int i = 0; i < rowSignature.size(); i++) { writer.writeRowField( @@ -576,12 +587,16 @@ private Optional getSampleResults( ); return Optional.of(new ResultSetInformation( null, - // since the rows can be sampled, get the number of rows from counters - rowsAndSize.orElse(new Pair<>(null, null)).lhs, - rowsAndSize.orElse(new Pair<>(null, null)).rhs, dataSource, // only populate sample results in case a select query is successful - isSelectQuery ? SqlStatementResourceHelper.getResults(payload).orElse(null) : null + isSelectQuery ? SqlStatementResourceHelper.getResults(payload).orElse(null) : null, + ImmutableList.of( + new PageInformation( + rowsAndSize.orElse(new Pair<>(null, null)).lhs, + rowsAndSize.orElse(new Pair<>(null, null)).rhs, + 0 + ) + ) )); } else { return Optional.empty(); diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlMsqStatementResourcePostTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlMsqStatementResourcePostTest.java index ceae64dcf703..6f8fd82567f7 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlMsqStatementResourcePostTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlMsqStatementResourcePostTest.java @@ -30,6 +30,7 @@ import org.apache.druid.msq.indexing.error.InsertCannotBeEmptyFault; import org.apache.druid.msq.indexing.error.MSQException; import org.apache.druid.msq.sql.entity.ColumnNameAndTypes; +import org.apache.druid.msq.sql.entity.PageInformation; import org.apache.druid.msq.sql.entity.ResultSetInformation; import org.apache.druid.msq.sql.entity.SqlStatementResult; import org.apache.druid.msq.sql.resources.SqlStatementResource; @@ -113,8 +114,6 @@ public void testMSQSelectQueryTest() throws IOException MSQTestOverlordServiceClient.DURATION, new ResultSetInformation( null, - 6L, - 316L, MSQControllerTask.DUMMY_DATASOURCE_FOR_SELECT, objectMapper.readValue( objectMapper.writeValueAsString( @@ -122,7 +121,8 @@ public void testMSQSelectQueryTest() throws IOException new TypeReference>() { } - ) + ), + ImmutableList.of(new PageInformation(6L, 316L, 0)) ), null ); diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlStatementResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlStatementResourceTest.java index b5d5addecd78..f6bffb20b928 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlStatementResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlStatementResourceTest.java @@ -54,6 +54,7 @@ import org.apache.druid.msq.indexing.report.MSQTaskReportPayload; import org.apache.druid.msq.indexing.report.MSQTaskReportTest; import org.apache.druid.msq.sql.entity.ColumnNameAndTypes; +import org.apache.druid.msq.sql.entity.PageInformation; import org.apache.druid.msq.sql.entity.ResultSetInformation; import org.apache.druid.msq.sql.entity.SqlStatementResult; import org.apache.druid.msq.sql.resources.SqlStatementResource; @@ -650,7 +651,7 @@ public void testMSQSelectAcceptedQuery() ); assertExceptionMessage( - resource.doGetResults(ACCEPTED_SELECT_MSQ_QUERY, null, null, makeOkRequest()), + resource.doGetResults(ACCEPTED_SELECT_MSQ_QUERY, 0L, makeOkRequest()), StringUtils.format( "Query[%s] is currently in [%s] state. Please wait for it to complete.", ACCEPTED_SELECT_MSQ_QUERY, @@ -665,7 +666,6 @@ public void testMSQSelectAcceptedQuery() } @Test - public void testMSQSelectRunningQuery() { @@ -685,7 +685,7 @@ public void testMSQSelectRunningQuery() ); assertExceptionMessage( - resource.doGetResults(RUNNING_SELECT_MSQ_QUERY, null, null, makeOkRequest()), + resource.doGetResults(RUNNING_SELECT_MSQ_QUERY, 0L, makeOkRequest()), StringUtils.format( "Query[%s] is currently in [%s] state. Please wait for it to complete.", RUNNING_SELECT_MSQ_QUERY, @@ -711,18 +711,17 @@ public void testFinishedSelectMSQQuery() throws Exception COL_NAME_AND_TYPES, 100L, new ResultSetInformation( - null, - null, null, MSQControllerTask.DUMMY_DATASOURCE_FOR_SELECT, RESULT_ROWS.stream() .map(Arrays::asList) - .collect(Collectors.toList()) + .collect(Collectors.toList()), + ImmutableList.of(new PageInformation(null, null, 0L)) ), null ), response.getEntity()); - Response resultsResponse = resource.doGetResults(FINISHED_SELECT_MSQ_QUERY, null, null, makeOkRequest()); + Response resultsResponse = resource.doGetResults(FINISHED_SELECT_MSQ_QUERY, 0L, makeOkRequest()); Assert.assertEquals(Response.Status.OK.getStatusCode(), resultsResponse.getStatus()); List> rows = new ArrayList<>(); @@ -736,43 +735,23 @@ public void testFinishedSelectMSQQuery() throws Exception resource.deleteQuery(FINISHED_SELECT_MSQ_QUERY, makeOkRequest()).getStatus() ); - Assert.assertEquals( - rows.subList(1, 2), - getResultRowsFromResponse(resource.doGetResults( - FINISHED_SELECT_MSQ_QUERY, - 1L, - null, - makeOkRequest() - )) - ); - Assert.assertEquals( - rows.subList(0, 1), - getResultRowsFromResponse(resource.doGetResults( - FINISHED_SELECT_MSQ_QUERY, - 0L, - 1L, - makeOkRequest() - )) - ); Assert.assertEquals( rows, getResultRowsFromResponse(resource.doGetResults( FINISHED_SELECT_MSQ_QUERY, 0L, - 3L, makeOkRequest() )) ); Assert.assertEquals( Response.Status.BAD_REQUEST.getStatusCode(), - resource.doGetResults(FINISHED_SELECT_MSQ_QUERY, -1L, 3L, makeOkRequest()).getStatus() + resource.doGetResults(FINISHED_SELECT_MSQ_QUERY, -1L, makeOkRequest()).getStatus() ); Assert.assertEquals( Response.Status.BAD_REQUEST.getStatusCode(), - resource.doGetResults(FINISHED_SELECT_MSQ_QUERY, null, -1L, makeOkRequest()).getStatus() + resource.doGetResults(FINISHED_SELECT_MSQ_QUERY, null, makeOkRequest()).getStatus() ); - } @Test @@ -781,7 +760,7 @@ public void testFailedMSQQuery() for (String queryID : ImmutableList.of(ERRORED_SELECT_MSQ_QUERY, ERRORED_INSERT_MSQ_QUERY)) { assertExceptionMessage(resource.doGetStatus(queryID, makeOkRequest()), FAILURE_MSG, Response.Status.OK); assertExceptionMessage( - resource.doGetResults(queryID, null, null, makeOkRequest()), + resource.doGetResults(queryID, 0L, makeOkRequest()), StringUtils.format( "Query[%s] failed. Hit status api for more details.", queryID @@ -807,40 +786,24 @@ public void testFinishedInsertMSQQuery() throws Exception CREATED_TIME, null, 100L, - new ResultSetInformation(null, null, null, "test", null), + new ResultSetInformation(null, "test", null, ImmutableList.of(new PageInformation(null, null, 0))), null ), response.getEntity()); - Response resultsResponse = resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, null, null, makeOkRequest()); + Response resultsResponse = resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, 0L, makeOkRequest()); Assert.assertEquals(Response.Status.OK.getStatusCode(), resultsResponse.getStatus()); - - Assert.assertNull(getResultRowsFromResponse(resource.doGetResults( - FINISHED_INSERT_MSQ_QUERY, - 1L, - null, - makeOkRequest() - ))); - Assert.assertNull(getResultRowsFromResponse(resource.doGetResults( - FINISHED_INSERT_MSQ_QUERY, - 0L, - 1L, - makeOkRequest() - ))); - Assert.assertNull(getResultRowsFromResponse(resource.doGetResults( - FINISHED_INSERT_MSQ_QUERY, - 0L, - 3L, - makeOkRequest() - ))); - Assert.assertEquals( Response.Status.BAD_REQUEST.getStatusCode(), - resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, -1L, 3L, makeOkRequest()).getStatus() + resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, -1L, makeOkRequest()).getStatus() ); Assert.assertEquals( Response.Status.BAD_REQUEST.getStatusCode(), - resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, null, -1L, makeOkRequest()).getStatus() + resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, -1L, makeOkRequest()).getStatus() + ); + Assert.assertEquals( + Response.Status.BAD_REQUEST.getStatusCode(), + resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, null, makeOkRequest()).getStatus() ); } @@ -850,7 +813,7 @@ public void testNonMSQTasks() { for (String queryID : ImmutableList.of(RUNNING_NON_MSQ_TASK, FAILED_NON_MSQ_TASK, FINISHED_NON_MSQ_TASK)) { assertNullResponse(resource.doGetStatus(queryID, makeOkRequest()), Response.Status.NOT_FOUND); - assertNullResponse(resource.doGetResults(queryID, null, null, makeOkRequest()), Response.Status.NOT_FOUND); + assertNullResponse(resource.doGetResults(queryID, 0L, makeOkRequest()), Response.Status.NOT_FOUND); assertNullResponse(resource.deleteQuery(queryID, makeOkRequest()), Response.Status.NOT_FOUND); } } @@ -874,7 +837,7 @@ public void testMSQInsertAcceptedQuery() ); assertExceptionMessage( - resource.doGetResults(ACCEPTED_INSERT_MSQ_TASK, null, null, makeOkRequest()), + resource.doGetResults(ACCEPTED_INSERT_MSQ_TASK, 0L, makeOkRequest()), StringUtils.format( "Query[%s] is currently in [%s] state. Please wait for it to complete.", ACCEPTED_INSERT_MSQ_TASK, @@ -907,7 +870,7 @@ public void testMSQInsertRunningQuery() ); assertExceptionMessage( - resource.doGetResults(RUNNING_INSERT_MSQ_QUERY, null, null, makeOkRequest()), + resource.doGetResults(RUNNING_INSERT_MSQ_QUERY, 0L, makeOkRequest()), StringUtils.format( "Query[%s] is currently in [%s] state. Please wait for it to complete.", RUNNING_INSERT_MSQ_QUERY, @@ -924,15 +887,12 @@ public void testMSQInsertRunningQuery() @Test public void forbiddenTests() { - Assert.assertEquals(Response.Status.FORBIDDEN.getStatusCode(), resource.doGetStatus(RUNNING_SELECT_MSQ_QUERY, makeExpectedReq(CalciteTests.SUPER_USER_AUTH_RESULT)).getStatus()); - Assert.assertEquals(Response.Status.FORBIDDEN.getStatusCode(), resource.doGetResults(RUNNING_SELECT_MSQ_QUERY, - null, - null, + 1L, makeExpectedReq(CalciteTests.SUPER_USER_AUTH_RESULT)).getStatus()); Assert.assertEquals(Response.Status.FORBIDDEN.getStatusCode(), resource.deleteQuery(RUNNING_SELECT_MSQ_QUERY, diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/ResultSetInformationTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/ResultSetInformationTest.java index 14d04b1b76d5..bf29588219d2 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/ResultSetInformationTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/ResultSetInformationTest.java @@ -31,14 +31,17 @@ public class ResultSetInformationTest { public static final ObjectMapper MAPPER = new ObjectMapper(); - public static final ResultSetInformation RESULTS = new ResultSetInformation(ResultFormat.OBJECT, 1L, 1L, "ds", - ImmutableList.of( - ImmutableList.of("1"), - ImmutableList.of("2"), - ImmutableList.of("3") - ) + public static final ResultSetInformation RESULTS = new ResultSetInformation( + ResultFormat.OBJECT, + "ds", + ImmutableList.of( + ImmutableList.of("1"), + ImmutableList.of("2"), + ImmutableList.of("3") + ), + ImmutableList.of(new PageInformation(1L, 1L, 0)) ); - public static final String JSON_STRING = "{\"resultFormat\":\"object\",\"numRows\":1,\"sizeInBytes\":1,\"dataSource\":\"ds\",\"sampleRecords\":[[\"1\"],[\"2\"],[\"3\"]]}"; + public static final String JSON_STRING = "{\"resultFormat\":\"object\",\"dataSource\":\"ds\",\"sampleRecords\":[[\"1\"],[\"2\"],[\"3\"]],\"pageInformationList\":[{\"numRows\":1,\"sizeInBytes\":1,\"id\":0}]}"; @Test @@ -51,7 +54,7 @@ public void sanityTest() throws JsonProcessingException MAPPER.readValue(MAPPER.writeValueAsString(RESULTS), ResultSetInformation.class).hashCode() ); Assert.assertEquals( - "ResultSetInformation{totalRows=1, totalSize=1, resultFormat=object, records=[[1], [2], [3]], dataSource='ds'}", + "ResultSetInformation{resultFormat=object, records=[[1], [2], [3]], dataSource='ds', pageInformationList=[PageInformation{numRows=1, sizeInBytes=1, id=0}]}", RESULTS.toString() ); } From 5580e9abe652ef5a54836e616785c9cc8e079476 Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Sat, 1 Jul 2023 10:49:06 +0530 Subject: [PATCH 2/5] Fix tests --- .../apache/druid/msq/sql/entity/SqlStatementResultTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/SqlStatementResultTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/SqlStatementResultTest.java index a0be3afcf74a..02c0d8400e19 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/SqlStatementResultTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/SqlStatementResultTest.java @@ -43,7 +43,7 @@ public class SqlStatementResultTest + "\"createdAt\":\"2023-05-31T12:00:00.000Z\"," + "\"schema\":[{\"name\":\"_time\",\"type\":\"TIMESTAMP\",\"nativeType\":\"LONG\"},{\"name\":\"alias\",\"type\":\"VARCHAR\",\"nativeType\":\"STRING\"},{\"name\":\"market\",\"type\":\"VARCHAR\",\"nativeType\":\"STRING\"}]," + "\"durationMs\":100," - + "\"result\":{\"resultFormat\":\"object\",\"numRows\":1,\"sizeInBytes\":1,\"dataSource\":\"ds\",\"sampleRecords\":[[\"1\"],[\"2\"],[\"3\"]]}," + + "\"result\":{\"resultFormat\":\"object\",\"dataSource\":\"ds\",\"sampleRecords\":[[\"1\"],[\"2\"],[\"3\"]],\"pageInformationList\":[{\"numRows\":1,\"sizeInBytes\":1,\"id\":0}]}," + "\"errorDetails\":{\"error\":\"druidException\",\"errorCode\":\"QueryNotSupported\",\"persona\":\"USER\",\"category\":\"UNCATEGORIZED\",\"errorMessage\":\"QueryNotSupported\",\"context\":{}}}"; public static final SqlStatementResult SQL_STATEMENT_RESULT = new SqlStatementResult( @@ -87,7 +87,7 @@ public void sanityTest() throws JsonProcessingException + " createdAt=2023-05-31T12:00:00.000Z," + " sqlRowSignature=[ColumnNameAndTypes{colName='_time', sqlTypeName='TIMESTAMP', nativeTypeName='LONG'}, ColumnNameAndTypes{colName='alias', sqlTypeName='VARCHAR', nativeTypeName='STRING'}, ColumnNameAndTypes{colName='market', sqlTypeName='VARCHAR', nativeTypeName='STRING'}]," + " durationInMs=100," - + " resultSetInformation=ResultSetInformation{totalRows=1, totalSize=1, resultFormat=object, records=[[1], [2], [3]], dataSource='ds'}," + + " resultSetInformation=ResultSetInformation{resultFormat=object, records=[[1], [2], [3]], dataSource='ds', pageInformationList=[PageInformation{numRows=1, sizeInBytes=1, id=0}]}," + " errorResponse={error=druidException, errorCode=QueryNotSupported, persona=USER, category=UNCATEGORIZED, errorMessage=QueryNotSupported, context={}}}", SQL_STATEMENT_RESULT.toString() ); From 13e889c79e160979250e0e6cd52669dde3ccfaeb Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Sat, 1 Jul 2023 15:03:55 +0530 Subject: [PATCH 3/5] Fix builds --- .../java/org/apache/druid/msq/sql/SqlStatementResourceTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlStatementResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlStatementResourceTest.java index f6bffb20b928..664ad9474906 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlStatementResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlStatementResourceTest.java @@ -776,7 +776,7 @@ public void testFailedMSQQuery() } @Test - public void testFinishedInsertMSQQuery() throws Exception + public void testFinishedInsertMSQQuery() { Response response = resource.doGetStatus(FINISHED_INSERT_MSQ_QUERY, makeOkRequest()); Assert.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()); From 349cd47a51b4a9d65f967f06aeac10aa7d325b2c Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Mon, 3 Jul 2023 11:29:33 +0530 Subject: [PATCH 4/5] Address review comments --- .../druid/msq/sql/entity/PageInformation.java | 3 ++ .../msq/sql/entity/ResultSetInformation.java | 53 ++++++++++++++----- .../sql/resources/SqlStatementResource.java | 15 ++---- .../sql/SqlMsqStatementResourcePostTest.java | 2 + .../msq/sql/SqlStatementResourceTest.java | 28 +++++----- .../sql/entity/ResultSetInformationTest.java | 7 +-- 6 files changed, 67 insertions(+), 41 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/PageInformation.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/PageInformation.java index 81abb49486b2..2754c52f1fe1 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/PageInformation.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/PageInformation.java @@ -26,6 +26,9 @@ import javax.annotation.Nullable; import java.util.Objects; +/** + * Contains information about a single page in the results. + */ public class PageInformation { @Nullable diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/ResultSetInformation.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/ResultSetInformation.java index 64228523e741..e131fa85c737 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/ResultSetInformation.java +++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/entity/ResultSetInformation.java @@ -31,6 +31,10 @@ public class ResultSetInformation { + @Nullable + private final Long numTotalRows; + @Nullable + private final Long totalSizeInBytes; @Nullable private final ResultFormat resultFormat; @Nullable @@ -38,20 +42,40 @@ public class ResultSetInformation @Nullable private final String dataSource; @Nullable - private final List pageInformationList; + private final List pages; @JsonCreator public ResultSetInformation( + @JsonProperty("numTotalRows") @Nullable Long numTotalRows, + @JsonProperty("totalSizeInBytes") @Nullable Long totalSizeInBytes, @JsonProperty("resultFormat") @Nullable ResultFormat resultFormat, @JsonProperty("dataSource") @Nullable String dataSource, @JsonProperty("sampleRecords") @Nullable List records, - @JsonProperty("pageInformationList") @Nullable List pageInformationList + @JsonProperty("pages") @Nullable List pages ) { + this.numTotalRows = numTotalRows; + this.totalSizeInBytes = totalSizeInBytes; this.resultFormat = resultFormat; this.dataSource = dataSource; this.records = records; - this.pageInformationList = pageInformationList; + this.pages = pages; + } + + @JsonProperty + @Nullable + @JsonInclude(JsonInclude.Include.NON_NULL) + public Long getNumTotalRows() + { + return numTotalRows; + } + + @JsonProperty + @Nullable + @JsonInclude(JsonInclude.Include.NON_NULL) + public Long getTotalSizeInBytes() + { + return totalSizeInBytes; } @JsonProperty @@ -81,9 +105,9 @@ public List getRecords() @JsonProperty @Nullable @JsonInclude(JsonInclude.Include.NON_NULL) - public List getPageInformationList() + public List getPages() { - return pageInformationList; + return pages; } @Override @@ -96,26 +120,31 @@ public boolean equals(Object o) return false; } ResultSetInformation that = (ResultSetInformation) o; - return resultFormat == that.resultFormat - && Objects.equals(records, that.records) - && Objects.equals(dataSource, that.dataSource) - && Objects.equals(pageInformationList, that.pageInformationList); + return Objects.equals(numTotalRows, that.numTotalRows) && Objects.equals( + totalSizeInBytes, + that.totalSizeInBytes + ) && resultFormat == that.resultFormat && Objects.equals(records, that.records) && Objects.equals( + dataSource, + that.dataSource + ) && Objects.equals(pages, that.pages); } @Override public int hashCode() { - return Objects.hash(resultFormat, records, dataSource, pageInformationList); + return Objects.hash(numTotalRows, totalSizeInBytes, resultFormat, records, dataSource, pages); } @Override public String toString() { return "ResultSetInformation{" + - "resultFormat=" + resultFormat + + "numTotalRows=" + numTotalRows + + ", totalSizeInBytes=" + totalSizeInBytes + + ", resultFormat=" + resultFormat + ", records=" + records + ", dataSource='" + dataSource + '\'' + - ", pageInformationList=" + pageInformationList + + ", pages=" + pages + '}'; } } 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 770126311eb9..ce30284e5936 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 @@ -285,16 +285,7 @@ public Response doGetResults( } final AuthenticationResult authenticationResult = AuthorizationUtils.authenticationResultFromRequest(req); - if (page == null) { - return buildNonOkResponse( - DruidException.forPersona(DruidException.Persona.USER) - .ofCategory(DruidException.Category.INVALID_INPUT) - .build( - "Page cannot be null." - ) - ); - } - if (page < 0) { + if (page != null && page < 0) { return buildNonOkResponse( DruidException.forPersona(DruidException.Persona.USER) .ofCategory(DruidException.Category.INVALID_INPUT) @@ -342,7 +333,7 @@ public Response doGetResults( return Response.ok().build(); } - if (page > 0) { + if (page != null && page > 0) { // Results from task report are only present as one page. return buildNonOkResponse( DruidException.forPersona(DruidException.Persona.USER) @@ -586,6 +577,8 @@ private Optional getSampleResults( isSelectQuery ); return Optional.of(new ResultSetInformation( + rowsAndSize.orElse(new Pair<>(null, null)).lhs, + rowsAndSize.orElse(new Pair<>(null, null)).rhs, null, dataSource, // only populate sample results in case a select query is successful diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlMsqStatementResourcePostTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlMsqStatementResourcePostTest.java index 6f8fd82567f7..51e10a93b2d3 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlMsqStatementResourcePostTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlMsqStatementResourcePostTest.java @@ -113,6 +113,8 @@ public void testMSQSelectQueryTest() throws IOException ), MSQTestOverlordServiceClient.DURATION, new ResultSetInformation( + 6L, + 316L, null, MSQControllerTask.DUMMY_DATASOURCE_FOR_SELECT, objectMapper.readValue( diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlStatementResourceTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlStatementResourceTest.java index 664ad9474906..6a19f3f792ef 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlStatementResourceTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/SqlStatementResourceTest.java @@ -711,6 +711,8 @@ public void testFinishedSelectMSQQuery() throws Exception COL_NAME_AND_TYPES, 100L, new ResultSetInformation( + null, + null, null, MSQControllerTask.DUMMY_DATASOURCE_FOR_SELECT, RESULT_ROWS.stream() @@ -745,12 +747,17 @@ public void testFinishedSelectMSQQuery() throws Exception ); Assert.assertEquals( - Response.Status.BAD_REQUEST.getStatusCode(), - resource.doGetResults(FINISHED_SELECT_MSQ_QUERY, -1L, makeOkRequest()).getStatus() + rows, + getResultRowsFromResponse(resource.doGetResults( + FINISHED_SELECT_MSQ_QUERY, + null, + makeOkRequest() + )) ); + Assert.assertEquals( Response.Status.BAD_REQUEST.getStatusCode(), - resource.doGetResults(FINISHED_SELECT_MSQ_QUERY, null, makeOkRequest()).getStatus() + resource.doGetResults(FINISHED_SELECT_MSQ_QUERY, -1L, makeOkRequest()).getStatus() ); } @@ -786,26 +793,17 @@ public void testFinishedInsertMSQQuery() CREATED_TIME, null, 100L, - new ResultSetInformation(null, "test", null, ImmutableList.of(new PageInformation(null, null, 0))), + new ResultSetInformation(null, null, null, "test", null, ImmutableList.of(new PageInformation(null, null, 0))), null ), response.getEntity()); - Response resultsResponse = resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, 0L, makeOkRequest()); - Assert.assertEquals(Response.Status.OK.getStatusCode(), resultsResponse.getStatus()); + Assert.assertEquals(Response.Status.OK.getStatusCode(), resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, 0L, makeOkRequest()).getStatus()); + Assert.assertEquals(Response.Status.OK.getStatusCode(), resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, null, makeOkRequest()).getStatus()); Assert.assertEquals( Response.Status.BAD_REQUEST.getStatusCode(), resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, -1L, makeOkRequest()).getStatus() ); - Assert.assertEquals( - Response.Status.BAD_REQUEST.getStatusCode(), - resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, -1L, makeOkRequest()).getStatus() - ); - Assert.assertEquals( - Response.Status.BAD_REQUEST.getStatusCode(), - resource.doGetResults(FINISHED_INSERT_MSQ_QUERY, null, makeOkRequest()).getStatus() - ); - } @Test diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/ResultSetInformationTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/ResultSetInformationTest.java index bf29588219d2..8e40d8daf1a7 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/ResultSetInformationTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/ResultSetInformationTest.java @@ -32,6 +32,8 @@ public class ResultSetInformationTest public static final ObjectMapper MAPPER = new ObjectMapper(); public static final ResultSetInformation RESULTS = new ResultSetInformation( + 1L, + 1L, ResultFormat.OBJECT, "ds", ImmutableList.of( @@ -41,8 +43,7 @@ public class ResultSetInformationTest ), ImmutableList.of(new PageInformation(1L, 1L, 0)) ); - public static final String JSON_STRING = "{\"resultFormat\":\"object\",\"dataSource\":\"ds\",\"sampleRecords\":[[\"1\"],[\"2\"],[\"3\"]],\"pageInformationList\":[{\"numRows\":1,\"sizeInBytes\":1,\"id\":0}]}"; - + public static final String JSON_STRING = "{\"numTotalRows\":1,\"totalSizeInBytes\":1,\"resultFormat\":\"object\",\"dataSource\":\"ds\",\"sampleRecords\":[[\"1\"],[\"2\"],[\"3\"]],\"pages\":[{\"numRows\":1,\"sizeInBytes\":1,\"id\":0}]}"; @Test public void sanityTest() throws JsonProcessingException @@ -54,7 +55,7 @@ public void sanityTest() throws JsonProcessingException MAPPER.readValue(MAPPER.writeValueAsString(RESULTS), ResultSetInformation.class).hashCode() ); Assert.assertEquals( - "ResultSetInformation{resultFormat=object, records=[[1], [2], [3]], dataSource='ds', pageInformationList=[PageInformation{numRows=1, sizeInBytes=1, id=0}]}", + "ResultSetInformation{numTotalRows=1, totalSizeInBytes=1, resultFormat=object, records=[[1], [2], [3]], dataSource='ds', pages=[PageInformation{numRows=1, sizeInBytes=1, id=0}]}", RESULTS.toString() ); } From d04ffddd0bac6114925d34d34c7ed4da670ffc0c Mon Sep 17 00:00:00 2001 From: Adarsh Sanjeev Date: Mon, 3 Jul 2023 12:14:22 +0530 Subject: [PATCH 5/5] Fix builds --- .../apache/druid/msq/sql/entity/SqlStatementResultTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/SqlStatementResultTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/SqlStatementResultTest.java index 02c0d8400e19..1409450c6972 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/SqlStatementResultTest.java +++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/sql/entity/SqlStatementResultTest.java @@ -43,7 +43,7 @@ public class SqlStatementResultTest + "\"createdAt\":\"2023-05-31T12:00:00.000Z\"," + "\"schema\":[{\"name\":\"_time\",\"type\":\"TIMESTAMP\",\"nativeType\":\"LONG\"},{\"name\":\"alias\",\"type\":\"VARCHAR\",\"nativeType\":\"STRING\"},{\"name\":\"market\",\"type\":\"VARCHAR\",\"nativeType\":\"STRING\"}]," + "\"durationMs\":100," - + "\"result\":{\"resultFormat\":\"object\",\"dataSource\":\"ds\",\"sampleRecords\":[[\"1\"],[\"2\"],[\"3\"]],\"pageInformationList\":[{\"numRows\":1,\"sizeInBytes\":1,\"id\":0}]}," + + "\"result\":{\"numTotalRows\":1,\"totalSizeInBytes\":1,\"resultFormat\":\"object\",\"dataSource\":\"ds\",\"sampleRecords\":[[\"1\"],[\"2\"],[\"3\"]],\"pages\":[{\"numRows\":1,\"sizeInBytes\":1,\"id\":0}]}," + "\"errorDetails\":{\"error\":\"druidException\",\"errorCode\":\"QueryNotSupported\",\"persona\":\"USER\",\"category\":\"UNCATEGORIZED\",\"errorMessage\":\"QueryNotSupported\",\"context\":{}}}"; public static final SqlStatementResult SQL_STATEMENT_RESULT = new SqlStatementResult( @@ -87,7 +87,7 @@ public void sanityTest() throws JsonProcessingException + " createdAt=2023-05-31T12:00:00.000Z," + " sqlRowSignature=[ColumnNameAndTypes{colName='_time', sqlTypeName='TIMESTAMP', nativeTypeName='LONG'}, ColumnNameAndTypes{colName='alias', sqlTypeName='VARCHAR', nativeTypeName='STRING'}, ColumnNameAndTypes{colName='market', sqlTypeName='VARCHAR', nativeTypeName='STRING'}]," + " durationInMs=100," - + " resultSetInformation=ResultSetInformation{resultFormat=object, records=[[1], [2], [3]], dataSource='ds', pageInformationList=[PageInformation{numRows=1, sizeInBytes=1, id=0}]}," + + " resultSetInformation=ResultSetInformation{numTotalRows=1, totalSizeInBytes=1, resultFormat=object, records=[[1], [2], [3]], dataSource='ds', pages=[PageInformation{numRows=1, sizeInBytes=1, id=0}]}," + " errorResponse={error=druidException, errorCode=QueryNotSupported, persona=USER, category=UNCATEGORIZED, errorMessage=QueryNotSupported, context={}}}", SQL_STATEMENT_RESULT.toString() );