feat(build): Support compiling well-known protobuf types (#522)

tdyas · web-flow · commit 61555ff2b5b7 · 2021-02-23T16:21:03.000-05:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -14,6 +14,7 @@ members = [
     "tests/included_service",
     "tests/same_name",
     "tests/wellknown",
+    "tests/wellknown-compiled",
     "tests/extern_path/uuid",
     "tests/ambiguous_methods",
     "tests/extern_path/my_application",
diff --git a/tests/wellknown-compiled/Cargo.toml b/tests/wellknown-compiled/Cargo.toml
@@ -0,0 +1,19 @@
+[package]
+name = "wellknown-compiled"
+version = "0.1.0"
+authors = ["Lucio Franco <luciofranco14@gmail.com>"]
+edition = "2018"
+publish = false
+license = "MIT"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[lib]
+doctest = false
+
+[dependencies]
+tonic = { path = "../../tonic" }
+prost = "0.7"
+
+[build-dependencies]
+tonic-build = { path = "../../tonic-build" }
diff --git a/tests/wellknown-compiled/build.rs b/tests/wellknown-compiled/build.rs
@@ -0,0 +1,6 @@
+fn main() {
+    tonic_build::configure()
+        .compile_well_known_types(true)
+        .compile(&["proto/google.proto"], &["proto"])
+        .unwrap();
+}
diff --git a/tests/wellknown-compiled/proto/google.proto b/tests/wellknown-compiled/proto/google.proto
@@ -0,0 +1,16 @@
+syntax = "proto3";
+
+package google.protobuf;
+
+import "google/protobuf/any.proto";
+import "google/protobuf/api.proto";
+import "google/protobuf/descriptor.proto";
+import "google/protobuf/duration.proto";
+import "google/protobuf/empty.proto";
+import "google/protobuf/field_mask.proto";
+import "google/protobuf/source_context.proto";
+import "google/protobuf/struct.proto";
+import "google/protobuf/timestamp.proto";
+import "google/protobuf/type.proto";
+import "google/protobuf/wrappers.proto";
+
diff --git a/tests/wellknown-compiled/src/lib.rs b/tests/wellknown-compiled/src/lib.rs
@@ -0,0 +1,9 @@
+pub mod google {
+    pub mod protobuf {
+        tonic::include_proto!("google.protobuf");
+    }
+}
+
+pub fn grok() {
+    let _empty = crate::google::protobuf::Empty {};
+}
diff --git a/tonic-build/src/client.rs b/tonic-build/src/client.rs
@@ -7,10 +7,15 @@ use quote::{format_ident, quote};
 ///
 /// This takes some `Service` and will generate a `TokenStream` that contains
 /// a public module with the generated client.
-pub fn generate<T: Service>(service: &T, emit_package: bool, proto_path: &str) -> TokenStream {
+pub fn generate<T: Service>(
+    service: &T,
+    emit_package: bool,
+    proto_path: &str,
+    compile_well_known_types: bool,
+) -> TokenStream {
     let service_ident = quote::format_ident!("{}Client", service.name());
     let client_mod = quote::format_ident!("{}_client", naive_snake_case(&service.name()));
-    let methods = generate_methods(service, emit_package, proto_path);
+    let methods = generate_methods(service, emit_package, proto_path, compile_well_known_types);
 
     let connect = generate_connect(&service_ident);
     let service_doc = generate_doc_comments(service.comment());
@@ -87,7 +92,12 @@ fn generate_connect(_service_ident: &syn::Ident) -> TokenStream {
     TokenStream::new()
 }
 
-fn generate_methods<T: Service>(service: &T, emit_package: bool, proto_path: &str) -> TokenStream {
+fn generate_methods<T: Service>(
+    service: &T,
+    emit_package: bool,
+    proto_path: &str,
+    compile_well_known_types: bool,
+) -> TokenStream {
     let mut stream = TokenStream::new();
     let package = if emit_package { service.package() } else { "" };
 
@@ -103,10 +113,14 @@ fn generate_methods<T: Service>(service: &T, emit_package: bool, proto_path: &st
         stream.extend(generate_doc_comments(method.comment()));
 
         let method = match (method.client_streaming(), method.server_streaming()) {
-            (false, false) => generate_unary(method, proto_path, path),
-            (false, true) => generate_server_streaming(method, proto_path, path),
-            (true, false) => generate_client_streaming(method, proto_path, path),
-            (true, true) => generate_streaming(method, proto_path, path),
+            (false, false) => generate_unary(method, proto_path, compile_well_known_types, path),
+            (false, true) => {
+                generate_server_streaming(method, proto_path, compile_well_known_types, path)
+            }
+            (true, false) => {
+                generate_client_streaming(method, proto_path, compile_well_known_types, path)
+            }
+            (true, true) => generate_streaming(method, proto_path, compile_well_known_types, path),
         };
 
         stream.extend(method);
@@ -115,10 +129,15 @@ fn generate_methods<T: Service>(service: &T, emit_package: bool, proto_path: &st
     stream
 }
 
-fn generate_unary<T: Method>(method: &T, proto_path: &str, path: String) -> TokenStream {
+fn generate_unary<T: Method>(
+    method: &T,
+    proto_path: &str,
+    compile_well_known_types: bool,
+    path: String,
+) -> TokenStream {
     let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap();
     let ident = format_ident!("{}", method.name());
-    let (request, response) = method.request_response_name(proto_path);
+    let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
 
     quote! {
         pub async fn #ident(
@@ -135,11 +154,16 @@ fn generate_unary<T: Method>(method: &T, proto_path: &str, path: String) -> Toke
     }
 }
 
-fn generate_server_streaming<T: Method>(method: &T, proto_path: &str, path: String) -> TokenStream {
+fn generate_server_streaming<T: Method>(
+    method: &T,
+    proto_path: &str,
+    compile_well_known_types: bool,
+    path: String,
+) -> TokenStream {
     let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap();
     let ident = format_ident!("{}", method.name());
 
-    let (request, response) = method.request_response_name(proto_path);
+    let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
 
     quote! {
         pub async fn #ident(
@@ -156,11 +180,16 @@ fn generate_server_streaming<T: Method>(method: &T, proto_path: &str, path: Stri
     }
 }
 
-fn generate_client_streaming<T: Method>(method: &T, proto_path: &str, path: String) -> TokenStream {
+fn generate_client_streaming<T: Method>(
+    method: &T,
+    proto_path: &str,
+    compile_well_known_types: bool,
+    path: String,
+) -> TokenStream {
     let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap();
     let ident = format_ident!("{}", method.name());
 
-    let (request, response) = method.request_response_name(proto_path);
+    let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
 
     quote! {
         pub async fn #ident(
@@ -177,11 +206,16 @@ fn generate_client_streaming<T: Method>(method: &T, proto_path: &str, path: Stri
     }
 }
 
-fn generate_streaming<T: Method>(method: &T, proto_path: &str, path: String) -> TokenStream {
+fn generate_streaming<T: Method>(
+    method: &T,
+    proto_path: &str,
+    compile_well_known_types: bool,
+    path: String,
+) -> TokenStream {
     let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap();
     let ident = format_ident!("{}", method.name());
 
-    let (request, response) = method.request_response_name(proto_path);
+    let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
 
     quote! {
         pub async fn #ident(
diff --git a/tonic-build/src/lib.rs b/tonic-build/src/lib.rs
@@ -147,7 +147,11 @@ pub trait Method {
     /// Get comments about this item.
     fn comment(&self) -> &[Self::Comment];
     /// Type name of request and response.
-    fn request_response_name(&self, proto_path: &str) -> (TokenStream, TokenStream);
+    fn request_response_name(
+        &self,
+        proto_path: &str,
+        compile_well_known_types: bool,
+    ) -> (TokenStream, TokenStream);
 }
 
 /// Format files under the out_dir with rustfmt
diff --git a/tonic-build/src/prost.rs b/tonic-build/src/prost.rs
@@ -18,6 +18,7 @@ pub fn configure() -> Builder {
         field_attributes: Vec::new(),
         type_attributes: Vec::new(),
         proto_path: "super".to_string(),
+        compile_well_known_types: false,
         #[cfg(feature = "rustfmt")]
         format: true,
         emit_package: true,
@@ -94,8 +95,13 @@ impl crate::Method for Method {
         &self.comments.leading[..]
     }
 
-    fn request_response_name(&self, proto_path: &str) -> (TokenStream, TokenStream) {
-        let request = if self.input_proto_type.starts_with(".google.protobuf")
+    fn request_response_name(
+        &self,
+        proto_path: &str,
+        compile_well_known_types: bool,
+    ) -> (TokenStream, TokenStream) {
+        let request = if (self.input_proto_type.starts_with(".google.protobuf")
+            && !compile_well_known_types)
             || self.input_type.starts_with("::")
         {
             self.input_type.parse::<TokenStream>().unwrap()
@@ -105,7 +111,8 @@ impl crate::Method for Method {
                 .to_token_stream()
         };
 
-        let response = if self.output_proto_type.starts_with(".google.protobuf")
+        let response = if (self.output_proto_type.starts_with(".google.protobuf")
+            && !compile_well_known_types)
             || self.output_type.starts_with("::")
         {
             self.output_type.parse::<TokenStream>().unwrap()
@@ -142,6 +149,7 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
                 &service,
                 self.builder.emit_package,
                 &self.builder.proto_path,
+                self.builder.compile_well_known_types,
             );
             self.servers.extend(server);
         }
@@ -151,6 +159,7 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
                 &service,
                 self.builder.emit_package,
                 &self.builder.proto_path,
+                self.builder.compile_well_known_types,
             );
             self.clients.extend(client);
         }
@@ -196,6 +205,7 @@ pub struct Builder {
     pub(crate) type_attributes: Vec<(String, String)>,
     pub(crate) proto_path: String,
     pub(crate) emit_package: bool,
+    pub(crate) compile_well_known_types: bool,
 
     out_dir: Option<PathBuf>,
     #[cfg(feature = "rustfmt")]
@@ -285,6 +295,15 @@ impl Builder {
         self
     }
 
+    /// Enable or disable directing Prost to compile well-known protobuf types instead
+    /// of using the already-compiled versions available in the `prost-types` crate.
+    ///
+    /// This defaults to `false`.
+    pub fn compile_well_known_types(mut self, compile_well_known_types: bool) -> Self {
+        self.compile_well_known_types = compile_well_known_types;
+        self
+    }
+
     /// Compile the .proto files and execute code generation.
     pub fn compile<P>(self, protos: &[P], includes: &[P]) -> io::Result<()>
     where
@@ -326,6 +345,9 @@ impl Builder {
         for (prost_path, attr) in self.type_attributes.iter() {
             config.type_attribute(prost_path, attr);
         }
+        if self.compile_well_known_types {
+            config.compile_well_known_types();
+        }
         config.service_generator(Box::new(ServiceGenerator::new(self)));
 
         config.compile_protos(protos, includes)?;
diff --git a/tonic-build/src/server.rs b/tonic-build/src/server.rs
