This is an automated email from the ASF dual-hosted git repository. JackieTien97 pushed a commit to branch rest-row-limit in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 65f2e5f16ea18e1de7194672643ed10a719f43f8 Author: JackieTien97 <[email protected]> AuthorDate: Wed Jul 8 12:15:44 2026 +0800 Enforce REST query row limit --- .../src/main/openapi3/iotdb_rest_table_v1.yaml | 1 + .../src/main/openapi3/iotdb_rest_v1.yaml | 1 + .../src/main/openapi3/iotdb_rest_v2.yaml | 1 + .../org/apache/iotdb/rest/i18n/RestMessages.java | 3 + .../org/apache/iotdb/rest/i18n/RestMessages.java | 3 + .../rest/protocol/handler/QueryRowLimitUtils.java | 66 ++++++ .../table/v1/handler/QueryDataSetHandler.java | 44 +--- .../protocol/table/v1/impl/RestApiServiceImpl.java | 9 +- .../protocol/v1/handler/QueryDataSetHandler.java | 114 +++++------ .../protocol/v1/impl/GrafanaApiServiceImpl.java | 19 +- .../rest/protocol/v1/impl/RestApiServiceImpl.java | 8 +- .../protocol/v2/handler/QueryDataSetHandler.java | 114 +++++------ .../protocol/v2/impl/GrafanaApiServiceImpl.java | 19 +- .../rest/protocol/v2/impl/RestApiServiceImpl.java | 13 +- .../protocol/handler/QueryRowLimitUtilsTest.java | 52 +++++ .../v1/handler/QueryDataSetHandlerTest.java | 222 +++++++++++++++++++++ .../src/test/resources/iotdb-common.properties | 3 +- .../src/test/resources/iotdb-common.properties | 3 +- .../src/test/resources/iotdb-system.properties | 3 +- .../conf/iotdb-system.properties.template | 3 +- 20 files changed, 524 insertions(+), 177 deletions(-) diff --git a/external-service-impl/rest-openapi/src/main/openapi3/iotdb_rest_table_v1.yaml b/external-service-impl/rest-openapi/src/main/openapi3/iotdb_rest_table_v1.yaml index 36cec631b54..27e347881c5 100644 --- a/external-service-impl/rest-openapi/src/main/openapi3/iotdb_rest_table_v1.yaml +++ b/external-service-impl/rest-openapi/src/main/openapi3/iotdb_rest_table_v1.yaml @@ -98,6 +98,7 @@ components: row_limit: type: integer format: int32 + description: Maximum rows to return. The effective limit is capped by rest_query_default_row_size_limit. InsertTabletRequest: title: InsertTabletRequest diff --git a/external-service-impl/rest-openapi/src/main/openapi3/iotdb_rest_v1.yaml b/external-service-impl/rest-openapi/src/main/openapi3/iotdb_rest_v1.yaml index 18080c693ab..7149ddd1cd3 100644 --- a/external-service-impl/rest-openapi/src/main/openapi3/iotdb_rest_v1.yaml +++ b/external-service-impl/rest-openapi/src/main/openapi3/iotdb_rest_v1.yaml @@ -166,6 +166,7 @@ components: rowLimit: type: integer format: int32 + description: Maximum rows to return. The effective limit is capped by rest_query_default_row_size_limit. InsertTabletRequest: title: InsertTabletRequest diff --git a/external-service-impl/rest-openapi/src/main/openapi3/iotdb_rest_v2.yaml b/external-service-impl/rest-openapi/src/main/openapi3/iotdb_rest_v2.yaml index 0cae51bef23..0ebb430b1b9 100644 --- a/external-service-impl/rest-openapi/src/main/openapi3/iotdb_rest_v2.yaml +++ b/external-service-impl/rest-openapi/src/main/openapi3/iotdb_rest_v2.yaml @@ -201,6 +201,7 @@ components: row_limit: type: integer format: int32 + description: Maximum rows to return. The effective limit is capped by rest_query_default_row_size_limit. PrefixPathList: title: PrefixPathList diff --git a/external-service-impl/rest/src/main/i18n/en/org/apache/iotdb/rest/i18n/RestMessages.java b/external-service-impl/rest/src/main/i18n/en/org/apache/iotdb/rest/i18n/RestMessages.java index 1b4bf36432c..882bb16c558 100644 --- a/external-service-impl/rest/src/main/i18n/en/org/apache/iotdb/rest/i18n/RestMessages.java +++ b/external-service-impl/rest/src/main/i18n/en/org/apache/iotdb/rest/i18n/RestMessages.java @@ -34,6 +34,9 @@ public final class RestMessages { public static final String SQL_SHOULD_NOT_BE_NULL = "sql should not be null"; public static final String ROW_LIMIT_SHOULD_BE_POSITIVE = "row_limit should be positive"; public static final String ROW_LIMIT_CAMEL_SHOULD_BE_POSITIVE = "rowLimit should be positive"; + public static final String + MESSAGE_DATASET_ROW_SIZE_EXCEEDED_THE_GIVEN_MAX_ROW_SIZE_ARG_7A3F6452 = + "Dataset row size exceeded the given max row size (%d)"; public static final String PREFIX_PATHS_NOT_NULL = "prefix_paths should not be null"; public static final String TIMESTAMPS_NOT_NULL = "timestamps should not be null"; public static final String IS_ALIGNED_NOT_NULL = "is_aligned should not be null"; diff --git a/external-service-impl/rest/src/main/i18n/zh/org/apache/iotdb/rest/i18n/RestMessages.java b/external-service-impl/rest/src/main/i18n/zh/org/apache/iotdb/rest/i18n/RestMessages.java index 23f3475cdc4..cb443b6ec97 100644 --- a/external-service-impl/rest/src/main/i18n/zh/org/apache/iotdb/rest/i18n/RestMessages.java +++ b/external-service-impl/rest/src/main/i18n/zh/org/apache/iotdb/rest/i18n/RestMessages.java @@ -34,6 +34,9 @@ public final class RestMessages { public static final String SQL_SHOULD_NOT_BE_NULL = "sql 不能为空"; public static final String ROW_LIMIT_SHOULD_BE_POSITIVE = "row_limit 应为正数"; public static final String ROW_LIMIT_CAMEL_SHOULD_BE_POSITIVE = "rowLimit 应为正数"; + public static final String + MESSAGE_DATASET_ROW_SIZE_EXCEEDED_THE_GIVEN_MAX_ROW_SIZE_ARG_7A3F6452 = + "数据集行数超过给定的最大行数（%d）"; public static final String PREFIX_PATHS_NOT_NULL = "prefix_paths 不能为空"; public static final String TIMESTAMPS_NOT_NULL = "timestamps 不能为空"; public static final String IS_ALIGNED_NOT_NULL = "is_aligned 不能为空"; diff --git a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/handler/QueryRowLimitUtils.java b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/handler/QueryRowLimitUtils.java new file mode 100644 index 00000000000..4d8bd08b4f0 --- /dev/null +++ b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/handler/QueryRowLimitUtils.java @@ -0,0 +1,66 @@ +/* + * 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.iotdb.rest.protocol.handler; + +import org.apache.iotdb.rest.i18n.RestMessages; +import org.apache.iotdb.rest.protocol.model.ExecutionStatus; +import org.apache.iotdb.rpc.TSStatusCode; + +import jakarta.ws.rs.core.Response; + +public final class QueryRowLimitUtils { + + private static final int MIN_ROW_SIZE_LIMIT = 1; + + private QueryRowLimitUtils() {} + + public static int resolveActualRowSizeLimit( + Integer requestedRowSizeLimit, int configuredRowSizeLimit) { + int hardLimit = normalizeRowSizeLimit(configuredRowSizeLimit); + if (requestedRowSizeLimit == null) { + return hardLimit; + } + return normalizeRowSizeLimit(Math.min(requestedRowSizeLimit, hardLimit)); + } + + public static int normalizeRowSizeLimit(int rowSizeLimit) { + return Math.max(MIN_ROW_SIZE_LIMIT, rowSizeLimit); + } + + public static boolean exceedsLimit( + int fetchedRowCount, int incomingRowCount, int actualRowSizeLimit) { + return incomingRowCount > 0 + && (long) fetchedRowCount + incomingRowCount > normalizeRowSizeLimit(actualRowSizeLimit); + } + + public static Response buildRowSizeLimitExceededResponse(int actualRowSizeLimit) { + int rowSizeLimit = normalizeRowSizeLimit(actualRowSizeLimit); + return Response.ok() + .entity( + new ExecutionStatus() + .code(TSStatusCode.QUERY_PROCESS_ERROR.getStatusCode()) + .message( + String.format( + RestMessages + .MESSAGE_DATASET_ROW_SIZE_EXCEEDED_THE_GIVEN_MAX_ROW_SIZE_ARG_7A3F6452, + rowSizeLimit))) + .build(); + } +} diff --git a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/table/v1/handler/QueryDataSetHandler.java b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/table/v1/handler/QueryDataSetHandler.java index 4bc94553414..9067756d5c6 100644 --- a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/table/v1/handler/QueryDataSetHandler.java +++ b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/table/v1/handler/QueryDataSetHandler.java @@ -22,9 +22,8 @@ import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Query; import org.apache.iotdb.commons.queryengine.plan.relational.sql.ast.Statement; import org.apache.iotdb.db.queryengine.common.header.DatasetHeader; import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution; -import org.apache.iotdb.rest.protocol.model.ExecutionStatus; +import org.apache.iotdb.rest.protocol.handler.QueryRowLimitUtils; import org.apache.iotdb.rest.protocol.table.v1.model.QueryDataSet; -import org.apache.iotdb.rpc.TSStatusCode; import org.apache.tsfile.block.column.Column; import org.apache.tsfile.common.conf.TSFileConfig; @@ -45,7 +44,7 @@ public class QueryDataSetHandler { private QueryDataSetHandler() {} /** - * @param actualRowSizeLimit max number of rows to return. no limit when actualRowSizeLimit <= 0. + * @param actualRowSizeLimit max number of rows to return. */ public static Response fillQueryDataSet( IQueryExecution queryExecution, Statement statement, int actualRowSizeLimit) @@ -61,7 +60,6 @@ public class QueryDataSetHandler { IQueryExecution queryExecution, final int actualRowSizeLimit) throws IoTDBException { QueryDataSet targetDataSet = new QueryDataSet(); int fetched = 0; - int columnNum = queryExecution.getOutputValueColumnCount(); DatasetHeader header = queryExecution.getDatasetHeader(); List<String> resultColumns = header.getRespColumns(); @@ -73,17 +71,6 @@ public class QueryDataSetHandler { targetDataSet.addValuesItem(new ArrayList<>()); } while (true) { - if (0 < actualRowSizeLimit && actualRowSizeLimit <= fetched) { - return Response.ok() - .entity( - new ExecutionStatus() - .code(TSStatusCode.QUERY_PROCESS_ERROR.getStatusCode()) - .message( - String.format( - "Dataset row size exceeded the given max row size (%d)", - actualRowSizeLimit))) - .build(); - } Optional<TsBlock> optionalTsBlock = queryExecution.getBatchResult(); if (!optionalTsBlock.isPresent() || optionalTsBlock.get().isEmpty()) { if (fetched == 0) { @@ -94,22 +81,22 @@ public class QueryDataSetHandler { } TsBlock tsBlock = optionalTsBlock.get(); int currentCount = tsBlock.getPositionCount(); + if (QueryRowLimitUtils.exceedsLimit(fetched, currentCount, actualRowSizeLimit)) { + return QueryRowLimitUtils.buildRowSizeLimitExceededResponse(actualRowSizeLimit); + } for (int k = 0; k < resultColumns.size(); k++) { Column column = tsBlock.getColumn(headerMap.get(resultColumns.get(k))); List<Object> targetDataSetColumn = targetDataSet.getValues().get(k); for (int i = 0; i < currentCount; i++) { - fetched++; if (column.isNull(i)) { targetDataSetColumn.add(null); } else { addTypedValueToTarget(targetDataSet.getDataTypes(), k, i, targetDataSetColumn, column); } } - if (k != columnNum - 1) { - fetched -= currentCount; - } } + fetched += currentCount; } targetDataSet.setValues(convertColumnToRow(targetDataSet.getValues())); return Response.ok().entity(targetDataSet).build(); @@ -174,17 +161,6 @@ public class QueryDataSetHandler { int fetched = 0; int columnNum = queryExecution.getOutputValueColumnCount(); while (true) { - if (0 < actualRowSizeLimit && actualRowSizeLimit <= fetched) { - return Response.ok() - .entity( - new ExecutionStatus() - .code(TSStatusCode.QUERY_PROCESS_ERROR.getStatusCode()) - .message( - String.format( - "Dataset row size exceeded the given max row size (%d)", - actualRowSizeLimit))) - .build(); - } Optional<TsBlock> optionalTsBlock = queryExecution.getBatchResult(); if (!optionalTsBlock.isPresent()) { if (fetched == 0) { @@ -199,21 +175,21 @@ public class QueryDataSetHandler { targetDataSet.setValues(new ArrayList<>()); return Response.ok().entity(targetDataSet).build(); } + if (QueryRowLimitUtils.exceedsLimit(fetched, currentCount, actualRowSizeLimit)) { + return QueryRowLimitUtils.buildRowSizeLimitExceededResponse(actualRowSizeLimit); + } for (int k = 0; k < columnNum; k++) { Column column = tsBlock.getColumn(targetDataSetIndexToSourceDataSetIndex[k]); List<Object> targetDataSetColumn = targetDataSet.getValues().get(k); for (int i = 0; i < currentCount; i++) { - fetched++; if (column.isNull(i)) { targetDataSetColumn.add(null); } else { addTypedValueToTarget(targetDataSet.getDataTypes(), k, i, targetDataSetColumn, column); } } - if (k != columnNum - 1) { - fetched -= currentCount; - } } + fetched += currentCount; } targetDataSet.setValues(convertColumnToRow(targetDataSet.getValues())); return Response.ok().entity(targetDataSet).build(); diff --git a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/table/v1/impl/RestApiServiceImpl.java b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/table/v1/impl/RestApiServiceImpl.java index 77197c3c088..a04ae7a16b4 100644 --- a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/table/v1/impl/RestApiServiceImpl.java +++ b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/table/v1/impl/RestApiServiceImpl.java @@ -38,6 +38,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.StatementType; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertTabletStatement; import org.apache.iotdb.db.utils.CommonUtils; import org.apache.iotdb.db.utils.SetThreadName; +import org.apache.iotdb.rest.protocol.handler.QueryRowLimitUtils; import org.apache.iotdb.rest.protocol.table.v1.NotFoundException; import org.apache.iotdb.rest.protocol.table.v1.RestApiService; import org.apache.iotdb.rest.protocol.table.v1.handler.ExceptionHandler; @@ -63,11 +64,12 @@ public class RestApiServiceImpl extends RestApiService { private static final IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig(); - private final Integer defaultQueryRowLimit; + private final int defaultQueryRowLimit; public RestApiServiceImpl() { defaultQueryRowLimit = - IoTDBRestServiceDescriptor.getInstance().getConfig().getRestQueryDefaultRowSizeLimit(); + QueryRowLimitUtils.normalizeRowSizeLimit( + IoTDBRestServiceDescriptor.getInstance().getConfig().getRestQueryDefaultRowSizeLimit()); } public Response executeQueryInternal( @@ -103,7 +105,8 @@ public class RestApiServiceImpl extends RestApiService { QueryDataSetHandler.fillQueryDataSet( queryExecution, statement, - sql.getRowLimit() == null ? defaultQueryRowLimit : sql.getRowLimit()); + QueryRowLimitUtils.resolveActualRowSizeLimit( + sql.getRowLimit(), defaultQueryRowLimit)); if (queryExecution.getQueryType() == QueryType.READ_WRITE) { return responseGenerateHelper(result); } diff --git a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v1/handler/QueryDataSetHandler.java b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v1/handler/QueryDataSetHandler.java index 3918a07dd52..4cee0926e77 100644 --- a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v1/handler/QueryDataSetHandler.java +++ b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v1/handler/QueryDataSetHandler.java @@ -29,6 +29,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowAIDevic import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowLoadedModelsStatement; import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowModelsStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.AuthorStatement; +import org.apache.iotdb.rest.protocol.handler.QueryRowLimitUtils; import org.apache.iotdb.rest.protocol.v1.model.ExecutionStatus; import org.apache.iotdb.rpc.TSStatusCode; @@ -51,7 +52,7 @@ public class QueryDataSetHandler { private QueryDataSetHandler() {} /** - * @param actualRowSizeLimit max number of rows to return. no limit when actualRowSizeLimit <= 0. + * @param actualRowSizeLimit max number of rows to return. */ public static Response fillQueryDataSet( IQueryExecution queryExecution, Statement statement, int actualRowSizeLimit) @@ -152,7 +153,6 @@ public class QueryDataSetHandler { final long timePrecision) throws IoTDBException { int fetched = 0; - int columnNum = queryExecution.getOutputValueColumnCount(); DatasetHeader header = queryExecution.getDatasetHeader(); List<String> resultColumns = header.getRespColumns(); @@ -164,17 +164,6 @@ public class QueryDataSetHandler { } while (true) { - if (0 < actualRowSizeLimit && actualRowSizeLimit <= fetched) { - return Response.ok() - .entity( - new org.apache.iotdb.rest.protocol.model.ExecutionStatus() - .code(TSStatusCode.QUERY_PROCESS_ERROR.getStatusCode()) - .message( - String.format( - "Dataset row size exceeded the given max row size (%d)", - actualRowSizeLimit))) - .build(); - } Optional<TsBlock> optionalTsBlock = queryExecution.getBatchResult(); if (!optionalTsBlock.isPresent() || optionalTsBlock.get().isEmpty()) { if (fetched == 0) { @@ -186,6 +175,9 @@ public class QueryDataSetHandler { } TsBlock tsBlock = optionalTsBlock.get(); int currentCount = tsBlock.getPositionCount(); + if (QueryRowLimitUtils.exceedsLimit(fetched, currentCount, actualRowSizeLimit)) { + return QueryRowLimitUtils.buildRowSizeLimitExceededResponse(actualRowSizeLimit); + } // time column for (int i = 0; i < currentCount; i++) { targetDataSet.addTimestampsItem( @@ -197,17 +189,14 @@ public class QueryDataSetHandler { Column column = tsBlock.getColumn(headerMap.get(resultColumns.get(k))); List<Object> targetDataSetColumn = targetDataSet.getValues().get(k); for (int i = 0; i < currentCount; i++) { - fetched++; if (column.isNull(i)) { targetDataSetColumn.add(null); } else { addTypedValueToTarget(dataTypes, k, i, targetDataSetColumn, column); } } - if (k != columnNum - 1) { - fetched -= currentCount; - } } + fetched += currentCount; } return Response.ok().entity(targetDataSet).build(); } @@ -222,17 +211,6 @@ public class QueryDataSetHandler { int fetched = 0; int columnNum = queryExecution.getOutputValueColumnCount(); while (true) { - if (0 < actualRowSizeLimit && actualRowSizeLimit <= fetched) { - return Response.ok() - .entity( - new org.apache.iotdb.rest.protocol.model.ExecutionStatus() - .code(TSStatusCode.QUERY_PROCESS_ERROR.getStatusCode()) - .message( - String.format( - "Dataset row size exceeded the given max row size (%d)", - actualRowSizeLimit))) - .build(); - } Optional<TsBlock> optionalTsBlock = queryExecution.getBatchResult(); if (!optionalTsBlock.isPresent()) { if (fetched == 0) { @@ -247,21 +225,21 @@ public class QueryDataSetHandler { targetDataSet.setValues(new ArrayList<>()); return Response.ok().entity(targetDataSet).build(); } + if (QueryRowLimitUtils.exceedsLimit(fetched, currentCount, actualRowSizeLimit)) { + return QueryRowLimitUtils.buildRowSizeLimitExceededResponse(actualRowSizeLimit); + } for (int k = 0; k < columnNum; k++) { Column column = tsBlock.getColumn(targetDataSetIndexToSourceDataSetIndex[k]); List<Object> targetDataSetColumn = targetDataSet.getValues().get(k); for (int i = 0; i < currentCount; i++) { - fetched++; if (column.isNull(i)) { targetDataSetColumn.add(null); } else { addTypedValueToTarget(dataTypes, k, i, targetDataSetColumn, column); } } - if (k != columnNum - 1) { - fetched -= currentCount; - } } + fetched += currentCount; } return Response.ok().entity(targetDataSet).build(); } @@ -291,44 +269,60 @@ public class QueryDataSetHandler { } public static Response fillGrafanaVariablesResult( - IQueryExecution queryExecution, Statement statement) throws IoTDBException { + IQueryExecution queryExecution, Statement statement, int actualRowSizeLimit) + throws IoTDBException { List<String> results = new ArrayList<>(); - Optional<TsBlock> optionalTsBlock = queryExecution.getBatchResult(); - if (!optionalTsBlock.isPresent()) { - return Response.ok().entity(results).build(); - } - TsBlock tsBlock = optionalTsBlock.get(); - int currentCount = tsBlock.getPositionCount(); - Column column = tsBlock.getColumn(0); + int fetched = 0; + while (true) { + Optional<TsBlock> optionalTsBlock = queryExecution.getBatchResult(); + if (!optionalTsBlock.isPresent() || optionalTsBlock.get().isEmpty()) { + return Response.ok().entity(results).build(); + } + TsBlock tsBlock = optionalTsBlock.get(); + int currentCount = tsBlock.getPositionCount(); + if (QueryRowLimitUtils.exceedsLimit(fetched, currentCount, actualRowSizeLimit)) { + return QueryRowLimitUtils.buildRowSizeLimitExceededResponse(actualRowSizeLimit); + } + Column column = tsBlock.getColumn(0); - for (int i = 0; i < currentCount; i++) { - String nodePaths = column.getObject(i).toString(); - if (statement instanceof ShowChildPathsStatement) { - String[] nodeSubPath = nodePaths.split("\\."); - results.add(nodeSubPath[nodeSubPath.length - 1]); - } else { - results.add(nodePaths); + for (int i = 0; i < currentCount; i++) { + String nodePaths = column.getObject(i).toString(); + if (statement instanceof ShowChildPathsStatement) { + String[] nodeSubPath = nodePaths.split("\\."); + results.add(nodeSubPath[nodeSubPath.length - 1]); + } else { + results.add(nodePaths); + } } + fetched += currentCount; } - return Response.ok().entity(results).build(); } - public static Response fillGrafanaNodesResult(IQueryExecution queryExecution) - throws IoTDBException { + public static Response fillGrafanaNodesResult( + IQueryExecution queryExecution, int actualRowSizeLimit) throws IoTDBException { List<String> nodes = new ArrayList<>(); - Optional<TsBlock> optionalTsBlock = queryExecution.getBatchResult(); - if (!optionalTsBlock.isPresent()) { + if (queryExecution == null) { return Response.ok().entity(nodes).build(); } - TsBlock tsBlock = optionalTsBlock.get(); - int currentCount = tsBlock.getPositionCount(); - Column column = tsBlock.getColumn(0); + int fetched = 0; + while (true) { + Optional<TsBlock> optionalTsBlock = queryExecution.getBatchResult(); + if (!optionalTsBlock.isPresent() || optionalTsBlock.get().isEmpty()) { + return Response.ok().entity(nodes).build(); + } + TsBlock tsBlock = optionalTsBlock.get(); + int currentCount = tsBlock.getPositionCount(); + if (QueryRowLimitUtils.exceedsLimit(fetched, currentCount, actualRowSizeLimit)) { + return QueryRowLimitUtils.buildRowSizeLimitExceededResponse(actualRowSizeLimit); + } + Column column = tsBlock.getColumn(0); - for (int i = 0; i < currentCount; i++) { - String nodePaths = column.getObject(i).toString(); - String[] nodeSubPath = nodePaths.split("\\."); - nodes.add(nodeSubPath[nodeSubPath.length - 1]); + for (int i = 0; i < currentCount; i++) { + String nodePaths = column.getObject(i).toString(); + String[] nodeSubPath = nodePaths.split("\\."); + nodes.add(nodeSubPath[nodeSubPath.length - 1]); + } + fetched += currentCount; } - return Response.ok().entity(nodes).build(); } } diff --git a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v1/impl/GrafanaApiServiceImpl.java b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v1/impl/GrafanaApiServiceImpl.java index 19b8ab8b38e..dcc17c43a59 100644 --- a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v1/impl/GrafanaApiServiceImpl.java +++ b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v1/impl/GrafanaApiServiceImpl.java @@ -21,6 +21,7 @@ import org.apache.iotdb.commons.conf.CommonDescriptor; import org.apache.iotdb.commons.path.PartialPath; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.conf.rest.IoTDBRestServiceDescriptor; import org.apache.iotdb.db.protocol.session.IClientSession; import org.apache.iotdb.db.protocol.session.SessionManager; import org.apache.iotdb.db.queryengine.plan.Coordinator; @@ -36,6 +37,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.crud.QueryStatement; import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowStatement; import org.apache.iotdb.db.utils.SetThreadName; import org.apache.iotdb.rest.protocol.handler.AuthorizationHandler; +import org.apache.iotdb.rest.protocol.handler.QueryRowLimitUtils; import org.apache.iotdb.rest.protocol.v1.GrafanaApiService; import org.apache.iotdb.rest.protocol.v1.handler.ExceptionHandler; import org.apache.iotdb.rest.protocol.v1.handler.QueryDataSetHandler; @@ -68,11 +70,15 @@ public class GrafanaApiServiceImpl extends GrafanaApiService { private final AuthorizationHandler authorizationHandler; private final long timePrecision; // the default timestamp precision is ms + private final int defaultQueryRowLimit; public GrafanaApiServiceImpl() { partitionFetcher = ClusterPartitionFetcher.getInstance(); schemaFetcher = ClusterSchemaFetcher.getInstance(); authorizationHandler = new AuthorizationHandler(); + defaultQueryRowLimit = + QueryRowLimitUtils.normalizeRowSizeLimit( + IoTDBRestServiceDescriptor.getInstance().getConfig().getRestQueryDefaultRowSizeLimit()); switch (CommonDescriptor.getInstance().getConfig().getTimestampPrecision()) { case "ns": @@ -132,7 +138,8 @@ public class GrafanaApiServiceImpl extends GrafanaApiService { } IQueryExecution queryExecution = COORDINATOR.getQueryExecution(queryId); try (SetThreadName threadName = new SetThreadName(result.queryId.getId())) { - return QueryDataSetHandler.fillGrafanaVariablesResult(queryExecution, statement); + return QueryDataSetHandler.fillGrafanaVariablesResult( + queryExecution, statement, defaultQueryRowLimit); } } catch (Exception e) { return Response.ok().entity(ExceptionHandler.tryCatchException(e)).build(); @@ -204,9 +211,11 @@ public class GrafanaApiServiceImpl extends GrafanaApiService { try (SetThreadName threadName = new SetThreadName(result.queryId.getId())) { if (((QueryStatement) statement).isAggregationQuery() && !((QueryStatement) statement).isGroupByTime()) { - return QueryDataSetHandler.fillAggregationPlanDataSet(queryExecution, 0); + return QueryDataSetHandler.fillAggregationPlanDataSet( + queryExecution, defaultQueryRowLimit); } else { - return QueryDataSetHandler.fillDataSetWithTimestamps(queryExecution, 0, timePrecision); + return QueryDataSetHandler.fillDataSetWithTimestamps( + queryExecution, defaultQueryRowLimit, timePrecision); } } } catch (Exception e) { @@ -269,10 +278,10 @@ public class GrafanaApiServiceImpl extends GrafanaApiService { IQueryExecution queryExecution = COORDINATOR.getQueryExecution(queryId); try (SetThreadName threadName = new SetThreadName(result.queryId.getId())) { - return QueryDataSetHandler.fillGrafanaNodesResult(queryExecution); + return QueryDataSetHandler.fillGrafanaNodesResult(queryExecution, defaultQueryRowLimit); } } else { - return QueryDataSetHandler.fillGrafanaNodesResult(null); + return QueryDataSetHandler.fillGrafanaNodesResult(null, defaultQueryRowLimit); } } catch (Exception e) { return Response.ok().entity(ExceptionHandler.tryCatchException(e)).build(); diff --git a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v1/impl/RestApiServiceImpl.java b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v1/impl/RestApiServiceImpl.java index b5b57ac10e7..4939e10a412 100644 --- a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v1/impl/RestApiServiceImpl.java +++ b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v1/impl/RestApiServiceImpl.java @@ -39,6 +39,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.metadata.view.CreateTableV import org.apache.iotdb.db.utils.CommonUtils; import org.apache.iotdb.db.utils.SetThreadName; import org.apache.iotdb.rest.protocol.handler.AuthorizationHandler; +import org.apache.iotdb.rest.protocol.handler.QueryRowLimitUtils; import org.apache.iotdb.rest.protocol.utils.InsertTabletSortDataUtils; import org.apache.iotdb.rest.protocol.v1.RestApiService; import org.apache.iotdb.rest.protocol.v1.handler.ExceptionHandler; @@ -70,14 +71,15 @@ public class RestApiServiceImpl extends RestApiService { private final ISchemaFetcher schemaFetcher; private final AuthorizationHandler authorizationHandler; - private final Integer defaultQueryRowLimit; + private final int defaultQueryRowLimit; public RestApiServiceImpl() { partitionFetcher = ClusterPartitionFetcher.getInstance(); schemaFetcher = ClusterSchemaFetcher.getInstance(); authorizationHandler = new AuthorizationHandler(); defaultQueryRowLimit = - IoTDBRestServiceDescriptor.getInstance().getConfig().getRestQueryDefaultRowSizeLimit(); + QueryRowLimitUtils.normalizeRowSizeLimit( + IoTDBRestServiceDescriptor.getInstance().getConfig().getRestQueryDefaultRowSizeLimit()); } @Override @@ -231,7 +233,7 @@ public class RestApiServiceImpl extends RestApiService { return QueryDataSetHandler.fillQueryDataSet( queryExecution, statement, - sql.getRowLimit() == null ? defaultQueryRowLimit : sql.getRowLimit()); + QueryRowLimitUtils.resolveActualRowSizeLimit(sql.getRowLimit(), defaultQueryRowLimit)); } } catch (Exception e) { finish = true; diff --git a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v2/handler/QueryDataSetHandler.java b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v2/handler/QueryDataSetHandler.java index b188454e179..34d03eac6ec 100644 --- a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v2/handler/QueryDataSetHandler.java +++ b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v2/handler/QueryDataSetHandler.java @@ -29,6 +29,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowAIDevic import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowLoadedModelsStatement; import org.apache.iotdb.db.queryengine.plan.statement.metadata.model.ShowModelsStatement; import org.apache.iotdb.db.queryengine.plan.statement.sys.AuthorStatement; +import org.apache.iotdb.rest.protocol.handler.QueryRowLimitUtils; import org.apache.iotdb.rest.protocol.model.ExecutionStatus; import org.apache.iotdb.rpc.TSStatusCode; @@ -51,7 +52,7 @@ public class QueryDataSetHandler { private QueryDataSetHandler() {} /** - * @param actualRowSizeLimit max number of rows to return. no limit when actualRowSizeLimit <= 0. + * @param actualRowSizeLimit max number of rows to return. */ public static Response fillQueryDataSet( IQueryExecution queryExecution, Statement statement, int actualRowSizeLimit) @@ -145,7 +146,6 @@ public class QueryDataSetHandler { final long timePrecision) throws IoTDBException { int fetched = 0; - int columnNum = queryExecution.getOutputValueColumnCount(); DatasetHeader header = queryExecution.getDatasetHeader(); List<String> resultColumns = header.getRespColumns(); @@ -158,17 +158,6 @@ public class QueryDataSetHandler { } while (true) { - if (0 < actualRowSizeLimit && actualRowSizeLimit <= fetched) { - return Response.ok() - .entity( - new ExecutionStatus() - .code(TSStatusCode.QUERY_PROCESS_ERROR.getStatusCode()) - .message( - String.format( - "Dataset row size exceeded the given max row size (%d)", - actualRowSizeLimit))) - .build(); - } Optional<TsBlock> optionalTsBlock = queryExecution.getBatchResult(); if (!optionalTsBlock.isPresent() || optionalTsBlock.get().isEmpty()) { if (fetched == 0) { @@ -180,6 +169,9 @@ public class QueryDataSetHandler { } TsBlock tsBlock = optionalTsBlock.get(); int currentCount = tsBlock.getPositionCount(); + if (QueryRowLimitUtils.exceedsLimit(fetched, currentCount, actualRowSizeLimit)) { + return QueryRowLimitUtils.buildRowSizeLimitExceededResponse(actualRowSizeLimit); + } // time column for (int i = 0; i < currentCount; i++) { targetDataSet.addTimestampsItem( @@ -191,17 +183,14 @@ public class QueryDataSetHandler { Column column = tsBlock.getColumn(headerMap.get(resultColumns.get(k))); List<Object> targetDataSetColumn = targetDataSet.getValues().get(k); for (int i = 0; i < currentCount; i++) { - fetched++; if (column.isNull(i)) { targetDataSetColumn.add(null); } else { addTypedValueToTarget(targetDataSet.getDataTypes(), k, i, targetDataSetColumn, column); } } - if (k != columnNum - 1) { - fetched -= currentCount; - } } + fetched += currentCount; } return Response.ok().entity(targetDataSet).build(); } @@ -216,17 +205,6 @@ public class QueryDataSetHandler { int columnNum = queryExecution.getOutputValueColumnCount(); List<String> dataTypes = targetDataSet.getDataTypes(); while (true) { - if (0 < actualRowSizeLimit && actualRowSizeLimit <= fetched) { - return Response.ok() - .entity( - new ExecutionStatus() - .code(TSStatusCode.QUERY_PROCESS_ERROR.getStatusCode()) - .message( - String.format( - "Dataset row size exceeded the given max row size (%d)", - actualRowSizeLimit))) - .build(); - } Optional<TsBlock> optionalTsBlock = queryExecution.getBatchResult(); if (!optionalTsBlock.isPresent()) { if (fetched == 0) { @@ -241,21 +219,21 @@ public class QueryDataSetHandler { targetDataSet.setValues(new ArrayList<>()); return Response.ok().entity(targetDataSet).build(); } + if (QueryRowLimitUtils.exceedsLimit(fetched, currentCount, actualRowSizeLimit)) { + return QueryRowLimitUtils.buildRowSizeLimitExceededResponse(actualRowSizeLimit); + } for (int k = 0; k < columnNum; k++) { Column column = tsBlock.getColumn(targetDataSetIndexToSourceDataSetIndex[k]); List<Object> targetDataSetColumn = targetDataSet.getValues().get(k); for (int i = 0; i < currentCount; i++) { - fetched++; if (column.isNull(i)) { targetDataSetColumn.add(null); } else { addTypedValueToTarget(dataTypes, k, i, targetDataSetColumn, column); } } - if (k != columnNum - 1) { - fetched -= currentCount; - } } + fetched += currentCount; } return Response.ok().entity(targetDataSet).build(); } @@ -285,44 +263,60 @@ public class QueryDataSetHandler { } public static Response fillGrafanaVariablesResult( - IQueryExecution queryExecution, Statement statement) throws IoTDBException { + IQueryExecution queryExecution, Statement statement, int actualRowSizeLimit) + throws IoTDBException { List<String> results = new ArrayList<>(); - Optional<TsBlock> optionalTsBlock = queryExecution.getBatchResult(); - if (!optionalTsBlock.isPresent()) { - return Response.ok().entity(results).build(); - } - TsBlock tsBlock = optionalTsBlock.get(); - int currentCount = tsBlock.getPositionCount(); - Column column = tsBlock.getColumn(0); + int fetched = 0; + while (true) { + Optional<TsBlock> optionalTsBlock = queryExecution.getBatchResult(); + if (!optionalTsBlock.isPresent() || optionalTsBlock.get().isEmpty()) { + return Response.ok().entity(results).build(); + } + TsBlock tsBlock = optionalTsBlock.get(); + int currentCount = tsBlock.getPositionCount(); + if (QueryRowLimitUtils.exceedsLimit(fetched, currentCount, actualRowSizeLimit)) { + return QueryRowLimitUtils.buildRowSizeLimitExceededResponse(actualRowSizeLimit); + } + Column column = tsBlock.getColumn(0); - for (int i = 0; i < currentCount; i++) { - String nodePaths = column.getObject(i).toString(); - if (statement instanceof ShowChildPathsStatement) { - String[] nodeSubPath = nodePaths.split("\\."); - results.add(nodeSubPath[nodeSubPath.length - 1]); - } else { - results.add(nodePaths); + for (int i = 0; i < currentCount; i++) { + String nodePaths = column.getObject(i).toString(); + if (statement instanceof ShowChildPathsStatement) { + String[] nodeSubPath = nodePaths.split("\\."); + results.add(nodeSubPath[nodeSubPath.length - 1]); + } else { + results.add(nodePaths); + } } + fetched += currentCount; } - return Response.ok().entity(results).build(); } - public static Response fillGrafanaNodesResult(IQueryExecution queryExecution) - throws IoTDBException { + public static Response fillGrafanaNodesResult( + IQueryExecution queryExecution, int actualRowSizeLimit) throws IoTDBException { List<String> nodes = new ArrayList<>(); - Optional<TsBlock> optionalTsBlock = queryExecution.getBatchResult(); - if (!optionalTsBlock.isPresent()) { + if (queryExecution == null) { return Response.ok().entity(nodes).build(); } - TsBlock tsBlock = optionalTsBlock.get(); - int currentCount = tsBlock.getPositionCount(); - Column column = tsBlock.getColumn(0); + int fetched = 0; + while (true) { + Optional<TsBlock> optionalTsBlock = queryExecution.getBatchResult(); + if (!optionalTsBlock.isPresent() || optionalTsBlock.get().isEmpty()) { + return Response.ok().entity(nodes).build(); + } + TsBlock tsBlock = optionalTsBlock.get(); + int currentCount = tsBlock.getPositionCount(); + if (QueryRowLimitUtils.exceedsLimit(fetched, currentCount, actualRowSizeLimit)) { + return QueryRowLimitUtils.buildRowSizeLimitExceededResponse(actualRowSizeLimit); + } + Column column = tsBlock.getColumn(0); - for (int i = 0; i < currentCount; i++) { - String nodePaths = column.getObject(i).toString(); - String[] nodeSubPath = nodePaths.split("\\."); - nodes.add(nodeSubPath[nodeSubPath.length - 1]); + for (int i = 0; i < currentCount; i++) { + String nodePaths = column.getObject(i).toString(); + String[] nodeSubPath = nodePaths.split("\\."); + nodes.add(nodeSubPath[nodeSubPath.length - 1]); + } + fetched += currentCount; } - return Response.ok().entity(nodes).build(); } } diff --git a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v2/impl/GrafanaApiServiceImpl.java b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v2/impl/GrafanaApiServiceImpl.java index e8701f89347..32834a5bced 100644 --- a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v2/impl/GrafanaApiServiceImpl.java +++ b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v2/impl/GrafanaApiServiceImpl.java @@ -21,6 +21,7 @@ import org.apache.iotdb.commons.conf.CommonDescriptor; import org.apache.iotdb.commons.path.PartialPath; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.conf.rest.IoTDBRestServiceDescriptor; import org.apache.iotdb.db.protocol.session.IClientSession; import org.apache.iotdb.db.protocol.session.SessionManager; import org.apache.iotdb.db.queryengine.plan.Coordinator; @@ -36,6 +37,7 @@ import org.apache.iotdb.db.queryengine.plan.statement.crud.QueryStatement; import org.apache.iotdb.db.queryengine.plan.statement.metadata.ShowStatement; import org.apache.iotdb.db.utils.SetThreadName; import org.apache.iotdb.rest.protocol.handler.AuthorizationHandler; +import org.apache.iotdb.rest.protocol.handler.QueryRowLimitUtils; import org.apache.iotdb.rest.protocol.v2.GrafanaApiService; import org.apache.iotdb.rest.protocol.v2.handler.ExceptionHandler; import org.apache.iotdb.rest.protocol.v2.handler.QueryDataSetHandler; @@ -68,11 +70,15 @@ public class GrafanaApiServiceImpl extends GrafanaApiService { private final AuthorizationHandler authorizationHandler; private final long timePrecision; // the default timestamp precision is ms + private final int defaultQueryRowLimit; public GrafanaApiServiceImpl() { partitionFetcher = ClusterPartitionFetcher.getInstance(); schemaFetcher = ClusterSchemaFetcher.getInstance(); authorizationHandler = new AuthorizationHandler(); + defaultQueryRowLimit = + QueryRowLimitUtils.normalizeRowSizeLimit( + IoTDBRestServiceDescriptor.getInstance().getConfig().getRestQueryDefaultRowSizeLimit()); switch (CommonDescriptor.getInstance().getConfig().getTimestampPrecision()) { case "ns": @@ -132,7 +138,8 @@ public class GrafanaApiServiceImpl extends GrafanaApiService { } IQueryExecution queryExecution = COORDINATOR.getQueryExecution(queryId); try (SetThreadName threadName = new SetThreadName(result.queryId.getId())) { - return QueryDataSetHandler.fillGrafanaVariablesResult(queryExecution, statement); + return QueryDataSetHandler.fillGrafanaVariablesResult( + queryExecution, statement, defaultQueryRowLimit); } } catch (Exception e) { return Response.ok().entity(ExceptionHandler.tryCatchException(e)).build(); @@ -204,9 +211,11 @@ public class GrafanaApiServiceImpl extends GrafanaApiService { try (SetThreadName threadName = new SetThreadName(result.queryId.getId())) { if (((QueryStatement) statement).isAggregationQuery() && !((QueryStatement) statement).isGroupByTime()) { - return QueryDataSetHandler.fillAggregationPlanDataSet(queryExecution, 0); + return QueryDataSetHandler.fillAggregationPlanDataSet( + queryExecution, defaultQueryRowLimit); } else { - return QueryDataSetHandler.fillDataSetWithTimestamps(queryExecution, 0, timePrecision); + return QueryDataSetHandler.fillDataSetWithTimestamps( + queryExecution, defaultQueryRowLimit, timePrecision); } } } catch (Exception e) { @@ -269,10 +278,10 @@ public class GrafanaApiServiceImpl extends GrafanaApiService { IQueryExecution queryExecution = COORDINATOR.getQueryExecution(queryId); try (SetThreadName threadName = new SetThreadName(result.queryId.getId())) { - return QueryDataSetHandler.fillGrafanaNodesResult(queryExecution); + return QueryDataSetHandler.fillGrafanaNodesResult(queryExecution, defaultQueryRowLimit); } } else { - return QueryDataSetHandler.fillGrafanaNodesResult(null); + return QueryDataSetHandler.fillGrafanaNodesResult(null, defaultQueryRowLimit); } } catch (Exception e) { return Response.ok().entity(ExceptionHandler.tryCatchException(e)).build(); diff --git a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v2/impl/RestApiServiceImpl.java b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v2/impl/RestApiServiceImpl.java index 6e2cef45b7f..c1dfd3e35d2 100644 --- a/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v2/impl/RestApiServiceImpl.java +++ b/external-service-impl/rest/src/main/java/org/apache/iotdb/rest/protocol/v2/impl/RestApiServiceImpl.java @@ -47,6 +47,7 @@ import org.apache.iotdb.db.schemaengine.schemaregion.ISchemaRegion; import org.apache.iotdb.db.utils.CommonUtils; import org.apache.iotdb.db.utils.SetThreadName; import org.apache.iotdb.rest.protocol.handler.AuthorizationHandler; +import org.apache.iotdb.rest.protocol.handler.QueryRowLimitUtils; import org.apache.iotdb.rest.protocol.model.ExecutionStatus; import org.apache.iotdb.rest.protocol.utils.InsertTabletSortDataUtils; import org.apache.iotdb.rest.protocol.v2.NotFoundException; @@ -99,14 +100,15 @@ public class RestApiServiceImpl extends RestApiService { private final ISchemaFetcher schemaFetcher; private final AuthorizationHandler authorizationHandler; - private final Integer defaultQueryRowLimit; + private final int defaultQueryRowLimit; public RestApiServiceImpl() { partitionFetcher = ClusterPartitionFetcher.getInstance(); schemaFetcher = ClusterSchemaFetcher.getInstance(); authorizationHandler = new AuthorizationHandler(); defaultQueryRowLimit = - IoTDBRestServiceDescriptor.getInstance().getConfig().getRestQueryDefaultRowSizeLimit(); + QueryRowLimitUtils.normalizeRowSizeLimit( + IoTDBRestServiceDescriptor.getInstance().getConfig().getRestQueryDefaultRowSizeLimit()); } @Override @@ -187,6 +189,7 @@ public class RestApiServiceImpl extends RestApiService { List<Object> timeseries = new ArrayList<>(); List<Object> valueList = new ArrayList<>(); List<Object> dataTypeList = new ArrayList<>(); + int fetched = 0; for (final Map.Entry<TableId, Map<IDeviceID, Map<String, Pair<TSDataType, TimeValuePair>>>> result : resultMap.entrySet()) { @@ -201,10 +204,14 @@ public class RestApiServiceImpl extends RestApiService { || tvPair.getValue() == null) { continue; } + if (QueryRowLimitUtils.exceedsLimit(fetched, 1, defaultQueryRowLimit)) { + return QueryRowLimitUtils.buildRowSizeLimitExceededResponse(defaultQueryRowLimit); + } valueList.add(tvPair.getValue().getStringValue()); dataTypeList.add(tvPair.getValue().getDataType().name()); targetDataSet.addTimestampsItem(tvPair.getTimestamp()); timeseries.add(deviceWithSeparator + measurementLastEntry.getKey()); + fetched++; } } } @@ -398,7 +405,7 @@ public class RestApiServiceImpl extends RestApiService { return QueryDataSetHandler.fillQueryDataSet( queryExecution, statement, - sql.getRowLimit() == null ? defaultQueryRowLimit : sql.getRowLimit()); + QueryRowLimitUtils.resolveActualRowSizeLimit(sql.getRowLimit(), defaultQueryRowLimit)); } } catch (Exception e) { finish = true; diff --git a/external-service-impl/rest/src/test/java/org/apache/iotdb/rest/protocol/handler/QueryRowLimitUtilsTest.java b/external-service-impl/rest/src/test/java/org/apache/iotdb/rest/protocol/handler/QueryRowLimitUtilsTest.java new file mode 100644 index 00000000000..02bdaf1c7fe --- /dev/null +++ b/external-service-impl/rest/src/test/java/org/apache/iotdb/rest/protocol/handler/QueryRowLimitUtilsTest.java @@ -0,0 +1,52 @@ +/* + * 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.iotdb.rest.protocol.handler; + +import org.junit.Test; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public class QueryRowLimitUtilsTest { + + @Test + public void resolveActualRowSizeLimitShouldUseConfiguredLimitAsHardLimit() { + assertEquals(10, QueryRowLimitUtils.resolveActualRowSizeLimit(null, 10)); + assertEquals(10, QueryRowLimitUtils.resolveActualRowSizeLimit(100, 10)); + assertEquals(5, QueryRowLimitUtils.resolveActualRowSizeLimit(5, 10)); + } + + @Test + public void resolveActualRowSizeLimitShouldKeepHardLimitPositive() { + assertEquals(1, QueryRowLimitUtils.resolveActualRowSizeLimit(null, 0)); + assertEquals(1, QueryRowLimitUtils.resolveActualRowSizeLimit(100, 0)); + assertEquals(1, QueryRowLimitUtils.resolveActualRowSizeLimit(null, -1)); + } + + @Test + public void exceedsLimitShouldRejectOnlyRowsBeyondTheHardLimit() { + assertFalse(QueryRowLimitUtils.exceedsLimit(0, 2, 2)); + assertTrue(QueryRowLimitUtils.exceedsLimit(2, 1, 2)); + assertTrue(QueryRowLimitUtils.exceedsLimit(0, 2, 1)); + assertFalse(QueryRowLimitUtils.exceedsLimit(0, 0, 1)); + assertTrue(QueryRowLimitUtils.exceedsLimit(0, 2, 0)); + } +} diff --git a/external-service-impl/rest/src/test/java/org/apache/iotdb/rest/protocol/v1/handler/QueryDataSetHandlerTest.java b/external-service-impl/rest/src/test/java/org/apache/iotdb/rest/protocol/v1/handler/QueryDataSetHandlerTest.java new file mode 100644 index 00000000000..36a2ec1292e --- /dev/null +++ b/external-service-impl/rest/src/test/java/org/apache/iotdb/rest/protocol/v1/handler/QueryDataSetHandlerTest.java @@ -0,0 +1,222 @@ +/* + * 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.iotdb.rest.protocol.v1.handler; + +import org.apache.iotdb.commons.exception.IoTDBException; +import org.apache.iotdb.commons.queryengine.common.SqlDialect; +import org.apache.iotdb.commons.schema.column.ColumnHeader; +import org.apache.iotdb.db.queryengine.common.header.DatasetHeader; +import org.apache.iotdb.db.queryengine.plan.analyze.QueryType; +import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult; +import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution; +import org.apache.iotdb.rest.protocol.model.ExecutionStatus; +import org.apache.iotdb.rest.protocol.v1.model.QueryDataSet; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.read.common.block.TsBlock; +import org.apache.tsfile.read.common.block.TsBlockBuilder; +import org.junit.Test; + +import jakarta.ws.rs.core.Response; + +import java.nio.ByteBuffer; +import java.util.ArrayDeque; +import java.util.Arrays; +import java.util.Collections; +import java.util.Optional; +import java.util.Queue; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +public class QueryDataSetHandlerTest { + + @Test + public void fillDataSetWithTimestampsShouldAllowResultAtExactLimit() throws Exception { + Response response = + QueryDataSetHandler.fillDataSetWithTimestamps( + new TestQueryExecution(newTsBlock(1L, 11L, 2L, 22L)), 2, 1); + + assertTrue(response.getEntity() instanceof QueryDataSet); + QueryDataSet dataSet = (QueryDataSet) response.getEntity(); + assertEquals(Arrays.asList(1L, 2L), dataSet.getTimestamps()); + assertEquals(Collections.singletonList(Arrays.asList(11L, 22L)), dataSet.getValues()); + } + + @Test + public void fillDataSetWithTimestampsShouldRejectRowsBeyondLimit() throws Exception { + Response response = + QueryDataSetHandler.fillDataSetWithTimestamps( + new TestQueryExecution(newTsBlock(1L, 11L, 2L, 22L)), 1, 1); + + assertTrue(response.getEntity() instanceof ExecutionStatus); + ExecutionStatus status = (ExecutionStatus) response.getEntity(); + assertEquals(TSStatusCode.QUERY_PROCESS_ERROR.getStatusCode(), status.getCode().intValue()); + } + + private static TsBlock newTsBlock(long... timeAndValues) { + TsBlockBuilder builder = new TsBlockBuilder(Collections.singletonList(TSDataType.INT64)); + for (int i = 0; i < timeAndValues.length; i += 2) { + builder.getTimeColumnBuilder().writeLong(timeAndValues[i]); + builder.getColumnBuilder(0).writeLong(timeAndValues[i + 1]); + builder.declarePosition(); + } + return builder.build(); + } + + private static final class TestQueryExecution implements IQueryExecution { + + private static final String COLUMN_NAME = "root.sg.d1.s1"; + + private final Queue<Optional<TsBlock>> batches = new ArrayDeque<>(); + private final DatasetHeader datasetHeader; + + private TestQueryExecution(TsBlock... tsBlocks) { + for (TsBlock tsBlock : tsBlocks) { + batches.add(Optional.of(tsBlock)); + } + batches.add(Optional.empty()); + + datasetHeader = + new DatasetHeader( + Collections.singletonList(new ColumnHeader(COLUMN_NAME, TSDataType.INT64)), false); + datasetHeader.setTreeColumnToTsBlockIndexMap(Collections.singletonList(COLUMN_NAME)); + } + + @Override + public void start() {} + + @Override + public void stop(Throwable t) {} + + @Override + public void stopAndCleanup(Throwable t) {} + + @Override + public void cancel() {} + + @Override + public ExecutionResult getStatus() { + return null; + } + + @Override + public Optional<TsBlock> getBatchResult() throws IoTDBException { + return batches.remove(); + } + + @Override + public Optional<ByteBuffer> getByteBufferBatchResult() { + return Optional.empty(); + } + + @Override + public boolean hasNextResult() { + return !batches.isEmpty(); + } + + @Override + public int getOutputValueColumnCount() { + return 1; + } + + @Override + public DatasetHeader getDatasetHeader() { + return datasetHeader; + } + + @Override + public QueryType getQueryType() { + return QueryType.READ; + } + + @Override + public boolean isQuery() { + return true; + } + + @Override + public boolean isUserQuery() { + return true; + } + + @Override + public String getQueryId() { + return null; + } + + @Override + public long getStartExecutionTime() { + return 0; + } + + @Override + public void recordExecutionTime(long executionTime) {} + + @Override + public void updateCurrentRpcStartTime(long startTime) {} + + @Override + public boolean isActive() { + return false; + } + + @Override + public long getTotalExecutionTime() { + return 0; + } + + @Override + public long getTimeout() { + return 0; + } + + @Override + public Optional<String> getExecuteSQL() { + return Optional.empty(); + } + + @Override + public String getStatementType() { + return null; + } + + @Override + public SqlDialect getSQLDialect() { + return SqlDialect.TREE; + } + + @Override + public String getUser() { + return null; + } + + @Override + public String getClientHostname() { + return null; + } + + @Override + public boolean isDebug() { + return false; + } + } +} diff --git a/iotdb-client/session/src/test/resources/iotdb-common.properties b/iotdb-client/session/src/test/resources/iotdb-common.properties index 2ec58b5db32..5c2150a3789 100644 --- a/iotdb-client/session/src/test/resources/iotdb-common.properties +++ b/iotdb-client/session/src/test/resources/iotdb-common.properties @@ -27,7 +27,8 @@ enable_rest_service=true # the binding port of the REST service # rest_service_port=18080 -# the default row limit to a REST query response when the rowSize parameter is not given in request +# The maximum row limit for REST and Grafana query responses. +# The request rowLimit/row_limit value cannot exceed this limit. # rest_query_default_row_size_limit=10000 # Whether to display rest service interface information through OpenAPI. eg: http://ip:port/openapi.json diff --git a/iotdb-core/datanode/src/test/resources/iotdb-common.properties b/iotdb-core/datanode/src/test/resources/iotdb-common.properties index b753b0519ee..d22cfaa6372 100644 --- a/iotdb-core/datanode/src/test/resources/iotdb-common.properties +++ b/iotdb-core/datanode/src/test/resources/iotdb-common.properties @@ -30,7 +30,8 @@ enable_rest_service=true # Whether to display rest service interface information through OpenAPI. eg: http://ip:port/openapi.json # enable_swagger=false -# the default row limit to a REST query response when the rowSize parameter is not given in request +# The maximum row limit for REST and Grafana query responses. +# The request rowLimit/row_limit value cannot exceed this limit. # rest_query_default_row_size_limit=10000 # is SSL enabled diff --git a/iotdb-core/datanode/src/test/resources/iotdb-system.properties b/iotdb-core/datanode/src/test/resources/iotdb-system.properties index a7a2a08bd1f..4cde9826f07 100644 --- a/iotdb-core/datanode/src/test/resources/iotdb-system.properties +++ b/iotdb-core/datanode/src/test/resources/iotdb-system.properties @@ -48,7 +48,8 @@ enable_rest_service=true # Whether to display rest service interface information through OpenAPI. eg: http://ip:port/openapi.json # enable_swagger=false -# the default row limit to a REST query response when the rowSize parameter is not given in request +# The maximum row limit for REST and Grafana query responses. +# The request rowLimit/row_limit value cannot exceed this limit. # rest_query_default_row_size_limit=10000 # is SSL enabled diff --git a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template index fdf79943255..91e84e89a0e 100644 --- a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template +++ b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template @@ -607,7 +607,8 @@ rest_service_port=18080 # Datatype: boolean enable_swagger=false -# the default row limit to a REST query response when the rowSize parameter is not given in request +# The maximum row limit for REST and Grafana query responses. +# The request rowLimit/row_limit value cannot exceed this limit. # effectiveMode: restart # Datatype: int rest_query_default_row_size_limit=10000
