load protobuf from schema registry

Author: xiangjinwu

src/connector/src/schema/protobuf.rs

@@ -17,28 +17,46 @@ use std::collections::BTreeMap;
 use itertools::Itertools as _;
 use prost_reflect::{DescriptorPool, FileDescriptor, MessageDescriptor};
 
-use super::loader::LoadedSchema;
+use super::loader::{LoadedSchema, SchemaLoader};
 use super::schema_registry::Subject;
 use super::{
     invalid_option_error, InvalidOptionError, SchemaFetchError, MESSAGE_NAME_KEY,
-    SCHEMA_LOCATION_KEY,
+    SCHEMA_LOCATION_KEY, SCHEMA_REGISTRY_KEY,
 };
 use crate::common::AwsAuthProps;
 use crate::parser::{EncodingProperties, ProtobufParserConfig, ProtobufProperties};
 
 /// `aws_auth_props` is only required when reading `s3://` URL.
 pub async fn fetch_descriptor(
     format_options: &BTreeMap<String, String>,
+    topic: &str,
     aws_auth_props: Option<&AwsAuthProps>,
-) -> Result<MessageDescriptor, SchemaFetchError> {
-    let row_schema_location = format_options
-        .get(SCHEMA_LOCATION_KEY)
-        .ok_or_else(|| invalid_option_error!("{SCHEMA_LOCATION_KEY} required"))?
-        .clone();
+) -> Result<(MessageDescriptor, Option<i32>), SchemaFetchError> {
     let message_name = format_options
         .get(MESSAGE_NAME_KEY)
         .ok_or_else(|| invalid_option_error!("{MESSAGE_NAME_KEY} required"))?
         .clone();
+    let schema_location = format_options.get(SCHEMA_LOCATION_KEY);
+    let schema_registry = format_options.get(SCHEMA_REGISTRY_KEY);
+    let row_schema_location = match (schema_location, schema_registry) {
+        (Some(_), Some(_)) => {
+            return Err(invalid_option_error!(
+                "cannot use {SCHEMA_LOCATION_KEY} and {SCHEMA_REGISTRY_KEY} together"
+            )
+            .into())
+        }
+        (None, None) => {
+            return Err(invalid_option_error!(
+                "requires one of {SCHEMA_LOCATION_KEY} or {SCHEMA_REGISTRY_KEY}"
+            )
+            .into())
+        }
+        (None, Some(_)) => {
+            let (md, sid) = fetch_from_registry(&message_name, format_options, topic).await?;
+            return Ok((md, Some(sid)));
+        }
+        (Some(url), None) => url.clone(),
+    };
 
     if row_schema_location.starts_with("s3") && aws_auth_props.is_none() {
         return Err(invalid_option_error!("s3 URL not supported yet").into());
@@ -59,7 +77,23 @@ pub async fn fetch_descriptor(
     let conf = ProtobufParserConfig::new(enc)
         .await
         .map_err(SchemaFetchError::YetToMigrate)?;
-    Ok(conf.message_descriptor)
+    Ok((conf.message_descriptor, None))
+}
+
+pub async fn fetch_from_registry(
+    message_name: &str,
+    format_options: &BTreeMap<String, String>,
+    topic: &str,
+) -> Result<(MessageDescriptor, i32), SchemaFetchError> {
+    let loader = SchemaLoader::from_format_options(topic, format_options)?;
+
+    let (vid, vpb) = loader.load_val_schema::<FileDescriptor>().await?;
+
+    // TODO: why file rather than pool or message?
+    Ok((
+        vpb.parent_pool().get_message_by_name(message_name).unwrap(),
+        vid,
+    ))
 }
 
 impl LoadedSchema for FileDescriptor {

src/connector/src/sink/formatter/mod.rs

@@ -123,12 +123,18 @@ impl SinkFormatterImpl {
                     }
                     SinkEncode::Protobuf => {
                         // By passing `None` as `aws_auth_props`, reading from `s3://` not supported yet.
-                        let descriptor =
-                            crate::schema::protobuf::fetch_descriptor(&format_desc.options, None)
-                                .await
-                                .map_err(|e| SinkError::Config(anyhow!(e)))?;
-                        let val_encoder =
-                            ProtoEncoder::new(schema, None, descriptor, ProtoHeader::None)?;
+                        let (descriptor, sid) = crate::schema::protobuf::fetch_descriptor(
+                            &format_desc.options,
+                            topic,
+                            None,
+                        )
+                        .await
+                        .map_err(|e| SinkError::Config(anyhow!(e)))?;
+                        let header = match sid {
+                            None => ProtoHeader::None,
+                            Some(sid) => ProtoHeader::ConfluentSchemaRegistry(sid),
+                        };
+                        let val_encoder = ProtoEncoder::new(schema, None, descriptor, header)?;
                         let formatter = AppendOnlyFormatter::new(key_encoder, val_encoder);
                         Ok(SinkFormatterImpl::AppendOnlyProto(formatter))
                     }