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        let sub = properties.subscription_client().await?;
50        if !sub
51            .exists(None)
52            .await
53            .context("error checking subscription validity")?
54        {
55            bail!("subscription {} does not exist", &sub.id())
56        }
57
58        let seek_to = match (properties.start_offset, properties.start_snapshot) {
59            (None, None) => None,
60            (Some(start_offset), None) => {
61                let ts = start_offset
62                    .parse::<i64>()
63                    .context("error parsing start_offset")
64                    .map(|nanos| Utc.timestamp_nanos(nanos).into())?;
65                Some(SeekTo::Timestamp(ts))
66            }
67            (None, Some(snapshot)) => Some(SeekTo::Snapshot(snapshot)),
68            (Some(_), Some(_)) => {
69                bail!("specify at most one of start_offset or start_snapshot")
70            }
71        };
72
73        if let Some(seek_to) = seek_to {
74            sub.seek(seek_to, None)
75                .await
76                .context("error seeking subscription")?;
77        }
78
79        Ok(Self {
80            subscription: properties.subscription,
81        })
82    }
83
84    async fn list_splits(&mut self) -> ConnectorResult<Vec<PubsubSplit>> {
85        tracing::debug!("enumerating pubsub splits (adaptive mode, returning 1 template split)");
86        Ok(vec![PubsubSplit {
87            index: 0,
88            subscription: self.subscription.clone(),
89            __deprecated_start_offset: None,
90            __deprecated_stop_offset: None,
91        }])
92    }
93}
risingwave_connector/source/google_pubsub/enumerator/client.rs

risingwave_connector/source/google_pubsub/enumerator/
client.rs
