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!("Failed to poll mqtt eventloop: {}", err.as_report());
311                            tokio::time::sleep(std::time::Duration::from_secs(1)).await;
312                        }
313                    },
314                }
315            }
316        });
317
318        let payload_writer = MqttSinkPayloadWriter {
319            topic: config.topic.clone(),
320            client,
321            qos,
322            retain: config.retain,
323            topic_index_path,
324        };
325
326        Ok::<_, SinkError>(Self {
327            config: config.clone(),
328            payload_writer,
329            schema: schema.clone(),
330            stopped,
331            encoder,
332        })
333    }
334}
335
336impl AsyncTruncateSinkWriter for MqttSinkWriter {
337    async fn write_chunk<'a>(
338        &'a mut self,
339        chunk: StreamChunk,
340        _add_future: DeliveryFutureManagerAddFuture<'a, Self::DeliveryFuture>,
341    ) -> Result<()> {
342        self.payload_writer.write_chunk(chunk, &self.encoder).await
343    }
344}
345
346impl Drop for MqttSinkWriter {
347    fn drop(&mut self) {
348        self.stopped
349            .store(true, std::sync::atomic::Ordering::Relaxed);
350    }
351}
352
353struct MqttSinkPayloadWriter {
354    // connection to mqtt, one per executor
355    client: rumqttc::v5::AsyncClient,
356    topic: Option<String>,
357    qos: QoS,
358    retain: bool,
359    topic_index_path: Vec<usize>,
360}
361
362impl MqttSinkPayloadWriter {
363    async fn write_chunk(&mut self, chunk: StreamChunk, encoder: &RowEncoderWrapper) -> Result<()> {
364        for (op, row) in chunk.rows() {
365            if op != Op::Insert {
366                continue;
367            }
368
369            let topic = match get_topic_from_index_path(
370                &self.topic_index_path,
371                self.topic.as_deref(),
372                &row,
373            ) {
374                Some(s) => s,
375                None => {
376                    tracing::error!("topic field not found in row, skipping: {:?}", row);
377                    return Ok(());
378                }
379            };
380
381            let v = encoder.encode(row)?;
382
383            self.client
384                .publish(topic, self.qos, self.retain, v)
385                .await
386                .context("mqtt sink error")
387                .map_err(SinkError::Mqtt)?;
388        }
389
390        Ok(())
391    }
392}
393
394fn get_topic_from_index_path<'s>(
395    path: &[usize],
396    default_topic: Option<&'s str>,
397    row: &'s RowRef<'s>,
398) -> Option<&'s str> {
399    if let Some(topic) = default_topic
400        && path.is_empty()
401    {
402        Some(topic)
403    } else {
404        let mut iter = path.iter();
405        let scalar = iter
406            .next()
407            .and_then(|pos| row.datum_at(*pos))
408            .and_then(|d| {
409                iter.try_fold(d, |d, pos| match d {
410                    ScalarRefImpl::Struct(struct_ref) => {
411                        struct_ref.iter_fields_ref().nth(*pos).flatten()
412                    }
413                    _ => None,
414                })
415            });
416        match scalar {
417            Some(ScalarRefImpl::Utf8(s)) => Some(s),
418            _ => {
419                if let Some(topic) = default_topic {
420                    Some(topic)
421                } else {
422                    None
423                }
424            }
425        }
426    }
427}
428
429// This function returns the index path to the topic field in the schema, validating that the field exists and is of type string
430// 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
431// to the topic field.
432fn get_topic_field_index_path(schema: &Schema, topic_field: &str) -> Result<Vec<usize>> {
433    let mut iter = topic_field.split('.');
434    let mut path = vec![];
435    let dt =
436        iter.next()
437            .and_then(|field| {
438                // Extract the field from the schema
439                schema
440                    .fields()
441                    .iter()
442                    .enumerate()
443                    .find(|(_, f)| f.name == field)
444                    .map(|(pos, f)| {
445                        path.push(pos);
446                        &f.data_type
447                    })
448            })
449            .and_then(|dt| {
450                // Iterate over the next fields to extract the fields from the nested structs
451                iter.try_fold(dt, |dt, field| match dt {
452                    DataType::Struct(st) => {
453                        st.iter().enumerate().find(|(_, (s, _))| *s == field).map(
454                            |(pos, (_, dt))| {
455                                path.push(pos);
456                                dt
457                            },
458                        )
459                    }
460                    _ => None,
461                })
462            });
463
464    match dt {
465        Some(DataType::Varchar) => Ok(path),
466        Some(dt) => Err(SinkError::Config(anyhow!(
467            "topic field `{}` must be of type string but got {:?}",
468            topic_field,
469            dt
470        ))),
471        None => Err(SinkError::Config(anyhow!(
472            "topic field `{}`  not found",
473            topic_field
474        ))),
475    }
476}
477
478#[cfg(test)]
479mod test {
480    use risingwave_common::array::{DataChunk, DataChunkTestExt, RowRef};
481    use risingwave_common::catalog::{Field, Schema};
482    use risingwave_common::types::{DataType, StructType};
483
484    use super::{get_topic_field_index_path, get_topic_from_index_path};
485
486    #[test]
487    fn test_single_field_extraction() {
488        let schema = Schema::new(vec![Field::with_name(DataType::Varchar, "topic")]);
489        let path = get_topic_field_index_path(&schema, "topic").unwrap();
490        assert_eq!(path, vec![0]);
491
492        let chunk = DataChunk::from_pretty(
493            "T
494            test",
495        );
496
497        let row = RowRef::new(&chunk, 0);
498
499        assert_eq!(get_topic_from_index_path(&path, None, &row), Some("test"));
500
501        let result = get_topic_field_index_path(&schema, "other_field");
502        assert!(result.is_err());
503    }
504
505    #[test]
506    fn test_nested_field_extraction() {
507        let schema = Schema::new(vec![Field::with_name(
508            DataType::Struct(StructType::new(vec![
509                ("field", DataType::Int32),
510                ("subtopic", DataType::Varchar),
511            ])),
512            "topic",
513        )]);
514        let path = get_topic_field_index_path(&schema, "topic.subtopic").unwrap();
515        assert_eq!(path, vec![0, 1]);
516
517        let chunk = DataChunk::from_pretty(
518            "<i,T>
519            (1,test)",
520        );
521
522        let row = RowRef::new(&chunk, 0);
523
524        assert_eq!(get_topic_from_index_path(&path, None, &row), Some("test"));
525
526        let result = get_topic_field_index_path(&schema, "topic.other_field");
527        assert!(result.is_err());
528    }
529
530    #[test]
531    fn test_null_values_extraction() {
532        let path = vec![0];
533        let chunk = DataChunk::from_pretty(
534            "T
535            .",
536        );
537        let row = RowRef::new(&chunk, 0);
538        assert_eq!(
539            get_topic_from_index_path(&path, Some("default"), &row),
540            Some("default")
541        );
542        assert_eq!(get_topic_from_index_path(&path, None, &row), None);
543
544        let path = vec![0, 1];
545        let chunk = DataChunk::from_pretty(
546            "<i,T>
547            (1,)",
548        );
549        let row = RowRef::new(&chunk, 0);
550        assert_eq!(
551            get_topic_from_index_path(&path, Some("default"), &row),
552            Some("default")
553        );
554        assert_eq!(get_topic_from_index_path(&path, None, &row), None);
555    }
556
557    #[test]
558    fn test_multiple_levels() {
559        let schema = Schema::new(vec![
560            Field::with_name(
561                DataType::Struct(StructType::new(vec![
562                    ("field", DataType::Int32),
563                    (
564                        "subtopic",
565                        DataType::Struct(StructType::new(vec![
566                            ("int_field", DataType::Int32),
567                            ("boolean_field", DataType::Boolean),
568                            ("string_field", DataType::Varchar),
569                        ])),
570                    ),
571                ])),
572                "topic",
573            ),
574            Field::with_name(DataType::Varchar, "other_field"),
575        ]);
576
577        let path = get_topic_field_index_path(&schema, "topic.subtopic.string_field").unwrap();
578        assert_eq!(path, vec![0, 1, 2]);
579
580        assert!(get_topic_field_index_path(&schema, "topic.subtopic.boolean_field").is_err());
581
582        assert!(get_topic_field_index_path(&schema, "topic.subtopic.int_field").is_err());
583
584        assert!(get_topic_field_index_path(&schema, "topic.field").is_err());
585
586        let path = get_topic_field_index_path(&schema, "other_field").unwrap();
587        assert_eq!(path, vec![1]);
588
589        let chunk = DataChunk::from_pretty(
590            "<i,<T>> T
591            (1,(test)) other",
592        );
593
594        let row = RowRef::new(&chunk, 0);
595
596        // topic.subtopic.string_field
597        assert_eq!(
598            get_topic_from_index_path(&[0, 1, 0], None, &row),
599            Some("test")
600        );
601
602        // topic.field
603        assert_eq!(get_topic_from_index_path(&[0, 0], None, &row), None);
604
605        // other_field
606        assert_eq!(get_topic_from_index_path(&[1], None, &row), Some("other"));
607    }
608}
risingwave_connector/sink/mqtt.rs

risingwave_connector/sink/
mqtt.rs
