risingwave_connector/source/google_pubsub/
mod.rs

1// Copyright 2022 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::HashMap;
16
17use anyhow::Context;
18use google_cloud_pubsub::client::{Client, ClientConfig};
19use google_cloud_pubsub::subscriber::SubscriberConfig;
20use google_cloud_pubsub::subscription::Subscription;
21use risingwave_common::bail;
22use serde::Deserialize;
23
24pub mod enumerator;
25pub mod source;
26pub mod split;
27
28pub use enumerator::*;
29use phf::{Set, phf_set};
30use serde_with::{DisplayFromStr, serde_as};
31pub use source::*;
32pub use split::*;
33use with_options::WithOptions;
34
35use crate::enforce_secret::EnforceSecret;
36use crate::error::ConnectorResult;
37use crate::source::SourceProperties;
38
39pub const GOOGLE_PUBSUB_CONNECTOR: &str = "google_pubsub";
40
41const DEFAULT_ACK_DEADLINE_SECONDS: i32 = 60;
42// Pub/Sub messages are acknowledged only after a checkpoint. The upstream client default of 50
43// can therefore stall each reader between checkpoints and severely limit throughput.
44const DEFAULT_MAX_OUTSTANDING_MESSAGES: i64 = 1024;
45const DEFAULT_MAX_OUTSTANDING_BYTES: i64 = 1_000_000_000;
46
47/// # Implementation Notes
48/// Pub/Sub does not rely on persisted state (`SplitImpl`) to start from a position.
49/// It rely on Pub/Sub to load-balance messages between all Readers.
50/// We `ack` received messages after checkpoint (see `WaitCheckpointWorker`) to achieve at-least-once delivery.
51#[serde_as]
52#[derive(Clone, Debug, Deserialize, WithOptions)]
53pub struct PubsubProperties {
54    /// Pub/Sub subscription to consume messages from.
55    ///
56    /// Note that we rely on Pub/Sub to load-balance messages between all Readers pulling from
57    /// the same subscription. So one `subscription` (i.e., one `Source`) can only used for one MV
58    /// (shared between the actors of its fragment).
59    /// Otherwise, different MVs on the same Source will both receive part of the messages.
60    /// TODO: check and enforce this on Meta.
61    #[serde(rename = "pubsub.subscription")]
62    pub subscription: String,
63
64    /// use the connector with a pubsub emulator
65    /// <https://cloud.google.com/pubsub/docs/emulator>
66    #[serde(rename = "pubsub.emulator_host")]
67    pub emulator_host: Option<String>,
68
69    /// `credentials` is a JSON string containing the service account credentials.
70    /// See the [service-account credentials guide](https://developers.google.com/workspace/guides/create-credentials#create_credentials_for_a_service_account).
71    /// The service account must have the `pubsub.subscriber` [role](https://cloud.google.com/pubsub/docs/access-control#roles).
72    #[serde(rename = "pubsub.credentials")]
73    pub credentials: Option<String>,
74
75    /// `start_offset` is a numeric timestamp, ideally the publish timestamp of a message
76    /// in the subscription. If present, the connector will attempt to seek the subscription
77    /// to the timestamp and start consuming from there. Note that the seek operation is
78    /// subject to limitations around the message retention policy of the subscription. See
79    /// [Seeking to a timestamp](https://cloud.google.com/pubsub/docs/replay-overview#seeking_to_a_timestamp) for
80    /// more details.
81    #[serde(rename = "pubsub.start_offset.nanos")]
82    pub start_offset: Option<String>,
83
84    /// `start_snapshot` is a named pub/sub snapshot. If present, the connector will first seek
85    /// to the snapshot before starting consumption. Snapshots are the preferred seeking mechanism
86    /// in pub/sub because they guarantee retention of:
87    /// - All unacknowledged messages at the time of their creation.
88    /// - All messages created after their creation.
89    /// Besides retention guarantees, snapshots are also more precise than timestamp-based seeks.
90    /// See [Seeking to a snapshot](https://cloud.google.com/pubsub/docs/replay-overview#seeking_to_a_timestamp) for
91    /// more details.
92    #[serde(rename = "pubsub.start_snapshot")]
93    pub start_snapshot: Option<String>,
94
95    /// Deprecated: ignored since adaptive split mode was introduced.
96    /// Split count now adapts automatically to the number of actors.
97    /// Kept for backward compatibility with existing DDL.
98    #[serde_as(as = "Option<DisplayFromStr>")]
99    #[serde(rename = "pubsub.parallelism")]
100    pub parallelism: Option<u32>,
101
102    /// The ack deadline in seconds for the streaming pull subscriber.
103    /// This is the maximum time the server will wait for an ack before redelivering the message.
104    /// Must be between 10 and 600 seconds. Defaults to 60.
105    #[serde_as(as = "Option<DisplayFromStr>")]
106    #[serde(rename = "pubsub.ack_deadline_seconds")]
107    #[with_option(allow_alter_on_fly)]
108    pub ack_deadline_seconds: Option<i32>,
109
110    /// The maximum number of unacknowledged messages delivered to each streaming pull reader.
111    /// Pub/Sub pauses delivery to a reader when this limit is reached. Must be greater than 0.
112    /// Defaults to 1024.
113    #[serde_as(as = "Option<DisplayFromStr>")]
114    #[serde(rename = "pubsub.max_outstanding_messages")]
115    #[with_option(allow_alter_on_fly)]
116    pub max_outstanding_messages: Option<i64>,
117
118    /// The maximum total size of unacknowledged messages delivered to each streaming pull reader.
119    /// Pub/Sub pauses delivery to a reader when this limit is reached. Must be greater than 0.
120    /// Defaults to 1 GB.
121    #[serde_as(as = "Option<DisplayFromStr>")]
122    #[serde(rename = "pubsub.max_outstanding_bytes")]
123    #[with_option(allow_alter_on_fly)]
124    pub max_outstanding_bytes: Option<i64>,
125
126    #[serde(flatten)]
127    pub unknown_fields: HashMap<String, String>,
128}
129
130impl EnforceSecret for PubsubProperties {
131    const ENFORCE_SECRET_PROPERTIES: Set<&'static str> = phf_set! {
132        "pubsub.credentials",
133    };
134}
135
136impl SourceProperties for PubsubProperties {
137    type Split = PubsubSplit;
138    type SplitEnumerator = PubsubSplitEnumerator;
139    type SplitReader = PubsubSplitReader;
140
141    const SOURCE_NAME: &'static str = GOOGLE_PUBSUB_CONNECTOR;
142}
143
144impl crate::source::UnknownFields for PubsubProperties {
145    fn unknown_fields(&self) -> HashMap<String, String> {
146        self.unknown_fields.clone()
147    }
148}
149
150impl PubsubProperties {
151    pub(crate) fn subscriber_config(&self) -> ConnectorResult<SubscriberConfig> {
152        let stream_ack_deadline_seconds = self
153            .ack_deadline_seconds
154            .unwrap_or(DEFAULT_ACK_DEADLINE_SECONDS);
155        if !(10..=600).contains(&stream_ack_deadline_seconds) {
156            bail!("pubsub.ack_deadline_seconds must be between 10 and 600");
157        }
158
159        let max_outstanding_messages = self
160            .max_outstanding_messages
161            .unwrap_or(DEFAULT_MAX_OUTSTANDING_MESSAGES);
162        if max_outstanding_messages <= 0 {
163            bail!("pubsub.max_outstanding_messages must be greater than 0");
164        }
165
166        let max_outstanding_bytes = self
167            .max_outstanding_bytes
168            .unwrap_or(DEFAULT_MAX_OUTSTANDING_BYTES);
169        if max_outstanding_bytes <= 0 {
170            bail!("pubsub.max_outstanding_bytes must be greater than 0");
171        }
172
173        Ok(SubscriberConfig {
174            stream_ack_deadline_seconds,
175            max_outstanding_messages,
176            max_outstanding_bytes,
177            ..Default::default()
178        })
179    }
180
181    pub(crate) async fn subscription_client(&self) -> ConnectorResult<Subscription> {
182        // initialize env
183        {
184            tracing::debug!("setting pubsub environment variables");
185            if let Some(emulator_host) = &self.emulator_host {
186                // safety: only read in the same thread below in with_auth
187                unsafe { std::env::set_var("PUBSUB_EMULATOR_HOST", emulator_host) };
188            }
189            if let Some(credentials) = &self.credentials {
190                // safety: only read in the same thread below in with_auth
191                unsafe { std::env::set_var("GOOGLE_APPLICATION_CREDENTIALS_JSON", credentials) };
192            }
193        };
194
195        // Validate config
196        let config = ClientConfig::default().with_auth().await?;
197        let client = Client::new(config)
198            .await
199            .context("error initializing pubsub client")?;
200
201        Ok(client.subscription(&self.subscription))
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use serde_json::json;
208
209    use super::*;
210
211    fn parse_pubsub_properties(extra: serde_json::Value) -> PubsubProperties {
212        let mut value = json!({
213            "pubsub.subscription": "projects/test/subscriptions/test",
214            "pubsub.emulator_host": "localhost:8900",
215        });
216        value
217            .as_object_mut()
218            .unwrap()
219            .extend(extra.as_object().unwrap().clone());
220        serde_json::from_value(value).unwrap()
221    }
222
223    #[test]
224    fn test_subscriber_config_defaults() {
225        let config = parse_pubsub_properties(json!({}))
226            .subscriber_config()
227            .unwrap();
228
229        assert_eq!(config.stream_ack_deadline_seconds, 60);
230        assert_eq!(config.max_outstanding_messages, 1024);
231        assert_eq!(config.max_outstanding_bytes, 1_000_000_000);
232    }
233
234    #[test]
235    fn test_subscriber_config_overrides() {
236        let config = parse_pubsub_properties(json!({
237            "pubsub.ack_deadline_seconds": "120",
238            "pubsub.max_outstanding_messages": "2048",
239            "pubsub.max_outstanding_bytes": "1048576",
240        }))
241        .subscriber_config()
242        .unwrap();
243
244        assert_eq!(config.stream_ack_deadline_seconds, 120);
245        assert_eq!(config.max_outstanding_messages, 2048);
246        assert_eq!(config.max_outstanding_bytes, 1_048_576);
247    }
248
249    #[test]
250    fn test_subscriber_config_validation() {
251        let invalid_values = [
252            (
253                json!({"pubsub.ack_deadline_seconds": "9"}),
254                "pubsub.ack_deadline_seconds must be between 10 and 600",
255            ),
256            (
257                json!({"pubsub.max_outstanding_messages": "0"}),
258                "pubsub.max_outstanding_messages must be greater than 0",
259            ),
260            (
261                json!({"pubsub.max_outstanding_bytes": "0"}),
262                "pubsub.max_outstanding_bytes must be greater than 0",
263            ),
264        ];
265
266        for (value, expected_error) in invalid_values {
267            let error = parse_pubsub_properties(value)
268                .subscriber_config()
269                .unwrap_err();
270            assert!(error.to_string().contains(expected_error));
271        }
272    }
273}
risingwave_connector/source/google_pubsub/mod.rs

risingwave_connector/source/google_pubsub/
mod.rs
