risingwave_connector/parser/protobuf/
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::collections::HashSet;
16
17use anyhow::Context;
18use prost_reflect::{DescriptorPool, DynamicMessage, FileDescriptor, MessageDescriptor};
19use risingwave_common::catalog::Field;
20use risingwave_common::{bail, try_match_expand};
21use risingwave_connector_codec::decoder::protobuf::ProtobufAccess;
22pub use risingwave_connector_codec::decoder::protobuf::parser::{PROTOBUF_MESSAGES_AS_JSONB, *};
23
24use crate::error::ConnectorResult;
25use crate::parser::unified::AccessImpl;
26use crate::parser::utils::bytes_from_url;
27use crate::parser::{AccessBuilder, EncodingProperties, SchemaLocation};
28use crate::schema::schema_registry::{Client, WireFormatError, extract_schema_id, handle_sr_list};
29use crate::schema::{
30    ConfluentSchemaLoader, InvalidOptionError, SchemaLoader, bail_invalid_option_error,
31};
32
33#[derive(Debug)]
34pub struct ProtobufAccessBuilder {
35    wire_type: WireType,
36    message_descriptor: MessageDescriptor,
37
38    // A HashSet containing protobuf message type full names (e.g. "google.protobuf.Any")
39    // that should be mapped to JSONB type when storing in RisingWave
40    messages_as_jsonb: HashSet<String>,
41}
42
43impl AccessBuilder for ProtobufAccessBuilder {
44    async fn generate_accessor(
45        &mut self,
46        payload: Vec<u8>,
47        _: &crate::source::SourceMeta,
48    ) -> ConnectorResult<AccessImpl<'_>> {
49        let payload = match self.wire_type {
50            WireType::Confluent => resolve_pb_header(&payload)?,
51            WireType::None => &payload,
52        };
53
54        let message = DynamicMessage::decode(self.message_descriptor.clone(), payload)
55            .context("failed to parse message")?;
56
57        Ok(AccessImpl::Protobuf(ProtobufAccess::new(
58            message,
59            &self.messages_as_jsonb,
60        )))
61    }
62}
63
64impl ProtobufAccessBuilder {
65    pub fn new(config: ProtobufParserConfig) -> ConnectorResult<Self> {
66        let ProtobufParserConfig {
67            wire_type,
68            message_descriptor,
69            messages_as_jsonb,
70        } = config;
71
72        Ok(Self {
73            wire_type,
74            message_descriptor,
75            messages_as_jsonb,
76        })
77    }
78}
79
80#[derive(Debug, Clone)]
81enum WireType {
82    None,
83    Confluent,
84    // Glue,
85    // Pulsar,
86}
87
88impl TryFrom<&SchemaLocation> for WireType {
89    type Error = InvalidOptionError;
90
91    fn try_from(value: &SchemaLocation) -> Result<Self, Self::Error> {
92        match value {
93            SchemaLocation::File { .. } => Ok(Self::None),
94            SchemaLocation::Confluent { .. } => Ok(Self::Confluent),
95            SchemaLocation::Glue { .. } => bail_invalid_option_error!(
96                "encode protobuf from aws glue schema registry not supported yet"
97            ),
98            SchemaLocation::Pulsar { .. } => bail_invalid_option_error!(
99                "encode protobuf from Pulsar schema registry not supported yet"
100            ),
101        }
102    }
103}
104
105#[derive(Debug, Clone)]
106pub struct ProtobufParserConfig {
107    wire_type: WireType,
108    pub(crate) message_descriptor: MessageDescriptor,
109    messages_as_jsonb: HashSet<String>,
110}
111
112impl ProtobufParserConfig {
113    pub async fn new(encoding_properties: EncodingProperties) -> ConnectorResult<Self> {
114        let protobuf_config = try_match_expand!(encoding_properties, EncodingProperties::Protobuf)?;
115        let message_name = &protobuf_config.message_name;
116
117        let wire_type = (&protobuf_config.schema_location).try_into()?;
118        if protobuf_config.key_message_name.is_some() {
119            // https://docs.confluent.io/platform/7.5/control-center/topics/schema.html#c3-schemas-best-practices-key-value-pairs
120            bail!("protobuf key is not supported");
121        }
122        let pool = match protobuf_config.schema_location {
123            SchemaLocation::Confluent {
124                urls,
125                client_config,
126                name_strategy,
127                topic,
128            } => {
129                let url = handle_sr_list(urls.as_str())?;
130                let client = Client::new(url, &client_config)?;
131                let loader = SchemaLoader::Confluent(ConfluentSchemaLoader {
132                    client,
133                    name_strategy,
134                    topic,
135                    key_record_name: None,
136                    val_record_name: Some(message_name.clone()),
137                });
138                let (_schema_id, root_file_descriptor) = loader
139                    .load_val_schema::<FileDescriptor>()
140                    .await
141                    .context("load schema failed")?;
142                root_file_descriptor.parent_pool().clone()
143            }
144            SchemaLocation::File {
145                url,
146                aws_auth_props,
147            } => {
148                let url = handle_sr_list(url.as_str())?;
149                let url = url.first().unwrap();
150                let schema_bytes = bytes_from_url(url, aws_auth_props.as_ref()).await?;
151                DescriptorPool::decode(schema_bytes.as_slice())
152                    .with_context(|| format!("cannot build descriptor pool from schema `{url}`"))?
153            }
154            SchemaLocation::Glue { .. } => bail_invalid_option_error!(
155                "encode protobuf from aws glue schema registry not supported yet"
156            ),
157            SchemaLocation::Pulsar { .. } => bail_invalid_option_error!(
158                "encode protobuf from Pulsar schema registry not supported yet"
159            ),
160        };
161
162        let message_descriptor = pool
163            .get_message_by_name(message_name)
164            .with_context(|| format!("cannot find message `{message_name}` in schema"))?;
165
166        Ok(Self {
167            message_descriptor,
168            wire_type,
169            messages_as_jsonb: protobuf_config.messages_as_jsonb,
170        })
171    }
172
173    /// Maps the protobuf schema to relational schema.
174    pub fn map_to_columns(&self) -> ConnectorResult<Vec<Field>> {
175        pb_schema_to_fields(&self.message_descriptor, &self.messages_as_jsonb).map_err(|e| e.into())
176    }
177}
178
179/// A port from the implementation of confluent's Varint Zig-zag deserialization.
180/// See `ReadVarint` in <https://github.com/apache/kafka/blob/trunk/clients/src/main/java/org/apache/kafka/common/utils/ByteUtils.java>
181fn decode_varint_zigzag(buffer: &[u8]) -> ConnectorResult<(i32, usize)> {
182    // We expect the decoded number to be 4 bytes.
183    let mut value = 0u32;
184    let mut shift = 0;
185    let mut len = 0usize;
186
187    for &byte in buffer {
188        len += 1;
189        // The Varint encoding is limited to 5 bytes.
190        if len > 5 {
191            break;
192        }
193        // The byte is cast to u32 to avoid shifting overflow.
194        let byte_ext = byte as u32;
195        // In Varint encoding, the lowest 7 bits are used to represent number,
196        // while the highest zero bit indicates the end of the number with Varint encoding.
197        value |= (byte_ext & 0x7F) << shift;
198        if byte_ext & 0x80 == 0 {
199            return Ok((((value >> 1) as i32) ^ -((value & 1) as i32), len));
200        }
201
202        shift += 7;
203    }
204
205    Err(WireFormatError::ParseMessageIndexes.into())
206}
207
208/// Reference: <https://docs.confluent.io/platform/current/schema-registry/fundamentals/serdes-develop/index.html#wire-format>
209/// Wire format for Confluent pb header is:
210/// | 0          | 1-4        | 5-x             | x+1-end
211/// | magic-byte | schema-id  | message-indexes | protobuf-payload
212pub(crate) fn resolve_pb_header(payload: &[u8]) -> ConnectorResult<&[u8]> {
213    // there's a message index array at the front of payload
214    // if it is the first message in proto def, the array is just and `0`
215    let (_, remained) = extract_schema_id(payload)?;
216    // The message indexes are encoded as int using variable-length zig-zag encoding,
217    // prefixed by the length of the array.
218    // Note that if the first byte is 0, it is equivalent to (1, 0) as an optimization.
219    match remained.first() {
220        Some(0) => Ok(&remained[1..]),
221        Some(_) => {
222            let (index_len, mut offset) = decode_varint_zigzag(remained)?;
223            for _ in 0..index_len {
224                offset += decode_varint_zigzag(&remained[offset..])?.1;
225            }
226            Ok(&remained[offset..])
227        }
228        None => bail!("The proto payload is empty"),
229    }
230}
231
232#[cfg(test)]
233mod test {
234    use super::*;
235
236    #[test]
237    fn test_decode_varint_zigzag() {
238        // 1. Positive number
239        let buffer = vec![0x02];
240        let (value, len) = decode_varint_zigzag(&buffer).unwrap();
241        assert_eq!(value, 1);
242        assert_eq!(len, 1);
243
244        // 2. Negative number
245        let buffer = vec![0x01];
246        let (value, len) = decode_varint_zigzag(&buffer).unwrap();
247        assert_eq!(value, -1);
248        assert_eq!(len, 1);
249
250        // 3. Larger positive number
251        let buffer = vec![0x9E, 0x03];
252        let (value, len) = decode_varint_zigzag(&buffer).unwrap();
253        assert_eq!(value, 207);
254        assert_eq!(len, 2);
255
256        // 4. Larger negative number
257        let buffer = vec![0xBF, 0x07];
258        let (value, len) = decode_varint_zigzag(&buffer).unwrap();
259        assert_eq!(value, -480);
260        assert_eq!(len, 2);
261
262        // 5. Maximum positive number
263        let buffer = vec![0xFE, 0xFF, 0xFF, 0xFF, 0x0F];
264        let (value, len) = decode_varint_zigzag(&buffer).unwrap();
265        assert_eq!(value, i32::MAX);
266        assert_eq!(len, 5);
267
268        // 6. Maximum negative number
269        let buffer = vec![0xFF, 0xFF, 0xFF, 0xFF, 0x0F];
270        let (value, len) = decode_varint_zigzag(&buffer).unwrap();
271        assert_eq!(value, i32::MIN);
272        assert_eq!(len, 5);
273
274        // 7. More than 32 bits
275        let buffer = vec![0xFF, 0xFF, 0xFF, 0xFF, 0x7F];
276        let (value, len) = decode_varint_zigzag(&buffer).unwrap();
277        assert_eq!(value, i32::MIN);
278        assert_eq!(len, 5);
279
280        // 8. Invalid input (more than 5 bytes)
281        let buffer = vec![0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF];
282        let result = decode_varint_zigzag(&buffer);
283        assert!(result.is_err());
284    }
285}
risingwave_connector/parser/protobuf/parser.rs

risingwave_connector/parser/protobuf/
parser.rs
