feat(grpc): preserve comments from protobuf when generating config (t…

…ailcallhq#2023)

src/core/generator/from_proto.rs

@@ -3,10 +3,13 @@ use std::collections::{BTreeSet, HashSet};
 use anyhow::{bail, Result};
 use derive_setters::Setters;
 use prost_reflect::prost_types::{
-    DescriptorProto, EnumDescriptorProto, FileDescriptorSet, ServiceDescriptorProto,
+    DescriptorProto, EnumDescriptorProto, FileDescriptorSet, ServiceDescriptorProto, SourceCodeInfo,
 };
 
 use super::graphql_type::{GraphQLType, Unparsed};
+use super::proto::comments_builder::CommentsBuilder;
+use super::proto::path_builder::PathBuilder;
+use super::proto::path_field::PathField;
 use crate::core::config::{Arg, Config, Enum, Field, Grpc, Tag, Type};
 
 /// Assists in the mapping and retrieval of proto type names to custom formatted
@@ -24,6 +27,10 @@ struct Context {
 
     /// Set of visited map types
     map_types: HashSet<String>,
+
+    /// Optional field to store source code information, including comments, for
+    /// each entity.
+    comments_builder: CommentsBuilder,
 }
 
 impl Context {
@@ -33,40 +40,84 @@ impl Context {
             namespace: Default::default(),
             config: Default::default(),
             map_types: Default::default(),
+            comments_builder: CommentsBuilder::new(None),
         }
     }
 
+    /// Sets source code information for preservation of comments.
+    fn with_source_code_info(mut self, source_code_info: SourceCodeInfo) -> Self {
+        self.comments_builder = CommentsBuilder::new(Some(source_code_info));
+        self
+    }
+
     /// Resolves the actual name and inserts the type.
     fn insert_type(mut self, name: String, ty: Type) -> Self {
         self.config.types.insert(name.to_string(), ty);
         self
     }
 
     /// Processes proto enum types.
-    fn append_enums(mut self, enums: &Vec<EnumDescriptorProto>) -> Self {
-        for enum_ in enums {
+    fn append_enums(
+        mut self,
+        enums: &[EnumDescriptorProto],
+        parent_path: &PathBuilder,
+        is_nested: bool,
+    ) -> Self {
+        for (index, enum_) in enums.iter().enumerate() {
             let enum_name = enum_.name();
 
-            let variants = enum_
-                .value
-                .iter()
-                .map(|v| GraphQLType::new(v.name()).into_enum_variant().to_string())
-                .collect::<BTreeSet<String>>();
+            let enum_type_path = if is_nested {
+                parent_path.extend(PathField::NestedEnum, index as i32)
+            } else {
+                parent_path.extend(PathField::EnumType, index as i32)
+            };
+
+            let mut variants_with_comments = BTreeSet::new();
+
+            for (value_index, v) in enum_.value.iter().enumerate() {
+                let variant_name = GraphQLType::new(v.name()).into_enum_variant().to_string();
+
+                // Path to the enum value's comments
+                let value_path = PathBuilder::new(&enum_type_path)
+                    .extend(PathField::EnumValue, value_index as i32); // 2: value field
+
+                // Get comments for the enum value
+                let comment = self.comments_builder.get_comments(&value_path);
+
+                // Format the variant with its comment as description
+                if let Some(comment) = comment {
+                    // TODO: better support for enum variant descriptions [There is no way to define
+                    // description for enum variant in current config structure]
+                    let variant_with_comment =
+                        format!("\"\"\n  {}\n  \"\"\n  {}", comment, variant_name);
+                    variants_with_comments.insert(variant_with_comment);
+                } else {
+                    variants_with_comments.insert(variant_name);
+                }
+            }
 
             let type_name = GraphQLType::new(enum_name)
                 .extend(self.namespace.as_slice())
                 .into_enum()
                 .to_string();
+
+            let doc = self.comments_builder.get_comments(&enum_type_path);
+
             self.config
                 .enums
-                .insert(type_name, Enum { variants, doc: None });
+                .insert(type_name, Enum { variants: variants_with_comments, doc });
         }
         self
     }
 
     /// Processes proto message types.
-    fn append_msg_type(mut self, messages: &Vec<DescriptorProto>) -> Result<Self> {
-        for message in messages {
+    fn append_msg_type(
+        mut self,
+        messages: &[DescriptorProto],
+        parent_path: &PathBuilder,
+        is_nested: bool,
+    ) -> Result<Self> {
+        for (index, message) in messages.iter().enumerate() {
             let msg_name = message.name();
 
             let msg_type = GraphQLType::new(msg_name)
@@ -87,15 +138,26 @@ impl Context {
                 continue;
             }
 
+            let msg_path = if is_nested {
+                parent_path.extend(PathField::NestedType, index as i32)
+            } else {
+                parent_path.extend(PathField::MessageType, index as i32)
+            };
+
             // first append the name of current message as namespace
             self.namespace.push(msg_name.to_string());
-            self = self.append_enums(&message.enum_type);
-            self = self.append_msg_type(&message.nested_type)?;
+            self = self.append_enums(&message.enum_type, &PathBuilder::new(&msg_path), true);
+            self =
+                self.append_msg_type(&message.nested_type, &PathBuilder::new(&msg_path), true)?;
             // then drop it after handling nested types
             self.namespace.pop();
 
-            let mut ty = Type::default();
-            for field in message.field.iter() {
+            let mut ty = Type {
+                doc: self.comments_builder.get_comments(&msg_path),
+                ..Default::default()
+            };
+
+            for (field_index, field) in message.field.iter().enumerate() {
                 let field_name = GraphQLType::new(field.name())
                     .extend(self.namespace.as_slice())
                     .into_field();
@@ -131,6 +193,10 @@ impl Context {
                     cfg_field.type_of = type_of;
                 }
 
+                let field_path =
+                    PathBuilder::new(&msg_path).extend(PathField::Field, field_index as i32);
+                cfg_field.doc = self.comments_builder.get_comments(&field_path);
+
                 ty.fields.insert(field_name.to_string(), cfg_field);
             }
 
@@ -142,14 +208,20 @@ impl Context {
     }
 
     /// Processes proto service definitions and their methods.
-    fn append_query_service(mut self, services: &Vec<ServiceDescriptorProto>) -> Result<Self> {
+    fn append_query_service(
+        mut self,
+        services: &[ServiceDescriptorProto],
+        parent_path: &PathBuilder,
+    ) -> Result<Self> {
         if services.is_empty() {
             return Ok(self);
         }
 
-        for service in services {
+        for (index, service) in services.iter().enumerate() {
             let service_name = service.name();
-            for method in &service.method {
+            let path = parent_path.extend(PathField::Service, index as i32);
+
+            for (method_index, method) in service.method.iter().enumerate() {
                 let field_name = GraphQLType::new(method.name())
                     .extend(self.namespace.as_slice())
                     .push(service_name)
@@ -190,6 +262,10 @@ impl Context {
                     method: field_name.id(),
                 });
 
+                let method_path =
+                    PathBuilder::new(&path).extend(PathField::Method, method_index as i32);
+                cfg_field.doc = self.comments_builder.get_comments(&method_path);
+
                 let ty = self
                     .config
                     .types
@@ -263,10 +339,16 @@ pub fn from_proto(descriptor_sets: &[FileDescriptorSet], query: &str) -> Result<
         for file_descriptor in descriptor_set.file.iter() {
             ctx.namespace = vec![file_descriptor.package().to_string()];
 
+            if let Some(source_code_info) = &file_descriptor.source_code_info {
+                ctx = ctx.with_source_code_info(source_code_info.clone());
+            }
+
+            let root_path = PathBuilder::new(&[]);
+
             ctx = ctx
-                .append_enums(&file_descriptor.enum_type)
-                .append_msg_type(&file_descriptor.message_type)?
-                .append_query_service(&file_descriptor.service)?;
+                .append_enums(&file_descriptor.enum_type, &root_path, false)
+                .append_msg_type(&file_descriptor.message_type, &root_path, false)?
+                .append_query_service(&file_descriptor.service, &root_path)?;
         }
     }
 

src/core/generator/mod.rs

@@ -1,6 +1,7 @@
 mod from_proto;
 mod generator;
 mod graphql_type;
+mod proto;
 mod source;
 pub use generator::Generator;
 pub use source::Source;

src/core/generator/proto/comments_builder.rs

@@ -0,0 +1,28 @@
+use prost_reflect::prost_types::SourceCodeInfo;
+
+pub struct CommentsBuilder {
+    source_code_info: Option<SourceCodeInfo>,
+}
+
+impl CommentsBuilder {
+    pub fn new(source_code_info: Option<SourceCodeInfo>) -> Self {
+        Self { source_code_info }
+    }
+
+    pub fn get_comments(&self, path: &[i32]) -> Option<String> {
+        self.source_code_info.as_ref().and_then(|info| {
+            info.location
+                .iter()
+                .find(|loc| loc.path == path)
+                .and_then(|loc| {
+                    loc.leading_comments.as_ref().map(|c| {
+                        c.lines()
+                            .map(|line| line.trim_start_matches('*').trim())
+                            .filter(|line| !line.is_empty())
+                            .collect::<Vec<_>>()
+                            .join("\n  ")
+                    })
+                })
+        })
+    }
+}

src/core/generator/proto/mod.rs

@@ -0,0 +1,3 @@
+pub mod comments_builder;
+pub mod path_builder;
+pub mod path_field;

src/core/generator/proto/path_builder.rs

@@ -0,0 +1,18 @@
+use super::path_field::PathField;
+
+pub struct PathBuilder {
+    base_path: Vec<i32>,
+}
+
+impl PathBuilder {
+    pub fn new(base_path: &[i32]) -> Self {
+        Self { base_path: base_path.to_vec() }
+    }
+
+    pub fn extend(&self, field: PathField, index: i32) -> Vec<i32> {
+        let mut extended_path = self.base_path.clone();
+        extended_path.push(field.value());
+        extended_path.push(index);
+        extended_path
+    }
+}

src/core/generator/proto/path_field.rs

@@ -0,0 +1,25 @@
+pub enum PathField {
+    EnumType,
+    MessageType,
+    Service,
+    Field,
+    Method,
+    EnumValue,
+    NestedType,
+    NestedEnum,
+}
+
+impl PathField {
+    pub fn value(&self) -> i32 {
+        match self {
+            PathField::EnumType => 5,
+            PathField::MessageType => 4,
+            PathField::Service => 6,
+            PathField::Field => 2,
+            PathField::Method => 2,
+            PathField::EnumValue => 2,
+            PathField::NestedType => 3,
+            PathField::NestedEnum => 4,
+        }
+    }
+}

src/core/generator/snapshots/tailcall__core__generator__from_proto__test__from_proto.snap

@@ -10,6 +10,9 @@ input greetings__HelloRequest @tag(id: "greetings.HelloRequest") {
   name: String
 }
 
+  """
+  The request message containing the user's name.
+  """
 input greetings_a__b__HelloRequest @tag(id: "greetings_a.b.HelloRequest") {
   name: String
 }
@@ -37,7 +40,13 @@ enum news__Status {
 }
 
 type Query {
+  """
+  Sends a greeting
+  """
   greetings_a__b__Greeter__SayHello(helloRequest: greetings_a__b__HelloRequest!): greetings_a__b__HelloReply! @grpc(body: "{{.args.helloRequest}}", method: "greetings_a.b.Greeter.SayHello")
+  """
+  Sends a greeting
+  """
   greetings_b__c__Greeter__SayHello(helloRequest: greetings__HelloRequest!): greetings__HelloReply! @grpc(body: "{{.args.helloRequest}}", method: "greetings_b.c.Greeter.SayHello")
   news__NewsService__AddNews(news: news__NewsInput!): news__News! @grpc(body: "{{.args.news}}", method: "news.NewsService.AddNews")
   news__NewsService__DeleteNews(newsId: news__NewsId!): Empty! @grpc(body: "{{.args.newsId}}", method: "news.NewsService.DeleteNews")
@@ -51,6 +60,9 @@ type greetings__HelloReply @tag(id: "greetings.HelloReply") {
   message: String
 }
 
+  """
+  The response message containing the greetings
+  """
 type greetings_a__b__HelloReply @tag(id: "greetings_a.b.HelloReply") {
   message: String
 }