Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

enhancement(codecs)!: Use DescriptorPool from ProtobufDeserializer #901

Draft
wants to merge 4 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 20 additions & 6 deletions src/protobuf/descriptor.rs
Original file line number Diff line number Diff line change
@@ -1,17 +1,31 @@
use prost_reflect::{DescriptorPool, MessageDescriptor};
use std::path::Path;

pub fn get_message_descriptor(
pub fn get_message_pool_descriptor(
descriptor_set_path: &Path,
message_type: &str,
) -> std::result::Result<MessageDescriptor, String> {
) -> std::result::Result<DescriptorPool, String> {
let b = std::fs::read(descriptor_set_path).map_err(|e| {
format!("Failed to open protobuf desc file '{descriptor_set_path:?}': {e}",)
})?;
let pool = DescriptorPool::decode(b.as_slice()).map_err(|e| {
format!("Failed to parse protobuf desc file '{descriptor_set_path:?}': {e}")
})?;
DescriptorPool::decode(b.as_slice())
.map_err(|e| format!("Failed to parse protobuf desc file '{descriptor_set_path:?}': {e}"))
}

pub fn get_message_descriptor(
descriptor_set_path: &Path,
message_type: &str,
) -> std::result::Result<MessageDescriptor, String> {
let pool = get_message_pool_descriptor(descriptor_set_path)?;
pool.get_message_by_name(message_type).ok_or_else(|| {
format!("The message type '{message_type}' could not be found in '{descriptor_set_path:?}'")
})
}

pub fn get_message_descriptor_from_pool(
pool: &DescriptorPool,
message_type: &str,
) -> std::result::Result<MessageDescriptor, String> {
pool.get_message_by_name(message_type).ok_or_else(|| {
format!("The message type '{message_type}' could not be found in the descriptor pool")
})
}
2 changes: 2 additions & 0 deletions src/protobuf/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@ mod encode;
mod parse;

pub use descriptor::get_message_descriptor;
pub use descriptor::get_message_descriptor_from_pool;
pub use descriptor::get_message_pool_descriptor;

pub use encode::encode_message;
pub(crate) use encode::encode_proto;
Expand Down
49 changes: 46 additions & 3 deletions src/protobuf/parse.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,12 @@
use crate::compiler::prelude::*;
use crate::protobuf::get_message_descriptor_from_pool;
use prost_reflect::ReflectMessage;
use prost_reflect::{DynamicMessage, MessageDescriptor};

use tracing::debug;

pub fn proto_to_value(
descriptor_pool: Option<&prost_reflect::DescriptorPool>,
prost_reflect_value: &prost_reflect::Value,
field_descriptor: Option<&prost_reflect::FieldDescriptor>,
) -> std::result::Result<Value, String> {
Expand Down Expand Up @@ -42,11 +46,49 @@ pub fn proto_to_value(
}
}
prost_reflect::Value::Message(v) => {
if let Some(descriptor_pool) = descriptor_pool {
if let Some(type_url_cow) = v.get_field_by_name("type_url") {
if let Some(value_cow) = v.get_field_by_name("value") {
if let prost_reflect::Value::String(type_url) = &*type_url_cow {
if let prost_reflect::Value::Bytes(value) = &*value_cow {
let type_name = type_url.trim_start_matches("type.googleapis.com/");
match get_message_descriptor_from_pool(descriptor_pool, type_name) {
Comment on lines +49 to +55
Copy link
Contributor

@pront pront Jul 2, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a lot of nesting, can you explain the motivation and the use case for this? I am also worried about the hardcoded values in here.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was a naive implementation for checking the presence of fields.

The FieldDescriptor actually have a type_name. If we can check if the given field of this type, we can decode it straight away:

    field {
      name: "anydata"
      number: 3
      label: LABEL_OPTIONAL
      type: TYPE_MESSAGE
      type_name: ".google.protobuf.Any"
      json_name: "anydata"
    }

Ok(message_descriptor) => {
match DynamicMessage::decode(
message_descriptor,
value.clone(),
) {
Ok(dynamic_message) => {
return proto_to_value(
Some(descriptor_pool),
&prost_reflect::Value::Message(dynamic_message),
None,
);
}
Err(error) => {
return Err(format!(
"Error parsing embedded protobuf message: {:?}",
error
));
}
}
}
Err(_) => {
debug!("Message type '{}' not found in the descriptor pool, ignoring message.", type_name);
}
}
}
}
}
}
}

let mut obj_map = ObjectMap::new();
for field_desc in v.descriptor().fields() {
if v.has_field(&field_desc) {
let field_value = v.get_field(&field_desc);
let out = proto_to_value(field_value.as_ref(), Some(&field_desc))?;
let out =
proto_to_value(descriptor_pool, field_value.as_ref(), Some(&field_desc))?;
obj_map.insert(field_desc.name().into(), out);
}
}
Expand All @@ -55,7 +97,7 @@ pub fn proto_to_value(
prost_reflect::Value::List(v) => {
let vec = v
.iter()
.map(|o| proto_to_value(o, field_descriptor))
.map(|o| proto_to_value(descriptor_pool, o, field_descriptor))
.collect::<Result<Vec<_>, String>>()?;
Value::from(vec)
}
Expand All @@ -80,7 +122,7 @@ pub fn proto_to_value(
)
})?
.into(),
proto_to_value(kv.1, Some(&message_desc.map_entry_value_field()))?,
proto_to_value(descriptor_pool, kv.1, Some(&message_desc.map_entry_value_field()))?,
))
})
.collect::<std::result::Result<ObjectMap, String>>()?,
Expand All @@ -99,6 +141,7 @@ pub(crate) fn parse_proto(descriptor: &MessageDescriptor, value: Value) -> Resol
let dynamic_message = DynamicMessage::decode(descriptor.clone(), bytes)
.map_err(|error| format!("Error parsing protobuf: {:?}", error))?;
Ok(proto_to_value(
None,
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we declare a singleton method to retrieve global configuration with loaded descriptor pool, we can use it here (when specified)

&prost_reflect::Value::Message(dynamic_message),
None,
)?)
Expand Down
Loading