[Improvement-18714][Task-Sql] Avoid redundant serialization in alert content limiting and log result rows incrementally in query result processing in SqlTask

dolphinscheduler-task-plugin/dolphinscheduler-task-sql/pom.xml

@@ -44,6 +44,14 @@
             <artifactId>dolphinscheduler-datasource-api</artifactId>
             <version>${project.version}</version>
             <scope>provided</scope>
+            <exclusions>
+                <!-- keep logback (brought by dolphinscheduler-common) the only slf4j binding on this
+                     module's classpath, so tests can capture logs through the logback api -->
+                <exclusion>
+                    <groupId>org.slf4j</groupId>
+                    <artifactId>slf4j-reload4j</artifactId>
+                </exclusion>
+            </exclusions>
         </dependency>
 
         <dependency>

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

@@ -302,22 +302,21 @@ private String resultProcess(ResultSet resultSet) throws Exception {
                 }
             }
 
+            int displayRows = sqlParameters.getDisplayRows() > 0 ? sqlParameters.getDisplayRows()
+                    : TaskConstants.DEFAULT_DISPLAY_ROWS;
+            log.info("display the first {} rows of sql result as follows:", displayRows);
             while (resultSet.next()) {
                 ObjectNode mapOfColValues = JSONUtils.createObjectNode();
                 for (int i = 1; i <= num; i++) {
                     mapOfColValues.set(columnLabels[i - 1], JSONUtils.toJsonNode(resultSet.getObject(i)));
                 }
                 resultJSONArray.add(mapOfColValues);
+                if (resultJSONArray.size() <= displayRows) {
+                    log.info("row {} : {}", resultJSONArray.size(), JSONUtils.toJsonString(mapOfColValues));
+                }
             }
-
-            int displayRows = sqlParameters.getDisplayRows() > 0 ? sqlParameters.getDisplayRows()
-                    : TaskConstants.DEFAULT_DISPLAY_ROWS;
-            displayRows = Math.min(displayRows, resultJSONArray.size());
-            log.info("display sql result {} rows as follows:", displayRows);
-            for (int i = 0; i < displayRows; i++) {
-                String row = JSONUtils.toJsonString(resultJSONArray.get(i));
-                log.info("row {} : {}", i + 1, row);
-            }
+            log.info("displayed {} of {} rows", Math.min(displayRows, resultJSONArray.size()),
+                    resultJSONArray.size());
         }
 
         String result;
@@ -385,31 +384,65 @@ private void prepareTaskResultAlert(ArrayNode resultJSONArray) {
      */
     private String limitAlertContent(ArrayNode resultJSONArray) {
         int totalRows = resultJSONArray.size();
-        ArrayNode limited = JSONUtils.createArrayNode();
-        int rowsToKeep = Math.min(totalRows, alertMaxRows);
-        for (int i = 0; i < rowsToKeep; i++) {
-            limited.add(limitRowValues(resultJSONArray.get(i)));
+
+        // Serialize each kept row once and track the running UTF-8 byte size, so the budget
+        // check on every dropped row is O(1) instead of rebuilding and re-encoding the whole
+        // content (which would turn the drop loop into O(rows^2)).
+        List<String> serializedRows = new ArrayList<>(Math.min(totalRows, alertMaxRows));
+        int rowsBytes = 0;
+        for (int i = 0; i < totalRows && i < alertMaxRows; i++) {
+            String rowJson = JSONUtils.toJsonString(limitRowValues(resultJSONArray.get(i)));
+            serializedRows.add(rowJson);
+            rowsBytes += rowJson.getBytes(StandardCharsets.UTF_8).length;
         }
 
-        int omittedRows = totalRows - rowsToKeep;
-        while (true) {
-            if (omittedRows > 0) {
-                limited.add(createOmissionMarker(omittedRows));
-            }
-            String content = JSONUtils.toJsonString(limited);
-            if (content.getBytes(StandardCharsets.UTF_8).length <= alertMaxContentBytes) {
-                return content;
-            }
-            if (omittedRows > 0) {
-                limited.remove(limited.size() - 1);
+        int omittedRows = totalRows - serializedRows.size();
+        while (!serializedRows.isEmpty()
+                && contentBytes(rowsBytes, serializedRows.size(), omittedRows) > alertMaxContentBytes) {
+            int last = serializedRows.size() - 1;
+            rowsBytes -= serializedRows.remove(last).getBytes(StandardCharsets.UTF_8).length;
+            omittedRows++;
+        }
+        return buildAlertContent(serializedRows, omittedRows);
+    }
+
+    /**
+     * UTF-8 byte size of the alert content that buildAlertContent produces from the kept rows
+     * and the omission marker, i.e. "[" + rows joined by "," + ("," + marker)? + "]".
+     */
+    private int contentBytes(int rowsBytes, int rowCount, int omittedRows) {
+        int bytes = 2; // '[' and ']'
+        if (rowCount > 0) {
+            bytes += rowsBytes + rowCount - 1; // the rows plus their ',' separators
+        }
+        if (omittedRows > 0) {
+            bytes += (rowCount > 0 ? 1 : 0)
+                    + JSONUtils.toJsonString(createOmissionMarker(omittedRows)).getBytes(StandardCharsets.UTF_8).length;
+        }
+        return bytes;
+    }
+
+    /**
+     * Assemble the alert content from the kept rows and the omission marker,
+     * i.e. "[" + rows joined by "," + ("," + marker)? + "]".
+     */
+    private String buildAlertContent(List<String> serializedRows, int omittedRows) {
+        StringBuilder content = new StringBuilder();
+        content.append('[');
+        for (int i = 0; i < serializedRows.size(); i++) {
+            if (i > 0) {
+                content.append(',');
             }
-            if (limited.isEmpty()) {
-                // Even a lone marker exceeds the budget (pathological budget); return the minimal notice.
-                return JSONUtils.toJsonString(createOmissionMarkerArray(totalRows));
+            content.append(serializedRows.get(i));
+        }
+        if (omittedRows > 0) {
+            if (!serializedRows.isEmpty()) {
+                content.append(',');
             }
-            limited.remove(limited.size() - 1);
-            omittedRows++;
+            content.append(JSONUtils.toJsonString(createOmissionMarker(omittedRows)));
         }
+        content.append(']');
+        return content.toString();
     }
 
     private ObjectNode createOmissionMarker(int omittedRows) {
@@ -418,12 +451,6 @@ private ObjectNode createOmissionMarker(int omittedRows) {
         return omitted;
     }
 
-    private ArrayNode createOmissionMarkerArray(int totalRows) {
-        ArrayNode markerArray = JSONUtils.createArrayNode();
-        markerArray.add(createOmissionMarker(totalRows));
-        return markerArray;
-    }
-
     private JsonNode limitRowValues(JsonNode row) {
         if (!row.isObject()) {
             return row;

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

@@ -48,14 +48,22 @@
 import java.sql.ResultSet;
 import java.sql.ResultSetMetaData;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
 
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
 import org.mockito.MockedStatic;
 import org.mockito.Mockito;
+import org.slf4j.LoggerFactory;
+
+import ch.qos.logback.classic.Logger;
+import ch.qos.logback.classic.spi.ILoggingEvent;
+import ch.qos.logback.core.read.ListAppender;
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ArrayNode;
@@ -384,6 +392,145 @@ void testResultProcess_DuplicateColumnLabels_ThrowsTaskException() throws Except
                 "TaskException message should mention duplicate column name");
     }
 
+    @Test
+    void testResultProcessLogsRowsWhileFetchingAndOnlyUpToDisplayRows() throws Exception {
+        String taskParams = "{\"type\":\"HIVE\",\"datasource\":1,\"sql\":\"select 1\",\"displayRows\":2}";
+        TaskExecutionContext ctx = new TaskExecutionContext();
+        ctx.setTaskParams(taskParams);
+        ctx.setTaskName("display_rows_task");
+        ctx.setResourceParametersHelper(getResourceParametersHelperWithDatasourceType(DbType.HIVE));
+
+        SqlTask task = new SqlTask(ctx);
+
+        ListAppender<ILoggingEvent> appender = attachLogAppender();
+        try {
+            ResultSet mockResultSet = mock(ResultSet.class);
+            ResultSetMetaData mockMetaData = mock(ResultSetMetaData.class);
+            when(mockResultSet.getMetaData()).thenReturn(mockMetaData);
+            when(mockMetaData.getColumnCount()).thenReturn(1);
+            when(mockMetaData.getColumnLabel(1)).thenReturn("id");
+            AtomicInteger fetchedRows = new AtomicInteger();
+            when(mockResultSet.next()).thenAnswer(inv -> fetchedRows.incrementAndGet() <= 5);
+            when(mockResultSet.getObject(1)).thenAnswer(inv -> {
+                int currentRow = fetchedRows.get();
+                if (currentRow > 1) {
+                    long loggedRows = appender.list.stream()
+                            .map(ILoggingEvent::getFormattedMessage)
+                            .filter(message -> message.startsWith("row "))
+                            .count();
+                    Assertions.assertEquals(Math.min(currentRow - 1, 2), loggedRows, "before fetching row "
+                            + currentRow + ", the first " + Math.min(currentRow - 1, 2)
+                            + " rows should already be logged");
+                }
+                return "v" + currentRow;
+            });
+
+            Method resultProcessMethod = SqlTask.class.getDeclaredMethod("resultProcess", ResultSet.class);
+            resultProcessMethod.setAccessible(true);
+            String result = (String) resultProcessMethod.invoke(task, mockResultSet);
+
+            List<String> messages = appender.list.stream()
+                    .map(ILoggingEvent::getFormattedMessage)
+                    .collect(Collectors.toList());
+            Assertions.assertTrue(messages.contains("display the first 2 rows of sql result as follows:"));
+            Assertions.assertTrue(messages.contains("row 1 : {\"id\":\"v1\"}"));
+            Assertions.assertTrue(messages.contains("row 2 : {\"id\":\"v2\"}"));
+            Assertions.assertEquals(2, messages.stream().filter(message -> message.startsWith("row ")).count(),
+                    "only the first displayRows=2 rows should be logged");
+            Assertions.assertTrue(messages.contains("displayed 2 of 5 rows"));
+            // the full result set is kept regardless of the display limit
+            Assertions.assertTrue(result.contains("v3"));
+            Assertions.assertTrue(result.contains("v5"));
+        } finally {
+            detachLogAppender(appender);
+        }
+    }
+
+    @Test
+    void testResultProcessLogsAllRowsWhenResultSmallerThanDefaultDisplayRows() throws Exception {
+        String taskParams = "{\"type\":\"HIVE\",\"datasource\":1,\"sql\":\"select 1\"}";
+        TaskExecutionContext ctx = new TaskExecutionContext();
+        ctx.setTaskParams(taskParams);
+        ctx.setTaskName("default_display_rows_task");
+        ctx.setResourceParametersHelper(getResourceParametersHelperWithDatasourceType(DbType.HIVE));
+
+        SqlTask task = new SqlTask(ctx);
+
+        ListAppender<ILoggingEvent> appender = attachLogAppender();
+        try {
+            ResultSet mockResultSet = mock(ResultSet.class);
+            ResultSetMetaData mockMetaData = mock(ResultSetMetaData.class);
+            when(mockResultSet.getMetaData()).thenReturn(mockMetaData);
+            when(mockMetaData.getColumnCount()).thenReturn(1);
+            when(mockMetaData.getColumnLabel(1)).thenReturn("id");
+            when(mockResultSet.next()).thenReturn(true, true, false);
+            when(mockResultSet.getObject(1)).thenReturn("v");
+
+            Method resultProcessMethod = SqlTask.class.getDeclaredMethod("resultProcess", ResultSet.class);
+            resultProcessMethod.setAccessible(true);
+            resultProcessMethod.invoke(task, mockResultSet);
+
+            List<String> messages = appender.list.stream()
+                    .map(ILoggingEvent::getFormattedMessage)
+                    .collect(Collectors.toList());
+            Assertions.assertTrue(messages.contains("display the first 10 rows of sql result as follows:"));
+            Assertions.assertTrue(messages.contains("row 1 : {\"id\":\"v\"}"));
+            Assertions.assertTrue(messages.contains("row 2 : {\"id\":\"v\"}"));
+            Assertions.assertEquals(2, messages.stream().filter(message -> message.startsWith("row ")).count());
+            Assertions.assertTrue(messages.contains("displayed 2 of 2 rows"));
+        } finally {
+            detachLogAppender(appender);
+        }
+    }
+
+    @Test
+    void testResultProcessLogsNoRowForEmptyResultSet() throws Exception {
+        String taskParams = "{\"type\":\"HIVE\",\"datasource\":1,\"sql\":\"select 1\",\"displayRows\":3}";
+        TaskExecutionContext ctx = new TaskExecutionContext();
+        ctx.setTaskParams(taskParams);
+        ctx.setTaskName("empty_result_display_rows_task");
+        ctx.setResourceParametersHelper(getResourceParametersHelperWithDatasourceType(DbType.HIVE));
+
+        SqlTask task = new SqlTask(ctx);
+
+        ListAppender<ILoggingEvent> appender = attachLogAppender();
+        try {
+            ResultSet mockResultSet = mock(ResultSet.class);
+            ResultSetMetaData mockMetaData = mock(ResultSetMetaData.class);
+            when(mockResultSet.getMetaData()).thenReturn(mockMetaData);
+            when(mockMetaData.getColumnCount()).thenReturn(1);
+            when(mockMetaData.getColumnLabel(1)).thenReturn("id");
+            when(mockResultSet.next()).thenReturn(false);
+
+            Method resultProcessMethod = SqlTask.class.getDeclaredMethod("resultProcess", ResultSet.class);
+            resultProcessMethod.setAccessible(true);
+            resultProcessMethod.invoke(task, mockResultSet);
+
+            List<String> messages = appender.list.stream()
+                    .map(ILoggingEvent::getFormattedMessage)
+                    .collect(Collectors.toList());
+            Assertions.assertTrue(messages.contains("display the first 3 rows of sql result as follows:"));
+            Assertions.assertEquals(0, messages.stream().filter(message -> message.startsWith("row ")).count());
+            Assertions.assertTrue(messages.contains("displayed 0 of 0 rows"));
+        } finally {
+            detachLogAppender(appender);
+        }
+    }
+
+    private ListAppender<ILoggingEvent> attachLogAppender() {
+        Logger logger = (Logger) LoggerFactory.getLogger(SqlTask.class);
+        ListAppender<ILoggingEvent> appender = new ListAppender<>();
+        appender.start();
+        logger.addAppender(appender);
+        return appender;
+    }
+
+    private void detachLogAppender(ListAppender<ILoggingEvent> appender) {
+        Logger logger = (Logger) LoggerFactory.getLogger(SqlTask.class);
+        logger.detachAppender(appender);
+        appender.stop();
+    }
+
     @Test
     void testGetSqlAndSqlParamsMap_nullPrepareParamsMap_replacesScheduleTimeAndTitlePlaceholder() throws Exception {
         long scheduleTimeMillis = 1700000000000L;
@@ -749,6 +896,53 @@ void testPrepareTaskResultAlertReplacesRowExceedingByteBudgetAlone() throws Exce
         Assertions.assertEquals("1 more rows omitted", only.get("...").asText());
     }
 
+    @Test
+    void testPrepareTaskResultAlertKeepsMaxRowsFittingExactByteBudget() throws Exception {
+        // each row {"id":"N"} serializes to 10 UTF-8 bytes and the marker row to 29 bytes,
+        // so keeping 2 rows costs exactly 53 bytes: brackets(2) + rows(20) + separator(1) + marker(30)
+        ArrayNode resultArray = JSONUtils.createArrayNode();
+        for (int i = 0; i < 5; i++) {
+            resultArray.add(JSONUtils.parseObject("{\"id\":\"" + i + "\"}"));
+        }
+
+        try (MockedStatic<PropertyUtils> propertyUtils = Mockito.mockStatic(PropertyUtils.class)) {
+            propertyUtils.when(() -> PropertyUtils.getInt("sql.task.alert.max.rows", 30)).thenReturn(5);
+            propertyUtils.when(() -> PropertyUtils.getInt("sql.task.alert.max.row.value.length", 100))
+                    .thenReturn(100);
+            propertyUtils.when(() -> PropertyUtils.getInt("sql.task.alert.max.content.bytes", 60000)).thenReturn(53);
+
+            TaskExecutionContext ctx = createSqlTaskWithAlert(resultArray, 10, "exact byte budget test", 1);
+
+            TaskAlertInfo alertInfo = ctx.getTaskAlertInfo();
+            Assertions.assertEquals("[{\"id\":\"0\"},{\"id\":\"1\"},{\"...\":\"3 more rows omitted\"}]",
+                    alertInfo.getContent());
+        }
+    }
+
+    @Test
+    void testPrepareTaskResultAlertMeasuresUtf8BytesNotChars() throws Exception {
+        // each row {"id":"中中中"} is 12 chars but 18 UTF-8 bytes, so a char-based budget
+        // check would keep 2 rows while the byte-based check keeps only 1
+        ArrayNode resultArray = JSONUtils.createArrayNode();
+        for (int i = 0; i < 5; i++) {
+            resultArray.add(JSONUtils.parseObject("{\"id\":\"" + repeatChar('\u4e2d', 3) + "\"}"));
+        }
+
+        try (MockedStatic<PropertyUtils> propertyUtils = Mockito.mockStatic(PropertyUtils.class)) {
+            propertyUtils.when(() -> PropertyUtils.getInt("sql.task.alert.max.rows", 30)).thenReturn(5);
+            propertyUtils.when(() -> PropertyUtils.getInt("sql.task.alert.max.row.value.length", 100))
+                    .thenReturn(100);
+            propertyUtils.when(() -> PropertyUtils.getInt("sql.task.alert.max.content.bytes", 60000)).thenReturn(60);
+
+            TaskExecutionContext ctx = createSqlTaskWithAlert(resultArray, 10, "utf8 bytes test", 1);
+
+            TaskAlertInfo alertInfo = ctx.getTaskAlertInfo();
+            Assertions.assertEquals(
+                    "[{\"id\":\"\u4e2d\u4e2d\u4e2d\"},{\"...\":\"4 more rows omitted\"}]",
+                    alertInfo.getContent());
+        }
+    }
+
     private String repeatChar(char c, int count) {
         StringBuilder builder = new StringBuilder(count);
         for (int i = 0; i < count; i++) {