risingwave_connector/sink/
google_pubsub.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 std::collections::BTreeMap;
16
17use anyhow::anyhow;
18use google_cloud_gax::conn::Environment;
19use google_cloud_gax::grpc::Status;
20use google_cloud_googleapis::pubsub::v1::PubsubMessage;
21use google_cloud_pubsub::apiv1;
22use google_cloud_pubsub::client::google_cloud_auth::credentials::CredentialsFile;
23use google_cloud_pubsub::client::google_cloud_auth::project;
24use google_cloud_pubsub::client::google_cloud_auth::token::DefaultTokenSourceProvider;
25use google_cloud_pubsub::client::{Client, ClientConfig};
26use google_cloud_pubsub::publisher::Publisher;
27use risingwave_common::array::StreamChunk;
28use risingwave_common::catalog::Schema;
29use serde::Deserialize;
30use serde_with::serde_as;
31use with_options::WithOptions;
32
33use super::catalog::SinkFormatDesc;
34use super::formatter::SinkFormatterImpl;
35use super::log_store::DeliveryFutureManagerAddFuture;
36use super::writer::{
37    AsyncTruncateLogSinkerOf, AsyncTruncateSinkWriter, AsyncTruncateSinkWriterExt, FormattedSink,
38};
39use super::{Result, Sink, SinkError, SinkParam, SinkWriterParam};
40use crate::dispatch_sink_formatter_str_key_impl;
41use crate::enforce_secret::EnforceSecret;
42
43pub const PUBSUB_SINK: &str = "google_pubsub";
44const PUBSUB_SEND_FUTURE_BUFFER_MAX_SIZE: usize = 65536;
45
46mod delivery_future {
47    use anyhow::Context;
48    use futures::future::try_join_all;
49    use futures::{FutureExt, TryFuture, TryFutureExt};
50    use google_cloud_pubsub::publisher::Awaiter;
51
52    use crate::sink::SinkError;
53
54    pub type GooglePubSubSinkDeliveryFuture =
55        impl TryFuture<Ok = (), Error = SinkError> + Unpin + 'static;
56
57    #[define_opaque(GooglePubSubSinkDeliveryFuture)]
58    pub(super) fn may_delivery_future(awaiter: Vec<Awaiter>) -> GooglePubSubSinkDeliveryFuture {
59        try_join_all(awaiter.into_iter().map(|awaiter| {
60            awaiter.get().map(|result| {
61                result
62                    .context("Google Pub/Sub sink error")
63                    .map_err(SinkError::GooglePubSub)
64                    .map(|_| ())
65            })
66        }))
67        .map_ok(|_: Vec<()>| ())
68        .boxed()
69    }
70}
71
72use delivery_future::*;
73
74#[serde_as]
75#[derive(Clone, Debug, Deserialize, WithOptions)]
76pub struct GooglePubSubConfig {
77    /// The Google Pub/Sub Project ID
78    #[serde(rename = "pubsub.project_id")]
79    pub project_id: String,
80
81    /// Specifies the Pub/Sub topic to publish messages
82    #[serde(rename = "pubsub.topic")]
83    pub topic: String,
84
85    /// The Google Pub/Sub endpoint URL
86    #[serde(rename = "pubsub.endpoint")]
87    pub endpoint: String,
88
89    /// use the connector with a pubsub emulator
90    /// <https://cloud.google.com/pubsub/docs/emulator>
91    #[serde(rename = "pubsub.emulator_host")]
92    pub emulator_host: Option<String>,
93
94    /// A JSON string containing the service account credentials for authorization,
95    /// see the [service-account](https://developers.google.com/workspace/guides/create-credentials#create_credentials_for_a_service_account) credentials guide.
96    /// The provided account credential must have the
97    /// `pubsub.publisher` [role](https://cloud.google.com/pubsub/docs/access-control#roles)
98    #[serde(rename = "pubsub.credentials")]
99    pub credentials: Option<String>,
100
101    #[serde(flatten)]
102    pub unknown_fields: std::collections::HashMap<String, String>,
103}
104
105crate::impl_sink_unknown_fields!(GooglePubSubConfig);
106
107impl EnforceSecret for GooglePubSubConfig {
108    const ENFORCE_SECRET_PROPERTIES: phf::Set<&'static str> = phf::phf_set! {
109        "pubsub.credentials",
110    };
111}
112
113impl GooglePubSubConfig {
114    fn from_btreemap(values: BTreeMap<String, String>) -> Result<Self> {
115        serde_json::from_value::<GooglePubSubConfig>(serde_json::to_value(values).unwrap())
116            .map_err(|e| SinkError::Config(anyhow!(e)))
117    }
118}
119
120#[derive(Clone, Debug)]
121pub struct GooglePubSubSink {
122    pub config: GooglePubSubConfig,
123    is_append_only: bool,
124
125    schema: Schema,
126    pk_indices: Vec<usize>,
127    format_desc: SinkFormatDesc,
128    db_name: String,
129    sink_from_name: String,
130}
131
132impl EnforceSecret for GooglePubSubSink {
133    fn enforce_secret<'a>(
134        prop_iter: impl Iterator<Item = &'a str>,
135    ) -> crate::error::ConnectorResult<()> {
136        for prop in prop_iter {
137            GooglePubSubConfig::enforce_one(prop)?;
138        }
139        Ok(())
140    }
141}
142impl Sink for GooglePubSubSink {
143    type LogSinker = AsyncTruncateLogSinkerOf<GooglePubSubSinkWriter>;
144
145    const SINK_NAME: &'static str = PUBSUB_SINK;
146
147    crate::impl_validate_sink_unknown_fields!();
148
149    async fn validate(&self) -> Result<()> {
150        if !self.is_append_only {
151            return Err(SinkError::GooglePubSub(anyhow!(
152                "Google Pub/Sub sink only support append-only mode"
153            )));
154        }
155
156        let conf = &self.config;
157        if matches!((&conf.emulator_host, &conf.credentials), (None, None)) {
158            return Err(SinkError::GooglePubSub(anyhow!(
159                "Configure at least one of `pubsub.emulator_host` and `pubsub.credentials` in the Google Pub/Sub sink"
160            )));
161        }
162
163        Ok(())
164    }
165
166    async fn new_log_sinker(&self, _writer_param: SinkWriterParam) -> Result<Self::LogSinker> {
167        Ok(GooglePubSubSinkWriter::new(
168            self.config.clone(),
169            self.schema.clone(),
170            self.pk_indices.clone(),
171            &self.format_desc,
172            self.db_name.clone(),
173            self.sink_from_name.clone(),
174        )
175        .await?
176        .into_log_sinker(PUBSUB_SEND_FUTURE_BUFFER_MAX_SIZE))
177    }
178}
179
180impl TryFrom<SinkParam> for GooglePubSubSink {
181    type Error = SinkError;
182
183    fn try_from(param: SinkParam) -> std::result::Result<Self, Self::Error> {
184        let schema = param.schema();
185        let pk_indices = param.downstream_pk_or_empty();
186        let config = GooglePubSubConfig::from_btreemap(param.properties)?;
187        let format_desc = param
188            .format_desc
189            .ok_or_else(|| SinkError::Config(anyhow!("missing FORMAT ... ENCODE ...")))?;
190        Ok(Self {
191            config,
192            is_append_only: param.sink_type.is_append_only(),
193            schema,
194            pk_indices,
195            format_desc,
196            db_name: param.db_name,
197            sink_from_name: param.sink_from_name,
198        })
199    }
200}
201
202struct GooglePubSubPayloadWriter<'w> {
203    publisher: &'w mut Publisher,
204    message_vec: Vec<PubsubMessage>,
205    add_future: DeliveryFutureManagerAddFuture<'w, GooglePubSubSinkDeliveryFuture>,
206}
207
208impl GooglePubSubSinkWriter {
209    pub async fn new(
210        config: GooglePubSubConfig,
211        schema: Schema,
212        pk_indices: Vec<usize>,
213        format_desc: &SinkFormatDesc,
214        db_name: String,
215        sink_from_name: String,
216    ) -> Result<Self> {
217        let environment = if let Some(ref cred) = config.credentials {
218            let mut auth_config = project::Config::default();
219            auth_config = auth_config.with_audience(apiv1::conn_pool::AUDIENCE);
220            auth_config = auth_config.with_scopes(&apiv1::conn_pool::SCOPES);
221            let cred_file = CredentialsFile::new_from_str(cred).await.map_err(|e| {
222                SinkError::GooglePubSub(
223                    anyhow!(e).context("Failed to create Google Cloud Pub/Sub credentials file"),
224                )
225            })?;
226            let provider =
227                DefaultTokenSourceProvider::new_with_credentials(auth_config, Box::new(cred_file))
228                    .await
229                    .map_err(|e| {
230                        SinkError::GooglePubSub(
231                            anyhow!(e).context(
232                                "Failed to create Google Cloud Pub/Sub token source provider",
233                            ),
234                        )
235                    })?;
236            Environment::GoogleCloud(Box::new(provider))
237        } else if let Some(emu_host) = config.emulator_host {
238            Environment::Emulator(emu_host)
239        } else {
240            return Err(SinkError::GooglePubSub(anyhow!(
241                "Missing emulator_host or credentials in Google Pub/Sub sink"
242            )));
243        };
244
245        let client_config = ClientConfig {
246            endpoint: config.endpoint,
247            project_id: Some(config.project_id),
248            environment,
249            ..Default::default()
250        };
251        let client = Client::new(client_config)
252            .await
253            .map_err(|e| SinkError::GooglePubSub(anyhow!(e)))?;
254
255        let topic = async {
256            let topic = client.topic(&config.topic);
257            if !topic.exists(None).await? {
258                topic.create(None, None).await?;
259            }
260            Ok(topic)
261        }
262        .await
263        .map_err(|e: Status| SinkError::GooglePubSub(anyhow!(e)))?;
264
265        let formatter = SinkFormatterImpl::new(
266            format_desc,
267            schema,
268            pk_indices,
269            db_name,
270            sink_from_name,
271            topic.fully_qualified_name(),
272        )
273        .await?;
274
275        let publisher = topic.new_publisher(None);
276
277        Ok(Self {
278            formatter,
279            publisher,
280        })
281    }
282}
283
284pub struct GooglePubSubSinkWriter {
285    formatter: SinkFormatterImpl,
286    publisher: Publisher,
287}
288
289impl AsyncTruncateSinkWriter for GooglePubSubSinkWriter {
290    type DeliveryFuture = GooglePubSubSinkDeliveryFuture;
291
292    async fn write_chunk<'a>(
293        &'a mut self,
294        chunk: StreamChunk,
295        add_future: DeliveryFutureManagerAddFuture<'a, Self::DeliveryFuture>,
296    ) -> Result<()> {
297        let mut payload_writer = GooglePubSubPayloadWriter {
298            publisher: &mut self.publisher,
299            message_vec: Vec::with_capacity(chunk.cardinality()),
300            add_future,
301        };
302        dispatch_sink_formatter_str_key_impl!(&self.formatter, formatter, {
303            payload_writer.write_chunk(chunk, formatter).await
304        })?;
305        payload_writer.finish().await
306    }
307}
308
309impl GooglePubSubPayloadWriter<'_> {
310    pub async fn finish(&mut self) -> Result<()> {
311        let message_vec = std::mem::take(&mut self.message_vec);
312        let awaiters = self.publisher.publish_bulk(message_vec).await;
313        self.add_future
314            .add_future_may_await(may_delivery_future(awaiters))
315            .await?;
316        Ok(())
317    }
318}
319
320impl FormattedSink for GooglePubSubPayloadWriter<'_> {
321    type K = String;
322    type V = Vec<u8>;
323
324    async fn write_one(&mut self, k: Option<Self::K>, v: Option<Self::V>) -> Result<()> {
325        let ordering_key = k.unwrap_or_default();
326        match v {
327            Some(data) => {
328                let msg = PubsubMessage {
329                    data,
330                    ordering_key,
331                    ..Default::default()
332                };
333                self.message_vec.push(msg);
334                Ok(())
335            }
336            None => Err(SinkError::GooglePubSub(anyhow!(
337                "Google Pub/Sub sink error: missing value to publish"
338            ))),
339        }
340    }
341}
risingwave_connector/sink/google_pubsub.rs

risingwave_connector/sink/
google_pubsub.rs
