Skip to main content

risingwave_connector/schema/schema_registry/
client.rs

1// Copyright 2023 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::HashSet;
16use std::fmt::Debug;
17use std::sync::Arc;
18use std::time::Duration;
19
20use futures::future::select_all;
21use itertools::Itertools;
22use reqwest::{Method, Url};
23use risingwave_common::util::retry::exponential_backoff;
24use serde::Deserialize;
25use serde::de::DeserializeOwned;
26use thiserror_ext::AsReport as _;
27use tokio_retry::Retry;
28use tokio_retry::strategy::jitter;
29
30use super::util::*;
31use crate::connector_common::ConfluentSchemaRegistryConnection;
32use crate::schema::{InvalidOptionError, invalid_option_error};
33use crate::with_options::Get;
34
35pub const SCHEMA_REGISTRY_USERNAME: &str = "schema.registry.username";
36pub const SCHEMA_REGISTRY_PASSWORD: &str = "schema.registry.password";
37pub const SCHEMA_REGISTRY_CA_PEM_PATH: &str = "schema.registry.ca_pem_path";
38
39pub const SCHEMA_REGISTRY_MAX_DELAY_KEY: &str = "schema.registry.max.delay.sec";
40pub const SCHEMA_REGISTRY_BACKOFF_DURATION_KEY: &str = "schema.registry.backoff.duration.ms";
41pub const SCHEMA_REGISTRY_BACKOFF_FACTOR_KEY: &str = "schema.registry.backoff.factor";
42pub const SCHEMA_REGISTRY_RETRIES_MAX_KEY: &str = "schema.registry.retries.max";
43
44const DEFAULT_MAX_DELAY_SEC: u32 = 3;
45const DEFAULT_BACKOFF_DURATION_MS: u64 = 100;
46const DEFAULT_BACKOFF_FACTOR: u64 = 2;
47const DEFAULT_RETRIES_MAX: usize = 3;
48
49#[derive(Debug, Clone)]
50struct SchemaRegistryRetryConfig {
51    pub max_delay_sec: u32,
52    pub backoff_duration_ms: u64,
53    pub backoff_factor: u64,
54    pub retries_max: usize,
55}
56
57impl Default for SchemaRegistryRetryConfig {
58    fn default() -> Self {
59        Self {
60            max_delay_sec: DEFAULT_MAX_DELAY_SEC,
61            backoff_duration_ms: DEFAULT_BACKOFF_DURATION_MS,
62            backoff_factor: DEFAULT_BACKOFF_FACTOR,
63            retries_max: DEFAULT_RETRIES_MAX,
64        }
65    }
66}
67
68#[derive(Debug, Clone, Default)]
69pub struct SchemaRegistryConfig {
70    username: Option<String>,
71    password: Option<String>,
72    ca_pem_path: Option<String>,
73
74    retry_config: SchemaRegistryRetryConfig,
75}
76
77impl<T: Get> From<&T> for SchemaRegistryConfig {
78    fn from(props: &T) -> Self {
79        SchemaRegistryConfig {
80            username: props.get(SCHEMA_REGISTRY_USERNAME).cloned(),
81            password: props.get(SCHEMA_REGISTRY_PASSWORD).cloned(),
82            ca_pem_path: props.get(SCHEMA_REGISTRY_CA_PEM_PATH).cloned(),
83
84            retry_config: SchemaRegistryRetryConfig {
85                max_delay_sec: props
86                    .get(SCHEMA_REGISTRY_MAX_DELAY_KEY)
87                    .and_then(|v| v.parse::<u32>().ok())
88                    .unwrap_or(DEFAULT_MAX_DELAY_SEC),
89                backoff_duration_ms: props
90                    .get(SCHEMA_REGISTRY_BACKOFF_DURATION_KEY)
91                    .and_then(|v| v.parse::<u64>().ok())
92                    .unwrap_or(DEFAULT_BACKOFF_DURATION_MS),
93                backoff_factor: props
94                    .get(SCHEMA_REGISTRY_BACKOFF_FACTOR_KEY)
95                    .and_then(|v| v.parse::<u64>().ok())
96                    .unwrap_or(DEFAULT_BACKOFF_FACTOR),
97                retries_max: props
98                    .get(SCHEMA_REGISTRY_RETRIES_MAX_KEY)
99                    .and_then(|v| v.parse::<usize>().ok())
100                    .unwrap_or(DEFAULT_RETRIES_MAX),
101            },
102        }
103    }
104}
105
106/// An client for communication with schema registry
107#[derive(Debug)]
108pub struct Client {
109    inner: reqwest::Client,
110    url: Vec<Url>,
111    username: Option<String>,
112    password: Option<String>,
113
114    retry_config: SchemaRegistryRetryConfig,
115}
116
117#[derive(Debug, thiserror::Error)]
118#[error("all request confluent registry all timeout, {context}\n{}", errs.iter().map(|e| format!("\t{}", e.as_report())).join("\n"))]
119pub struct ConcurrentRequestError {
120    errs: Vec<itertools::Either<RequestError, tokio::task::JoinError>>,
121    context: String,
122}
123
124type SrResult<T> = Result<T, ConcurrentRequestError>;
125
126#[derive(thiserror::Error, Debug)]
127pub enum SchemaRegistryClientError {
128    #[error(transparent)]
129    InvalidOption(#[from] InvalidOptionError),
130    #[error("read ca file error: {0}")]
131    ReadFile(#[source] std::io::Error),
132    #[error("parse ca file error: {0}")]
133    ParsePem(#[source] reqwest::Error),
134    #[error("build schema registry client error: {0}")]
135    Build(#[source] reqwest::Error),
136}
137
138impl TryFrom<&ConfluentSchemaRegistryConnection> for Client {
139    type Error = InvalidOptionError;
140
141    fn try_from(value: &ConfluentSchemaRegistryConnection) -> Result<Self, Self::Error> {
142        let urls = handle_sr_list(value.url.as_str())?;
143
144        Client::new(
145            urls,
146            &SchemaRegistryConfig {
147                username: value.username.clone(),
148                password: value.password.clone(),
149                ..Default::default()
150            },
151        )
152        .map_err(|e| match e {
153            SchemaRegistryClientError::InvalidOption(e) => e,
154            e => {
155                invalid_option_error!("failed to create schema registry client: {}", e.as_report())
156            }
157        })
158    }
159}
160
161impl Client {
162    pub(crate) fn new(
163        url: Vec<Url>,
164        client_config: &SchemaRegistryConfig,
165    ) -> Result<Self, SchemaRegistryClientError> {
166        let valid_urls = url
167            .iter()
168            .map(|url| (url.cannot_be_a_base(), url))
169            .filter(|(x, _)| !*x)
170            .map(|(_, url)| url.clone())
171            .collect_vec();
172        if valid_urls.is_empty() {
173            return Err(SchemaRegistryClientError::InvalidOption(
174                invalid_option_error!(
175                    "the following schema registry URLs are not valid base URLs: {}",
176                    url.iter().join(" ")
177                ),
178            ));
179        } else {
180            tracing::debug!(
181                "schema registry client will use url {:?} to connect",
182                valid_urls
183            );
184        }
185
186        let mut client_builder = reqwest::Client::builder();
187        if let Some(ca_path) = client_config.ca_pem_path.as_ref() {
188            if ca_path.eq_ignore_ascii_case("ignore") {
189                client_builder = client_builder.danger_accept_invalid_certs(true);
190            } else {
191                client_builder = client_builder.add_root_certificate(
192                    reqwest::Certificate::from_pem(
193                        &std::fs::read(ca_path).map_err(SchemaRegistryClientError::ReadFile)?,
194                    )
195                    .map_err(SchemaRegistryClientError::ParsePem)?,
196                );
197            }
198        }
199
200        let inner = client_builder
201            .build()
202            .map_err(SchemaRegistryClientError::Build)?;
203
204        Ok(Client {
205            inner,
206            url: valid_urls,
207            username: client_config.username.clone(),
208            password: client_config.password.clone(),
209            retry_config: client_config.retry_config.clone(),
210        })
211    }
212
213    async fn concurrent_req<'a, T>(
214        &'a self,
215        method: Method,
216        path: &'a [&'a (impl AsRef<str> + ?Sized + Debug + ToString)],
217    ) -> SrResult<T>
218    where
219        T: DeserializeOwned + Send + Sync + 'static,
220    {
221        let mut fut_req = Vec::with_capacity(self.url.len());
222        let mut errs = Vec::with_capacity(self.url.len());
223        let ctx = Arc::new(SchemaRegistryCtx {
224            username: self.username.clone(),
225            password: self.password.clone(),
226            client: self.inner.clone(),
227            path: path.iter().map(|p| p.to_string()).collect_vec(),
228        });
229        tracing::debug!("retry config: {:?}", self.retry_config);
230
231        let retry_strategy = exponential_backoff(
232            Duration::from_millis(self.retry_config.backoff_duration_ms),
233            self.retry_config.backoff_factor,
234            Duration::from_secs(self.retry_config.max_delay_sec as u64),
235        )
236        .take(self.retry_config.retries_max)
237        .map(jitter);
238
239        for url in &self.url {
240            let url_clone = url.clone();
241            let ctx_clone = ctx.clone();
242            let method_clone = method.clone();
243
244            let retry_future = Retry::spawn(retry_strategy.clone(), move || {
245                let ctx = ctx_clone.clone();
246                let url = url_clone.clone();
247                let method = method_clone.clone();
248                async move { req_inner(ctx, url, method).await }
249            });
250
251            fut_req.push(tokio::spawn(retry_future));
252        }
253
254        while !fut_req.is_empty() {
255            let (result, _index, remaining) = select_all(fut_req).await;
256            match result {
257                Ok(Ok(res)) => {
258                    let _ = remaining.iter().map(|ele| ele.abort());
259                    return Ok(res);
260                }
261                Ok(Err(e)) => errs.push(itertools::Either::Left(e)),
262                Err(e) => errs.push(itertools::Either::Right(e)),
263            }
264            fut_req = remaining;
265        }
266
267        Err(ConcurrentRequestError {
268            errs,
269            context: format!("req path {:?}, urls {}", path, self.url.iter().join(" ")),
270        })
271    }
272
273    /// get schema by id
274    pub async fn get_schema_by_id(&self, id: i32) -> SrResult<ConfluentSchema> {
275        let res: GetByIdResp = self
276            .concurrent_req(Method::GET, &["schemas", "ids", &id.to_string()])
277            .await?;
278        Ok(ConfluentSchema {
279            id,
280            content: res.schema,
281        })
282    }
283
284    /// get the latest schema of the subject
285    pub async fn get_schema_by_subject(&self, subject: &str) -> SrResult<ConfluentSchema> {
286        self.get_subject(subject).await.map(|s| s.schema)
287    }
288
289    // used for connection validate, just check if request is ok
290    pub async fn validate_connection(&self) -> SrResult<()> {
291        #[derive(Debug, Deserialize)]
292        struct GetConfigResp {
293            #[serde(rename = "compatibilityLevel")]
294            _compatibility_level: String,
295        }
296
297        let _: GetConfigResp = self.concurrent_req(Method::GET, &["config"]).await?;
298        Ok(())
299    }
300
301    /// get the latest version of the subject
302    pub async fn get_subject(&self, subject: &str) -> SrResult<Subject> {
303        let res: GetBySubjectResp = self
304            .concurrent_req(Method::GET, &["subjects", subject, "versions", "latest"])
305            .await?;
306        tracing::debug!("update schema: {:?}", res);
307        Ok(Subject {
308            schema: ConfluentSchema {
309                id: res.id,
310                content: res.schema,
311            },
312            version: res.version,
313            name: res.subject,
314        })
315    }
316
317    /// get the latest version of the subject and all it's references(deps)
318    pub async fn get_subject_and_references(
319        &self,
320        subject: &str,
321    ) -> SrResult<(Subject, Vec<Subject>)> {
322        let mut subjects = vec![];
323        let mut visited = HashSet::new();
324        let mut queue = vec![(subject.to_owned(), "latest".to_owned())];
325        // use bfs to get all references
326        while let Some((subject, version)) = queue.pop() {
327            let res: GetBySubjectResp = self
328                .concurrent_req(Method::GET, &["subjects", &subject, "versions", &version])
329                .await?;
330            let ref_subject = Subject {
331                schema: ConfluentSchema {
332                    id: res.id,
333                    content: res.schema,
334                },
335                version: res.version,
336                name: res.subject.clone(),
337            };
338            subjects.push(ref_subject);
339            visited.insert(res.subject);
340            queue.extend(
341                res.references
342                    .into_iter()
343                    .filter(|r| !visited.contains(&r.subject))
344                    .map(|r| (r.subject, r.version.to_string())),
345            );
346        }
347        let origin_subject = subjects.remove(0);
348
349        Ok((origin_subject, subjects))
350    }
351}