tonic-build: add missing APIs (#172)

Signed-off-by: Runji Wang <[email protected]>

CHANGELOG.md

@@ -6,6 +6,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## tonic-build [0.4.2] - 2023-10-08
+
+### Added
+
+- Add missing APIs for `Builder` in v0.10.0.
+
 ## etcd [0.4.0] - 2023-10-07
 
 ### Changed

madsim-tonic-build/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "madsim-tonic-build"
-version = "0.4.1+0.10.0"
+version = "0.4.2+0.10.0"
 edition = "2021"
 authors = [
     "Lucio Franco <[email protected]>",

madsim-tonic-build/src/prost.rs

@@ -2,6 +2,7 @@ use super::{client, server, Attributes};
 use proc_macro2::TokenStream;
 use prost_build::Config;
 use std::{
+    collections::HashSet,
     ffi::OsString,
     io,
     path::{Path, PathBuf},
@@ -14,18 +15,29 @@ pub fn configure() -> Builder {
     Builder {
         build_client: true,
         build_server: true,
+        build_transport: true,
         file_descriptor_set_path: None,
+        skip_protoc_run: false,
         out_dir: None,
         extern_path: Vec::new(),
         field_attributes: Vec::new(),
+        message_attributes: Vec::new(),
+        enum_attributes: Vec::new(),
         type_attributes: Vec::new(),
+        boxed: Vec::new(),
+        btree_map: None,
+        bytes: None,
         server_attributes: Attributes::default(),
         client_attributes: Attributes::default(),
         proto_path: "super".to_string(),
         compile_well_known_types: false,
         emit_package: true,
         protoc_args: Vec::new(),
         include_file: None,
+        emit_rerun_if_changed: std::env::var_os("CARGO").is_some(),
+        disable_comments: HashSet::default(),
+        use_arc_self: false,
+        generate_default_stubs: false,
         builder: tonic_build::configure(),
     }
 }
@@ -124,17 +136,28 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
 pub struct Builder {
     pub(crate) build_client: bool,
     pub(crate) build_server: bool,
+    pub(crate) build_transport: bool,
     pub(crate) file_descriptor_set_path: Option<PathBuf>,
+    pub(crate) skip_protoc_run: bool,
     pub(crate) extern_path: Vec<(String, String)>,
     pub(crate) field_attributes: Vec<(String, String)>,
     pub(crate) type_attributes: Vec<(String, String)>,
+    pub(crate) message_attributes: Vec<(String, String)>,
+    pub(crate) enum_attributes: Vec<(String, String)>,
+    pub(crate) boxed: Vec<String>,
+    pub(crate) btree_map: Option<Vec<String>>,
+    pub(crate) bytes: Option<Vec<String>>,
     pub(crate) server_attributes: Attributes,
     pub(crate) client_attributes: Attributes,
     pub(crate) proto_path: String,
     pub(crate) emit_package: bool,
     pub(crate) compile_well_known_types: bool,
     pub(crate) protoc_args: Vec<OsString>,
     pub(crate) include_file: Option<PathBuf>,
+    pub(crate) emit_rerun_if_changed: bool,
+    pub(crate) disable_comments: HashSet<String>,
+    pub(crate) use_arc_self: bool,
+    pub(crate) generate_default_stubs: bool,
 
     out_dir: Option<PathBuf>,
 
@@ -157,13 +180,31 @@ impl Builder {
         self
     }
 
+    /// Enable or disable generated clients and servers to have built-in tonic
+    /// transport features.
+    ///
+    /// When the `transport` feature is disabled this does nothing.
+    pub fn build_transport(mut self, enable: bool) -> Self {
+        self.builder = self.builder.build_transport(enable);
+        self.build_transport = enable;
+        self
+    }
+
     /// Generate a file containing the encoded `prost_types::FileDescriptorSet` for protocol buffers
     /// modules. This is required for implementing gRPC Server Reflection.
     pub fn file_descriptor_set_path(mut self, path: impl AsRef<Path>) -> Self {
         self.file_descriptor_set_path = Some(path.as_ref().to_path_buf());
         self
     }
 
+    /// In combination with with file_descriptor_set_path, this can be used to provide a file
+    /// descriptor set as an input file, rather than having prost-build generate the file by
+    /// calling protoc.
+    pub fn skip_protoc_run(mut self) -> Self {
+        self.skip_protoc_run = true;
+        self
+    }
+
     /// Set the output directory to generate code to.
     ///
     /// Defaults to the `OUT_DIR` environment variable.
@@ -204,6 +245,76 @@ impl Builder {
         self
     }
 
+    /// Add additional attribute to matched messages.
+    ///
+    /// Passed directly to `prost_build::Config.message_attribute`.
+    pub fn message_attribute<P: AsRef<str>, A: AsRef<str>>(
+        mut self,
+        path: P,
+        attribute: A,
+    ) -> Self {
+        self.message_attributes
+            .push((path.as_ref().to_string(), attribute.as_ref().to_string()));
+        self
+    }
+
+    /// Add additional attribute to matched enums.
+    ///
+    /// Passed directly to `prost_build::Config.enum_attribute`.
+    pub fn enum_attribute<P: AsRef<str>, A: AsRef<str>>(mut self, path: P, attribute: A) -> Self {
+        self.enum_attributes
+            .push((path.as_ref().to_string(), attribute.as_ref().to_string()));
+        self
+    }
+
+    /// Add additional boxed fields.
+    ///
+    /// Passed directly to `prost_build::Config.boxed`.
+    pub fn boxed<P: AsRef<str>>(mut self, path: P) -> Self {
+        self.boxed.push(path.as_ref().to_string());
+        self
+    }
+
+    /// Configure the code generator to generate Rust `BTreeMap` fields for Protobuf `map` type
+    /// fields.
+    ///
+    /// Passed directly to `prost_build::Config.btree_map`.
+    ///
+    /// Note: previous configurated paths for `btree_map` will be cleared.
+    pub fn btree_map<I, S>(mut self, paths: I) -> Self
+    where
+        I: IntoIterator<Item = S>,
+        S: AsRef<str>,
+    {
+        self.btree_map = Some(
+            paths
+                .into_iter()
+                .map(|path| path.as_ref().to_string())
+                .collect(),
+        );
+        self
+    }
+
+    /// Configure the code generator to generate Rust `bytes::Bytes` fields for Protobuf `bytes`
+    /// type fields.
+    ///
+    /// Passed directly to `prost_build::Config.bytes`.
+    ///
+    /// Note: previous configurated paths for `bytes` will be cleared.
+    pub fn bytes<I, S>(mut self, paths: I) -> Self
+    where
+        I: IntoIterator<Item = S>,
+        S: AsRef<str>,
+    {
+        self.bytes = Some(
+            paths
+                .into_iter()
+                .map(|path| path.as_ref().to_string())
+                .collect(),
+        );
+        self
+    }
+
     /// Add additional attribute to matched server `mod`s. Matches on the package name.
     pub fn server_mod_attribute<P: AsRef<str>, A: AsRef<str>>(
         mut self,
@@ -270,6 +381,20 @@ impl Builder {
         self
     }
 
+    /// Disable service and rpc comments emission.
+    pub fn disable_comments(mut self, path: impl AsRef<str>) -> Self {
+        self.builder = self.builder.disable_comments(path.as_ref());
+        self.disable_comments.insert(path.as_ref().to_string());
+        self
+    }
+
+    /// Emit `Arc<Self>` receiver type in server traits instead of `&self`.
+    pub fn use_arc_self(mut self, enable: bool) -> Self {
+        self.builder = self.builder.use_arc_self(enable);
+        self.use_arc_self = enable;
+        self
+    }
+
     /// Emits GRPC endpoints with no attached package. Effectively ignores protofile package declaration from grpc context.
     ///
     /// This effectively sets prost's exported package to an empty string.
@@ -299,6 +424,36 @@ impl Builder {
         self
     }
 
+    /// Enable or disable emitting
+    /// [`cargo:rerun-if-changed=PATH`](https://doc.rust-lang.org/cargo/reference/build-scripts.html#rerun-if-changed)
+    /// instructions for Cargo.
+    ///
+    /// If set, writes instructions to `stdout` for Cargo so that it understands
+    /// when to rerun the build script. By default, this setting is enabled if
+    /// the `CARGO` environment variable is set. The `CARGO` environment
+    /// variable is set by Cargo for build scripts. Therefore, this setting
+    /// should be enabled automatically when run from a build script. However,
+    /// the method of detection is not completely reliable since the `CARGO`
+    /// environment variable can have been set by anything else. If writing the
+    /// instructions to `stdout` is undesireable, you can disable this setting
+    /// explicitly.
+    pub fn emit_rerun_if_changed(mut self, enable: bool) -> Self {
+        self.emit_rerun_if_changed = enable;
+        self
+    }
+
+    /// Enable or disable directing service generation to providing a default implementation for service methods.
+    /// When this is false all gRPC methods must be explicitly implemented.
+    /// When this is true any unimplemented service methods will return 'unimplemented' gRPC error code.
+    /// When this is true all streaming server request RPC types explicitly use tonic::codegen::BoxStream type.
+    ///
+    /// This defaults to `false`.
+    pub fn generate_default_stubs(mut self, enable: bool) -> Self {
+        self.builder = self.builder.generate_default_stubs(enable);
+        self.generate_default_stubs = enable;
+        self
+    }
+
     /// Compile the .proto files and execute code generation.
     pub fn compile(
         self,
@@ -331,6 +486,9 @@ impl Builder {
         if let Some(path) = self.file_descriptor_set_path.as_ref() {
             config.file_descriptor_set_path(path);
         }
+        if self.skip_protoc_run {
+            config.skip_protoc_run();
+        }
         for (proto_path, rust_path) in self.extern_path.iter() {
             config.extern_path(proto_path, rust_path);
         }
@@ -340,6 +498,21 @@ impl Builder {
         for (prost_path, attr) in self.type_attributes.iter() {
             config.type_attribute(prost_path, attr);
         }
+        for (prost_path, attr) in self.message_attributes.iter() {
+            config.message_attribute(prost_path, attr);
+        }
+        for (prost_path, attr) in self.enum_attributes.iter() {
+            config.enum_attribute(prost_path, attr);
+        }
+        for prost_path in self.boxed.iter() {
+            config.boxed(prost_path);
+        }
+        if let Some(ref paths) = self.btree_map {
+            config.btree_map(paths);
+        }
+        if let Some(ref paths) = self.bytes {
+            config.bytes(paths);
+        }
         if self.compile_well_known_types {
             config.compile_well_known_types();
         }
@@ -351,6 +524,19 @@ impl Builder {
             config.protoc_arg(arg);
         }
 
+        if self.emit_rerun_if_changed {
+            for path in protos.iter() {
+                println!("cargo:rerun-if-changed={}", path.as_ref().display())
+            }
+
+            for path in includes.iter() {
+                // Cargo will watch the **entire** directory recursively. If we
+                // could figure out which files are imported by our protos we
+                // could specify only those files instead.
+                println!("cargo:rerun-if-changed={}", path.as_ref().display())
+            }
+        }
+
         config.service_generator(self.service_generator());
 
         config.compile_protos(protos, includes)?;