Merge pull request #1 from NickLarsenNZ/optional_return_overrides

feat(build): allow optional service method return overrides

tonic-build/README.md

@@ -40,6 +40,25 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
    Ok(())
 }
 ```
+
+### Override service method return values
+
+```rust
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+   tonic_build::configure()
+        .service_return_overrides(tonic_build::ServiceReturnOverrideMap::from([
+            (String::from("Echo.unary_echo"), ServiceReturnOverride::None), // implied if no mapping is made for the service. trait method implementation therefore required.
+            (String::from("Echo.server_streaming_echo"), ServiceReturnOverride::OkDefault), // return the default value for the response message.
+            (String::from("Echo.client_streaming_echo"), ServiceReturnOverride::ErrUnimplemented), // return the unimplemented error status.
+            // any other service not mapped here will require an explicit trait method implementation.
+        ]))
+        .compile(
+            &["proto/echo/echo.proto"],
+            &["proto/echo"],
+        )?;
+   Ok(())
+}
+```
 See [more examples here](https://github.com/hyperium/tonic/tree/master/examples)
 
 ### Google APIs example

tonic-build/src/code_gen.rs

@@ -2,7 +2,7 @@ use std::collections::HashSet;
 
 use proc_macro2::TokenStream;
 
-use crate::{Attributes, Service};
+use crate::{prost::ServiceReturnOverrideMap, Attributes, Service};
 
 /// Builder for the generic code generation of server and clients.
 #[derive(Debug)]
@@ -12,7 +12,7 @@ pub struct CodeGenBuilder {
     attributes: Attributes,
     build_transport: bool,
     disable_comments: HashSet<String>,
-    default_impl: bool
+    service_return_overrides: ServiceReturnOverrideMap,
 }
 
 impl CodeGenBuilder {
@@ -59,8 +59,11 @@ impl CodeGenBuilder {
     }
 
     /// Enable or disable returning automatic unimplemented gRPC error code for generated traits.
-    pub fn default_impl(&mut self, default_impl: bool) -> &mut Self {
-        self.default_impl = default_impl;
+    pub fn service_return_overrides(
+        &mut self,
+        service_return_overrides: ServiceReturnOverrideMap,
+    ) -> &mut Self {
+        self.service_return_overrides = service_return_overrides;
         self
     }
 
@@ -92,7 +95,7 @@ impl CodeGenBuilder {
             self.compile_well_known_types,
             &self.attributes,
             &self.disable_comments,
-            self.default_impl
+            &self.service_return_overrides,
         )
     }
 }
@@ -105,7 +108,7 @@ impl Default for CodeGenBuilder {
             attributes: Attributes::default(),
             build_transport: true,
             disable_comments: HashSet::default(),
-            default_impl: false
+            service_return_overrides: ServiceReturnOverrideMap::default(),
         }
     }
 }

tonic-build/src/lib.rs

@@ -85,7 +85,7 @@ mod prost;
 
 #[cfg(feature = "prost")]
 #[cfg_attr(docsrs, doc(cfg(feature = "prost")))]
-pub use prost::{compile_protos, configure, Builder};
+pub use prost::{compile_protos, configure, Builder, ServiceReturnOverrideMap};
 
 pub mod manual;
 

tonic-build/src/prost.rs

@@ -5,7 +5,7 @@ use proc_macro2::TokenStream;
 use prost_build::{Config, Method, Service};
 use quote::ToTokens;
 use std::{
-    collections::HashSet,
+    collections::{HashMap, HashSet},
     ffi::OsString,
     io,
     path::{Path, PathBuf},
@@ -36,7 +36,7 @@ pub fn configure() -> Builder {
         include_file: None,
         emit_rerun_if_changed: std::env::var_os("CARGO").is_some(),
         disable_comments: HashSet::default(),
-        default_impl: false
+        service_return_overrides: ServiceReturnOverrideMap::default(),
     }
 }
 
@@ -170,7 +170,7 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
                 .compile_well_known_types(self.builder.compile_well_known_types)
                 .attributes(self.builder.server_attributes.clone())
                 .disable_comments(self.builder.disable_comments.clone())
-                .default_impl(self.builder.default_impl)
+                .service_return_overrides(self.builder.service_return_overrides.clone())
                 .generate_server(&service, &self.builder.proto_path);
 
             self.servers.extend(server);
@@ -242,7 +242,7 @@ pub struct Builder {
     pub(crate) include_file: Option<PathBuf>,
     pub(crate) emit_rerun_if_changed: bool,
     pub(crate) disable_comments: HashSet<String>,
-    pub(crate) default_impl: bool,
+    pub(crate) service_return_overrides: ServiceReturnOverrideMap,
 
     out_dir: Option<PathBuf>,
 }
@@ -450,12 +450,19 @@ impl Builder {
         self
     }
 
-    /// When generating services enables or disables the default implementation stubs of returning 'unimplemented' gRPC error code.
-    /// When this is false all gRPC methods must be explicitly implemented.
+    /// When generating services allows for specifying default service trait
+    /// method implementations. By default, services will require explicit
+    /// implementations.
     ///
-    /// This defaults to `false`.
-    pub fn default_impl(mut self, enable: bool) -> Self {
-        self.default_impl = enable;
+    /// Example:
+    /// ```rs
+    /// todo!()
+    /// ```
+    pub fn service_return_overrides(
+        mut self,
+        service_return_overrides: HashMap<String, ServiceReturnOverride>,
+    ) -> Self {
+        self.service_return_overrides = service_return_overrides;
         self
     }
 
@@ -541,3 +548,20 @@ impl Builder {
         Box::new(ServiceGenerator::new(self))
     }
 }
+
+/// An alias to abstract the underlying type and ensure its usages in the library are correct.
+pub type ServiceReturnOverrideMap = HashMap<String, ServiceReturnOverride>;
+
+#[derive(Clone, Copy, Debug, Default)]
+pub enum ServiceReturnOverride {
+    /// Leave trait method unimplemented. This forces explicit implementation.
+    /// Default bahaviour.
+    #[default]
+    None,
+
+    /// Return a `Result::Ok` with a `tonic::Response` of the default for the type.
+    OkDefault,
+
+    /// Return a `Result::Error` with a `tonic::Status::unimplemented` and the method name included in the friendly error returned to the client.
+    ErrUnimplemented,
+}

tonic-build/src/server.rs

@@ -3,7 +3,7 @@ use std::collections::HashSet;
 use super::{Attributes, Method, Service};
 use crate::{
     format_method_name, format_method_path, format_service_name, generate_doc_comment,
-    generate_doc_comments, naive_snake_case,
+    generate_doc_comments, naive_snake_case, prost::ServiceReturnOverrideMap,
 };
 use proc_macro2::{Span, TokenStream};
 use quote::quote;
@@ -28,7 +28,7 @@ pub fn generate<T: Service>(
         compile_well_known_types,
         attributes,
         &HashSet::default(),
-        false
+        &ServiceReturnOverrideMap::default(),
     )
 }
 
@@ -39,7 +39,7 @@ pub(crate) fn generate_internal<T: Service>(
     compile_well_known_types: bool,
     attributes: &Attributes,
     disable_comments: &HashSet<String>,
-    default_impl: bool
+    service_return_overrides: &ServiceReturnOverrideMap,
 ) -> TokenStream {
     let methods = generate_methods(service, emit_package, proto_path, compile_well_known_types);
 
@@ -53,7 +53,7 @@ pub(crate) fn generate_internal<T: Service>(
         compile_well_known_types,
         server_trait.clone(),
         disable_comments,
-        default_impl
+        service_return_overrides,
     );
     let package = if emit_package { service.package() } else { "" };
     // Transport based implementations
@@ -230,15 +230,15 @@ fn generate_trait<T: Service>(
     compile_well_known_types: bool,
     server_trait: Ident,
     disable_comments: &HashSet<String>,
-    default_impl: bool
+    service_return_overrides: &ServiceReturnOverrideMap,
 ) -> TokenStream {
     let methods = generate_trait_methods(
         service,
         emit_package,
         proto_path,
         compile_well_known_types,
         disable_comments,
-        default_impl
+        service_return_overrides,
     );
     let trait_doc = generate_doc_comment(format!(
         " Generated trait containing gRPC methods that should be implemented for use with {}Server.",
@@ -260,7 +260,7 @@ fn generate_trait_methods<T: Service>(
     proto_path: &str,
     compile_well_known_types: bool,
     disable_comments: &HashSet<String>,
-    default_impl: bool
+    service_return_overrides: &ServiceReturnOverrideMap,
 ) -> TokenStream {
     let mut stream = TokenStream::new();
 
@@ -277,40 +277,64 @@ fn generate_trait_methods<T: Service>(
                 generate_doc_comments(method.comment())
             };
 
-        let method = match (method.client_streaming(), method.server_streaming(), default_impl) {
-            (false, false, true) => {
+        // let return_value = match service_return_overrides {};
+        let override_lookup_key = format!("{}.{}", service.name(), method.name());
+        let return_value = service_return_overrides
+            .get(&override_lookup_key)
+            .and_then(|v| {
+                return {
+                    match v {
+                        crate::prost::ServiceReturnOverride::None => None,
+                        crate::prost::ServiceReturnOverride::OkDefault => {
+                            Some(quote! { Ok(tonic::Response::new(#res_message::default())) })
+                        }
+
+                        crate::prost::ServiceReturnOverride::ErrUnimplemented => {
+                            let not_implemented_error = format!("{} is not implemented", method.name());
+                            Some(quote! { Err(tonic::Status::unimplemented(#not_implemented_error)) })
+                        },
+                    }
+                };
+            });
+
+        let method = match (
+            method.client_streaming(),
+            method.server_streaming(),
+            return_value,
+        ) {
+            (false, false, Some(return_value)) => {
                 quote! {
                     #method_doc
                     async fn #name(&self, request: tonic::Request<#req_message>)
                         -> std::result::Result<tonic::Response<#res_message>, tonic::Status> {
-                        Err(tonic::Status::unimplemented("Not yet implemented"))
+                        #return_value
                     }
                 }
-            },
-            (false, false, false) => {
+            }
+            (false, false, None) => {
                 quote! {
                     #method_doc
                     async fn #name(&self, request: tonic::Request<#req_message>)
                         -> std::result::Result<tonic::Response<#res_message>, tonic::Status>;
                 }
-            },
-            (true, false, true) => {
+            }
+            (true, false, Some(return_value)) => {
                 quote! {
                     #method_doc
                     async fn #name(&self, request: tonic::Request<tonic::Streaming<#req_message>>)
                         -> std::result::Result<tonic::Response<#res_message>, tonic::Status> {
-                        Err(tonic::Status::unimplemented("Not yet implemented"))
+                        #return_value
                     }
                 }
-            },
-            (true, false, false) => {
+            }
+            (true, false, None) => {
                 quote! {
                     #method_doc
                     async fn #name(&self, request: tonic::Request<tonic::Streaming<#req_message>>)
                         -> std::result::Result<tonic::Response<#res_message>, tonic::Status>;
                 }
-            },
-            (false, true, true) => {
+            }
+            (false, true, Some(return_value)) => {
                 let stream = quote::format_ident!("{}Stream", method.identifier());
                 let stream_doc = generate_doc_comment(format!(
                     " Server streaming response type for the {} method.",
@@ -324,11 +348,11 @@ fn generate_trait_methods<T: Service>(
                     #method_doc
                     async fn #name(&self, request: tonic::Request<#req_message>)
                         -> std::result::Result<tonic::Response<Self::#stream>, tonic::Status> {
-                        Err(tonic::Status::unimplemented("Not yet implemented"))
+                        #return_value
                     }
                 }
-            },
-            (false, true, false) => {
+            }
+            (false, true, None) => {
                 let stream = quote::format_ident!("{}Stream", method.identifier());
                 let stream_doc = generate_doc_comment(format!(
                     " Server streaming response type for the {} method.",
@@ -343,8 +367,8 @@ fn generate_trait_methods<T: Service>(
                     async fn #name(&self, request: tonic::Request<#req_message>)
                         -> std::result::Result<tonic::Response<Self::#stream>, tonic::Status>;
                 }
-            },
-            (true, true, true) => {
+            }
+            (true, true, Some(return_value)) => {
                 let stream = quote::format_ident!("{}Stream", method.identifier());
                 let stream_doc = generate_doc_comment(format!(
                     " Server streaming response type for the {} method.",
@@ -358,11 +382,11 @@ fn generate_trait_methods<T: Service>(
                     #method_doc
                     async fn #name(&self, request: tonic::Request<#req_message>)
                         -> std::result::Result<tonic::Response<Self::#stream>, tonic::Status> {
-                        Err(tonic::Status::unimplemented("Not yet implemented"))
+                        #return_value
                     }
                 }
-            },
-            (true, true, false) => {
+            }
+            (true, true, None) => {
                 let stream = quote::format_ident!("{}Stream", method.identifier());
                 let stream_doc = generate_doc_comment(format!(
                     " Server streaming response type for the {} method.",