risingwave_connector/source/google_pubsub/enumerator/
client.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 anyhow::Context;
16use async_trait::async_trait;
17use chrono::{TimeZone, Utc};
18use google_cloud_pubsub::subscription::SeekTo;
19use risingwave_common::bail;
20
21use crate::error::ConnectorResult;
22use crate::source::SourceEnumeratorContextRef;
23use crate::source::base::SplitEnumerator;
24use crate::source::google_pubsub::PubsubProperties;
25use crate::source::google_pubsub::split::PubsubSplit;
26
27pub struct PubsubSplitEnumerator {
28    subscription: String,
29}
30
31#[async_trait]
32impl SplitEnumerator for PubsubSplitEnumerator {
33    type Properties = PubsubProperties;
34    type Split = PubsubSplit;
35
36    async fn new(
37        properties: Self::Properties,
38        _context: SourceEnumeratorContextRef,
39    ) -> ConnectorResult<PubsubSplitEnumerator> {
40        if properties.parallelism.is_some() {
41            tracing::warn!(
42                "pubsub.parallelism is deprecated and will be ignored. \
43                 Split count now adapts automatically to the number of actors."
44            );
45        }
46
47        properties.subscriber_config()?;
48
49        if properties.credentials.is_none() && properties.emulator_host.is_none() {
50            bail!("credentials must be set if not using the pubsub emulator")
51        }
52
53        let sub = properties.subscription_client().await?;
54        if !sub
55            .exists(None)
56            .await
57            .context("error checking subscription validity")?
58        {
59            bail!("subscription {} does not exist", &sub.id())
60        }
61
62        let seek_to = match (properties.start_offset, properties.start_snapshot) {
63            (None, None) => None,
64            (Some(start_offset), None) => {
65                let ts = start_offset
66                    .parse::<i64>()
67                    .context("error parsing start_offset")
68                    .map(|nanos| Utc.timestamp_nanos(nanos).into())?;
69                Some(SeekTo::Timestamp(ts))
70            }
71            (None, Some(snapshot)) => Some(SeekTo::Snapshot(snapshot)),
72            (Some(_), Some(_)) => {
73                bail!("specify at most one of start_offset or start_snapshot")
74            }
75        };
76
77        if let Some(seek_to) = seek_to {
78            sub.seek(seek_to, None)
79                .await
80                .context("error seeking subscription")?;
81        }
82
83        Ok(Self {
84            subscription: properties.subscription,
85        })
86    }
87
88    async fn list_splits(&mut self) -> ConnectorResult<Vec<PubsubSplit>> {
89        tracing::debug!("enumerating pubsub splits (adaptive mode, returning 1 template split)");
90        Ok(vec![PubsubSplit {
91            index: 0,
92            subscription: self.subscription.clone(),
93            __deprecated_start_offset: None,
94            __deprecated_stop_offset: None,
95        }])
96    }
97}
risingwave_connector/source/google_pubsub/enumerator/client.rs

risingwave_connector/source/google_pubsub/enumerator/
client.rs
