[To dev/1.3] Enforce REST query row limit as hard cap (#18149)

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 swagger. eg: http://ip:port/swagger.json

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/handler/QueryRowLimitUtils.java

@@ -0,0 +1,68 @@
+/*
+ * 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.db.protocol.rest.handler;
+
+import org.apache.iotdb.db.protocol.rest.model.ExecutionStatus;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+import javax.ws.rs.core.Response;
+
+public final class QueryRowLimitUtils {
+
+  /**
+   * Fallback used when {@code rest_query_default_row_size_limit} is missing or non-positive.
+   * Matches the built-in default in {@code IoTDBRestServiceConfig}; a non-positive configured value
+   * used to mean "unlimited" and is now treated as this cap instead of being clamped down to 1.
+   */
+  private static final int DEFAULT_ROW_SIZE_LIMIT = 10000;
+
+  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 rowSizeLimit > 0 ? rowSizeLimit : DEFAULT_ROW_SIZE_LIMIT;
+  }
+
+  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(
+                        "Dataset row size exceeded the given max row size (%d)", rowSizeLimit)))
+        .build();
+  }
+}

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v1/handler/QueryDataSetHandler.java

@@ -18,6 +18,7 @@
 package org.apache.iotdb.db.protocol.rest.v1.handler;
 
 import org.apache.iotdb.commons.exception.IoTDBException;
+import org.apache.iotdb.db.protocol.rest.handler.QueryRowLimitUtils;
 import org.apache.iotdb.db.protocol.rest.v1.model.ExecutionStatus;
 import org.apache.iotdb.db.queryengine.common.header.DatasetHeader;
 import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution;
@@ -47,7 +48,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)
@@ -136,7 +137,6 @@ private static Response fillQueryDataSetWithTimestamps(
       final long timePrecision)
       throws IoTDBException {
     int fetched = 0;
-    int columnNum = queryExecution.getOutputValueColumnCount();
 
     DatasetHeader header = queryExecution.getDatasetHeader();
     List<String> resultColumns = header.getRespColumns();
@@ -147,17 +147,6 @@ private static Response fillQueryDataSetWithTimestamps(
     }
 
     while (true) {
-      if (0 < actualRowSizeLimit && actualRowSizeLimit <= fetched) {
-        return Response.ok()
-            .entity(
-                new org.apache.iotdb.db.protocol.rest.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) {
@@ -169,6 +158,9 @@ private static Response fillQueryDataSetWithTimestamps(
       }
       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(
@@ -180,7 +172,6 @@ private static Response fillQueryDataSetWithTimestamps(
         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 {
@@ -190,10 +181,8 @@ private static Response fillQueryDataSetWithTimestamps(
                     : column.getObject(i));
           }
         }
-        if (k != columnNum - 1) {
-          fetched -= currentCount;
-        }
       }
+      fetched += currentCount;
     }
     return Response.ok().entity(targetDataSet).build();
   }
@@ -207,17 +196,6 @@ private static Response fillQueryDataSetWithoutTimestamps(
     int fetched = 0;
     int columnNum = queryExecution.getOutputValueColumnCount();
     while (true) {
-      if (0 < actualRowSizeLimit && actualRowSizeLimit <= fetched) {
-        return Response.ok()
-            .entity(
-                new org.apache.iotdb.db.protocol.rest.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) {
@@ -232,11 +210,13 @@ private static Response fillQueryDataSetWithoutTimestamps(
         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 {
@@ -246,53 +226,67 @@ private static Response fillQueryDataSetWithoutTimestamps(
                     : column.getObject(i));
           }
         }
-        if (k != columnNum - 1) {
-          fetched -= currentCount;
-        }
       }
+      fetched += currentCount;
     }
     return Response.ok().entity(targetDataSet).build();
   }
 
   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();
   }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v1/impl/GrafanaApiServiceImpl.java

@@ -21,7 +21,9 @@
 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.rest.handler.AuthorizationHandler;
+import org.apache.iotdb.db.protocol.rest.handler.QueryRowLimitUtils;
 import org.apache.iotdb.db.protocol.rest.v1.GrafanaApiService;
 import org.apache.iotdb.db.protocol.rest.v1.handler.ExceptionHandler;
 import org.apache.iotdb.db.protocol.rest.v1.handler.QueryDataSetHandler;
@@ -67,11 +69,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":
@@ -130,7 +136,8 @@ public Response variables(SQL sql, SecurityContext securityContext) {
       }
       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();
@@ -200,9 +207,11 @@ public Response expression(ExpressionRequest expressionRequest, SecurityContext
       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) {
@@ -262,10 +271,10 @@ public Response node(List<String> requestBody, SecurityContext securityContext)
         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();

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/rest/v1/impl/RestApiServiceImpl.java

@@ -21,6 +21,7 @@
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.conf.rest.IoTDBRestServiceDescriptor;
 import org.apache.iotdb.db.protocol.rest.handler.AuthorizationHandler;
+import org.apache.iotdb.db.protocol.rest.handler.QueryRowLimitUtils;
 import org.apache.iotdb.db.protocol.rest.utils.InsertTabletSortDataUtils;
 import org.apache.iotdb.db.protocol.rest.v1.RestApiService;
 import org.apache.iotdb.db.protocol.rest.v1.handler.ExceptionHandler;
@@ -66,14 +67,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
@@ -208,7 +210,7 @@ public Response executeQueryStatement(SQL sql, SecurityContext securityContext)
         return QueryDataSetHandler.fillQueryDataSet(
             queryExecution,
             statement,
-            sql.getRowLimit() == null ? defaultQueryRowLimit : sql.getRowLimit());
+            QueryRowLimitUtils.resolveActualRowSizeLimit(sql.getRowLimit(), defaultQueryRowLimit));
       }
     } catch (Exception e) {
       finish = true;