From 3f9ab801f7ee50ec04ab0f73cd457898dc687e61 Mon Sep 17 00:00:00 2001
From: Derek Strickland <1111455+DerekStrickland@users.noreply.github.com>
Date: Wed, 20 Oct 2021 10:11:01 -0400
Subject: [PATCH] feat(build): Support prost's include_file option (#774)

---
 tonic-build/Cargo.toml   |  2 +-
 tonic-build/src/prost.rs | 16 ++++++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/tonic-build/Cargo.toml b/tonic-build/Cargo.toml
index 92ecfce7e..433c4b6d5 100644
--- a/tonic-build/Cargo.toml
+++ b/tonic-build/Cargo.toml
@@ -29,4 +29,4 @@ prost = ["prost-build"]
 compression = []
 
 [package.metadata.docs.rs]
-all-features = true
+all-features = true
\ No newline at end of file
diff --git a/tonic-build/src/prost.rs b/tonic-build/src/prost.rs
index 251dace7a..2a346f395 100644
--- a/tonic-build/src/prost.rs
+++ b/tonic-build/src/prost.rs
@@ -26,6 +26,7 @@ pub fn configure() -> Builder {
         format: true,
         emit_package: true,
         protoc_args: Vec::new(),
+        include_file: None,
     }
 }
 
@@ -225,6 +226,7 @@ pub struct Builder {
     pub(crate) emit_package: bool,
     pub(crate) compile_well_known_types: bool,
     pub(crate) protoc_args: Vec<OsString>,
+    pub(crate) include_file: Option<PathBuf>,
 
     out_dir: Option<PathBuf>,
     #[cfg(feature = "rustfmt")]
@@ -367,6 +369,17 @@ impl Builder {
         self
     }
 
+    /// Configures the optional module filename for easy inclusion of all generated Rust files
+    ///
+    /// If set, generates a file (inside the `OUT_DIR` or `out_dir()` as appropriate) which contains
+    /// a set of `pub mod XXX` statements combining to load all Rust files generated.  This can allow
+    /// for a shortcut where multiple related proto files have been compiled together resulting in
+    /// a semi-complex set of includes.
+    pub fn include_file(mut self, path: impl AsRef<Path>) -> Self {
+        self.include_file = Some(path.as_ref().to_path_buf());
+        self
+    }
+
     /// Compile the .proto files and execute code generation.
     pub fn compile<P>(self, protos: &[P], includes: &[P]) -> io::Result<()>
     where
@@ -411,6 +424,9 @@ impl Builder {
         if self.compile_well_known_types {
             config.compile_well_known_types();
         }
+        if let Some(path) = self.include_file.as_ref() {
+            config.include_file(path);
+        }
 
         for arg in self.protoc_args.iter() {
             config.protoc_arg(arg);