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

This commit is contained in:
Tom Dyas
2021-02-23 16:21:03 -05:00
committed by GitHub
parent 49d517fac9
commit 61555ff2b5
9 changed files with 198 additions and 39 deletions
+1
View File
@@ -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",
+19
View File
@@ -0,0 +1,19 @@
[package]
name = "wellknown-compiled"
version = "0.1.0"
authors = ["Lucio Franco <[email protected]>"]
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" }
+6
View File
@@ -0,0 +1,6 @@
fn main() {
tonic_build::configure()
.compile_well_known_types(true)
.compile(&["proto/google.proto"], &["proto"])
.unwrap();
}
@@ -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";
+9
View File
@@ -0,0 +1,9 @@
pub mod google {
pub mod protobuf {
tonic::include_proto!("google.protobuf");
}
}
pub fn grok() {
let _empty = crate::google::protobuf::Empty {};
}
+49 -15
View File
@@ -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(
+5 -1
View File
@@ -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
+25 -3
View File
@@ -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)?;
+68 -20
View File
@@ -8,13 +8,23 @@ use syn::{Ident, Lit, LitStr};
///
/// This takes some `Service` and will generate a `TokenStream` that contains
/// a public module containing the server service and handler trait.
pub fn generate<T: Service>(service: &T, emit_package: bool, proto_path: &str) -> TokenStream {
let methods = generate_methods(service, proto_path);
pub fn generate<T: Service>(
service: &T,
emit_package: bool,
proto_path: &str,
compile_well_known_types: bool,
) -> TokenStream {
let methods = generate_methods(service, proto_path, compile_well_known_types);
let server_service = quote::format_ident!("{}Server", service.name());
let server_trait = quote::format_ident!("{}", service.name());
let server_mod = quote::format_ident!("{}_server", naive_snake_case(&service.name()));
let generated_trait = generate_trait(service, proto_path, server_trait.clone());
let generated_trait = generate_trait(
service,
proto_path,
compile_well_known_types,
server_trait.clone(),
);
let service_doc = generate_doc_comments(service.comment());
let package = if emit_package { service.package() } else { "" };
// Transport based implementations
@@ -112,8 +122,13 @@ pub fn generate<T: Service>(service: &T, emit_package: bool, proto_path: &str) -
}
}
fn generate_trait<T: Service>(service: &T, proto_path: &str, server_trait: Ident) -> TokenStream {
let methods = generate_trait_methods(service, proto_path);
fn generate_trait<T: Service>(
service: &T,
proto_path: &str,
compile_well_known_types: bool,
server_trait: Ident,
) -> TokenStream {
let methods = generate_trait_methods(service, proto_path, compile_well_known_types);
let trait_doc = generate_doc_comment(&format!(
"Generated trait containing gRPC methods that should be implemented for use with {}Server.",
service.name()
@@ -128,13 +143,18 @@ fn generate_trait<T: Service>(service: &T, proto_path: &str, server_trait: Ident
}
}
fn generate_trait_methods<T: Service>(service: &T, proto_path: &str) -> TokenStream {
fn generate_trait_methods<T: Service>(
service: &T,
proto_path: &str,
compile_well_known_types: bool,
) -> TokenStream {
let mut stream = TokenStream::new();
for method in service.methods() {
let name = quote::format_ident!("{}", method.name());
let (req_message, res_message) = method.request_response_name(proto_path);
let (req_message, res_message) =
method.request_response_name(proto_path, compile_well_known_types);
let method_doc = generate_doc_comments(method.comment());
@@ -217,7 +237,11 @@ fn generate_transport(
TokenStream::new()
}
fn generate_methods<T: Service>(service: &T, proto_path: &str) -> TokenStream {
fn generate_methods<T: Service>(
service: &T,
proto_path: &str,
compile_well_known_types: bool,
) -> TokenStream {
let mut stream = TokenStream::new();
for method in service.methods() {
@@ -237,16 +261,36 @@ fn generate_methods<T: Service>(service: &T, proto_path: &str) -> TokenStream {
let server_trait = quote::format_ident!("{}", service.name());
let method_stream = match (method.client_streaming(), method.server_streaming()) {
(false, false) => generate_unary(method, proto_path, ident, server_trait),
(false, false) => generate_unary(
method,
proto_path,
compile_well_known_types,
ident,
server_trait,
),
(false, true) => {
generate_server_streaming(method, proto_path, ident.clone(), server_trait)
}
(true, false) => {
generate_client_streaming(method, proto_path, ident.clone(), server_trait)
}
(false, true) => generate_server_streaming(
method,
proto_path,
compile_well_known_types,
ident.clone(),
server_trait,
),
(true, false) => generate_client_streaming(
method,
proto_path,
compile_well_known_types,
ident.clone(),
server_trait,
),
(true, true) => generate_streaming(method, proto_path, ident.clone(), server_trait),
(true, true) => generate_streaming(
method,
proto_path,
compile_well_known_types,
ident.clone(),
server_trait,
),
};
let method = quote! {
@@ -263,6 +307,7 @@ fn generate_methods<T: Service>(service: &T, proto_path: &str) -> TokenStream {
fn generate_unary<T: Method>(
method: &T,
proto_path: &str,
compile_well_known_types: bool,
method_ident: Ident,
server_trait: Ident,
) -> TokenStream {
@@ -270,7 +315,7 @@ fn generate_unary<T: Method>(
let service_ident = quote::format_ident!("{}Svc", method.identifier());
let (request, response) = method.request_response_name(proto_path);
let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
quote! {
#[allow(non_camel_case_types)]
@@ -313,6 +358,7 @@ fn generate_unary<T: Method>(
fn generate_server_streaming<T: Method>(
method: &T,
proto_path: &str,
compile_well_known_types: bool,
method_ident: Ident,
server_trait: Ident,
) -> TokenStream {
@@ -320,7 +366,7 @@ fn generate_server_streaming<T: Method>(
let service_ident = quote::format_ident!("{}Svc", method.identifier());
let (request, response) = method.request_response_name(proto_path);
let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
let response_stream = quote::format_ident!("{}Stream", method.identifier());
@@ -367,12 +413,13 @@ fn generate_server_streaming<T: Method>(
fn generate_client_streaming<T: Method>(
method: &T,
proto_path: &str,
compile_well_known_types: bool,
method_ident: Ident,
server_trait: Ident,
) -> TokenStream {
let service_ident = quote::format_ident!("{}Svc", method.identifier());
let (request, response) = method.request_response_name(proto_path);
let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
let codec_name = syn::parse_str::<syn::Path>(T::CODEC_PATH).unwrap();
quote! {
@@ -418,6 +465,7 @@ fn generate_client_streaming<T: Method>(
fn generate_streaming<T: Method>(
method: &T,
proto_path: &str,
compile_well_known_types: bool,
method_ident: Ident,
server_trait: Ident,
) -> TokenStream {
@@ -425,7 +473,7 @@ fn generate_streaming<T: Method>(
let service_ident = quote::format_ident!("{}Svc", method.identifier());
let (request, response) = method.request_response_name(proto_path);
let (request, response) = method.request_response_name(proto_path, compile_well_known_types);
let response_stream = quote::format_ident!("{}Stream", method.identifier());