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