compile_build_api_proto: Add SDK protobuf compilation.

In the interest of avoiding more custom protobuf stories, we're adding
a compiled version of the protobuf bindings compatible with the version
of the protobuf in the SDK.

BUG=chromium:1181505
TEST=./compile_build_api_proto

Change-Id: Id8381d84c4a7369938551515608b62539960a2a9
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/2728425
Tested-by: Alex Klein <saklein@chromium.org>
Reviewed-by: LaMont Jones <lamontjones@chromium.org>
Commit-Queue: Alex Klein <saklein@chromium.org>
diff --git a/api/compile_build_api_proto.py b/api/compile_build_api_proto.py
index 331d0fd..35bf467 100644
--- a/api/compile_build_api_proto.py
+++ b/api/compile_build_api_proto.py
@@ -10,6 +10,7 @@
 
 from __future__ import print_function
 
+import enum
 import os
 import sys
 
@@ -19,15 +20,11 @@
 from chromite.lib import cros_logging as logging
 from chromite.lib import osutils
 
-
 assert sys.version_info >= (3, 6), 'This module requires Python 3.6+'
 
-
-_API_DIR = os.path.join(constants.CHROMITE_DIR, 'api')
 _CIPD_ROOT = os.path.join(constants.CHROMITE_DIR, '.cipd_bin')
-_PROTOC = os.path.join(_CIPD_ROOT, 'protoc')
-_PROTO_DIR = os.path.join(constants.CHROMITE_DIR, 'infra', 'proto')
 
+# Chromite's protobuf library version (third_party/google/protobuf).
 PROTOC_VERSION = '3.13.0'
 
 
@@ -39,8 +36,45 @@
   """A failure we can't recover from."""
 
 
-def _InstallProtoc():
+@enum.unique
+class ProtocVersion(enum.Enum):
+  """Enum for possible protoc versions."""
+  # The SDK version of the bindings use the protoc in the SDK, and so is
+  # compatible with the protobuf library in the SDK, i.e. the one installed
+  # via the ebuild.
+  SDK = enum.auto()
+  # The Chromite version of the bindings uses a protoc binary downloaded from
+  # CIPD that matches the version of the protobuf library in
+  # chromite/third_party/google/protobuf.
+  CHROMITE = enum.auto()
+
+
+def _get_gen_dir(protoc_version: ProtocVersion):
+  """Get the chromite/api directory path."""
+  if protoc_version is ProtocVersion.SDK:
+    return os.path.join(constants.CHROMITE_DIR, 'api', 'gen_sdk')
+  else:
+    return os.path.join(constants.CHROMITE_DIR, 'api', 'gen')
+
+
+def _get_protoc_command(protoc_version: ProtocVersion):
+  """Get the protoc command for the target protoc."""
+  if protoc_version is ProtocVersion.SDK:
+    return 'protoc'
+  else:
+    return os.path.join(_CIPD_ROOT, 'protoc')
+
+
+def _get_proto_dir(_protoc_version):
+  """Get the proto directory for the target protoc."""
+  return os.path.join(constants.CHROMITE_DIR, 'infra', 'proto')
+
+
+def _InstallProtoc(protoc_version: ProtocVersion):
   """Install protoc from CIPD."""
+  if protoc_version is not ProtocVersion.CHROMITE:
+    return
+
   logging.info('Installing protoc.')
   cmd = ['cipd', 'ensure']
   # Clean up the output.
@@ -71,7 +105,7 @@
   Args:
     directory (str): Path to be cleaned up.
   """
-  logging.info('Cleaning old files.')
+  logging.info('Cleaning old files from %s.', directory)
   for dirpath, _dirnames, filenames in os.walk(directory):
     old = [os.path.join(dirpath, f) for f in filenames if f.endswith('_pb2.py')]
     # Remove empty init files to clean up otherwise empty directories.
@@ -84,14 +118,17 @@
       osutils.SafeUnlink(current)
 
 
-def _GenerateFiles(source, output):
+def _GenerateFiles(source: str, output: str, protoc_version: ProtocVersion):
   """Generate the proto files from the |source| tree into |output|.
 
   Args:
-    source (str): Path to the proto source root directory.
-    output (str): Path to the output root directory.
+    source: Path to the proto source root directory.
+    output: Path to the output root directory.
+    protoc_version: Which protoc to use.
   """
-  logging.info('Generating files.')
+  logging.info('Generating files to %s.', output)
+  osutils.SafeMakedirs(output)
+
   targets = []
 
   # Only compile the subset we need for the API.
@@ -110,9 +147,21 @@
           # We have a match, add the file.
           targets.append(os.path.join(dirpath, filename))
 
-  cmd = [_PROTOC, '--python_out', output, '--proto_path', source] + targets
+  cmd = [
+      _get_protoc_command(protoc_version),
+      '--python_out',
+      output,
+      '--proto_path',
+      source,
+  ]
+  cmd.extend(targets)
+
   result = cros_build_lib.run(
-      cmd, cwd=source, print_cmd=False, check=False)
+      cmd,
+      cwd=source,
+      print_cmd=False,
+      check=False,
+      enter_chroot=protoc_version is ProtocVersion.SDK)
 
   if result.returncode:
     raise GenerationError('Error compiling the proto. See the output for a '
@@ -121,20 +170,21 @@
 
 def _InstallMissingInits(directory):
   """Add any __init__.py files not present in the generated protobuf folders."""
-  logging.info('Adding missing __init__.py files.')
+  logging.info('Adding missing __init__.py files in %s.', directory)
   for dirpath, _dirnames, filenames in os.walk(directory):
     if '__init__.py' not in filenames:
       osutils.Touch(os.path.join(dirpath, '__init__.py'))
 
 
-def _PostprocessFiles(directory):
+def _PostprocessFiles(directory: str, protoc_version: ProtocVersion):
   """Do postprocessing on the generated files.
 
   Args:
-    directory (str): The root directory containing the generated files that are
+    directory: The root directory containing the generated files that are
       to be processed.
+    protoc_version: Which protoc is being used to generate the files.
   """
-  logging.info('Postprocessing: Fix imports.')
+  logging.info('Postprocessing: Fix imports in %s.', directory)
   # We are using a negative address here (the /address/! portion of the sed
   # command) to make sure we don't change any imports from protobuf itself.
   address = '^from google.protobuf'
@@ -145,8 +195,12 @@
   #   - [^ ] = Not a space. The [:space:] character set is too broad, but would
   #       technically work too.
   find = r'^from \([^ ]*\) import \([^ ]*\)_pb2 as \([^ ]*\)$'
-  # Substitute: 'from chromite.api.gen.x import y_pb2 as x_dot_y_pb2'.
-  sub = 'from chromite.api.gen.\\1 import \\2_pb2 as \\3'
+  # Substitute: 'from chromite.api.gen[_sdk].x import y_pb2 as x_dot_y_pb2'.
+  if protoc_version is ProtocVersion.SDK:
+    sub = 'from chromite.api.gen_sdk.\\1 import \\2_pb2 as \\3'
+  else:
+    sub = 'from chromite.api.gen.\\1 import \\2_pb2 as \\3'
+
   from_sed = [
       'sed', '-i',
       '/%(address)s/!s/%(find)s/%(sub)s/g' % {
@@ -157,14 +211,14 @@
   ]
 
   for dirpath, _dirnames, filenames in os.walk(directory):
-    # Update the
+    # Update the imports in the generated files.
     pb2 = [os.path.join(dirpath, f) for f in filenames if f.endswith('_pb2.py')]
     if pb2:
       cmd = from_sed + pb2
       cros_build_lib.run(cmd, print_cmd=False)
 
 
-def CompileProto(output=None):
+def CompileProto(output: str, protoc_version: ProtocVersion):
   """Compile the Build API protobuf files.
 
   By default this will compile from infra/proto/src to api/gen. The output
@@ -172,26 +226,60 @@
   in the default location.
 
   Args:
-    output (str|None): The output directory.
+    output: The output directory.
+    protoc_version: Which protoc to use for the compile.
   """
-  source = os.path.join(_PROTO_DIR, 'src')
-  output = output or os.path.join(_API_DIR, 'gen')
+  source = os.path.join(_get_proto_dir(protoc_version), 'src')
+  protoc_version = protoc_version or ProtocVersion.CHROMITE
 
-  _InstallProtoc()
+  _InstallProtoc(protoc_version)
   _CleanTargetDirectory(output)
-  _GenerateFiles(source, output)
+  _GenerateFiles(source, output, protoc_version)
   _InstallMissingInits(output)
-  _PostprocessFiles(output)
+  _PostprocessFiles(output, protoc_version)
 
 
 def GetParser():
   """Build the argument parser."""
   parser = commandline.ArgumentParser(description=__doc__)
-  parser.add_argument(
+  standard_group = parser.add_argument_group(
+      'Committed Bindings',
+      description='Options for generating the bindings in chromite/api/.')
+  standard_group.add_argument(
+      '--chromite',
+      dest='protoc_version',
+      action='append_const',
+      const=ProtocVersion.CHROMITE,
+      help='Generate only the chromite bindings. Generates all by default. The '
+           'chromite bindings are compatible with the version of protobuf in '
+           'chromite/third_party.')
+  standard_group.add_argument(
+      '--sdk',
+      dest='protoc_version',
+      action='append_const',
+      const=ProtocVersion.SDK,
+      help='Generate only the SDK bindings. Generates all by default. The SDK '
+           'bindings are compiled by protoc in the SDK, and is compatible '
+           'with the version of protobuf in the SDK (i.e. the one installed by '
+           'the ebuild).')
+
+  dest_group = parser.add_argument_group(
+      'Out of Tree Bindings',
+      description='Options for generating bindings in a custom location.')
+  dest_group.add_argument(
       '--destination',
       type='path',
-      help='The directory where the proto should be generated. Defaults to '
-           'the correct directory for the API.')
+      help='A directory where a single version of the proto should be '
+           'generated. When not given, the proto generates in all default '
+           'locations instead.')
+  dest_group.add_argument(
+      '--dest-sdk',
+      action='store_const',
+      dest='dest_protoc',
+      default=ProtocVersion.CHROMITE,
+      const=ProtocVersion.SDK,
+      help='Generate the SDK version of the protos in --destination instead of '
+           'the chromite version.')
   return parser
 
 
@@ -200,6 +288,9 @@
   parser = GetParser()
   opts = parser.parse_args(argv)
 
+  if not opts.protoc_version:
+    opts.protoc_version = [ProtocVersion.CHROMITE, ProtocVersion.SDK]
+
   opts.Freeze()
   return opts
 
@@ -207,8 +298,40 @@
 def main(argv):
   opts = _ParseArguments(argv)
 
-  try:
-    CompileProto(output=opts.destination)
-  except Error as e:
-    logging.error(e)
-    return 1
+  if opts.destination:
+    # Destination set, only compile a single version in the destination.
+    try:
+      CompileProto(output=opts.destination, protoc_version=opts.dest_protoc)
+    except Error as e:
+      cros_build_lib.Die('Error compiling bindings to destination: %s', str(e))
+    else:
+      return 0
+
+  if ProtocVersion.CHROMITE in opts.protoc_version:
+    # Compile the chromite bindings.
+    try:
+      CompileProto(
+          output=_get_gen_dir(ProtocVersion.CHROMITE),
+          protoc_version=ProtocVersion.CHROMITE)
+    except Error as e:
+      cros_build_lib.Die('Error compiling chromite bindings: %s', str(e))
+
+  if ProtocVersion.SDK in opts.protoc_version:
+    # Compile the SDK bindings.
+    if not cros_build_lib.IsInsideChroot():
+      # Rerun inside of the SDK instead of trying to map all of the paths.
+      cmd = [
+          os.path.join(constants.CHROOT_SOURCE_ROOT, 'chromite', 'api',
+                       'compile_build_api_proto'),
+          '--sdk',
+      ]
+      result = cros_build_lib.run(
+          cmd, print_cmd=False, enter_chroot=True, check=False)
+      return result.returncode
+    else:
+      try:
+        CompileProto(
+            output=_get_gen_dir(ProtocVersion.SDK),
+            protoc_version=ProtocVersion.SDK)
+      except Error as e:
+        cros_build_lib.Die('Error compiling SDK bindings: %s', str(e))