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

risingwave_connector/sink/
google_pubsub.rs
