1 files changed, 26 insertions, 9 deletions
diff --git a/tools/distrib/python/make_grpcio_tools.py b/tools/distrib/python/make_grpcio_tools.py
index 50fbdbb14c..fd9b38b084 100755
--- a/tools/distrib/python/make_grpcio_tools.py
+++ b/tools/distrib/python/make_grpcio_tools.py
@@ -67,11 +67,16 @@ DEPS_FILE_CONTENT="""
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 # AUTO-GENERATED BY make_grpcio_tools.py!
-CC_FILES={}
+CC_FILES={cc_files}
+PROTO_FILES={proto_files}
+
+CC_INCLUDE={cc_include}
+PROTO_INCLUDE={proto_include}
 """
 
 # Bazel query result prefix for expected source files in protobuf.
 PROTOBUF_CC_PREFIX = '//:src/'
+PROTOBUF_PROTO_PREFIX = '//:src/'
 
 GRPC_ROOT = os.path.abspath(
     os.path.join(os.path.dirname(os.path.abspath(__file__)),
@@ -79,7 +84,8 @@ GRPC_ROOT = os.path.abspath(
 
 GRPC_PYTHON_ROOT = os.path.join(GRPC_ROOT, 'tools/distrib/python/grpcio_tools')
 
-GRPC_PROTOBUF = os.path.join(GRPC_ROOT, 'third_party/protobuf/src')
+GRPC_PYTHON_PROTOBUF_RELATIVE_ROOT = 'third_party/protobuf/src'
+GRPC_PROTOBUF = os.path.join(GRPC_ROOT, GRPC_PYTHON_PROTOBUF_RELATIVE_ROOT)
 GRPC_PROTOC_PLUGINS = os.path.join(GRPC_ROOT, 'src/compiler')
 GRPC_PYTHON_PROTOBUF = os.path.join(GRPC_PYTHON_ROOT,
                                     'third_party/protobuf/src')
@@ -93,18 +99,29 @@ GRPC_PYTHON_INCLUDE = os.path.join(GRPC_PYTHON_ROOT, 'grpc_root/include')
 
 BAZEL_DEPS = os.path.join(GRPC_ROOT, 'tools/distrib/python/bazel_deps.sh')
 BAZEL_DEPS_PROTOC_LIB_QUERY = '//:protoc_lib'
+BAZEL_DEPS_COMMON_PROTOS_QUERY = '//:well_known_protos'
+
 
+def bazel_query(query):
+  output = subprocess.check_output([BAZEL_DEPS, query])
+  return output.splitlines()
 
-def get_deps(query):
+def get_deps():
   """Write the result of the bazel query `query` against protobuf to
      `out_file`."""
-  output = subprocess.check_output([BAZEL_DEPS, query])
-  output = output.splitlines()
+  cc_files_output = bazel_query(BAZEL_DEPS_PROTOC_LIB_QUERY)
   cc_files = [
-      name for name in output
+      name[len(PROTOBUF_CC_PREFIX):] for name in cc_files_output
       if name.endswith('.cc') and name.startswith(PROTOBUF_CC_PREFIX)]
-  cc_files = [cc_file[len(PROTOBUF_CC_PREFIX):] for cc_file in cc_files]
-  deps_file_content = DEPS_FILE_CONTENT.format(cc_files)
+  proto_files_output = bazel_query(BAZEL_DEPS_COMMON_PROTOS_QUERY)
+  proto_files = [
+      name[len(PROTOBUF_PROTO_PREFIX):] for name in proto_files_output
+      if name.endswith('.proto') and name.startswith(PROTOBUF_PROTO_PREFIX)]
+  deps_file_content = DEPS_FILE_CONTENT.format(
+      cc_files=cc_files,
+      proto_files=proto_files,
+      cc_include=repr(GRPC_PYTHON_PROTOBUF_RELATIVE_ROOT),
+      proto_include=repr(GRPC_PYTHON_PROTOBUF_RELATIVE_ROOT))
   return deps_file_content
 
 
@@ -123,7 +140,7 @@ def main():
   shutil.copytree(GRPC_INCLUDE, GRPC_PYTHON_INCLUDE)
 
   try:
-    protoc_lib_deps_content = get_deps(BAZEL_DEPS_PROTOC_LIB_QUERY)
+    protoc_lib_deps_content = get_deps()
   except Exception as error:
     # We allow this script to succeed even if we couldn't get the dependencies,
     # as then we can assume that even without a successful bazel run the