api: add type database rule and use it from protoxform (envoyproxy#9276)

This PR bazelifies the type_db generation in type_database rule, and defines default_type_db for protoxform (and other api_proto_plugin) to take them.

This makes type_db cacheable and protoxform remote runnable because Bazel is now aware of full dependency graph. This should make format CI faster too.

Risk Level: Low
Testing: CI

Signed-off-by: Lizan Zhou <[email protected]>

Author: lizan

tools/api_proto_plugin/BUILD

@@ -1,5 +1,8 @@
 licenses(["notice"])  # Apache 2
 
+load("@rules_python//python:defs.bzl", "py_library")
+load("//tools/type_whisperer:type_database.bzl", "type_database")
+
 py_library(
     name = "api_proto_plugin",
     srcs = [
@@ -22,3 +25,16 @@ py_library(
     srcs_version = "PY3",
     visibility = ["//visibility:public"],
 )
+
+label_flag(
+    name = "default_type_db_target",
+    # While this is not completely empty but type_db_gen generates nothing on this target.
+    build_setting_default = "@com_google_protobuf//:empty_proto",
+    visibility = ["//visibility:public"],
+)
+
+type_database(
+    name = "default_type_db",
+    targets = [":default_type_db_target"],
+    visibility = ["//visibility:public"],
+)

tools/api_proto_plugin/plugin.bzl

@@ -1,3 +1,5 @@
+load("@rules_proto//proto:defs.bzl", "ProtoInfo")
+
 # Borrowed from https://github.com/grpc/grpc-java/blob/v1.24.1/java_grpc_library.bzl#L61
 def _path_ignoring_repository(f):
     # Bazel creates a _virtual_imports directory in case the .proto source files
@@ -43,16 +45,24 @@ def api_proto_plugin_impl(target, ctx, output_group, mnemonic, output_suffixes):
                                              output_suffix) for f in proto_sources]
 
     # Create the protoc command-line args.
+    inputs = target[ProtoInfo].transitive_sources
     ctx_path = ctx.label.package + "/" + ctx.label.name
     output_path = outputs[0].root.path + "/" + outputs[0].owner.workspace_root + "/" + ctx_path
     args = ["-I./" + ctx.label.workspace_root]
     args += ["-I" + import_path for import_path in import_paths]
     args += ["--plugin=protoc-gen-api_proto_plugin=" + ctx.executable._api_proto_plugin.path, "--api_proto_plugin_out=" + output_path]
+    if hasattr(ctx.attr, "_type_db"):
+        inputs = depset(transitive = [inputs] + [ctx.attr._type_db.files])
+        if len(ctx.attr._type_db.files.to_list()) != 1:
+            fail("{} must have one type database file".format(ctx.attr._type_db))
+        args += ["--api_proto_plugin_opt=type_db_path=" + ctx.attr._type_db.files.to_list()[0].path]
     args += [src.path for src in target[ProtoInfo].direct_sources]
+    env = {}
+
     ctx.actions.run(
         executable = ctx.executable._protoc,
         arguments = args,
-        inputs = target[ProtoInfo].transitive_sources,
+        inputs = inputs,
         tools = [ctx.executable._api_proto_plugin],
         outputs = outputs,
         mnemonic = mnemonic,
@@ -62,20 +72,25 @@ def api_proto_plugin_impl(target, ctx, output_group, mnemonic, output_suffixes):
     transitive_outputs = depset(outputs, transitive = [transitive_outputs])
     return [OutputGroupInfo(**{output_group: transitive_outputs})]
 
-def api_proto_plugin_aspect(tool_label, aspect_impl):
+def api_proto_plugin_aspect(tool_label, aspect_impl, use_type_db = False):
+    _attrs = {
+        "_protoc": attr.label(
+            default = Label("@com_google_protobuf//:protoc"),
+            executable = True,
+            cfg = "exec",
+        ),
+        "_api_proto_plugin": attr.label(
+            default = Label(tool_label),
+            executable = True,
+            cfg = "exec",
+        ),
+    }
+    if use_type_db:
+        _attrs["_type_db"] = attr.label(
+            default = Label("@envoy//tools/api_proto_plugin:default_type_db"),
+        )
     return aspect(
         attr_aspects = ["deps"],
-        attrs = {
-            "_protoc": attr.label(
-                default = Label("@com_google_protobuf//:protoc"),
-                executable = True,
-                cfg = "exec",
-            ),
-            "_api_proto_plugin": attr.label(
-                default = Label(tool_label),
-                executable = True,
-                cfg = "exec",
-            ),
-        },
+        attrs = _attrs,
         implementation = aspect_impl,
     )

tools/api_proto_plugin/plugin.py

@@ -17,9 +17,9 @@
         # Output files are generated alongside their corresponding input .proto,
         # with the output_suffix appended.
         'output_suffix',
-        # The visitor is a visitor.Visitor defining the business logic of the plugin
-        # for the specific output descriptor.
-        'visitor',
+        # The visitor is a function to create a visitor.Visitor defining the business
+        # logic of the plugin for the specific output descriptor.
+        'visitor_factory',
         # FileDescriptorProto transformer; this is applied to the input
         # before any output generation.
         'xform',
@@ -30,7 +30,7 @@ def DirectOutputDescriptor(output_suffix, visitor):
   return OutputDescriptor(output_suffix, visitor, lambda x: x)
 
 
-def Plugin(output_descriptors):
+def Plugin(output_descriptors, parameter_callback=None):
   """Protoc plugin entry point.
 
   This defines protoc plugin and manages the stdin -> stdout flow. An
@@ -48,6 +48,9 @@ def Plugin(output_descriptors):
   response = plugin_pb2.CodeGeneratorResponse()
   cprofile_enabled = os.getenv('CPROFILE_ENABLED')
 
+  if request.HasField("parameter") and parameter_callback:
+    parameter_callback(request.parameter)
+
   # We use request.file_to_generate rather than request.file_proto here since we
   # are invoked inside a Bazel aspect, each node in the DAG will be visited once
   # by the aspect and we only want to generate docs for the current node.
@@ -61,7 +64,8 @@ def Plugin(output_descriptors):
       f = response.file.add()
       f.name = file_proto.name + od.output_suffix
       xformed_proto = od.xform(file_proto)
-      f.content = traverse.TraverseFile(od.xform(file_proto), od.visitor) if xformed_proto else ''
+      f.content = traverse.TraverseFile(xformed_proto,
+                                        od.visitor_factory()) if xformed_proto else ''
     if cprofile_enabled:
       pr.disable()
       stats_stream = io.StringIO()

tools/proto_format.sh

@@ -14,21 +14,15 @@ rm -rf bazel-bin/external/envoy_api
 declare -r PROTO_TARGETS=$(bazel query "labels(srcs, labels(deps, @envoy_api//docs:protos))")
 
 # This is for local RBE setup, should be no-op for builds without RBE setting in bazelrc files.
-BAZEL_BUILD_OPTIONS+=" --remote_download_outputs=all --strategy=protoxform=sandboxed,local"
+BAZEL_BUILD_OPTIONS+=" --remote_download_outputs=all"
 
 # TODO(htuch): This script started life by cloning docs/build.sh. It depends on
 # the @envoy_api//docs:protos target in a few places as a result. This is not
 # guaranteed to be the precise set of protos we want to format, but as a
 # starting place it seems reasonable. In the future, we should change the logic
 # here.
-bazel build ${BAZEL_BUILD_OPTIONS} @envoy_api//docs:protos --aspects \
-  tools/type_whisperer/type_whisperer.bzl%type_whisperer_aspect --output_groups=types_pb_text \
-  --host_force_python=PY3
-declare -x -r TYPE_DB_PATH="${PWD}"/source/common/config/api_type_db.generated.pb_text
-bazel run ${BAZEL_BUILD_OPTIONS} //tools/type_whisperer:typedb_gen -- \
-  ${PWD} ${TYPE_DB_PATH} ${PROTO_TARGETS}
-bazel build ${BAZEL_BUILD_OPTIONS} @envoy_api//docs:protos --aspects \
-  tools/protoxform/protoxform.bzl%protoxform_aspect --output_groups=proto --action_env=CPROFILE_ENABLED=1 \
-  --action_env=TYPE_DB_PATH --host_force_python=PY3
+bazel build ${BAZEL_BUILD_OPTIONS} --//tools/api_proto_plugin:default_type_db_target=@envoy_api//docs:protos \
+  @envoy_api//docs:protos --aspects //tools/protoxform:protoxform.bzl%protoxform_aspect --output_groups=proto \
+  --action_env=CPROFILE_ENABLED=1 --host_force_python=PY3
 
 ./tools/proto_sync.py "$1" ${PROTO_TARGETS}

tools/protodoc/protodoc.py

@@ -562,7 +562,7 @@ def VisitFile(self, file_proto, type_context, services, msgs, enums):
 
 
 def Main():
-  plugin.Plugin([plugin.DirectOutputDescriptor('.rst', RstFormatVisitor())])
+  plugin.Plugin([plugin.DirectOutputDescriptor('.rst', RstFormatVisitor)])
 
 
 if __name__ == '__main__':

tools/protoxform/migrate.py

@@ -153,7 +153,7 @@ def V3MigrationXform(file_proto):
     v3 FileDescriptorProto message.
   """
   # Load type database.
-  typedb = utils.LoadTypeDb()
+  typedb = utils.GetTypeDb()
   # If this isn't a proto in an upgraded package, return None.
   if file_proto.package not in typedb.next_version_packages or not typedb.next_version_packages[
       file_proto.package]:

tools/protoxform/protoxform.bzl

@@ -9,4 +9,4 @@ def _protoxform_impl(target, ctx):
 #
 #   bazel build //api --aspects tools/protoxform/protoxform.bzl%protoxform_aspect \
 #       --output_groups=proto
-protoxform_aspect = api_proto_plugin_aspect("//tools/protoxform", _protoxform_impl)
+protoxform_aspect = api_proto_plugin_aspect("//tools/protoxform", _protoxform_impl, use_type_db = True)

tools/protoxform/protoxform.py

@@ -143,7 +143,7 @@ def FormatHeaderFromFile(source_code_info, file_proto):
     Formatted proto header as a string.
   """
   # Load the type database.
-  typedb = utils.LoadTypeDb()
+  typedb = utils.GetTypeDb()
   # Figure out type dependencies in this .proto.
   types = Types()
   text_format.Merge(traverse.TraverseFile(file_proto, type_whisperer.TypeWhispererVisitor()), types)
@@ -524,11 +524,17 @@ def VisitFile(self, file_proto, type_context, services, msgs, enums):
     return ClangFormat(header + formatted_services + formatted_enums + formatted_msgs)
 
 
+def ParameterCallback(parameter):
+  params = dict(param.split('=') for param in parameter.split(','))
+  if params["type_db_path"]:
+    utils.LoadTypeDb(params["type_db_path"])
+
+
 def Main():
   plugin.Plugin([
-      plugin.DirectOutputDescriptor('.v2.proto', ProtoFormatVisitor()),
-      plugin.OutputDescriptor('.v3alpha.proto', ProtoFormatVisitor(), migrate.V3MigrationXform)
-  ])
+      plugin.DirectOutputDescriptor('.v2.proto', ProtoFormatVisitor),
+      plugin.OutputDescriptor('.v3alpha.proto', ProtoFormatVisitor, migrate.V3MigrationXform)
+  ], ParameterCallback)
 
 
 if __name__ == '__main__':

tools/protoxform/utils.py

@@ -4,9 +4,16 @@
 
 from google.protobuf import text_format
 
+_typedb = None
 
-def LoadTypeDb():
-  typedb = TypeDb()
-  with open(os.getenv('TYPE_DB_PATH'), 'r') as f:
-    text_format.Merge(f.read(), typedb)
-  return typedb
+
+def GetTypeDb():
+  assert _typedb != None
+  return _typedb
+
+
+def LoadTypeDb(type_db_path):
+  global _typedb
+  _typedb = TypeDb()
+  with open(type_db_path, 'r') as f:
+    text_format.Merge(f.read(), _typedb)

tools/type_whisperer/type_database.bzl

@@ -0,0 +1,44 @@
+load(":type_whisperer.bzl", "type_whisperer_aspect")
+
+def _type_database_impl(ctx):
+    type_db_deps = []
+    for target in ctx.attr.targets:
+        type_db_deps.append(target[OutputGroupInfo].types_pb_text)
+
+    type_db_deps = depset(transitive = type_db_deps)
+
+    args = [ctx.outputs.pb_text.path]
+    for dep in type_db_deps.to_list():
+        if dep.owner.workspace_name in ctx.attr.proto_repositories:
+            args.append(dep.path)
+
+    ctx.actions.run(
+        executable = ctx.executable._type_db_gen,
+        arguments = args,
+        inputs = type_db_deps,
+        outputs = [ctx.outputs.pb_text],
+        mnemonic = "TypeDbGen",
+        use_default_shell_env = True,
+    )
+
+type_database = rule(
+    attrs = {
+        "targets": attr.label_list(
+            aspects = [type_whisperer_aspect],
+            doc = "List of all proto_library target to be included.",
+        ),
+        "proto_repositories": attr.string_list(
+            default = ["envoy_api"],
+            allow_empty = False,
+        ),
+        "_type_db_gen": attr.label(
+            default = Label("//tools/type_whisperer:typedb_gen"),
+            executable = True,
+            cfg = "exec",
+        ),
+    },
+    outputs = {
+        "pb_text": "%{name}.pb_text",
+    },
+    implementation = _type_database_impl,
+)

tools/type_whisperer/type_whisperer.py

@@ -46,7 +46,7 @@ def VisitFile(self, file_proto, type_context, services, msgs, enums):
 
 def Main():
   plugin.Plugin([
-      plugin.DirectOutputDescriptor('.types.pb_text', TypeWhispererVisitor()),
+      plugin.DirectOutputDescriptor('.types.pb_text', TypeWhispererVisitor),
   ])
 
 

tools/type_whisperer/typedb_gen.py

@@ -7,7 +7,6 @@
 
 from google.protobuf import text_format
 
-from tools.api_proto_plugin.utils import BazelBinPathForOutputArtifact
 from tools.type_whisperer.api_type_db_pb2 import TypeDb
 from tools.type_whisperer.types_pb2 import Types, TypeDescription
 
@@ -118,18 +117,13 @@ def NextVersionUpgrade(type_name, type_map, next_version_upgrade_memo, visited=N
 
 
 if __name__ == '__main__':
-  # Root of source tree.
-  src_root = sys.argv[1]
   # Output path for type database.
-  out_path = sys.argv[2]
-  # Bazel labels for source .proto.
-  src_labels = sys.argv[3:]
-
-  # Load type descriptors for each .proto.
-  type_desc_paths = [
-      BazelBinPathForOutputArtifact(label, '.types.pb_text', root=src_root) for label in src_labels
-  ]
+  out_path = sys.argv[1]
+
+  # Load type descriptors for each type whisper
+  type_desc_paths = sys.argv[2:]
   type_whispers = map(LoadTypes, type_desc_paths)
+
   # Aggregate type descriptors to a single type map.
   type_map = dict(sum([list(t.types.items()) for t in type_whispers], []))
   all_pkgs = set([type_desc.qualified_package for type_desc in type_map.values()])