risingwave_connector/sink/
nats.rs

1// Copyright 2023 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 core::future::IntoFuture;
17use std::collections::BTreeMap;
18use std::sync::Arc;
19
20use anyhow::{Context as _, anyhow};
21use async_nats::jetstream::context::Context;
22use futures::FutureExt;
23use futures::prelude::TryFuture;
24use risingwave_common::array::StreamChunk;
25use risingwave_common::catalog::Schema;
26use serde::Deserialize;
27use serde_with::serde_as;
28use tokio_retry::Retry;
29use tokio_retry::strategy::{ExponentialBackoff, jitter};
30use with_options::WithOptions;
31
32use super::SinkWriterParam;
33use super::encoder::{
34    DateHandlingMode, JsonbHandlingMode, TimeHandlingMode, TimestamptzHandlingMode,
35};
36use super::utils::chunk_to_json;
37use crate::connector_common::NatsCommon;
38use crate::enforce_secret::EnforceSecret;
39use crate::sink::encoder::{JsonEncoder, TimestampHandlingMode};
40use crate::sink::log_store::DeliveryFutureManagerAddFuture;
41use crate::sink::writer::{
42    AsyncTruncateLogSinkerOf, AsyncTruncateSinkWriter, AsyncTruncateSinkWriterExt,
43};
44use crate::sink::{Result, SINK_TYPE_APPEND_ONLY, Sink, SinkError, SinkParam};
45
46pub const NATS_SINK: &str = "nats";
47const NATS_SEND_FUTURE_BUFFER_MAX_SIZE: usize = 65536;
48
49#[serde_as]
50#[derive(Clone, Debug, Deserialize, WithOptions)]
51pub struct NatsConfig {
52    #[serde(flatten)]
53    pub common: NatsCommon,
54    // accept "append-only"
55    pub r#type: String,
56
57    #[serde(flatten)]
58    pub unknown_fields: std::collections::HashMap<String, String>,
59}
60
61crate::impl_sink_unknown_fields!(NatsConfig);
62
63#[derive(Clone, Debug)]
64pub struct NatsSink {
65    pub config: NatsConfig,
66    schema: Schema,
67    is_append_only: bool,
68}
69
70impl EnforceSecret for NatsSink {
71    fn enforce_secret<'a>(
72        prop_iter: impl Iterator<Item = &'a str>,
73    ) -> crate::error::ConnectorResult<()> {
74        for prop in prop_iter {
75            NatsCommon::enforce_one(prop)?;
76        }
77        Ok(())
78    }
79}
80
81// sink write
82pub struct NatsSinkWriter {
83    pub config: NatsConfig,
84    context: Context,
85    /// Hold the client Arc to keep it alive. This allows the shared client cache to reuse
86    /// the connection while we're still using it.
87    #[expect(dead_code)]
88    client: Arc<async_nats::Client>,
89    #[expect(dead_code)]
90    schema: Schema,
91    json_encoder: JsonEncoder,
92}
93
94pub type NatsSinkDeliveryFuture = impl TryFuture<Ok = (), Error = SinkError> + Unpin + 'static;
95
96/// Basic data types for use with the nats interface
97impl NatsConfig {
98    pub fn from_btreemap(values: BTreeMap<String, String>) -> Result<Self> {
99        let config = serde_json::from_value::<NatsConfig>(serde_json::to_value(values).unwrap())
100            .map_err(|e| SinkError::Config(anyhow!(e)))?;
101        if config.r#type != SINK_TYPE_APPEND_ONLY {
102            Err(SinkError::Config(anyhow!(
103                "NATS sink only supports append-only mode"
104            )))
105        } else {
106            Ok(config)
107        }
108    }
109}
110
111impl TryFrom<SinkParam> for NatsSink {
112    type Error = SinkError;
113
114    fn try_from(param: SinkParam) -> std::result::Result<Self, Self::Error> {
115        let schema = param.schema();
116        let config = NatsConfig::from_btreemap(param.properties)?;
117        Ok(Self {
118            config,
119            schema,
120            is_append_only: param.sink_type.is_append_only(),
121        })
122    }
123}
124
125impl Sink for NatsSink {
126    type LogSinker = AsyncTruncateLogSinkerOf<NatsSinkWriter>;
127
128    const SINK_NAME: &'static str = NATS_SINK;
129
130    crate::impl_validate_sink_unknown_fields!();
131
132    async fn validate(&self) -> Result<()> {
133        if !self.is_append_only {
134            return Err(SinkError::Nats(anyhow!(
135                "NATS sink only supports append-only mode"
136            )));
137        }
138        let _client = (self.config.common.build_client().await)
139            .context("validate nats sink error")
140            .map_err(SinkError::Nats)?;
141        Ok(())
142    }
143
144    async fn new_log_sinker(&self, _writer_param: SinkWriterParam) -> Result<Self::LogSinker> {
145        Ok(
146            NatsSinkWriter::new(self.config.clone(), self.schema.clone())
147                .await?
148                .into_log_sinker(NATS_SEND_FUTURE_BUFFER_MAX_SIZE),
149        )
150    }
151}
152
153impl NatsSinkWriter {
154    pub async fn new(config: NatsConfig, schema: Schema) -> Result<Self> {
155        let client = config
156            .common
157            .build_client()
158            .await
159            .map_err(|e| SinkError::Nats(anyhow!(e)))?;
160        let context = NatsCommon::build_context_from_client(&client);
161        Ok::<_, SinkError>(Self {
162            config: config.clone(),
163            context,
164            client,
165            schema: schema.clone(),
166            json_encoder: JsonEncoder::new(
167                schema,
168                None,
169                DateHandlingMode::FromCe,
170                TimestampHandlingMode::Milli,
171                TimestamptzHandlingMode::UtcWithoutSuffix,
172                TimeHandlingMode::Milli,
173                JsonbHandlingMode::String,
174            ),
175        })
176    }
177}
178
179impl AsyncTruncateSinkWriter for NatsSinkWriter {
180    type DeliveryFuture = NatsSinkDeliveryFuture;
181
182    #[define_opaque(NatsSinkDeliveryFuture)]
183    async fn write_chunk<'a>(
184        &'a mut self,
185        chunk: StreamChunk,
186        mut add_future: DeliveryFutureManagerAddFuture<'a, Self::DeliveryFuture>,
187    ) -> Result<()> {
188        let mut data = chunk_to_json(chunk, &self.json_encoder)?;
189        for item in &mut data {
190            let publish_ack_future = Retry::spawn(
191                ExponentialBackoff::from_millis(100).map(jitter).take(3),
192                || async {
193                    self.context
194                        .publish(self.config.common.subject.clone(), item.clone().into())
195                        .await
196                        .context("nats sink error")
197                        .map_err(SinkError::Nats)
198                },
199            )
200            .await
201            .context("nats sink error")
202            .map_err(SinkError::Nats)?;
203            let future = publish_ack_future.into_future().map(|result| {
204                result
205                    .context("Nats sink error")
206                    .map_err(SinkError::Nats)
207                    .map(|_| ())
208            });
209            add_future.add_future_may_await(future).await?;
210        }
211        Ok(())
212    }
213}
risingwave_connector/sink/nats.rs

risingwave_connector/sink/
nats.rs
