grpc_build/
builder.rs

1use std::{
2    ffi::OsString,
3    path::{Path, PathBuf},
4};
5
6/// A mirror of [`tonic_build::Builder`] for our own control
7pub struct Builder {
8    pub(crate) tonic: tonic_build::Builder,
9    pub(crate) prost: prost_build::Config,
10    pub(crate) protoc_args: Vec<OsString>,
11    pub(crate) out_dir: Option<PathBuf>,
12    pub(crate) force: bool,
13    pub(crate) default_module_name: Option<String>,
14    pub(crate) follow_links: bool,
15    pub(crate) file_descriptor_set_path: Option<PathBuf>,
16}
17
18impl Default for Builder {
19    fn default() -> Self {
20        Self {
21            tonic: tonic_build::configure(),
22            prost: prost_build::Config::default(),
23            protoc_args: Vec::default(),
24            out_dir: None,
25            force: false,
26            default_module_name: None,
27            follow_links: false,
28            file_descriptor_set_path: None,
29        }
30    }
31}
32
33impl Builder {
34    pub(crate) fn get_out_dir(&self) -> Result<PathBuf, anyhow::Error> {
35        if let Some(out_dir) = &self.out_dir {
36            return Ok(out_dir.clone());
37        }
38
39        std::env::var_os("OUT_DIR")
40            .ok_or_else(|| anyhow::anyhow!("could not determine $OUT_DIR"))
41            .map(Into::into)
42    }
43
44    pub fn new() -> Self {
45        Self::default()
46    }
47
48    pub fn force(mut self, force: bool) -> Self {
49        self.force = force;
50        self
51    }
52
53    /// Follow symbolic links when finding .proto files.
54    ///
55    /// This defaults to `false`.
56    pub fn follow_links(mut self, follow_links: bool) -> Self {
57        self.follow_links = follow_links;
58        self
59    }
60
61    pub fn out_dir(mut self, out_dir: impl AsRef<Path>) -> Self {
62        self.out_dir = Some(out_dir.as_ref().to_owned());
63        self
64    }
65
66    /// Configures what filename protobufs with no package definition are written to.
67    pub fn default_module_name(mut self, name: impl AsRef<str>) -> Self {
68        self.default_module_name = Some(name.as_ref().to_string());
69        self
70    }
71
72    /// Enable or disable gRPC client code generation.
73    pub fn build_client(mut self, enable: bool) -> Self {
74        self.tonic = self.tonic.build_client(enable);
75        self
76    }
77
78    /// Enable or disable gRPC server code generation.
79    pub fn build_server(mut self, enable: bool) -> Self {
80        self.tonic = self.tonic.build_server(enable);
81        self
82    }
83
84    /// Declare externally provided Protobuf package or type.
85    ///
86    /// Passed directly to `prost_build::Config.extern_path`.
87    /// Note that both the Protobuf path and the rust package paths should both be fully qualified.
88    /// i.e. Protobuf paths should start with "." and rust paths should start with "::"
89    pub fn extern_path(mut self, proto_path: impl AsRef<str>, rust_path: impl AsRef<str>) -> Self {
90        self.prost.extern_path(
91            proto_path.as_ref().to_string(),
92            rust_path.as_ref().to_string(),
93        );
94        self
95    }
96
97    /// Add additional attribute to matched messages, enums, and one-offs.
98    ///
99    /// Passed directly to `prost_build::Config.field_attribute`.
100    pub fn field_attribute<P: AsRef<str>, A: AsRef<str>>(mut self, path: P, attribute: A) -> Self {
101        self.prost
102            .field_attribute(path.as_ref(), attribute.as_ref());
103        self
104    }
105
106    /// Add additional attribute to matched messages, enums, and one-offs.
107    ///
108    /// Passed directly to `prost_build::Config.type_attribute`.
109    pub fn type_attribute<P: AsRef<str>, A: AsRef<str>>(mut self, path: P, attribute: A) -> Self {
110        self.prost.type_attribute(path.as_ref(), attribute.as_ref());
111        self
112    }
113
114    /// Add additional attribute to matched server `mod`s. Matches on the package name.
115    pub fn server_mod_attribute<P: AsRef<str>, A: AsRef<str>>(
116        mut self,
117        path: P,
118        attribute: A,
119    ) -> Self {
120        self.tonic = self.tonic.server_mod_attribute(path, attribute);
121        self
122    }
123
124    /// Add additional attribute to matched service servers. Matches on the service name.
125    pub fn server_attribute<P: AsRef<str>, A: AsRef<str>>(mut self, path: P, attribute: A) -> Self {
126        self.tonic = self.tonic.server_attribute(path, attribute);
127        self
128    }
129
130    /// Add additional attribute to matched client `mod`s. Matches on the package name.
131    pub fn client_mod_attribute<P: AsRef<str>, A: AsRef<str>>(
132        mut self,
133        path: P,
134        attribute: A,
135    ) -> Self {
136        self.tonic = self.tonic.client_mod_attribute(path, attribute);
137        self
138    }
139
140    /// Add additional attribute to matched service clients. Matches on the service name.
141    pub fn client_attribute<P: AsRef<str>, A: AsRef<str>>(mut self, path: P, attribute: A) -> Self {
142        self.tonic = self.tonic.client_attribute(path, attribute);
143        self
144    }
145
146    /// Configure Prost `protoc_args` build arguments.
147    ///
148    /// Note: Enabling `--experimental_allow_proto3_optional` requires protobuf >= 3.12.
149    pub fn protoc_arg<A: AsRef<str>>(mut self, arg: A) -> Self {
150        self.protoc_args.push(arg.as_ref().into());
151        self
152    }
153
154    /// Enable or disable directing Prost to compile well-known protobuf types instead
155    /// of using the already-compiled versions available in the `prost-types` crate.
156    ///
157    /// This defaults to `false`.
158    pub fn compile_well_known_types(mut self, compile_well_known_types: bool) -> Self {
159        if compile_well_known_types {
160            self.prost.compile_well_known_types();
161        };
162        self
163    }
164
165    /// Configures the optional module filename for easy inclusion of all generated Rust files
166    ///
167    /// If set, generates a file (inside the `OUT_DIR` or `out_dir()` as appropriate) which contains
168    /// a set of `pub mod XXX` statements combining to load all Rust files generated.  This can allow
169    /// for a shortcut where multiple related proto files have been compiled together resulting in
170    /// a semi-complex set of includes.
171    pub fn include_file(mut self, path: impl AsRef<Path>) -> Self {
172        self.prost.include_file(path.as_ref());
173        self
174    }
175
176    /// When set, the `FileDescriptorSet` generated by protoc is written to the provided filesystem path.
177    ///
178    /// This option can be used in conjunction with the `include_bytes!` macro and the types in the `prost-types` crate
179    /// for implementing reflection capabilities, among other things.
180    pub fn file_descriptor_set_path(mut self, path: impl Into<PathBuf>) -> Self {
181        self.file_descriptor_set_path = Some(path.into());
182        self
183    }
184}