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

risingwave_connector/sink/
nats.rs
