[Improve][Connector-V2] RocketMQ Source add message tag config (#8825)

Author: wuchunfu

docs/en/connector-v2/source/RocketMQ.md

@@ -27,10 +27,11 @@ Source connector for Apache RocketMQ.
 
 ## Source Options
 
-|                Name                 |  Type   | Required |          Default           |                                                                                                    Description                                                                                                     |
+| Name                                |  Type   | Required |          Default           | Description                                                                                                                                                                                                        |
 |-------------------------------------|---------|----------|----------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
 | topics                              | String  | yes      | -                          | `RocketMQ topic` name. If there are multiple `topics`, use `,` to split, for example: `"tpc1,tpc2"`.                                                                                                               |
 | name.srv.addr                       | String  | yes      | -                          | `RocketMQ` name server cluster address.                                                                                                                                                                            |
+| tags                                | String  | no       | -                          | `RocketMQ tag` name. If there are multiple `tags`, use `,` to split, for example: `"tag1,tag2"`.                                                                                                                   |
 | acl.enabled                         | Boolean | no       | false                      | If true, access control is enabled, and access key and secret key need to be configured.                                                                                                                           |
 | access.key                          | String  | no       |                            |                                                                                                                                                                                                                    |
 | secret.key                          | String  | no       |                            | When ACL_ENABLED is true, secret key cannot be empty.                                                                                                                                                              |
@@ -218,3 +219,62 @@ sink {
 }
 ```
 
+### Specified tag example:
+
+> Here you can specify a tag to consume data. If there are multiple tags, use `,` to separate them, for example: "tag1,tag2"
+
+```hocon
+env {
+  parallelism = 1
+  job.mode = "BATCH"
+  
+  # You can set spark configuration here
+  spark.app.name = "SeaTunnel"
+  spark.executor.instances = 2
+  spark.executor.cores = 1
+  spark.executor.memory = "1g"
+  spark.master = local
+}
+
+source {
+  Rocketmq {
+    plugin_output = "rocketmq_table"
+    name.srv.addr = "localhost:9876"
+    topics = "test_topic"
+    format = text
+    # The default field delimiter is ","
+    field_delimiter = ","
+    tags = "test_tag"
+    schema = {
+      fields {
+        id = bigint
+        c_map = "map<string, smallint>"
+        c_array = "array<tinyint>"
+        c_string = string
+        c_boolean = boolean
+        c_tinyint = tinyint
+        c_smallint = smallint
+        c_int = int
+        c_bigint = bigint
+        c_float = float
+        c_double = double
+        c_decimal = "decimal(2, 1)"
+        c_bytes = bytes
+        c_date = date
+        c_timestamp = timestamp
+      }
+    }
+  }
+}
+
+transform {
+  # If you would like to get more information about how to configure seatunnel and see full list of transform plugins,
+  # please go to https://seatunnel.apache.org/docs/category/transform
+}
+
+sink {
+  Console {
+    plugin_input = "rocketmq_table"
+  }
+}
+```

seatunnel-connectors-v2/connector-rocketmq/src/main/java/org/apache/seatunnel/connectors/seatunnel/rocketmq/config/ConsumerConfig.java

@@ -29,8 +29,17 @@ public class ConsumerConfig extends Config {
                     .stringType()
                     .noDefaultValue()
                     .withDescription(
-                            "RocketMq topic name. If there are multiple topics, use , to split, for example: "
+                            "RocketMq topic name. If there are multiple topics, use `,` to split, for example: "
                                     + "\"tpc1,tpc2\".");
+
+    public static final Option<String> TAGS =
+            Options.key("tags")
+                    .stringType()
+                    .noDefaultValue()
+                    .withDescription(
+                            "RocketMq tag name. If there are multiple tags, use `,` to split, for example: "
+                                    + "\"tag1,tag2\".");
+
     public static final Option<String> CONSUMER_GROUP =
             Options.key("consumer.group")
                     .stringType()

seatunnel-connectors-v2/connector-rocketmq/src/main/java/org/apache/seatunnel/connectors/seatunnel/rocketmq/source/ConsumerMetadata.java

@@ -33,6 +33,7 @@
 public class ConsumerMetadata implements Serializable {
     private RocketMqBaseConfiguration baseConfig = RocketMqBaseConfiguration.newBuilder().build();
     private List<String> topics;
+    private List<String> tags;
     private boolean enabledCommitCheckpoint = false;
     private StartMode startMode;
     private Map<MessageQueue, Long> specificStartOffsets;

seatunnel-connectors-v2/connector-rocketmq/src/main/java/org/apache/seatunnel/connectors/seatunnel/rocketmq/source/RocketMqSource.java

@@ -55,8 +55,10 @@
 import com.google.auto.service.AutoService;
 
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 import static org.apache.seatunnel.connectors.seatunnel.rocketmq.config.Config.ACCESS_KEY;
 import static org.apache.seatunnel.connectors.seatunnel.rocketmq.config.Config.ACL_ENABLED;
@@ -112,6 +114,20 @@ public void prepare(Config config) throws PrepareFailException {
                         config.getString(ConsumerConfig.TOPICS.key())
                                 .split(DEFAULT_FIELD_DELIMITER)));
 
+        if (config.hasPath(ConsumerConfig.TAGS.key())) {
+            String tags = config.getString(ConsumerConfig.TAGS.key());
+            if (tags != null && !tags.trim().isEmpty()) {
+                this.metadata.setTags(
+                        Arrays.stream(tags.split(DEFAULT_FIELD_DELIMITER))
+                                .map(String::trim)
+                                .filter(tag -> !tag.isEmpty())
+                                .distinct()
+                                .collect(Collectors.toList()));
+            } else {
+                this.metadata.setTags(Collections.emptyList());
+            }
+        }
+
         RocketMqBaseConfiguration.Builder baseConfigBuilder =
                 RocketMqBaseConfiguration.newBuilder()
                         .consumer()

seatunnel-connectors-v2/connector-rocketmq/src/main/java/org/apache/seatunnel/connectors/seatunnel/rocketmq/source/RocketMqSourceFactory.java

@@ -41,6 +41,7 @@ public OptionRule optionRule() {
                 .required(ConsumerConfig.TOPICS, Config.NAME_SRV_ADDR)
                 .optional(
                         Config.FORMAT,
+                        ConsumerConfig.TAGS,
                         ConsumerConfig.START_MODE,
                         ConsumerConfig.CONSUMER_GROUP,
                         ConsumerConfig.COMMIT_ON_CHECKPOINT,

seatunnel-connectors-v2/connector-rocketmq/src/main/java/org/apache/seatunnel/connectors/seatunnel/rocketmq/source/RocketMqSourceReader.java

@@ -146,9 +146,20 @@ record ->
                                                                 .collect(Collectors.toList());
                                                 long lastOffset = -1;
                                                 for (MessageExt record : messages) {
-                                                    deserializationSchema.deserialize(
-                                                            record.getBody(), output);
-                                                    lastOffset = record.getQueueOffset();
+                                                    // Check if the tags are specified and match the
+                                                    // record's tag
+                                                    boolean shouldProcess =
+                                                            metadata.getTags() == null
+                                                                    || metadata.getTags().isEmpty()
+                                                                    || metadata.getTags()
+                                                                            .contains(
+                                                                                    record
+                                                                                            .getTags());
+                                                    if (shouldProcess) {
+                                                        deserializationSchema.deserialize(
+                                                                record.getBody(), output);
+                                                        lastOffset = record.getQueueOffset();
+                                                    }
                                                     if (Boundedness.BOUNDED.equals(
                                                                     context.getBoundedness())
                                                             && record.getQueueOffset()

seatunnel-e2e/seatunnel-connector-v2-e2e/connector-rocketmq-e2e/src/test/java/org/apache/seatunnel/e2e/connector/rocketmq/RocketMqIT.java

@@ -49,7 +49,10 @@
 import org.apache.rocketmq.common.message.Message;
 import org.apache.rocketmq.common.message.MessageExt;
 import org.apache.rocketmq.common.message.MessageQueue;
+import org.apache.rocketmq.common.protocol.route.QueueData;
+import org.apache.rocketmq.common.protocol.route.TopicRouteData;
 import org.apache.rocketmq.remoting.protocol.LanguageCode;
+import org.apache.rocketmq.tools.admin.DefaultMQAdminExt;
 
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;
@@ -77,6 +80,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
+import java.util.stream.Collectors;
 
 import static org.apache.seatunnel.e2e.connector.rocketmq.RocketMqContainer.NAMESRV_PORT;
 
@@ -204,6 +208,42 @@ public void testTextFormatSinkRocketMq(TestContainer container)
         Assertions.assertEquals(10, data.size());
     }
 
+    @TestTemplate
+    public void testSourceRocketMqTextTagToConsole(TestContainer container)
+            throws IOException, InterruptedException {
+        String topic = "test_topic_text_tag";
+        String tag = "tag_test";
+
+        // delete topic if exist
+        deleteTopicIfExist(topic);
+
+        DefaultSeaTunnelRowSerializer serializer =
+                new DefaultSeaTunnelRowSerializer(
+                        topic, tag, SEATUNNEL_ROW_TYPE, SchemaFormat.TEXT, DEFAULT_FIELD_DELIMITER);
+        generateTestData(serializer::serializeRow, topic, 0, 32);
+        Container.ExecResult execResult =
+                container.executeJob("/rocketmq-source_text_tag_to_console.conf");
+        Assertions.assertEquals(0, execResult.getExitCode(), execResult.getStderr());
+    }
+
+    @TestTemplate
+    public void testSourceRocketMqTextErrorTagToConsole(TestContainer container)
+            throws IOException, InterruptedException {
+        String topic = "test_topic_text_error_tag";
+        String tag = "test_error_tag";
+
+        // delete topic if exist
+        deleteTopicIfExist(topic);
+
+        DefaultSeaTunnelRowSerializer serializer =
+                new DefaultSeaTunnelRowSerializer(
+                        topic, tag, SEATUNNEL_ROW_TYPE, SchemaFormat.TEXT, DEFAULT_FIELD_DELIMITER);
+        generateTestData(serializer::serializeRow, topic, 0, 32);
+        Container.ExecResult execResult =
+                container.executeJob("/rocketmq-source_text_error_tag_to_console.conf");
+        Assertions.assertEquals(0, execResult.getExitCode(), execResult.getStderr());
+    }
+
     @TestTemplate
     public void testSourceRocketMqTextToConsole(TestContainer container)
             throws IOException, InterruptedException {
@@ -458,4 +498,32 @@ public RocketMqBaseConfiguration newConfiguration() {
     interface ProducerRecordConverter {
         Message convert(SeaTunnelRow row);
     }
+
+    private void deleteTopicIfExist(String topicName) {
+        DefaultMQAdminExt admin = new DefaultMQAdminExt();
+        admin.setInstanceName(UUID.randomUUID().toString());
+        try {
+            admin.start();
+            TopicRouteData topicRouteData = admin.examineTopicRouteInfo(topicName);
+            if (topicRouteData != null
+                    && topicRouteData.getQueueDatas() != null
+                    && !topicRouteData.getQueueDatas().isEmpty()) {
+                Set<String> brokerNames =
+                        topicRouteData.getQueueDatas().stream()
+                                .map(QueueData::getBrokerName)
+                                .collect(Collectors.toSet());
+                admin.deleteTopicInBroker(brokerNames, topicName);
+                admin.deleteTopicInNameServer(brokerNames, topicName, "delete_topic");
+                log.info("Deleted topic: {}", topicName);
+            } else {
+                log.info("Topic {} does not exist", topicName);
+            }
+        } catch (Exception e) {
+            log.warn("Failed to delete topic {}: {}", topicName, e.getMessage());
+        } finally {
+            if (admin != null) {
+                admin.shutdown();
+            }
+        }
+    }
 }

seatunnel-e2e/seatunnel-connector-v2-e2e/connector-rocketmq-e2e/src/test/resources/rocketmq-source_text_error_tag_to_console.conf

@@ -0,0 +1,94 @@
+#
+# 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.
+#
+######
+###### This config file is a demonstration of streaming processing in seatunnel config
+######
+
+env {
+  parallelism = 1
+  job.mode = "BATCH"
+
+  # You can set spark configuration here
+  spark.app.name = "SeaTunnel"
+  spark.executor.instances = 1
+  spark.executor.cores = 1
+  spark.executor.memory = "1g"
+  spark.master = local
+}
+
+source {
+  Rocketmq {
+    plugin_output = "rocketmq_table"
+    name.srv.addr = "rocketmq-e2e:9876"
+    topics = "test_topic_text_error_tag"
+    format = text
+    # The default field delimiter is ","
+    field_delimiter = ","
+    tags = "error_tag_test"
+    schema = {
+      fields {
+        id = bigint
+        c_map = "map<string, smallint>"
+        c_array = "array<tinyint>"
+        c_string = string
+        c_boolean = boolean
+        c_tinyint = tinyint
+        c_smallint = smallint
+        c_int = int
+        c_bigint = bigint
+        c_float = float
+        c_double = double
+        c_decimal = "decimal(2, 1)"
+        c_bytes = bytes
+        c_date = date
+        c_timestamp = timestamp
+      }
+    }
+  }
+}
+
+transform {
+}
+
+sink {
+  Assert {
+    plugin_input = "rocketmq_table"
+    rules = {
+      row_rules = [
+        {
+          rule_type = MIN_ROW
+          rule_value = 0
+        },
+        {
+          rule_type = MAX_ROW
+          rule_value = 0
+        }
+      ],
+      field_rules = [
+        {
+          field_name = c_string
+          field_type = string
+          field_value = [
+            {
+              rule_type = NULL
+            }
+          ]
+        }
+      ]
+    }
+  }
+}