Skip to content

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

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.

Sign up for GitHub

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

Draft
wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

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

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
10cf878
enhancement(codecs): Use DescriptorPool from ProtobufDeserializer
nb-mouse Jun 21, 2024
a3ed24f
Do not fail on missed MessageDescriptor
nb-mouse Jun 21, 2024
4f4ea92
Merge remote-tracking branch 'origin' into protobuf-descriptor-pool
pront Jul 2, 2024
e19891d
cargo fmt
pront Jul 2, 2024
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
25 changes: 20 additions & 5 deletions src/protobuf/descriptor.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,17 +1,32 @@
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| {
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
View file
Open in desktop
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_pool_descriptor;
pub use descriptor::get_message_descriptor_from_pool;

pub use encode::encode_message;
pub(crate) use encode::encode_proto;
Expand Down
38 changes: 35 additions & 3 deletions src/protobuf/parse.rs
View file
Open in desktop
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,38 @@ 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
Member

@pront pront Jul 2, 2024 •
edited
Loading

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))?;
Copy link
Author

Choose a reason for hiding this comment

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

There is a possible improvement if fields in descriptor is ordered and type_url comes before value field then logic from above can be moved here.

obj_map.insert(field_desc.name().into(), out);
}
}
Expand All @@ -55,7 +86,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 +111,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 +130,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