rocketmq_model/common/message/
message_builder.rs

1// Copyright 2023 The RocketMQ Rust Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use bytes::Bytes;
16use cheetah_string::CheetahString;
17use rocketmq_error::Result;
18
19use crate::common::message::message_body::MessageBody;
20use crate::common::message::message_flag::MessageFlag;
21use crate::common::message::message_property::MessageProperties;
22use crate::common::message::message_property::MessagePropertyKey;
23use crate::common::message::MessageConst;
24
25/// Builder for constructing messages.
26///
27/// # Examples
28///
29/// ```
30/// use rocketmq_model::common::message::message_builder::MessageBuilder;
31///
32/// let msg = MessageBuilder::new()
33///     .topic("test-topic")
34///     .body_slice(b"hello world")
35///     .tags("important")
36///     .keys(vec!["key1".to_string(), "key2".to_string()])
37///     .build();
38/// ```
39#[derive(Default)]
40pub struct MessageBuilder {
41    topic: Option<CheetahString>,
42    body: MessageBody,
43    properties: MessageProperties,
44    flag: MessageFlag,
45    transaction_id: Option<CheetahString>,
46}
47
48impl MessageBuilder {
49    /// Creates a new message builder.
50    #[inline]
51    pub fn new() -> Self {
52        Self::default()
53    }
54
55    /// Sets the topic name (required).
56    #[inline]
57    pub fn topic(mut self, topic: impl Into<CheetahString>) -> Self {
58        self.topic = Some(topic.into());
59        self
60    }
61
62    /// Sets the message body from bytes.
63    #[inline]
64    pub fn body(mut self, body: impl Into<Bytes>) -> Self {
65        self.body = MessageBody::new(body);
66        self
67    }
68
69    /// Sets the message body from a byte slice.
70    #[inline]
71    pub fn body_slice(mut self, body: &[u8]) -> Self {
72        self.body = MessageBody::new(Bytes::copy_from_slice(body));
73        self
74    }
75
76    /// Sets an empty body.
77    #[inline]
78    pub fn empty_body(mut self) -> Self {
79        self.body = MessageBody::empty();
80        self
81    }
82
83    /// Sets the message tags.
84    #[inline]
85    pub fn tags(mut self, tags: impl Into<CheetahString>) -> Self {
86        self.properties.insert(MessagePropertyKey::Tags, tags.into());
87        self
88    }
89
90    /// Sets the message keys from a vector.
91    #[inline]
92    pub fn keys(mut self, keys: Vec<String>) -> Self {
93        if !keys.is_empty() {
94            let keys_str = keys.join(MessageConst::KEY_SEPARATOR);
95            self.properties.insert(MessagePropertyKey::Keys, keys_str);
96        }
97        self
98    }
99
100    /// Sets a single message key.
101    #[inline]
102    pub fn key(mut self, key: impl Into<CheetahString>) -> Self {
103        self.properties.insert(MessagePropertyKey::Keys, key.into());
104        self
105    }
106
107    /// Sets the message flag.
108    #[inline]
109    pub fn flag(mut self, flag: MessageFlag) -> Self {
110        self.flag = flag;
111        self
112    }
113
114    /// Sets the message flag from raw bits.
115    #[inline]
116    pub fn flag_bits(mut self, bits: i32) -> Self {
117        self.flag = MessageFlag::from_bits(bits);
118        self
119    }
120
121    /// Sets the delay time level (1-18 for predefined levels).
122    #[inline]
123    pub fn delay_level(mut self, level: i32) -> Self {
124        self.properties
125            .insert(MessagePropertyKey::DelayTimeLevel, level.to_string());
126        self
127    }
128
129    /// Sets the delay time in seconds.
130    #[inline]
131    pub fn delay_secs(mut self, secs: u64) -> Self {
132        self.properties
133            .insert(MessagePropertyKey::DelayTimeSec, secs.to_string());
134        self
135    }
136
137    /// Sets the delay time in milliseconds.
138    #[inline]
139    pub fn delay_millis(mut self, millis: u64) -> Self {
140        self.properties
141            .insert(MessagePropertyKey::DelayTimeMs, millis.to_string());
142        self
143    }
144
145    /// Sets the delivery time in milliseconds.
146    #[inline]
147    pub fn deliver_time_ms(mut self, time_ms: u64) -> Self {
148        self.properties
149            .insert(MessagePropertyKey::DeliverTimeMs, time_ms.to_string());
150        self
151    }
152
153    /// Sets the message priority.
154    #[inline]
155    pub fn priority(mut self, priority: u32) -> Self {
156        self.properties
157            .insert(MessagePropertyKey::Priority, priority.to_string());
158        self
159    }
160
161    /// Tries to set the message priority with Java-compatible non-negative validation.
162    pub fn try_priority(mut self, priority: i32) -> Result<Self> {
163        if priority < 0 {
164            return Err(crate::error::invalid_argument(
165                "The priority must be greater than or equal to 0",
166            ));
167        }
168        self.properties
169            .insert(MessagePropertyKey::Priority, priority.to_string());
170        Ok(self)
171    }
172
173    /// Sets whether to wait for store confirmation.
174    #[inline]
175    pub fn wait_store_msg_ok(mut self, wait: bool) -> Self {
176        self.properties
177            .insert(MessagePropertyKey::WaitStoreMsgOk, if wait { "true" } else { "false" });
178        self
179    }
180
181    /// Sets the transaction ID.
182    #[inline]
183    pub fn transaction_id(mut self, id: impl Into<CheetahString>) -> Self {
184        self.transaction_id = Some(id.into());
185        self
186    }
187
188    /// Sets the buyer ID.
189    #[inline]
190    pub fn buyer_id(mut self, buyer_id: impl Into<CheetahString>) -> Self {
191        self.properties.insert(MessagePropertyKey::BuyerId, buyer_id.into());
192        self
193    }
194
195    /// Sets the instance ID.
196    #[inline]
197    pub fn instance_id(mut self, instance_id: impl Into<CheetahString>) -> Self {
198        self.properties
199            .insert(MessagePropertyKey::InstanceId, instance_id.into());
200        self
201    }
202
203    /// Sets the correlation ID.
204    #[inline]
205    pub fn correlation_id(mut self, correlation_id: impl Into<CheetahString>) -> Self {
206        self.properties
207            .insert(MessagePropertyKey::CorrelationId, correlation_id.into());
208        self
209    }
210
211    /// Sets the sharding key.
212    #[inline]
213    pub fn sharding_key(mut self, sharding_key: impl Into<CheetahString>) -> Self {
214        self.properties
215            .insert(MessagePropertyKey::ShardingKey, sharding_key.into());
216        self
217    }
218
219    /// Sets the trace switch.
220    #[inline]
221    pub fn trace_switch(mut self, enabled: bool) -> Self {
222        self.properties
223            .insert(MessagePropertyKey::TraceSwitch, if enabled { "true" } else { "false" });
224        self
225    }
226
227    /// Sets a custom property directly using MessagePropertyKey.
228    #[inline]
229    pub fn property(mut self, key: MessagePropertyKey, value: impl Into<CheetahString>) -> Self {
230        self.properties.insert(key, value.into());
231        self
232    }
233
234    /// Sets a raw property with string key (for advanced use).
235    ///
236    /// # Errors
237    ///
238    /// Returns an error if the property name is reserved by the system.
239    pub fn raw_property(mut self, key: impl Into<CheetahString>, value: impl Into<CheetahString>) -> Result<Self> {
240        let key = key.into();
241        let value = value.into();
242
243        // Validate not a reserved key
244        if crate::common::message::STRING_HASH_SET.contains(key.as_str()) {
245            return Err(crate::error::invalid_property(format!(
246                "The Property<{key}> is used by system, input another please"
247            )));
248        }
249
250        self.properties.as_map_mut().insert(key, value);
251        Ok(self)
252    }
253
254    /// Builds the message.
255    ///
256    /// # Errors
257    ///
258    /// Returns an error if required fields (topic) are not set.
259    pub fn build(mut self) -> Result<super::message_single::Message> {
260        let topic = self
261            .topic
262            .ok_or_else(|| crate::error::invalid_property("Topic is required for message"))?;
263        if self.properties.get(MessagePropertyKey::WaitStoreMsgOk).is_none() {
264            self.properties.insert(MessagePropertyKey::WaitStoreMsgOk, "true");
265        }
266
267        Ok(super::message_single::Message::from_builder(
268            topic,
269            self.body,
270            self.properties,
271            self.flag,
272            self.transaction_id,
273        ))
274    }
275
276    /// Builds the message, panicking if validation fails.
277    ///
278    /// This is convenient for test code where you know the message is valid.
279    #[track_caller]
280    pub fn build_unchecked(self) -> super::message_single::Message {
281        self.build().expect("message validation failed")
282    }
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288
289    #[test]
290    fn test_builder_basic() {
291        let msg = MessageBuilder::new()
292            .topic("test-topic")
293            .body_slice(b"hello world")
294            .build_unchecked();
295
296        assert_eq!(msg.topic().as_str(), "test-topic");
297        assert_eq!(msg.body_slice(), b"hello world");
298        assert_eq!(msg.properties().get(MessagePropertyKey::WaitStoreMsgOk), Some("true"));
299    }
300
301    #[test]
302    fn test_builder_with_tags_and_keys() {
303        let msg = MessageBuilder::new()
304            .topic("test-topic")
305            .body_slice(b"test")
306            .tags("tag1")
307            .keys(vec!["key1".to_string(), "key2".to_string()])
308            .build_unchecked();
309
310        assert_eq!(msg.tags(), Some("tag1"));
311        assert_eq!(msg.keys(), Some(vec!["key1".to_string(), "key2".to_string()]));
312    }
313
314    #[test]
315    fn test_builder_with_delay() {
316        let msg = MessageBuilder::new()
317            .topic("test-topic")
318            .body_slice(b"test")
319            .delay_level(3)
320            .build_unchecked();
321
322        assert_eq!(msg.delay_time_level(), 3);
323    }
324
325    #[test]
326    fn test_builder_with_priority_matches_java_property() {
327        let msg = MessageBuilder::new()
328            .topic("test-topic")
329            .body_slice(b"test")
330            .priority(7)
331            .build_unchecked();
332
333        assert_eq!(msg.priority(), 7);
334        assert_eq!(msg.properties().get(MessagePropertyKey::Priority), Some("7"));
335    }
336
337    #[test]
338    fn test_builder_try_priority_rejects_negative_like_java() {
339        match MessageBuilder::new().topic("test-topic").try_priority(-1) {
340            Ok(_) => panic!("negative priority should be rejected"),
341            Err(error) => assert_eq!(error.descriptor(), &rocketmq_error::CORE_ARGUMENT_INVALID),
342        }
343    }
344
345    #[test]
346    fn test_builder_missing_topic() {
347        let result = MessageBuilder::new().body_slice(b"test").build();
348
349        assert!(result.is_err());
350    }
351
352    #[test]
353    fn test_builder_with_transaction_id() {
354        let msg = MessageBuilder::new()
355            .topic("test-topic")
356            .body_slice(b"test")
357            .transaction_id("tx-123")
358            .build_unchecked();
359
360        assert_eq!(msg.transaction_id(), Some("tx-123"));
361    }
362
363    #[test]
364    fn test_builder_wait_store_msg_ok_writes_java_boolean_property() {
365        let wait_msg = MessageBuilder::new()
366            .topic("test-topic")
367            .body_slice(b"test")
368            .wait_store_msg_ok(true)
369            .build_unchecked();
370        assert!(wait_msg.wait_store_msg_ok());
371        assert_eq!(
372            wait_msg.properties().get(MessagePropertyKey::WaitStoreMsgOk),
373            Some("true")
374        );
375
376        let no_wait_msg = MessageBuilder::new()
377            .topic("test-topic")
378            .body_slice(b"test")
379            .wait_store_msg_ok(false)
380            .build_unchecked();
381        assert!(!no_wait_msg.wait_store_msg_ok());
382        assert_eq!(
383            no_wait_msg.properties().get(MessagePropertyKey::WaitStoreMsgOk),
384            Some("false")
385        );
386    }
387}
rocketmq_model/common/message/message_builder.rs

rocketmq_model/common/message/
message_builder.rs
