risingwave_connector/schema/
loader.rs

1// Copyright 2024 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::BTreeMap;
16
17use risingwave_pb::catalog::PbSchemaRegistryNameStrategy;
18
19use super::schema_registry::{
20    Client, Subject, get_subject_by_strategy, handle_sr_list, name_strategy_from_str,
21};
22use super::{
23    AWS_GLUE_SCHEMA_ARN_KEY, InvalidOptionError, KEY_MESSAGE_NAME_KEY, MESSAGE_NAME_KEY,
24    MalformedResponseError, NAME_STRATEGY_KEY, SCHEMA_REGISTRY_KEY, SchemaFetchError,
25    invalid_option_error, malformed_response_error,
26};
27use crate::connector_common::AwsAuthProps;
28
29pub enum SchemaLoader {
30    Confluent(ConfluentSchemaLoader),
31    Glue(GlueSchemaLoader),
32}
33
34pub struct ConfluentSchemaLoader {
35    pub client: Client,
36    pub name_strategy: PbSchemaRegistryNameStrategy,
37    pub topic: String,
38    pub key_record_name: Option<String>,
39    pub val_record_name: Option<String>,
40}
41
42pub enum GlueSchemaLoader {
43    Real {
44        client: aws_sdk_glue::Client,
45        schema_arn: String,
46    },
47    Mock {
48        schema_version_id: uuid::Uuid,
49        definition: String,
50    },
51}
52
53pub enum SchemaVersion {
54    Confluent(i32),
55    Glue(uuid::Uuid),
56}
57
58impl ConfluentSchemaLoader {
59    pub fn from_format_options(
60        topic: &str,
61        format_options: &BTreeMap<String, String>,
62    ) -> Result<Self, SchemaFetchError> {
63        let schema_location = format_options
64            .get(SCHEMA_REGISTRY_KEY)
65            .ok_or_else(|| invalid_option_error!("`{SCHEMA_REGISTRY_KEY}` is required"))?;
66        let client_config = format_options.into();
67        let urls = handle_sr_list(schema_location)?;
68        let client = Client::new(urls, &client_config)?;
69
70        let name_strategy = format_options
71            .get(NAME_STRATEGY_KEY)
72            .map(|s| {
73                name_strategy_from_str(s).ok_or_else(|| {
74                    invalid_option_error!("unrecognized schema registry naming strategy: {s}")
75                })
76            })
77            .transpose()?
78            .unwrap_or_default();
79        let key_record_name = format_options.get(KEY_MESSAGE_NAME_KEY).cloned();
80        let val_record_name = format_options.get(MESSAGE_NAME_KEY).cloned();
81
82        Ok(Self {
83            client,
84            name_strategy,
85            topic: topic.into(),
86            key_record_name,
87            val_record_name,
88        })
89    }
90
91    async fn load_schema<Out: LoadedSchema, const IS_KEY: bool>(
92        &self,
93    ) -> Result<(SchemaVersion, Out), SchemaFetchError> {
94        let record = match IS_KEY {
95            true => self.key_record_name.as_deref(),
96            false => self.val_record_name.as_deref(),
97        };
98        let subject = get_subject_by_strategy(&self.name_strategy, &self.topic, record, IS_KEY)?;
99        let (primary_subject, dependency_subjects) =
100            self.client.get_subject_and_references(&subject).await?;
101        let schema_id = primary_subject.schema.id;
102        let out = Out::compile(primary_subject, dependency_subjects)?;
103        Ok((SchemaVersion::Confluent(schema_id), out))
104    }
105}
106
107impl GlueSchemaLoader {
108    pub async fn from_format_options(
109        schema_arn: &str,
110        format_options: &BTreeMap<String, String>,
111    ) -> Result<Self, SchemaFetchError> {
112        risingwave_common::license::Feature::GlueSchemaRegistry.check_available()?;
113        if let Some(mock_config) = format_options.get("aws.glue.mock_config") {
114            // Internal format for easy testing. See `MockGlueSchemaCache` for details.
115            let parsed: serde_json::Value =
116                serde_json::from_str(mock_config).expect("mock config shall be valid json");
117            let schema_version_id_str = parsed
118                .get("arn_to_latest_id")
119                .unwrap()
120                .as_object()
121                .unwrap()
122                .get(schema_arn)
123                .unwrap()
124                .as_str()
125                .unwrap();
126            let definition = parsed
127                .get("by_id")
128                .unwrap()
129                .as_object()
130                .unwrap()
131                .get(schema_version_id_str)
132                .unwrap()
133                .to_string();
134            return Ok(Self::Mock {
135                schema_version_id: schema_version_id_str.parse()?,
136                definition,
137            });
138        };
139        let aws_auth_props =
140            serde_json::from_value::<AwsAuthProps>(serde_json::to_value(format_options).unwrap())
141                .map_err(|_e| invalid_option_error!(""))?;
142        let client = aws_sdk_glue::Client::new(
143            &aws_auth_props
144                .build_config()
145                .await
146                .map_err(SchemaFetchError::YetToMigrate)?,
147        );
148        Ok(Self::Real {
149            client,
150            schema_arn: schema_arn.to_owned(),
151        })
152    }
153
154    async fn load_schema<Out: LoadedSchema, const IS_KEY: bool>(
155        &self,
156    ) -> Result<(SchemaVersion, Out), SchemaFetchError> {
157        if IS_KEY {
158            return Err(invalid_option_error!(
159                "GlueSchemaRegistry cannot be key. Specify `KEY ENCODE [TEXT | BYTES]` please."
160            )
161            .into());
162        }
163        let (schema_version_id, definition) = match self {
164            Self::Mock {
165                schema_version_id,
166                definition,
167            } => (*schema_version_id, definition.clone()),
168            Self::Real { client, schema_arn } => {
169                use aws_sdk_glue::types::{SchemaId, SchemaVersionNumber};
170
171                let res = client
172                    .get_schema_version()
173                    .schema_id(SchemaId::builder().schema_arn(schema_arn).build())
174                    .schema_version_number(
175                        SchemaVersionNumber::builder().latest_version(true).build(),
176                    )
177                    .send()
178                    .await
179                    .map_err(|e| Box::new(e.into_service_error()))?;
180                let schema_version_id = res
181                    .schema_version_id()
182                    .ok_or_else(|| malformed_response_error!("missing schema_version_id"))?
183                    .parse()?;
184                let definition = res
185                    .schema_definition()
186                    .ok_or_else(|| malformed_response_error!("missing schema_definition"))?
187                    .to_owned();
188                (schema_version_id, definition)
189            }
190        };
191
192        // https://github.com/awslabs/aws-glue-schema-registry/issues/32
193        // No references in AWS Glue Schema Registry yet
194        let primary = Subject {
195            version: 0,
196            name: "".to_owned(),
197            schema: super::schema_registry::ConfluentSchema {
198                id: 0,
199                content: definition,
200            },
201        };
202        let out = Out::compile(primary, vec![])?;
203        Ok((SchemaVersion::Glue(schema_version_id), out))
204    }
205}
206
207impl SchemaLoader {
208    pub async fn from_format_options(
209        topic: &str,
210        format_options: &BTreeMap<String, String>,
211    ) -> Result<Self, SchemaFetchError> {
212        if let Some(schema_arn) = format_options.get(AWS_GLUE_SCHEMA_ARN_KEY) {
213            Ok(Self::Glue(
214                GlueSchemaLoader::from_format_options(schema_arn, format_options).await?,
215            ))
216        } else {
217            Ok(Self::Confluent(ConfluentSchemaLoader::from_format_options(
218                topic,
219                format_options,
220            )?))
221        }
222    }
223
224    async fn load_schema<Out: LoadedSchema, const IS_KEY: bool>(
225        &self,
226    ) -> Result<(SchemaVersion, Out), SchemaFetchError> {
227        match self {
228            Self::Confluent(inner) => inner.load_schema::<Out, IS_KEY>().await,
229            Self::Glue(inner) => inner.load_schema::<Out, IS_KEY>().await,
230        }
231    }
232
233    pub async fn load_key_schema<Out: LoadedSchema>(
234        &self,
235    ) -> Result<(SchemaVersion, Out), SchemaFetchError> {
236        self.load_schema::<Out, true>().await
237    }
238
239    pub async fn load_val_schema<Out: LoadedSchema>(
240        &self,
241    ) -> Result<(SchemaVersion, Out), SchemaFetchError> {
242        self.load_schema::<Out, false>().await
243    }
244}
245
246pub trait LoadedSchema: Sized {
247    fn compile(primary: Subject, references: Vec<Subject>) -> Result<Self, SchemaFetchError>;
248}
risingwave_connector/schema/loader.rs

risingwave_connector/schema/
loader.rs
