Add task report fields in response of SQL statements endpoint (#16808)

If the optional query parameter detail is supplied, then the response also includes the following:

 * A stages object that summarizes information about the different stages being used for query execution, such as stage number, phase, start time, duration, input and output information, processing methods, and partitioning.
* A counters object that provides details on the rows, bytes, and files processed at various stages for each worker across different channels, along with sort progress.
* A warnings object that provides details about any warnings.
This commit is contained in:
Akshat Jain 2024-08-01 10:26:04 +05:30 committed by GitHub
parent 01f6cfcbf5
commit bb4d6cc001
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 682 additions and 90 deletions

View File

@ -629,10 +629,21 @@ Retrieves information about the query associated with the given query ID. The re
- `sizeInBytes`: the size of the page.
- `id`: the page number that you can use to reference a specific page when you get query results.
If the optional query parameter `detail` is supplied, then the response also includes the following:
- A `stages` object that summarizes information about the different stages being used for query execution, such as stage number, phase, start time, duration, input and output information, processing methods, and partitioning.
- A `counters` object that provides details on the rows, bytes, and files processed at various stages for each worker across different channels, along with sort progress.
- A `warnings` object that provides details about any warnings.
#### URL
`GET` `/druid/v2/sql/statements/{queryId}`
#### Query parameters
* `detail` (optional)
* Type: Boolean
* Default: false
* Fetch additional details about the query, which includes the information about different stages, counters for each stage, and any warnings.
#### Responses
<Tabs>
@ -672,7 +683,7 @@ The following example retrieves the status of a query with specified ID `query-9
```shell
curl "http://ROUTER_IP:ROUTER_PORT/druid/v2/sql/statements/query-9b93f6f7-ab0e-48f5-986a-3520f84f0804"
curl "http://ROUTER_IP:ROUTER_PORT/druid/v2/sql/statements/query-9b93f6f7-ab0e-48f5-986a-3520f84f0804?detail=true"
```
</TabItem>
@ -680,7 +691,7 @@ curl "http://ROUTER_IP:ROUTER_PORT/druid/v2/sql/statements/query-9b93f6f7-ab0e-4
```HTTP
GET /druid/v2/sql/statements/query-9b93f6f7-ab0e-48f5-986a-3520f84f0804 HTTP/1.1
GET /druid/v2/sql/statements/query-9b93f6f7-ab0e-48f5-986a-3520f84f0804?detail=true HTTP/1.1
Host: http://ROUTER_IP:ROUTER_PORT
```
@ -835,7 +846,422 @@ Host: http://ROUTER_IP:ROUTER_PORT
"sizeInBytes": 375
}
]
}
},
"stages": [
{
"stageNumber": 0,
"definition": {
"id": "query-9b93f6f7-ab0e-48f5-986a-3520f84f0804_0",
"input": [
{
"type": "table",
"dataSource": "wikipedia",
"intervals": [
"-146136543-09-08T08:23:32.096Z/146140482-04-24T15:36:27.903Z"
],
"filter": {
"type": "equals",
"column": "user",
"matchValueType": "STRING",
"matchValue": "BlueMoon2662"
},
"filterFields": [
"user"
]
}
],
"processor": {
"type": "scan",
"query": {
"queryType": "scan",
"dataSource": {
"type": "inputNumber",
"inputNumber": 0
},
"intervals": {
"type": "intervals",
"intervals": [
"-146136543-09-08T08:23:32.096Z/146140482-04-24T15:36:27.903Z"
]
},
"virtualColumns": [
{
"type": "expression",
"name": "v0",
"expression": "'BlueMoon2662'",
"outputType": "STRING"
}
],
"resultFormat": "compactedList",
"limit": 1001,
"filter": {
"type": "equals",
"column": "user",
"matchValueType": "STRING",
"matchValue": "BlueMoon2662"
},
"columns": [
"__time",
"added",
"channel",
"cityName",
"comment",
"commentLength",
"countryIsoCode",
"countryName",
"deleted",
"delta",
"deltaBucket",
"diffUrl",
"flags",
"isAnonymous",
"isMinor",
"isNew",
"isRobot",
"isUnpatrolled",
"metroCode",
"namespace",
"page",
"regionIsoCode",
"regionName",
"v0"
],
"context": {
"__resultFormat": "array",
"__user": "allowAll",
"enableWindowing": true,
"executionMode": "async",
"finalize": true,
"maxNumTasks": 2,
"maxParseExceptions": 0,
"queryId": "33b53acb-7533-4880-a81b-51c16c489eab",
"scanSignature": "[{\"name\":\"__time\",\"type\":\"LONG\"},{\"name\":\"added\",\"type\":\"LONG\"},{\"name\":\"channel\",\"type\":\"STRING\"},{\"name\":\"cityName\",\"type\":\"STRING\"},{\"name\":\"comment\",\"type\":\"STRING\"},{\"name\":\"commentLength\",\"type\":\"LONG\"},{\"name\":\"countryIsoCode\",\"type\":\"STRING\"},{\"name\":\"countryName\",\"type\":\"STRING\"},{\"name\":\"deleted\",\"type\":\"LONG\"},{\"name\":\"delta\",\"type\":\"LONG\"},{\"name\":\"deltaBucket\",\"type\":\"LONG\"},{\"name\":\"diffUrl\",\"type\":\"STRING\"},{\"name\":\"flags\",\"type\":\"STRING\"},{\"name\":\"isAnonymous\",\"type\":\"STRING\"},{\"name\":\"isMinor\",\"type\":\"STRING\"},{\"name\":\"isNew\",\"type\":\"STRING\"},{\"name\":\"isRobot\",\"type\":\"STRING\"},{\"name\":\"isUnpatrolled\",\"type\":\"STRING\"},{\"name\":\"metroCode\",\"type\":\"STRING\"},{\"name\":\"namespace\",\"type\":\"STRING\"},{\"name\":\"page\",\"type\":\"STRING\"},{\"name\":\"regionIsoCode\",\"type\":\"STRING\"},{\"name\":\"regionName\",\"type\":\"STRING\"},{\"name\":\"v0\",\"type\":\"STRING\"}]",
"sqlOuterLimit": 1001,
"sqlQueryId": "33b53acb-7533-4880-a81b-51c16c489eab",
"sqlStringifyArrays": false
},
"columnTypes": [
"LONG",
"LONG",
"STRING",
"STRING",
"STRING",
"LONG",
"STRING",
"STRING",
"LONG",
"LONG",
"LONG",
"STRING",
"STRING",
"STRING",
"STRING",
"STRING",
"STRING",
"STRING",
"STRING",
"STRING",
"STRING",
"STRING",
"STRING",
"STRING"
],
"granularity": {
"type": "all"
},
"legacy": false
}
},
"signature": [
{
"name": "__boost",
"type": "LONG"
},
{
"name": "__time",
"type": "LONG"
},
{
"name": "added",
"type": "LONG"
},
{
"name": "channel",
"type": "STRING"
},
{
"name": "cityName",
"type": "STRING"
},
{
"name": "comment",
"type": "STRING"
},
{
"name": "commentLength",
"type": "LONG"
},
{
"name": "countryIsoCode",
"type": "STRING"
},
{
"name": "countryName",
"type": "STRING"
},
{
"name": "deleted",
"type": "LONG"
},
{
"name": "delta",
"type": "LONG"
},
{
"name": "deltaBucket",
"type": "LONG"
},
{
"name": "diffUrl",
"type": "STRING"
},
{
"name": "flags",
"type": "STRING"
},
{
"name": "isAnonymous",
"type": "STRING"
},
{
"name": "isMinor",
"type": "STRING"
},
{
"name": "isNew",
"type": "STRING"
},
{
"name": "isRobot",
"type": "STRING"
},
{
"name": "isUnpatrolled",
"type": "STRING"
},
{
"name": "metroCode",
"type": "STRING"
},
{
"name": "namespace",
"type": "STRING"
},
{
"name": "page",
"type": "STRING"
},
{
"name": "regionIsoCode",
"type": "STRING"
},
{
"name": "regionName",
"type": "STRING"
},
{
"name": "v0",
"type": "STRING"
}
],
"shuffleSpec": {
"type": "mix"
},
"maxWorkerCount": 1
},
"phase": "FINISHED",
"workerCount": 1,
"partitionCount": 1,
"shuffle": "mix",
"output": "localStorage",
"startTime": "2024-07-31T15:20:21.255Z",
"duration": 103
},
{
"stageNumber": 1,
"definition": {
"id": "query-9b93f6f7-ab0e-48f5-986a-3520f84f0804_1",
"input": [
{
"type": "stage",
"stage": 0
}
],
"processor": {
"type": "limit",
"limit": 1001
},
"signature": [
{
"name": "__boost",
"type": "LONG"
},
{
"name": "__time",
"type": "LONG"
},
{
"name": "added",
"type": "LONG"
},
{
"name": "channel",
"type": "STRING"
},
{
"name": "cityName",
"type": "STRING"
},
{
"name": "comment",
"type": "STRING"
},
{
"name": "commentLength",
"type": "LONG"
},
{
"name": "countryIsoCode",
"type": "STRING"
},
{
"name": "countryName",
"type": "STRING"
},
{
"name": "deleted",
"type": "LONG"
},
{
"name": "delta",
"type": "LONG"
},
{
"name": "deltaBucket",
"type": "LONG"
},
{
"name": "diffUrl",
"type": "STRING"
},
{
"name": "flags",
"type": "STRING"
},
{
"name": "isAnonymous",
"type": "STRING"
},
{
"name": "isMinor",
"type": "STRING"
},
{
"name": "isNew",
"type": "STRING"
},
{
"name": "isRobot",
"type": "STRING"
},
{
"name": "isUnpatrolled",
"type": "STRING"
},
{
"name": "metroCode",
"type": "STRING"
},
{
"name": "namespace",
"type": "STRING"
},
{
"name": "page",
"type": "STRING"
},
{
"name": "regionIsoCode",
"type": "STRING"
},
{
"name": "regionName",
"type": "STRING"
},
{
"name": "v0",
"type": "STRING"
}
],
"shuffleSpec": {
"type": "maxCount",
"clusterBy": {
"columns": [
{
"columnName": "__boost",
"order": "ASCENDING"
}
]
},
"partitions": 1
},
"maxWorkerCount": 1
},
"phase": "FINISHED",
"workerCount": 1,
"partitionCount": 1,
"shuffle": "globalSort",
"output": "localStorage",
"startTime": "2024-07-31T15:20:21.355Z",
"duration": 10,
"sort": true
}
],
"counters": {
"0": {
"0": {
"input0": {
"type": "channel",
"rows": [
24433
],
"bytes": [
7393933
],
"files": [
22
],
"totalFiles": [
22
]
}
}
},
"1": {
"0": {
"sortProgress": {
"type": "sortProgress",
"totalMergingLevels": -1,
"levelToTotalBatches": {},
"levelToMergedBatches": {},
"totalMergersForUltimateLevel": -1,
"triviallyComplete": true,
"progressDigest": 1
}
}
}
},
"warnings": []
}
```
</details>

View File

@ -108,4 +108,14 @@ public class CounterSnapshotsTree
}
}
}
@Override
public String toString()
{
synchronized (snapshotsMap) {
return "CounterSnapshotsTree{" +
"snapshotsMap=" + snapshotsMap +
'}';
}
}
}

View File

@ -23,12 +23,14 @@ import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import org.apache.druid.error.ErrorResponse;
import org.apache.druid.msq.counters.CounterSnapshotsTree;
import org.apache.druid.msq.indexing.error.MSQErrorReport;
import org.apache.druid.msq.indexing.report.MSQStagesReport;
import org.apache.druid.msq.sql.SqlStatementState;
import org.joda.time.DateTime;
import javax.annotation.Nullable;
import java.util.List;
import java.util.Objects;
public class SqlStatementResult
{
@ -51,6 +53,27 @@ public class SqlStatementResult
@Nullable
private final ErrorResponse errorResponse;
@Nullable
private final MSQStagesReport stages;
@Nullable
private final CounterSnapshotsTree counters;
@Nullable
private final List<MSQErrorReport> warnings;
public SqlStatementResult(
String queryId,
SqlStatementState state,
DateTime createdAt,
List<ColumnNameAndTypes> sqlRowSignature,
Long durationMs,
ResultSetInformation resultSetInformation,
ErrorResponse errorResponse
)
{
this(queryId, state, createdAt, sqlRowSignature, durationMs, resultSetInformation, errorResponse, null, null, null);
}
@JsonCreator
public SqlStatementResult(
@ -67,8 +90,13 @@ public class SqlStatementResult
@Nullable @JsonProperty("result")
ResultSetInformation resultSetInformation,
@Nullable @JsonProperty("errorDetails")
ErrorResponse errorResponse
ErrorResponse errorResponse,
@Nullable @JsonProperty("stages")
MSQStagesReport stages,
@Nullable @JsonProperty("counters")
CounterSnapshotsTree counters,
@Nullable @JsonProperty("warnings")
List<MSQErrorReport> warnings
)
{
this.queryId = queryId;
@ -78,6 +106,9 @@ public class SqlStatementResult
this.durationMs = durationMs;
this.resultSetInformation = resultSetInformation;
this.errorResponse = errorResponse;
this.stages = stages;
this.counters = counters;
this.warnings = warnings;
}
@JsonProperty
@ -130,41 +161,28 @@ public class SqlStatementResult
return errorResponse;
}
@Override
public boolean equals(Object o)
@JsonProperty("stages")
@Nullable
@JsonInclude(JsonInclude.Include.NON_NULL)
public MSQStagesReport getStages()
{
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}
SqlStatementResult that = (SqlStatementResult) o;
return Objects.equals(queryId, that.queryId) && state == that.state && Objects.equals(
createdAt,
that.createdAt
) && Objects.equals(sqlRowSignature, that.sqlRowSignature) && Objects.equals(
durationMs,
that.durationMs
) && Objects.equals(resultSetInformation, that.resultSetInformation) && Objects.equals(
errorResponse == null ? null : errorResponse.getAsMap(),
that.errorResponse == null ? null : that.errorResponse.getAsMap()
);
return stages;
}
@Override
public int hashCode()
@JsonProperty("counters")
@Nullable
@JsonInclude(JsonInclude.Include.NON_NULL)
public CounterSnapshotsTree getCounters()
{
return Objects.hash(
queryId,
state,
createdAt,
sqlRowSignature,
durationMs,
resultSetInformation,
errorResponse == null ? null : errorResponse.getAsMap()
);
return counters;
}
@JsonProperty("warnings")
@Nullable
@JsonInclude(JsonInclude.Include.NON_NULL)
public List<MSQErrorReport> getWarnings()
{
return warnings;
}
@Override
@ -180,6 +198,9 @@ public class SqlStatementResult
", errorResponse=" + (errorResponse == null
? "{}"
: errorResponse.getAsMap().toString()) +
", stages=" + stages +
", counters=" + counters +
", warnings=" + warnings +
'}';
}
}

View File

@ -113,6 +113,7 @@ import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.function.Supplier;
import java.util.stream.Collectors;
@ -231,7 +232,9 @@ public class SqlStatementResource
@Path("/{id}")
@Produces(MediaType.APPLICATION_JSON)
public Response doGetStatus(
@PathParam("id") final String queryId, @Context final HttpServletRequest req
@PathParam("id") final String queryId,
@QueryParam("detail") boolean detail,
@Context final HttpServletRequest req
)
{
try {
@ -242,7 +245,8 @@ public class SqlStatementResource
queryId,
authenticationResult,
true,
Action.READ
Action.READ,
detail
);
if (sqlStatementResult.isPresent()) {
@ -369,7 +373,8 @@ public class SqlStatementResource
queryId,
authenticationResult,
false,
Action.WRITE
Action.WRITE,
false
);
if (sqlStatementResult.isPresent()) {
switch (sqlStatementResult.get().getState()) {
@ -479,7 +484,7 @@ public class SqlStatementResource
}
String taskId = String.valueOf(firstRow[0]);
Optional<SqlStatementResult> statementResult = getStatementStatus(taskId, authenticationResult, true, Action.READ);
Optional<SqlStatementResult> statementResult = getStatementStatus(taskId, authenticationResult, true, Action.READ, false);
if (statementResult.isPresent()) {
return Response.status(Response.Status.OK).entity(statementResult.get()).build();
@ -565,7 +570,8 @@ public class SqlStatementResource
String queryId,
AuthenticationResult authenticationResult,
boolean withResults,
Action forAction
Action forAction,
boolean detail
) throws DruidException
{
TaskStatusResponse taskResponse = contactOverlord(overlordClient.taskStatus(queryId), queryId);
@ -582,14 +588,29 @@ public class SqlStatementResource
MSQControllerTask msqControllerTask = getMSQControllerTaskAndCheckPermission(queryId, authenticationResult, forAction);
SqlStatementState sqlStatementState = SqlStatementResourceHelper.getSqlStatementState(statusPlus);
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,
contactOverlord(overlordClient.taskReportAsMap(queryId), queryId),
jsonMapper
msqTaskReportPayloadSupplier.get().orElse(null),
jsonMapper,
detail
);
} else {
Optional<List<ColumnNameAndTypes>> signature = SqlStatementResourceHelper.getSignature(msqControllerTask);
@ -605,7 +626,10 @@ public class SqlStatementResource
sqlStatementState,
msqControllerTask.getQuerySpec().getDestination()
).orElse(null) : null,
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
));
}
}

View File

@ -250,11 +250,12 @@ public class SqlStatementResourceHelper
TaskStatusResponse taskResponse,
TaskStatusPlus statusPlus,
SqlStatementState sqlStatementState,
TaskReport.ReportMap msqPayload,
ObjectMapper jsonMapper
MSQTaskReportPayload msqTaskReportPayload,
ObjectMapper jsonMapper,
boolean detail
)
{
final MSQErrorReport exceptionDetails = getQueryExceptionDetails(getPayload(msqPayload));
final MSQErrorReport exceptionDetails = getQueryExceptionDetails(msqTaskReportPayload);
final MSQFault fault = exceptionDetails == null ? null : exceptionDetails.getFault();
if (exceptionDetails == null || fault == null) {
return Optional.of(new SqlStatementResult(
@ -267,7 +268,10 @@ public class SqlStatementResourceHelper
DruidException.forPersona(DruidException.Persona.DEVELOPER)
.ofCategory(DruidException.Category.UNCATEGORIZED)
.build("%s", taskResponse.getStatus().getErrorMsg())
.toErrorResponse()
.toErrorResponse(),
detail ? getQueryStagesReport(msqTaskReportPayload) : null,
detail ? getQueryCounters(msqTaskReportPayload) : null,
detail ? getQueryWarningDetails(msqTaskReportPayload) : null
));
}
@ -293,7 +297,10 @@ public class SqlStatementResourceHelper
ex.withContext(exceptionContext);
return ex;
}
}).toErrorResponse()
}).toErrorResponse(),
detail ? getQueryStagesReport(msqTaskReportPayload) : null,
detail ? getQueryCounters(msqTaskReportPayload) : null,
detail ? getQueryWarningDetails(msqTaskReportPayload) : null
));
}
@ -353,7 +360,7 @@ public class SqlStatementResourceHelper
}
@Nullable
public static MSQStagesReport.Stage getFinalStage(MSQTaskReportPayload msqTaskReportPayload)
public static MSQStagesReport.Stage getFinalStage(@Nullable MSQTaskReportPayload msqTaskReportPayload)
{
if (msqTaskReportPayload == null || msqTaskReportPayload.getStages().getStages() == null) {
return null;
@ -369,11 +376,29 @@ public class SqlStatementResourceHelper
}
@Nullable
private static MSQErrorReport getQueryExceptionDetails(MSQTaskReportPayload payload)
private static MSQErrorReport getQueryExceptionDetails(@Nullable MSQTaskReportPayload payload)
{
return payload == null ? null : payload.getStatus().getErrorReport();
}
@Nullable
public static List<MSQErrorReport> getQueryWarningDetails(@Nullable MSQTaskReportPayload payload)
{
return payload == null ? null : new ArrayList<>(payload.getStatus().getWarningReports());
}
@Nullable
public static MSQStagesReport getQueryStagesReport(@Nullable MSQTaskReportPayload payload)
{
return payload == null ? null : payload.getStages();
}
@Nullable
public static CounterSnapshotsTree getQueryCounters(@Nullable MSQTaskReportPayload payload)
{
return payload == null ? null : payload.getCounters();
}
@Nullable
public static MSQTaskReportPayload getPayload(TaskReport.ReportMap reportMap)
{

View File

@ -71,14 +71,6 @@ public class SqlStatementResultTest
{
Assert.assertEquals(JSON_STRING, MAPPER.writeValueAsString(SQL_STATEMENT_RESULT));
Assert.assertEquals(
SQL_STATEMENT_RESULT,
MAPPER.readValue(MAPPER.writeValueAsString(SQL_STATEMENT_RESULT), SqlStatementResult.class)
);
Assert.assertEquals(
SQL_STATEMENT_RESULT.hashCode(),
MAPPER.readValue(MAPPER.writeValueAsString(SQL_STATEMENT_RESULT), SqlStatementResult.class).hashCode()
);
Assert.assertEquals(
"SqlStatementResult{"
+ "queryId='q1',"
@ -87,7 +79,10 @@ public class SqlStatementResultTest
+ " 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{numTotalRows=1, totalSizeInBytes=1, resultFormat=object, records=null, dataSource='ds', pages=[PageInformation{id=0, numRows=null, sizeInBytes=1, worker=null, partition=null}]},"
+ " errorResponse={error=druidException, errorCode=QueryNotSupported, persona=USER, category=UNCATEGORIZED, errorMessage=QueryNotSupported, context={}}}",
+ " errorResponse={error=druidException, errorCode=QueryNotSupported, persona=USER, category=UNCATEGORIZED, errorMessage=QueryNotSupported, context={}},"
+ " stages=null,"
+ " counters=null,"
+ " warnings=null}",
SQL_STATEMENT_RESULT.toString()
);
}

View File

@ -206,7 +206,7 @@ public class SqlMSQStatementResourcePostTest extends MSQTestBase
new ResultSetInformation(0L, 0L, null, "foo1", null, null),
null
);
Assert.assertEquals(expected, actual);
assertSqlStatementResult(expected, actual);
}
@Test
@ -236,7 +236,7 @@ public class SqlMSQStatementResourcePostTest extends MSQTestBase
new ResultSetInformation(0L, 0L, null, "foo1", null, null),
null
);
Assert.assertEquals(expected, actual);
assertSqlStatementResult(expected, actual);
}
@Test
@ -282,7 +282,7 @@ public class SqlMSQStatementResourcePostTest extends MSQTestBase
}
}).toErrorResponse()
);
Assert.assertEquals(expected, actual);
assertSqlStatementResult(expected, actual);
}
@Test
@ -687,11 +687,11 @@ public class SqlMSQStatementResourcePostTest extends MSQTestBase
new ResultSetInformation(NullHandling.sqlCompatible() ? 6L : 5L, 0L, null, "foo1", null, null),
null
);
Assert.assertEquals(expected, actual);
assertSqlStatementResult(expected, actual);
Response getResponse = resource.doGetStatus(actual.getQueryId(), SqlStatementResourceTest.makeOkRequest());
Response getResponse = resource.doGetStatus(actual.getQueryId(), false, SqlStatementResourceTest.makeOkRequest());
Assert.assertEquals(Response.Status.OK.getStatusCode(), getResponse.getStatus());
Assert.assertEquals(expected, getResponse.getEntity());
assertSqlStatementResult(expected, (SqlStatementResult) getResponse.getEntity());
Response resultsResponse = resource.doGetResults(
actual.getQueryId(),
@ -730,11 +730,11 @@ public class SqlMSQStatementResourcePostTest extends MSQTestBase
new ResultSetInformation(NullHandling.sqlCompatible() ? 6L : 5L, 0L, null, "foo1", null, null),
null
);
Assert.assertEquals(expected, actual);
assertSqlStatementResult(expected, actual);
Response getResponse = resource.doGetStatus(actual.getQueryId(), SqlStatementResourceTest.makeOkRequest());
Response getResponse = resource.doGetStatus(actual.getQueryId(), false, SqlStatementResourceTest.makeOkRequest());
Assert.assertEquals(Response.Status.OK.getStatusCode(), getResponse.getStatus());
Assert.assertEquals(expected, getResponse.getEntity());
assertSqlStatementResult(expected, (SqlStatementResult) getResponse.getEntity());
Response resultsResponse = resource.doGetResults(
actual.getQueryId(),
@ -754,4 +754,27 @@ public class SqlMSQStatementResourcePostTest extends MSQTestBase
return context;
}
private void assertSqlStatementResult(SqlStatementResult expected, SqlStatementResult actual)
{
Assert.assertEquals(expected.getQueryId(), actual.getQueryId());
Assert.assertEquals(expected.getCreatedAt(), actual.getCreatedAt());
Assert.assertEquals(expected.getSqlRowSignature(), actual.getSqlRowSignature());
Assert.assertEquals(expected.getDurationMs(), actual.getDurationMs());
Assert.assertEquals(expected.getStages(), actual.getStages());
Assert.assertEquals(expected.getState(), actual.getState());
Assert.assertEquals(expected.getWarnings(), actual.getWarnings());
Assert.assertEquals(expected.getResultSetInformation(), actual.getResultSetInformation());
if (actual.getCounters() == null || expected.getCounters() == null) {
Assert.assertEquals(expected.getCounters(), actual.getCounters());
} else {
Assert.assertEquals(expected.getCounters().toString(), actual.getCounters().toString());
}
if (actual.getErrorResponse() == null || expected.getErrorResponse() == null) {
Assert.assertEquals(expected.getErrorResponse(), actual.getErrorResponse());
} else {
Assert.assertEquals(expected.getErrorResponse().getAsMap(), actual.getErrorResponse().getAsMap());
}
}
}

View File

@ -97,6 +97,7 @@ import javax.ws.rs.core.Response;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
@ -449,7 +450,11 @@ public class SqlStatementResourceTest extends MSQTestBase
)));
Mockito.when(indexingServiceClient.taskReportAsMap(FINISHED_SELECT_MSQ_QUERY))
Mockito.when(indexingServiceClient.taskReportAsMap(ArgumentMatchers.eq(FINISHED_SELECT_MSQ_QUERY)))
.thenAnswer(inv -> Futures.immediateFuture(TaskReport.buildTaskReports(selectTaskReport.get())));
Mockito.when(indexingServiceClient.taskReportAsMap(ArgumentMatchers.eq(ACCEPTED_SELECT_MSQ_QUERY)))
.thenAnswer(inv -> Futures.immediateFuture(TaskReport.buildTaskReports(selectTaskReport.get())));
Mockito.when(indexingServiceClient.taskReportAsMap(ArgumentMatchers.eq(RUNNING_SELECT_MSQ_QUERY)))
.thenAnswer(inv -> Futures.immediateFuture(TaskReport.buildTaskReports(selectTaskReport.get())));
Mockito.when(indexingServiceClient.taskStatus(ArgumentMatchers.eq(ERRORED_SELECT_MSQ_QUERY)))
@ -584,6 +589,10 @@ public class SqlStatementResourceTest extends MSQTestBase
Mockito.when(indexingServiceClient.taskReportAsMap(ArgumentMatchers.eq(FINISHED_INSERT_MSQ_QUERY)))
.thenReturn(Futures.immediateFuture(TaskReport.buildTaskReports(MSQ_INSERT_TASK_REPORT)));
Mockito.when(indexingServiceClient.taskReportAsMap(ArgumentMatchers.eq(ACCEPTED_INSERT_MSQ_TASK)))
.thenReturn(Futures.immediateFuture(TaskReport.buildTaskReports(MSQ_INSERT_TASK_REPORT)));
Mockito.when(indexingServiceClient.taskReportAsMap(ArgumentMatchers.eq(RUNNING_INSERT_MSQ_QUERY)))
.thenReturn(Futures.immediateFuture(TaskReport.buildTaskReports(MSQ_INSERT_TASK_REPORT)));
Mockito.when(indexingServiceClient.taskPayload(FINISHED_INSERT_MSQ_QUERY))
.thenReturn(Futures.immediateFuture(new TaskPayloadResponse(
@ -690,9 +699,9 @@ public class SqlStatementResourceTest extends MSQTestBase
@Test
public void testMSQSelectAcceptedQuery()
{
Response response = resource.doGetStatus(ACCEPTED_SELECT_MSQ_QUERY, makeOkRequest());
Response response = resource.doGetStatus(ACCEPTED_SELECT_MSQ_QUERY, false, makeOkRequest());
Assert.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus());
Assert.assertEquals(
assertSqlStatementResult(
new SqlStatementResult(
ACCEPTED_SELECT_MSQ_QUERY,
SqlStatementState.ACCEPTED,
@ -702,7 +711,7 @@ public class SqlStatementResourceTest extends MSQTestBase
null,
null
),
response.getEntity()
(SqlStatementResult) response.getEntity()
);
assertExceptionMessage(
@ -724,9 +733,9 @@ public class SqlStatementResourceTest extends MSQTestBase
public void testMSQSelectRunningQuery()
{
Response response = resource.doGetStatus(RUNNING_SELECT_MSQ_QUERY, makeOkRequest());
Response response = resource.doGetStatus(RUNNING_SELECT_MSQ_QUERY, false, makeOkRequest());
Assert.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus());
Assert.assertEquals(
assertSqlStatementResult(
new SqlStatementResult(
RUNNING_SELECT_MSQ_QUERY,
SqlStatementState.RUNNING,
@ -736,7 +745,7 @@ public class SqlStatementResourceTest extends MSQTestBase
null,
null
),
response.getEntity()
(SqlStatementResult) response.getEntity()
);
assertExceptionMessage(
@ -754,10 +763,40 @@ public class SqlStatementResourceTest extends MSQTestBase
);
}
@Test
public void testMSQSelectRunningQueryWithDetail()
{
Response response = resource.doGetStatus(RUNNING_SELECT_MSQ_QUERY, true, makeOkRequest());
Assert.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus());
SqlStatementResult expectedSqlStatementResult = new SqlStatementResult(
RUNNING_SELECT_MSQ_QUERY,
SqlStatementState.RUNNING,
CREATED_TIME,
COL_NAME_AND_TYPES,
null,
null,
null,
selectTaskReport.get().getPayload().getStages(),
selectTaskReport.get().getPayload().getCounters(),
new ArrayList<>(selectTaskReport.get().getPayload().getStatus().getWarningReports())
);
assertSqlStatementResult(
expectedSqlStatementResult,
(SqlStatementResult) response.getEntity()
);
Assert.assertEquals(
Response.Status.ACCEPTED.getStatusCode(),
resource.deleteQuery(RUNNING_SELECT_MSQ_QUERY, makeOkRequest()).getStatus()
);
}
@Test
public void testFinishedSelectMSQQuery() throws Exception
{
Response response = resource.doGetStatus(FINISHED_SELECT_MSQ_QUERY, makeOkRequest());
Response response = resource.doGetStatus(FINISHED_SELECT_MSQ_QUERY, false, makeOkRequest());
Assert.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus());
Assert.assertEquals(objectMapper.writeValueAsString(new SqlStatementResult(
FINISHED_SELECT_MSQ_QUERY,
@ -825,7 +864,7 @@ public class SqlStatementResourceTest extends MSQTestBase
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.doGetStatus(queryID, false, makeOkRequest()), FAILURE_MSG, Response.Status.OK);
assertExceptionMessage(
resource.doGetResults(queryID, 0L, null, makeOkRequest()),
StringUtils.format(
@ -845,9 +884,9 @@ public class SqlStatementResourceTest extends MSQTestBase
@Test
public void testFinishedInsertMSQQuery()
{
Response response = resource.doGetStatus(FINISHED_INSERT_MSQ_QUERY, makeOkRequest());
Response response = resource.doGetStatus(FINISHED_INSERT_MSQ_QUERY, false, makeOkRequest());
Assert.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus());
Assert.assertEquals(new SqlStatementResult(
assertSqlStatementResult(new SqlStatementResult(
FINISHED_INSERT_MSQ_QUERY,
SqlStatementState.SUCCESS,
CREATED_TIME,
@ -855,7 +894,7 @@ public class SqlStatementResourceTest extends MSQTestBase
100L,
new ResultSetInformation(null, null, null, "test", null, null),
null
), response.getEntity());
), (SqlStatementResult) response.getEntity());
Assert.assertEquals(
Response.Status.OK.getStatusCode(),
@ -876,7 +915,7 @@ public class SqlStatementResourceTest extends MSQTestBase
public void testNonMSQTasks()
{
for (String queryID : ImmutableList.of(RUNNING_NON_MSQ_TASK, FAILED_NON_MSQ_TASK, FINISHED_NON_MSQ_TASK)) {
assertNotFound(resource.doGetStatus(queryID, makeOkRequest()), queryID);
assertNotFound(resource.doGetStatus(queryID, false, makeOkRequest()), queryID);
assertNotFound(resource.doGetResults(queryID, 0L, null, makeOkRequest()), queryID);
assertNotFound(resource.deleteQuery(queryID, makeOkRequest()), queryID);
}
@ -885,9 +924,9 @@ public class SqlStatementResourceTest extends MSQTestBase
@Test
public void testMSQInsertAcceptedQuery()
{
Response response = resource.doGetStatus(ACCEPTED_INSERT_MSQ_TASK, makeOkRequest());
Response response = resource.doGetStatus(ACCEPTED_INSERT_MSQ_TASK, false, makeOkRequest());
Assert.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus());
Assert.assertEquals(
assertSqlStatementResult(
new SqlStatementResult(
ACCEPTED_INSERT_MSQ_TASK,
SqlStatementState.ACCEPTED,
@ -897,7 +936,7 @@ public class SqlStatementResourceTest extends MSQTestBase
null,
null
),
response.getEntity()
(SqlStatementResult) response.getEntity()
);
assertExceptionMessage(
@ -918,9 +957,9 @@ public class SqlStatementResourceTest extends MSQTestBase
@Test
public void testMSQInsertRunningQuery()
{
Response response = resource.doGetStatus(RUNNING_INSERT_MSQ_QUERY, makeOkRequest());
Response response = resource.doGetStatus(RUNNING_INSERT_MSQ_QUERY, false, makeOkRequest());
Assert.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus());
Assert.assertEquals(
assertSqlStatementResult(
new SqlStatementResult(
RUNNING_INSERT_MSQ_QUERY,
SqlStatementState.RUNNING,
@ -930,7 +969,7 @@ public class SqlStatementResourceTest extends MSQTestBase
null,
null
),
response.getEntity()
(SqlStatementResult) response.getEntity()
);
assertExceptionMessage(
@ -955,6 +994,7 @@ public class SqlStatementResourceTest extends MSQTestBase
Response.Status.OK.getStatusCode(),
resource.doGetStatus(
RUNNING_SELECT_MSQ_QUERY,
false,
makeExpectedReq(makeAuthResultForUser(SUPERUSER))
).getStatus()
);
@ -984,6 +1024,7 @@ public class SqlStatementResourceTest extends MSQTestBase
Response.Status.FORBIDDEN.getStatusCode(),
resource.doGetStatus(
RUNNING_SELECT_MSQ_QUERY,
false,
makeExpectedReq(differentUserAuthResult)
).getStatus()
);
@ -1013,6 +1054,7 @@ public class SqlStatementResourceTest extends MSQTestBase
Response.Status.OK.getStatusCode(),
resource.doGetStatus(
RUNNING_SELECT_MSQ_QUERY,
false,
makeExpectedReq(differentUserAuthResult)
).getStatus()
);
@ -1042,6 +1084,7 @@ public class SqlStatementResourceTest extends MSQTestBase
Response.Status.FORBIDDEN.getStatusCode(),
resource.doGetStatus(
RUNNING_SELECT_MSQ_QUERY,
false,
makeExpectedReq(differentUserAuthResult)
).getStatus()
);
@ -1071,6 +1114,7 @@ public class SqlStatementResourceTest extends MSQTestBase
Response.Status.OK.getStatusCode(),
resource.doGetStatus(
RUNNING_SELECT_MSQ_QUERY,
false,
makeExpectedReq(differentUserAuthResult)
).getStatus()
);
@ -1107,7 +1151,7 @@ public class SqlStatementResourceTest extends MSQTestBase
Assert.assertEquals(
Response.Status.NOT_FOUND.getStatusCode(),
resource.doGetStatus(taskIdNotFound, makeOkRequest()).getStatus()
resource.doGetStatus(taskIdNotFound, false, makeOkRequest()).getStatus()
);
Assert.assertEquals(
Response.Status.NOT_FOUND.getStatusCode(),
@ -1124,4 +1168,28 @@ public class SqlStatementResourceTest extends MSQTestBase
{
Assert.assertEquals(Response.Status.OK.getStatusCode(), resource.isEnabled(makeOkRequest()).getStatus());
}
private void assertSqlStatementResult(SqlStatementResult expected, SqlStatementResult actual)
{
Assert.assertEquals(expected.getQueryId(), actual.getQueryId());
Assert.assertEquals(expected.getCreatedAt(), actual.getCreatedAt());
Assert.assertEquals(expected.getSqlRowSignature(), actual.getSqlRowSignature());
Assert.assertEquals(expected.getDurationMs(), actual.getDurationMs());
Assert.assertEquals(expected.getStages(), actual.getStages());
Assert.assertEquals(expected.getState(), actual.getState());
Assert.assertEquals(expected.getWarnings(), actual.getWarnings());
Assert.assertEquals(expected.getResultSetInformation(), actual.getResultSetInformation());
if (actual.getCounters() == null || expected.getCounters() == null) {
Assert.assertEquals(expected.getCounters(), actual.getCounters());
} else {
Assert.assertEquals(expected.getCounters().toString(), actual.getCounters().toString());
}
if (actual.getErrorResponse() == null || expected.getErrorResponse() == null) {
Assert.assertEquals(expected.getErrorResponse(), actual.getErrorResponse());
} else {
Assert.assertEquals(expected.getErrorResponse().getAsMap(), actual.getErrorResponse().getAsMap());
}
}
}