[Improvement][SQL] Query return number In SQL should be configurable (#5632)

* [Improvement][SQL] Query return number should be configurable #5580

* Update SqlTaskTest.java

* Update sql.vue

dolphinscheduler-common/src/main/java/org/apache/dolphinscheduler/common/task/sql/SqlParameters.java

@@ -94,6 +94,16 @@ public class SqlParameters extends AbstractParameters {
      */
     private String title;
 
+    private int limit;
+
+    public int getLimit() {
+        return limit;
+    }
+
+    public void setLimit(int limit) {
+        this.limit = limit;
+    }
+
     public String getType() {
         return type;
     }
@@ -217,6 +227,7 @@ public class SqlParameters extends AbstractParameters {
                 + ", sqlType=" + sqlType
                 + ", sendEmail=" + sendEmail
                 + ", displayRows=" + displayRows
+                + ", limit=" + limit
                 + ", udfs='" + udfs + '\''
                 + ", showType='" + showType + '\''
                 + ", connParams='" + connParams + '\''

dolphinscheduler-common/src/test/java/org/apache/dolphinscheduler/common/task/SqlParametersTest.java

@@ -35,6 +35,7 @@ public class SqlParametersTest {
     private final String showType = "TABLE";
     private final String title = "sql test";
     private final int groupId = 0;
+    private final int limit = 0;
 
     @Test
     public void testSqlParameters() {
@@ -51,6 +52,7 @@ public class SqlParametersTest {
         sqlParameters.setShowType(showType);
         sqlParameters.setTitle(title);
         sqlParameters.setGroupId(groupId);
+        sqlParameters.setLimit(limit);
 
         Assert.assertEquals(type, sqlParameters.getType());
         Assert.assertEquals(sql, sqlParameters.getSql());
@@ -62,6 +64,7 @@ public class SqlParametersTest {
         Assert.assertEquals(showType, sqlParameters.getShowType());
         Assert.assertEquals(title, sqlParameters.getTitle());
         Assert.assertEquals(groupId, sqlParameters.getGroupId());
+        Assert.assertEquals(limit, sqlParameters.getLimit());
 
         Assert.assertTrue(sqlParameters.checkParameters());
     }

dolphinscheduler-server/src/main/java/org/apache/dolphinscheduler/server/worker/task/sql/SqlTask.java

@@ -86,12 +86,6 @@ public class SqlTask extends AbstractTask {
      */
     private TaskExecutionContext taskExecutionContext;
 
-    /**
-     * default query sql limit
-     */
-    private static final int LIMIT = 10000;
-
-
     private AlertClientService alertClientService;
 
     public SqlTask(TaskExecutionContext taskExecutionContext, Logger logger, AlertClientService alertClientService) {
@@ -117,14 +111,16 @@ public class SqlTask extends AbstractTask {
         Thread.currentThread().setName(threadLoggerInfoName);
 
         logger.info("Full sql parameters: {}", sqlParameters);
-        logger.info("sql type : {}, datasource : {}, sql : {} , localParams : {},udfs : {},showType : {},connParams : {}",
+        logger.info("sql type : {}, datasource : {}, sql : {} , localParams : {}, udfs : {},showType : {}, "
+                        + "connParams : {}, query max result limit : {}",
                 sqlParameters.getType(),
                 sqlParameters.getDatasource(),
                 sqlParameters.getSql(),
                 sqlParameters.getLocalParams(),
                 sqlParameters.getUdfs(),
                 sqlParameters.getShowType(),
-                sqlParameters.getConnParams());
+                sqlParameters.getConnParams(),
+                sqlParameters.getLimit());
         try {
             SQLTaskExecutionContext sqlTaskExecutionContext = taskExecutionContext.getSqlTaskExecutionContext();
 
@@ -309,7 +305,7 @@ public class SqlTask extends AbstractTask {
 
             int rowCount = 0;
 
-            while (rowCount < LIMIT && resultSet.next()) {
+            while (rowCount < sqlParameters.getLimit() && resultSet.next()) {
                 ObjectNode mapOfColValues = JSONUtils.createObjectNode();
                 for (int i = 1; i <= num; i++) {
                     mapOfColValues.set(md.getColumnLabel(i), JSONUtils.toJsonNode(resultSet.getObject(i)));
@@ -326,12 +322,11 @@ public class SqlTask extends AbstractTask {
                 logger.info("row {} : {}", i + 1, row);
             }
         }
-
         String result = JSONUtils.toJsonString(resultJSONArray);
         if (sqlParameters.getSendEmail() == null || sqlParameters.getSendEmail()) {
             sendAttachment(sqlParameters.getGroupId(), StringUtils.isNotEmpty(sqlParameters.getTitle())
-                            ? sqlParameters.getTitle()
-                            : taskExecutionContext.getTaskName() + " query result sets", result);
+                    ? sqlParameters.getTitle()
+                    : taskExecutionContext.getTaskName() + " query result sets", result);
         }
         logger.debug("execute sql result : {}", result);
         return result;

dolphinscheduler-server/src/test/java/org/apache/dolphinscheduler/server/worker/task/sql/SqlTaskTest.java

@@ -19,6 +19,7 @@ package org.apache.dolphinscheduler.server.worker.task.sql;
 
 import org.apache.dolphinscheduler.common.Constants;
 import org.apache.dolphinscheduler.common.datasource.DatasourceUtil;
+import org.apache.dolphinscheduler.common.task.sql.SqlParameters;
 import org.apache.dolphinscheduler.common.utils.ParameterUtils;
 import org.apache.dolphinscheduler.dao.AlertDao;
 import org.apache.dolphinscheduler.remote.command.alert.AlertSendResponseCommand;
@@ -148,4 +149,60 @@ public class SqlTaskTest {
         String result = Whitebox.invokeMethod(sqlTask, "resultProcess", resultSet);
         Assert.assertNotNull(result);
     }
+
+    @Test
+    public void testQueryBySQLUsingLimit() throws Exception {
+        TaskExecutionContext localTaskExecutionContext;
+        TaskProps props = new TaskProps();
+        props.setExecutePath("/tmp");
+        props.setTaskAppId(String.valueOf(System.currentTimeMillis()));
+        props.setTaskInstanceId(1);
+        props.setTenantCode("1");
+        props.setEnvFile(".dolphinscheduler_env.sh");
+        props.setTaskStartTime(new Date());
+        props.setTaskTimeout(0);
+        props.setTaskParams(
+            "{\"localParams\":[{\"prop\":\"ret\", \"direct\":\"OUT\", \"type\":\"VARCHAR\", \"value\":\"\"}],"
+                + "\"type\":\"POSTGRESQL\",\"datasource\":1,\"sql\":\"SELECT * FROM tb_1\","
+                + "\"sqlType\":0, \"limit\":1, \"sendEmail\":\"false\"}");
+
+        localTaskExecutionContext = PowerMockito.mock(TaskExecutionContext.class);
+        PowerMockito.when(localTaskExecutionContext.getTaskParams()).thenReturn(props.getTaskParams());
+        PowerMockito.when(localTaskExecutionContext.getExecutePath()).thenReturn("/tmp");
+        PowerMockito.when(localTaskExecutionContext.getTaskAppId()).thenReturn("1");
+        PowerMockito.when(localTaskExecutionContext.getTenantCode()).thenReturn("root");
+        PowerMockito.when(localTaskExecutionContext.getStartTime()).thenReturn(new Date());
+        PowerMockito.when(localTaskExecutionContext.getTaskTimeout()).thenReturn(10000);
+        PowerMockito.when(localTaskExecutionContext.getLogPath()).thenReturn("/tmp/dx");
+
+        SQLTaskExecutionContext sqlTaskExecutionContext = new SQLTaskExecutionContext();
+        sqlTaskExecutionContext.setConnectionParams(CONNECTION_PARAMS);
+        PowerMockito.when(localTaskExecutionContext.getSqlTaskExecutionContext()).thenReturn(sqlTaskExecutionContext);
+
+        PowerMockito.mockStatic(SpringApplicationContext.class);
+        PowerMockito.when(SpringApplicationContext.getBean(Mockito.any())).thenReturn(new AlertDao());
+        AlertClientService localAlertClientService = PowerMockito.mock(AlertClientService.class);
+        SqlTask localSqlTask = new SqlTask(localTaskExecutionContext, logger, localAlertClientService);
+        localSqlTask.init();
+
+        ResultSet resultSet = PowerMockito.mock(ResultSet.class);
+        ResultSetMetaData mockResultMetaData = PowerMockito.mock(ResultSetMetaData.class);
+        PowerMockito.when(resultSet.getMetaData()).thenReturn(mockResultMetaData);
+        PowerMockito.when(mockResultMetaData.getColumnCount()).thenReturn(2);
+        PowerMockito.when(resultSet.next()).thenReturn(true);
+        PowerMockito.when(resultSet.getObject(Mockito.anyInt())).thenReturn(1);
+        PowerMockito.when(mockResultMetaData.getColumnLabel(Mockito.anyInt())).thenReturn("a");
+
+        AlertSendResponseCommand mockResponseCommand = PowerMockito.mock(AlertSendResponseCommand.class);
+        PowerMockito.when(mockResponseCommand.getResStatus()).thenReturn(true);
+        PowerMockito.when(localAlertClientService.sendAlert(Mockito.anyInt(), Mockito.anyString(), Mockito.anyString()))
+            .thenReturn(mockResponseCommand);
+
+        String result = Whitebox.invokeMethod(localSqlTask, "resultProcess", resultSet);
+        Assert.assertEquals(1, ((SqlParameters) localSqlTask.getParameters()).getLimit());
+
+        // In fact, the target table has 2 rows, as we set the limit to 1, if the limit works, the `resultProcess` method
+        // should return [{"a":1}] rather then [{"a":1},{"a":1}]
+        Assert.assertEquals("[{\"a\":1}]", result);
+    }
 }

dolphinscheduler-ui/src/js/conf/home/pages/dag/_source/formModel/tasks/sql.vue

@@ -43,6 +43,13 @@
         </div>
       </div>
     </m-list-box>
+    <m-list-box v-show="sqlType === '0'">
+        <div slot="text"><strong class='requiredIcon'>*</strong>{{$t('Max Numbers Return')}}</div>
+        <div slot="content">
+          <el-input type="input" :disabled="isDetails" size="medium" v-model="limit" :placeholder="$t('Max Numbers Return placeholder')">
+          </el-input>
+        </div>
+    </m-list-box>
     <template v-if="sqlType === '0' && sendEmail">
       <m-list-box>
         <div slot="text"><strong class='requiredIcon'>*</strong>{{$t('Title')}}</div>
@@ -178,6 +185,8 @@
         sendEmail: false,
         // Display rows
         displayRows: 10,
+        // Max returned rows
+        limit: 10000,
         // Email title
         title: '',
         // Sql parameter
@@ -197,6 +206,13 @@
       createNodeId: Number
     },
     methods: {
+      /**
+       * limit should't be empty;limit should be a non-negative number str;
+       * limit should be a number smaller or equal than Integer.MAX_VALUE in java.
+       */
+      isLimitInvalid () {
+        return !this.limit || !/^(0|[1-9]\d*)$/.test(this.limit) || parseInt(this.limit, 10) > 2147483647
+      },
       setEditorVal () {
         this.item = editor.getValue()
         this.scriptBoxDialog = true
@@ -258,6 +274,10 @@
           this.$message.warning(`${i18n.$t('Mail subject required')}`)
           return false
         }
+        if (this.sqlType === '0' && this.isLimitInvalid()) {
+          this.$message.warning(`${i18n.$t('Max Numbers Return required')}`)
+          return false
+        }
         if (this.sqlType === '0' && this.sendEmail && (this.groupId === '' || this.groupId === null)) {
           this.$message.warning(`${i18n.$t('Alarm group required')}`)
           return false
@@ -293,6 +313,7 @@
           sqlType: this.sqlType,
           sendEmail: this.sendEmail,
           displayRows: this.displayRows,
+          limit: this.limit,
           title: this.title,
           groupId: this.groupId,
           localParams: this.localParams,
@@ -344,6 +365,7 @@
           sqlType: this.sqlType,
           sendEmail: this.sendEmail,
           displayRows: this.displayRows,
+          limit: this.limit,
           title: this.title,
           groupId: this.groupId,
           localParams: this.localParams,
@@ -392,6 +414,7 @@
         this.sqlType = o.params.sqlType
         this.sendEmail = o.params.sendEmail || false
         this.displayRows = o.params.displayRows || 10
+        this.limit = o.params.limit || 10000
         this.connParams = o.params.connParams || ''
         this.localParams = o.params.localParams || []
         this.preStatements = o.params.preStatements || []
@@ -424,6 +447,7 @@
           sqlType: this.sqlType,
           sendEmail: this.sendEmail,
           displayRows: this.displayRows,
+          limit: this.limit,
           title: this.title,
           groupId: this.groupId,
           localParams: this.localParams,

dolphinscheduler-ui/src/js/module/i18n/locale/en_US.js

@@ -133,6 +133,9 @@ export default {
   'SQL Type': 'SQL Type',
   'Send Email': 'Send Email',
   'Log display': 'Log display',
+  'Max Numbers Return': 'Number of records to return',
+  'Max Numbers Return placeholder': 'Default is 10000, a large value may cause high pressure on the memory',
+  'Max Numbers Return required': 'Number of records to return parameter must be a number in the range of 0 - 2147483647',
   'rows of result': 'rows of result',
   Title: 'Title',
   'Please enter the title of email': 'Please enter the title of email',

dolphinscheduler-ui/src/js/module/i18n/locale/zh_CN.js

@@ -134,6 +134,9 @@ export default {
   'Send Email': '发送邮件',
   'Log display': '日志显示',
   'rows of result': '行查询结果',
+  'Max Numbers Return': '返回的记录行数',
+  'Max Numbers Return placeholder': '默认值10000,如果值过大可能会对内存造成较大压力',
+  'Max Numbers Return required': '返回的记录行数值必须是一个在0-2147483647范围内的整数',
   Title: '主题',
   'Please enter the title of email': '请输入邮件主题',
   Table: '表名',