risingwave_connector/sink/
mqtt.rs

1// Copyright 2024 RisingWave Labs
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 core::fmt::Debug;
16use std::collections::BTreeMap;
17use std::sync::Arc;
18use std::sync::atomic::AtomicBool;
19
20use anyhow::{Context as _, anyhow};
21use risingwave_common::array::{Op, RowRef, StreamChunk};
22use risingwave_common::catalog::Schema;
23use risingwave_common::id::ActorId;
24use risingwave_common::row::Row;
25use risingwave_common::types::{DataType, ScalarRefImpl};
26use rumqttc::v5::ConnectionError;
27use rumqttc::v5::mqttbytes::QoS;
28use serde::Deserialize;
29use serde_with::serde_as;
30use thiserror_ext::AsReport;
31use with_options::WithOptions;
32
33use super::SinkWriterParam;
34use super::catalog::{SinkEncode, SinkFormat, SinkFormatDesc, SinkId};
35use super::encoder::{
36    DateHandlingMode, JsonEncoder, JsonbHandlingMode, ProtoEncoder, ProtoHeader, RowEncoder, SerTo,
37    TimeHandlingMode, TimestampHandlingMode, TimestamptzHandlingMode,
38};
39use super::writer::AsyncTruncateSinkWriterExt;
40use crate::connector_common::MqttCommon;
41use crate::deserialize_bool_from_string;
42use crate::enforce_secret::EnforceSecret;
43use crate::sink::log_store::DeliveryFutureManagerAddFuture;
44use crate::sink::writer::{AsyncTruncateLogSinkerOf, AsyncTruncateSinkWriter};
45use crate::sink::{Result, SINK_TYPE_APPEND_ONLY, Sink, SinkError, SinkParam};
46
47pub const MQTT_SINK: &str = "mqtt";
48
49#[serde_as]
50#[derive(Clone, Debug, Deserialize, WithOptions)]
51pub struct MqttConfig {
52    #[serde(flatten)]
53    pub common: MqttCommon,
54
55    /// The topic name to subscribe or publish to. When subscribing, it can be a wildcard topic. e.g /topic/#
56    pub topic: Option<String>,
57
58    /// Whether the message should be retained by the broker
59    #[serde(default, deserialize_with = "deserialize_bool_from_string")]
60    pub retain: bool,
61
62    // accept "append-only"
63    pub r#type: String,
64
65    // if set, will use a field value as the topic name, if topic is also set it will be used as a fallback
66    #[serde(rename = "topic.field")]
67    pub topic_field: Option<String>,
68
69    #[serde(flatten)]
70    pub unknown_fields: std::collections::HashMap<String, String>,
71}
72
73crate::impl_sink_unknown_fields!(MqttConfig);
74
75impl EnforceSecret for MqttConfig {
76    fn enforce_one(prop: &str) -> crate::error::ConnectorResult<()> {
77        MqttCommon::enforce_one(prop)
78    }
79}
80
81pub enum RowEncoderWrapper {
82    Json(JsonEncoder),
83    Proto(ProtoEncoder),
84}
85
86impl RowEncoder for RowEncoderWrapper {
87    type Output = Vec<u8>;
88
89    fn encode_cols(
90        &self,
91        row: impl Row,
92        col_indices: impl Iterator<Item = usize>,
93    ) -> Result<Self::Output> {
94        match self {
95            RowEncoderWrapper::Json(json) => json.encode_cols(row, col_indices)?.ser_to(),
96            RowEncoderWrapper::Proto(proto) => proto.encode_cols(row, col_indices)?.ser_to(),
97        }
98    }
99
100    fn schema(&self) -> &Schema {
101        match self {
102            RowEncoderWrapper::Json(json) => json.schema(),
103            RowEncoderWrapper::Proto(proto) => proto.schema(),
104        }
105    }
106
107    fn col_indices(&self) -> Option<&[usize]> {
108        match self {
109            RowEncoderWrapper::Json(json) => json.col_indices(),
110            RowEncoderWrapper::Proto(proto) => proto.col_indices(),
111        }
112    }
113
114    fn encode(&self, row: impl Row) -> Result<Self::Output> {
115        match self {
116            RowEncoderWrapper::Json(json) => json.encode(row)?.ser_to(),
117            RowEncoderWrapper::Proto(proto) => proto.encode(row)?.ser_to(),
118        }
119    }
120}
121
122#[derive(Clone, Debug)]
123pub struct MqttSink {
124    pub config: MqttConfig,
125    schema: Schema,
126    format_desc: SinkFormatDesc,
127    is_append_only: bool,
128    name: String,
129}
130
131impl EnforceSecret for MqttSink {
132    fn enforce_secret<'a>(
133        prop_iter: impl Iterator<Item = &'a str>,
134    ) -> crate::error::ConnectorResult<()> {
135        for prop in prop_iter {
136            MqttConfig::enforce_one(prop)?;
137        }
138        Ok(())
139    }
140}
141
142// sink write
143pub struct MqttSinkWriter {
144    pub config: MqttConfig,
145    payload_writer: MqttSinkPayloadWriter,
146    #[expect(dead_code)]
147    schema: Schema,
148    encoder: RowEncoderWrapper,
149    stopped: Arc<AtomicBool>,
150}
151
152/// Basic data types for use with the mqtt interface
153impl MqttConfig {
154    pub fn from_btreemap(values: BTreeMap<String, String>) -> Result<Self> {
155        let config = serde_json::from_value::<MqttConfig>(serde_json::to_value(values).unwrap())
156            .map_err(|e| SinkError::Config(anyhow!(e)))?;
157        if config.r#type != SINK_TYPE_APPEND_ONLY {
158            Err(SinkError::Config(anyhow!(
159                "MQTT sink only supports append-only mode"
160            )))
161        } else {
162            Ok(config)
163        }
164    }
165}
166
167impl TryFrom<SinkParam> for MqttSink {
168    type Error = SinkError;
169
170    fn try_from(param: SinkParam) -> std::result::Result<Self, Self::Error> {
171        let schema = param.schema();
172        let config = MqttConfig::from_btreemap(param.properties)?;
173        Ok(Self {
174            config,
175            schema,
176            name: param.sink_name,
177            format_desc: param
178                .format_desc
179                .ok_or_else(|| SinkError::Config(anyhow!("missing FORMAT ... ENCODE ...")))?,
180            is_append_only: param.sink_type.is_append_only(),
181        })
182    }
183}
184
185impl Sink for MqttSink {
186    type LogSinker = AsyncTruncateLogSinkerOf<MqttSinkWriter>;
187
188    const SINK_NAME: &'static str = MQTT_SINK;
189
190    crate::impl_validate_sink_unknown_fields!();
191
192    async fn validate(&self) -> Result<()> {
193        if !self.is_append_only {
194            return Err(SinkError::Mqtt(anyhow!(
195                "MQTT sink only supports append-only mode"
196            )));
197        }
198
199        if let Some(field) = &self.config.topic_field {
200            let _ = get_topic_field_index_path(&self.schema, field.as_str())?;
201        } else if self.config.topic.is_none() {
202            return Err(SinkError::Config(anyhow!(
203                "either topic or topic.field must be set"
204            )));
205        }
206
207        let _client = (self.config.common.build_client(0.into(), 0))
208            .context("validate mqtt sink error")
209            .map_err(SinkError::Mqtt)?;
210
211        Ok(())
212    }
213
214    async fn new_log_sinker(&self, writer_param: SinkWriterParam) -> Result<Self::LogSinker> {
215        Ok(MqttSinkWriter::new(
216            self.config.clone(),
217            self.schema.clone(),
218            &self.format_desc,
219            &self.name,
220            writer_param.sink_id,
221            writer_param.actor_id,
222        )
223        .await?
224        .into_log_sinker(usize::MAX))
225    }
226}
227
228impl MqttSinkWriter {
229    pub async fn new(
230        config: MqttConfig,
231        schema: Schema,
232        format_desc: &SinkFormatDesc,
233        name: &str,
234        sink_id: SinkId,
235        actor_id: ActorId,
236    ) -> Result<Self> {
237        let mut topic_index_path = vec![];
238        if let Some(field) = &config.topic_field {
239            topic_index_path = get_topic_field_index_path(&schema, field.as_str())?;
240        }
241
242        let timestamptz_mode = TimestamptzHandlingMode::from_options(&format_desc.options)?;
243        let jsonb_handling_mode = JsonbHandlingMode::from_options(&format_desc.options)?;
244        let encoder = match format_desc.format {
245            SinkFormat::AppendOnly => match format_desc.encode {
246                SinkEncode::Json => RowEncoderWrapper::Json(JsonEncoder::new(
247                    schema.clone(),
248                    None,
249                    DateHandlingMode::FromCe,
250                    TimestampHandlingMode::Milli,
251                    timestamptz_mode,
252                    TimeHandlingMode::Milli,
253                    jsonb_handling_mode,
254                )),
255                SinkEncode::Protobuf => {
256                    let (descriptor, sid) = crate::schema::protobuf::fetch_descriptor(
257                        &format_desc.options,
258                        config.topic.as_deref().unwrap_or(name),
259                        None,
260                    )
261                    .await
262                    .map_err(|e| SinkError::Config(anyhow!(e)))?;
263                    let header = match sid {
264                        None => ProtoHeader::None,
265                        Some(sid) => ProtoHeader::ConfluentSchemaRegistry(sid),
266                    };
267                    RowEncoderWrapper::Proto(ProtoEncoder::new(
268                        schema.clone(),
269                        None,
270                        descriptor,
271                        header,
272                    )?)
273                }
274                _ => {
275                    return Err(SinkError::Config(anyhow!(
276                        "mqtt sink encode unsupported: {:?}",
277                        format_desc.encode,
278                    )));
279                }
280            },
281            _ => {
282                return Err(SinkError::Config(anyhow!(
283                    "MQTT sink only supports append-only mode"
284                )));
285            }
286        };
287        let qos = config.common.qos();
288
289        let (client, mut eventloop) = config
290            .common
291            .build_client(actor_id, sink_id.as_raw_id())
292            .map_err(|e| SinkError::Mqtt(anyhow!(e)))?;
293
294        let stopped = Arc::new(AtomicBool::new(false));
295        let stopped_clone = stopped.clone();
296        tokio::spawn(async move {
297            while !stopped_clone.load(std::sync::atomic::Ordering::Relaxed) {
298                match eventloop.poll().await {
299                    Ok(_) => (),
300                    Err(err) => match err {
301                        ConnectionError::Timeout(_) => (),
302                        ConnectionError::MqttState(rumqttc::v5::StateError::Io(err))
303                        | ConnectionError::Io(err)
304                            if err.kind() == std::io::ErrorKind::ConnectionAborted
305                                || err.kind() == std::io::ErrorKind::ConnectionReset =>
306                        {
307                            continue;
308                        }
309                        err => {
310                            tracing::error!(
311                                "failed to poll the MQTT event loop: {}",
312                                err.as_report()
313                            );
314                            tokio::time::sleep(std::time::Duration::from_secs(1)).await;
315                        }
316                    },
317                }
318            }
319        });
320
321        let payload_writer = MqttSinkPayloadWriter {
322            topic: config.topic.clone(),
323            client,
324            qos,
325            retain: config.retain,
326            topic_index_path,
327        };
328
329        Ok::<_, SinkError>(Self {
330            config: config.clone(),
331            payload_writer,
332            schema: schema.clone(),
333            stopped,
334            encoder,
335        })
336    }
337}
338
339impl AsyncTruncateSinkWriter for MqttSinkWriter {
340    async fn write_chunk<'a>(
341        &'a mut self,
342        chunk: StreamChunk,
343        _add_future: DeliveryFutureManagerAddFuture<'a, Self::DeliveryFuture>,
344    ) -> Result<()> {
345        self.payload_writer.write_chunk(chunk, &self.encoder).await
346    }
347}
348
349impl Drop for MqttSinkWriter {
350    fn drop(&mut self) {
351        self.stopped
352            .store(true, std::sync::atomic::Ordering::Relaxed);
353    }
354}
355
356struct MqttSinkPayloadWriter {
357    // connection to mqtt, one per executor
358    client: rumqttc::v5::AsyncClient,
359    topic: Option<String>,
360    qos: QoS,
361    retain: bool,
362    topic_index_path: Vec<usize>,
363}
364
365impl MqttSinkPayloadWriter {
366    async fn write_chunk(&mut self, chunk: StreamChunk, encoder: &RowEncoderWrapper) -> Result<()> {
367        for (op, row) in chunk.rows() {
368            if op != Op::Insert {
369                continue;
370            }
371
372            let topic = match get_topic_from_index_path(
373                &self.topic_index_path,
374                self.topic.as_deref(),
375                &row,
376            ) {
377                Some(s) => s,
378                None => {
379                    tracing::error!("topic field not found in row; skipping row: {:?}", row);
380                    return Ok(());
381                }
382            };
383
384            let v = encoder.encode(row)?;
385
386            self.client
387                .publish(topic, self.qos, self.retain, v)
388                .await
389                .context("mqtt sink error")
390                .map_err(SinkError::Mqtt)?;
391        }
392
393        Ok(())
394    }
395}
396
397fn get_topic_from_index_path<'s>(
398    path: &[usize],
399    default_topic: Option<&'s str>,
400    row: &'s RowRef<'s>,
401) -> Option<&'s str> {
402    if let Some(topic) = default_topic
403        && path.is_empty()
404    {
405        Some(topic)
406    } else {
407        let mut iter = path.iter();
408        let scalar = iter
409            .next()
410            .and_then(|pos| row.datum_at(*pos))
411            .and_then(|d| {
412                iter.try_fold(d, |d, pos| match d {
413                    ScalarRefImpl::Struct(struct_ref) => {
414                        struct_ref.iter_fields_ref().nth(*pos).flatten()
415                    }
416                    _ => None,
417                })
418            });
419        match scalar {
420            Some(ScalarRefImpl::Utf8(s)) => Some(s),
421            _ => {
422                if let Some(topic) = default_topic {
423                    Some(topic)
424                } else {
425                    None
426                }
427            }
428        }
429    }
430}
431
432// This function returns the index path to the topic field in the schema, validating that the field exists and is of type string
433// the returnent path can be used to extract the topic field from a row. The path is a list of indexes to be used to navigate the row
434// to the topic field.
435fn get_topic_field_index_path(schema: &Schema, topic_field: &str) -> Result<Vec<usize>> {
436    let mut iter = topic_field.split('.');
437    let mut path = vec![];
438    let dt =
439        iter.next()
440            .and_then(|field| {
441                // Extract the field from the schema
442                schema
443                    .fields()
444                    .iter()
445                    .enumerate()
446                    .find(|(_, f)| f.name == field)
447                    .map(|(pos, f)| {
448                        path.push(pos);
449                        &f.data_type
450                    })
451            })
452            .and_then(|dt| {
453                // Iterate over the next fields to extract the fields from the nested structs
454                iter.try_fold(dt, |dt, field| match dt {
455                    DataType::Struct(st) => {
456                        st.iter().enumerate().find(|(_, (s, _))| *s == field).map(
457                            |(pos, (_, dt))| {
458                                path.push(pos);
459                                dt
460                            },
461                        )
462                    }
463                    _ => None,
464                })
465            });
466
467    match dt {
468        Some(DataType::Varchar) => Ok(path),
469        Some(dt) => Err(SinkError::Config(anyhow!(
470            "topic field `{}` must be of type string but got {:?}",
471            topic_field,
472            dt
473        ))),
474        None => Err(SinkError::Config(anyhow!(
475            "topic field `{}`  not found",
476            topic_field
477        ))),
478    }
479}
480
481#[cfg(test)]
482mod test {
483    use risingwave_common::array::{DataChunk, DataChunkTestExt, RowRef};
484    use risingwave_common::catalog::{Field, Schema};
485    use risingwave_common::types::{DataType, StructType};
486
487    use super::{get_topic_field_index_path, get_topic_from_index_path};
488
489    #[test]
490    fn test_single_field_extraction() {
491        let schema = Schema::new(vec![Field::with_name(DataType::Varchar, "topic")]);
492        let path = get_topic_field_index_path(&schema, "topic").unwrap();
493        assert_eq!(path, vec![0]);
494
495        let chunk = DataChunk::from_pretty(
496            "T
497            test",
498        );
499
500        let row = RowRef::new(&chunk, 0);
501
502        assert_eq!(get_topic_from_index_path(&path, None, &row), Some("test"));
503
504        let result = get_topic_field_index_path(&schema, "other_field");
505        assert!(result.is_err());
506    }
507
508    #[test]
509    fn test_nested_field_extraction() {
510        let schema = Schema::new(vec![Field::with_name(
511            DataType::Struct(StructType::new(vec![
512                ("field", DataType::Int32),
513                ("subtopic", DataType::Varchar),
514            ])),
515            "topic",
516        )]);
517        let path = get_topic_field_index_path(&schema, "topic.subtopic").unwrap();
518        assert_eq!(path, vec![0, 1]);
519
520        let chunk = DataChunk::from_pretty(
521            "<i,T>
522            (1,test)",
523        );
524
525        let row = RowRef::new(&chunk, 0);
526
527        assert_eq!(get_topic_from_index_path(&path, None, &row), Some("test"));
528
529        let result = get_topic_field_index_path(&schema, "topic.other_field");
530        assert!(result.is_err());
531    }
532
533    #[test]
534    fn test_null_values_extraction() {
535        let path = vec![0];
536        let chunk = DataChunk::from_pretty(
537            "T
538            .",
539        );
540        let row = RowRef::new(&chunk, 0);
541        assert_eq!(
542            get_topic_from_index_path(&path, Some("default"), &row),
543            Some("default")
544        );
545        assert_eq!(get_topic_from_index_path(&path, None, &row), None);
546
547        let path = vec![0, 1];
548        let chunk = DataChunk::from_pretty(
549            "<i,T>
550            (1,)",
551        );
552        let row = RowRef::new(&chunk, 0);
553        assert_eq!(
554            get_topic_from_index_path(&path, Some("default"), &row),
555            Some("default")
556        );
557        assert_eq!(get_topic_from_index_path(&path, None, &row), None);
558    }
559
560    #[test]
561    fn test_multiple_levels() {
562        let schema = Schema::new(vec![
563            Field::with_name(
564                DataType::Struct(StructType::new(vec![
565                    ("field", DataType::Int32),
566                    (
567                        "subtopic",
568                        DataType::Struct(StructType::new(vec![
569                            ("int_field", DataType::Int32),
570                            ("boolean_field", DataType::Boolean),
571                            ("string_field", DataType::Varchar),
572                        ])),
573                    ),
574                ])),
575                "topic",
576            ),
577            Field::with_name(DataType::Varchar, "other_field"),
578        ]);
579
580        let path = get_topic_field_index_path(&schema, "topic.subtopic.string_field").unwrap();
581        assert_eq!(path, vec![0, 1, 2]);
582
583        assert!(get_topic_field_index_path(&schema, "topic.subtopic.boolean_field").is_err());
584
585        assert!(get_topic_field_index_path(&schema, "topic.subtopic.int_field").is_err());
586
587        assert!(get_topic_field_index_path(&schema, "topic.field").is_err());
588
589        let path = get_topic_field_index_path(&schema, "other_field").unwrap();
590        assert_eq!(path, vec![1]);
591
592        let chunk = DataChunk::from_pretty(
593            "<i,<T>> T
594            (1,(test)) other",
595        );
596
597        let row = RowRef::new(&chunk, 0);
598
599        // topic.subtopic.string_field
600        assert_eq!(
601            get_topic_from_index_path(&[0, 1, 0], None, &row),
602            Some("test")
603        );
604
605        // topic.field
606        assert_eq!(get_topic_from_index_path(&[0, 0], None, &row), None);
607
608        // other_field
609        assert_eq!(get_topic_from_index_path(&[1], None, &row), Some("other"));
610    }
611}
risingwave_connector/sink/mqtt.rs

risingwave_connector/sink/
mqtt.rs
