risingwave_connector/parser/avro/
parser.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::fmt::Debug;
16use std::sync::Arc;
17
18use anyhow::Context;
19use apache_avro::types::Value;
20use apache_avro::{Reader, Schema, from_avro_datum};
21use risingwave_common::catalog::Field;
22use risingwave_common::{bail, try_match_expand};
23use risingwave_connector_codec::decoder::avro::{
24    AvroAccess, AvroParseOptions, ResolvedAvroSchema, avro_schema_to_fields,
25};
26
27use super::{ConfluentSchemaCache, GlueSchemaCache as _, GlueSchemaCacheImpl, PulsarSchemaCache};
28use crate::error::ConnectorResult;
29use crate::parser::unified::AccessImpl;
30use crate::parser::utils::bytes_from_url;
31use crate::parser::{
32    AccessBuilder, AvroProperties, EncodingProperties, MapHandling, SchemaLocation,
33};
34use crate::schema::pulsar_schema::PulsarSchemaVersion;
35use crate::schema::schema_registry::{
36    Client, extract_schema_id, get_subject_by_strategy, handle_sr_list,
37};
38use crate::source::SourceMeta;
39
40// Default avro access builder
41#[derive(Debug)]
42pub struct AvroAccessBuilder {
43    schema: Arc<ResolvedAvroSchema>,
44    /// Refer to [`AvroParserConfig::writer_schema_cache`].
45    writer_schema_cache: WriterSchemaCache,
46    value: Option<Value>,
47}
48
49impl AccessBuilder for AvroAccessBuilder {
50    async fn generate_accessor(
51        &mut self,
52        payload: Vec<u8>,
53        source_meta: &SourceMeta,
54    ) -> ConnectorResult<AccessImpl<'_>> {
55        self.value = self.parse_avro_value(&payload, source_meta).await?;
56        Ok(AccessImpl::Avro(AvroAccess::new(
57            self.value.as_ref().unwrap(),
58            AvroParseOptions::create(&self.schema.original_schema),
59        )))
60    }
61}
62
63impl AvroAccessBuilder {
64    pub fn new(config: AvroParserConfig) -> ConnectorResult<Self> {
65        let AvroParserConfig {
66            schema,
67            writer_schema_cache,
68            ..
69        } = config;
70        Ok(Self {
71            schema,
72            writer_schema_cache,
73            value: None,
74        })
75    }
76
77    /// Note: we should use unresolved schema to parsing bytes into avro value.
78    /// Otherwise it's an invalid schema and parsing will fail. (Avro error: Two named schema defined for same fullname)
79    ///
80    /// # Notes about how Avro data looks like
81    ///
82    /// First, it has two [serialization encodings: binary and JSON](https://avro.apache.org/docs/1.11.1/specification/#encodings).
83    /// They don't have magic bytes and cannot be distinguished on their own.
84    ///
85    /// But in different cases, it starts with different headers, or magic bytes, which can be confusing.
86    ///
87    /// ## `apache_avro` API and headers
88    ///
89    /// - `apache_avro::Reader`: [Object Container Files](https://avro.apache.org/docs/1.11.1/specification/#object-container-files): contains file header, starting with 4 bytes `Obj1`. This is a batch file encoding. We don't use it.
90    /// - `apache_avro::GenericSingleObjectReader`: [Single-object encoding](https://avro.apache.org/docs/1.11.1/specification/#single-object-encoding): starts with 2 bytes `0xC301`. This is designed to be used in places like Kafka, but Confluent schema registry doesn't use it.
91    /// - `apache_avro::from_avro_datum`: no header, binary encoding. This is what we should use.
92    ///
93    /// ## Confluent schema registry
94    ///
95    /// - In Kafka ([Confluent schema registry wire format](https://docs.confluent.io/platform/7.6/schema-registry/fundamentals/serdes-develop/index.html#wire-format)):
96    ///   starts with 5 bytes`0x00{schema_id:08x}` followed by Avro binary encoding.
97    ///
98    /// ## Pulsar schema
99    ///
100    /// - In Pulsar, the payload is raw Avro binary encoding without a schema-registry header.
101    ///   The writer schema version is carried in Pulsar message metadata.
102    async fn parse_avro_value(
103        &self,
104        payload: &[u8],
105        source_meta: &SourceMeta,
106    ) -> ConnectorResult<Option<Value>> {
107        // Each schema backend has its own wire-format contract, so keep payload framing close to
108        // writer-schema lookup.
109        match &self.writer_schema_cache {
110            WriterSchemaCache::Confluent(resolver) => {
111                let (schema_id, mut raw_payload) = extract_schema_id(payload)?;
112                let writer_schema = resolver.get_by_id(schema_id).await?;
113                Ok(Some(from_avro_datum(
114                    writer_schema.as_ref(),
115                    &mut raw_payload,
116                    Some(&self.schema.original_schema),
117                )?))
118            }
119            WriterSchemaCache::File => {
120                // FIXME: we should not use `Reader` (file header) here. See comment above and https://github.com/risingwavelabs/risingwave/issues/12871
121                let mut reader = Reader::with_schema(&self.schema.original_schema, payload)?;
122                match reader.next() {
123                    Some(Ok(v)) => Ok(Some(v)),
124                    Some(Err(e)) => Err(e)?,
125                    None => bail!("avro parse unexpected eof"),
126                }
127            }
128            WriterSchemaCache::Glue(resolver) => {
129                // <https://github.com/awslabs/aws-glue-schema-registry/blob/v1.1.20/common/src/main/java/com/amazonaws/services/schemaregistry/utils/AWSSchemaRegistryConstants.java#L59-L61>
130                // byte 0:      header version = 3
131                // byte 1:      compression: 0 = no compression; 5 = zlib (unsupported)
132                // byte 2..=17: 16-byte UUID as schema version id
133                // byte 18..:   raw avro payload
134                if payload.len() < 18 {
135                    bail!("payload shorter than 18-byte glue header");
136                }
137                if payload[0] != 3 {
138                    bail!(
139                        "Only support glue header version 3 but found {}",
140                        payload[0]
141                    );
142                }
143                if payload[1] != 0 {
144                    bail!("Non-zero compression {} not supported", payload[1]);
145                }
146                let schema_version_id = uuid::Uuid::from_slice(&payload[2..18]).unwrap();
147                let writer_schema = resolver.get_by_id(schema_version_id).await?;
148                let mut raw_payload = &payload[18..];
149                Ok(Some(from_avro_datum(
150                    writer_schema.as_ref(),
151                    &mut raw_payload,
152                    Some(&self.schema.original_schema),
153                )?))
154            }
155            WriterSchemaCache::Pulsar(resolver) => {
156                let SourceMeta::Pulsar(meta) = source_meta else {
157                    bail!("Pulsar Avro parser received non-Pulsar source metadata");
158                };
159                let writer_schema = match meta
160                    .schema_version
161                    .as_deref()
162                    .map(PulsarSchemaVersion::try_from)
163                    .transpose()?
164                {
165                    Some(version) => resolver.get_by_version(version.0).await?,
166                    None => Arc::clone(&self.schema.original_schema),
167                };
168                let mut raw_payload = payload;
169                Ok(Some(from_avro_datum(
170                    writer_schema.as_ref(),
171                    &mut raw_payload,
172                    Some(&self.schema.original_schema),
173                )?))
174            }
175        }
176    }
177}
178
179#[derive(Debug, Clone)]
180pub struct AvroParserConfig {
181    schema: Arc<ResolvedAvroSchema>,
182    /// Writer schema is the schema used to write the data. When parsing Avro data, the exactly same schema
183    /// must be used to decode the message, and then convert it with the reader schema.
184    writer_schema_cache: WriterSchemaCache,
185
186    map_handling: Option<MapHandling>,
187}
188
189#[derive(Debug, Clone)]
190enum WriterSchemaCache {
191    Confluent(Arc<ConfluentSchemaCache>),
192    Glue(Arc<GlueSchemaCacheImpl>),
193    Pulsar(Arc<PulsarSchemaCache>),
194    File,
195}
196
197impl AvroParserConfig {
198    pub async fn new(encoding_properties: EncodingProperties) -> ConnectorResult<Self> {
199        let AvroProperties {
200            schema_location,
201            record_name,
202            key_record_name,
203            map_handling,
204        } = try_match_expand!(encoding_properties, EncodingProperties::Avro)?;
205        match schema_location {
206            SchemaLocation::Confluent {
207                urls: schema_location,
208                client_config,
209                name_strategy,
210                topic,
211            } => {
212                let url = handle_sr_list(schema_location.as_str())?;
213                let client = Client::new(url, &client_config)?;
214                let resolver = ConfluentSchemaCache::new(client);
215
216                if let Some(name) = &key_record_name {
217                    bail!("unused FORMAT ENCODE option: key.message='{name}'");
218                }
219                let subject_value = get_subject_by_strategy(
220                    &name_strategy,
221                    topic.as_str(),
222                    record_name.as_deref(),
223                    false,
224                )?;
225                tracing::debug!("value subject {subject_value}");
226
227                Ok(Self {
228                    schema: Arc::new(ResolvedAvroSchema::create(
229                        resolver.get_by_subject(&subject_value).await?,
230                    )?),
231                    writer_schema_cache: WriterSchemaCache::Confluent(Arc::new(resolver)),
232                    map_handling,
233                })
234            }
235            SchemaLocation::File {
236                url: schema_location,
237                aws_auth_props,
238            } => {
239                let url = handle_sr_list(schema_location.as_str())?;
240                let url = url.first().unwrap();
241                let schema_content = bytes_from_url(url, aws_auth_props.as_ref()).await?;
242                let schema = Schema::parse_reader(&mut schema_content.as_slice())
243                    .context("failed to parse avro schema")?;
244                Ok(Self {
245                    schema: Arc::new(ResolvedAvroSchema::create(Arc::new(schema))?),
246                    writer_schema_cache: WriterSchemaCache::File,
247                    map_handling,
248                })
249            }
250            SchemaLocation::Glue {
251                schema_arn,
252                aws_auth_props,
253                mock_config,
254            } => {
255                let resolver =
256                    GlueSchemaCacheImpl::new(&aws_auth_props, mock_config.as_deref()).await?;
257                let schema = resolver.get_by_name(&schema_arn).await?;
258                Ok(Self {
259                    schema: Arc::new(ResolvedAvroSchema::create(schema)?),
260                    writer_schema_cache: WriterSchemaCache::Glue(Arc::new(resolver)),
261                    map_handling,
262                })
263            }
264            SchemaLocation::Pulsar {
265                client_config,
266                topic,
267            } => {
268                let client = crate::schema::pulsar_schema::Client::new(&client_config)?;
269                let resolver = PulsarSchemaCache::new(client, &topic)?;
270                let schema = resolver.get_latest().await?;
271                Ok(Self {
272                    schema: Arc::new(ResolvedAvroSchema::create(schema)?),
273                    writer_schema_cache: WriterSchemaCache::Pulsar(Arc::new(resolver)),
274                    map_handling,
275                })
276            }
277        }
278    }
279
280    pub fn map_to_columns(&self) -> ConnectorResult<Vec<Field>> {
281        avro_schema_to_fields(&self.schema.original_schema, self.map_handling).map_err(Into::into)
282    }
283}
284
285#[cfg(test)]
286mod test {
287    use std::env;
288
289    use url::Url;
290
291    use super::*;
292    use crate::connector_common::AwsAuthProps;
293
294    fn test_data_path(file_name: &str) -> String {
295        let curr_dir = env::current_dir().unwrap().into_os_string();
296        curr_dir.into_string().unwrap() + "/src/test_data/" + file_name
297    }
298
299    #[tokio::test]
300    #[ignore]
301    async fn test_load_schema_from_s3() {
302        let schema_location = "s3://mingchao-schemas/complex-schema.avsc".to_owned();
303        let url = Url::parse(&schema_location).unwrap();
304        let aws_auth_config: AwsAuthProps =
305            serde_json::from_str(r#"region":"ap-southeast-1"#).unwrap();
306        let schema_content = bytes_from_url(&url, Some(&aws_auth_config)).await;
307        assert!(schema_content.is_ok());
308        let schema = Schema::parse_reader(&mut schema_content.unwrap().as_slice());
309        assert!(schema.is_ok());
310        println!("schema = {:?}", schema.unwrap());
311    }
312
313    #[tokio::test]
314    async fn test_load_schema_from_local() {
315        let schema_location = Url::from_file_path(test_data_path("complex-schema.avsc")).unwrap();
316        let schema_content = bytes_from_url(&schema_location, None).await;
317        assert!(schema_content.is_ok());
318        let schema = Schema::parse_reader(&mut schema_content.unwrap().as_slice());
319        assert!(schema.is_ok());
320        println!("schema = {:?}", schema.unwrap());
321    }
322
323    #[tokio::test]
324    #[ignore]
325    async fn test_load_schema_from_https() {
326        let schema_location =
327            "https://mingchao-schemas.s3.ap-southeast-1.amazonaws.com/complex-schema.avsc";
328        let url = Url::parse(schema_location).unwrap();
329        let schema_content = bytes_from_url(&url, None).await;
330        assert!(schema_content.is_ok());
331        let schema = Schema::parse_reader(&mut schema_content.unwrap().as_slice());
332        assert!(schema.is_ok());
333        println!("schema = {:?}", schema.unwrap());
334    }
335}
risingwave_connector/parser/avro/parser.rs

risingwave_connector/parser/avro/
parser.rs
