From ed4d01e3e611ef99b0517c65b530905ab4b3a5a4 Mon Sep 17 00:00:00 2001 From: Tom Dyas Date: Wed, 2 Mar 2022 18:55:26 -0500 Subject: [PATCH] [internal] Add Go Protobuf compiler (#13985) Add support for generating Go from protobuf sources using the [protoc-gen-go](https://pkg.go.dev/google.golang.org/protobuf/cmd/protoc-gen-go) and [protoc-gen-go-grpc](https://pkg.go.dev/google.golang.org/grpc/cmd/protoc-gen-go-grpc) plugins to `protoc`. This is not actually wired up yet to Go because we need to solve https://github.com/pantsbuild/pants/issues/14258, but it does give us the technology to generate the `.go` files. Note that this adds a new technique for us to install a Go tool deterministically. For now, the `go.mod` and `go.sum` are hardcoded, but we can choose to expose this through the options system in the future if need be. [ci skip-rust] --- .../pants/backend/codegen/protobuf/go/BUILD | 6 + .../backend/codegen/protobuf/go/__init__.py | 0 .../backend/codegen/protobuf/go/rules.py | 251 ++++++++++++++++++ .../protobuf/go/rules_integration_test.py | 201 ++++++++++++++ src/python/pants/backend/go/util_rules/sdk.py | 13 +- 5 files changed, 467 insertions(+), 4 deletions(-) create mode 100644 src/python/pants/backend/codegen/protobuf/go/BUILD create mode 100644 src/python/pants/backend/codegen/protobuf/go/__init__.py create mode 100644 src/python/pants/backend/codegen/protobuf/go/rules.py create mode 100644 src/python/pants/backend/codegen/protobuf/go/rules_integration_test.py diff --git a/src/python/pants/backend/codegen/protobuf/go/BUILD b/src/python/pants/backend/codegen/protobuf/go/BUILD new file mode 100644 index 00000000000..e54f5d10bf6 --- /dev/null +++ b/src/python/pants/backend/codegen/protobuf/go/BUILD @@ -0,0 +1,6 @@ +# Copyright 2021 Pants project contributors (see CONTRIBUTORS.md). +# Licensed under the Apache License, Version 2.0 (see LICENSE). + +python_sources() + +python_tests(name="tests") diff --git a/src/python/pants/backend/codegen/protobuf/go/__init__.py b/src/python/pants/backend/codegen/protobuf/go/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/src/python/pants/backend/codegen/protobuf/go/rules.py b/src/python/pants/backend/codegen/protobuf/go/rules.py new file mode 100644 index 00000000000..03b218368d6 --- /dev/null +++ b/src/python/pants/backend/codegen/protobuf/go/rules.py @@ -0,0 +1,251 @@ +# Copyright 2021 Pants project contributors (see CONTRIBUTORS.md). +# Licensed under the Apache License, Version 2.0 (see LICENSE). +from __future__ import annotations + +import os +from dataclasses import dataclass + +from pants.backend.codegen.protobuf.protoc import Protoc +from pants.backend.codegen.protobuf.target_types import ProtobufGrpcToggleField, ProtobufSourceField +from pants.backend.go.target_types import GoPackageSourcesField +from pants.backend.go.util_rules.sdk import GoSdkProcess +from pants.core.util_rules.external_tool import DownloadedExternalTool, ExternalToolRequest +from pants.core.util_rules.source_files import SourceFilesRequest +from pants.core.util_rules.stripped_source_files import StrippedSourceFiles +from pants.engine.fs import ( + AddPrefix, + CreateDigest, + Digest, + Directory, + FileContent, + MergeDigests, + RemovePrefix, + Snapshot, +) +from pants.engine.internals.native_engine import EMPTY_DIGEST +from pants.engine.internals.selectors import Get, MultiGet +from pants.engine.platform import Platform +from pants.engine.process import Process, ProcessResult +from pants.engine.rules import collect_rules, rule +from pants.engine.target import ( + GeneratedSources, + GenerateSourcesRequest, + TransitiveTargets, + TransitiveTargetsRequest, +) +from pants.engine.unions import UnionRule +from pants.source.source_root import SourceRoot, SourceRootRequest +from pants.util.logging import LogLevel + + +class GenerateGoFromProtobufRequest(GenerateSourcesRequest): + input = ProtobufSourceField + output = GoPackageSourcesField + + +@dataclass(frozen=True) +class SetupGoProtocPlugin: + digest: Digest + + +@rule(desc="Generate Go from Protobuf", level=LogLevel.DEBUG) +async def generate_go_from_protobuf( + request: GenerateGoFromProtobufRequest, + protoc: Protoc, + go_protoc_plugin: SetupGoProtocPlugin, +) -> GeneratedSources: + output_dir = "_generated_files" + protoc_relpath = "__protoc" + protoc_go_plugin_relpath = "__protoc_gen_go" + + downloaded_protoc_binary, empty_output_dir, transitive_targets = await MultiGet( + Get(DownloadedExternalTool, ExternalToolRequest, protoc.get_request(Platform.current)), + Get(Digest, CreateDigest([Directory(output_dir)])), + Get(TransitiveTargets, TransitiveTargetsRequest([request.protocol_target.address])), + ) + + # NB: By stripping the source roots, we avoid having to set the value `--proto_path` + # for Protobuf imports to be discoverable. + all_sources_stripped, target_sources_stripped = await MultiGet( + Get( + StrippedSourceFiles, + SourceFilesRequest( + tgt[ProtobufSourceField] + for tgt in transitive_targets.closure + if tgt.has_field(ProtobufSourceField) + ), + ), + Get( + StrippedSourceFiles, SourceFilesRequest([request.protocol_target[ProtobufSourceField]]) + ), + ) + + input_digest = await Get( + Digest, + MergeDigests( + [ + all_sources_stripped.snapshot.digest, + empty_output_dir, + ] + ), + ) + + maybe_grpc_plugin_args = [] + if request.protocol_target.get(ProtobufGrpcToggleField).value: + maybe_grpc_plugin_args = [ + f"--go-grpc_out={output_dir}", + "--go-grpc_opt=paths=source_relative", + ] + + result = await Get( + ProcessResult, + Process( + argv=[ + os.path.join(protoc_relpath, downloaded_protoc_binary.exe), + f"--plugin=go={os.path.join('.', protoc_go_plugin_relpath, 'protoc-gen-go')}", + f"--plugin=go-grpc={os.path.join('.', protoc_go_plugin_relpath, 'protoc-gen-go-grpc')}", + f"--go_out={output_dir}", + "--go_opt=paths=source_relative", + *maybe_grpc_plugin_args, + *target_sources_stripped.snapshot.files, + ], + # Note: Necessary or else --plugin option needs absolute path. + env={"PATH": protoc_go_plugin_relpath}, + input_digest=input_digest, + immutable_input_digests={ + protoc_relpath: downloaded_protoc_binary.digest, + protoc_go_plugin_relpath: go_protoc_plugin.digest, + }, + description=f"Generating Go sources from {request.protocol_target.address}.", + level=LogLevel.DEBUG, + output_directories=(output_dir,), + ), + ) + + normalized_digest, source_root = await MultiGet( + Get(Digest, RemovePrefix(result.output_digest, output_dir)), + Get(SourceRoot, SourceRootRequest, SourceRootRequest.for_target(request.protocol_target)), + ) + + source_root_restored = ( + await Get(Snapshot, AddPrefix(normalized_digest, source_root.path)) + if source_root.path != "." + else await Get(Snapshot, Digest, normalized_digest) + ) + return GeneratedSources(source_root_restored) + + +# Note: The versions of the Go protoc and gRPC plugins are hard coded in the following go.mod. To update, +# copy the following go.mod and go.sum contents to go.mod and go.sum files in a new directory. Then update the +# versions and run `go mod download all`. Copy the go.mod and go.sum contents back into these constants, +# making sure to replace tabs with `\t`. + +GO_PROTOBUF_GO_MOD = """\ +module org.pantsbuild.backend.go.protobuf + +go 1.17 + +require ( +\tgoogle.golang.org/grpc/cmd/protoc-gen-go-grpc v1.2.0 +\tgoogle.golang.org/protobuf v1.27.1 +) + +require ( +\tgithub.com/golang/protobuf v1.5.0 // indirect +\tgithub.com/google/go-cmp v0.5.5 // indirect +\tgolang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 // indirect +) +""" + +GO_PROTOBUF_GO_SUM = """\ +github.com/golang/protobuf v1.5.0 h1:LUVKkCeviFUMKqHa4tXIIij/lbhnMbP7Fn5wKdKkRh4= +github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= +github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU= +github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/grpc v1.2.0 h1:v8eFdETH8nqZHQ9x+0f2PLuU6W7zo5PFZuVEwH5126Y= +google.golang.org/grpc v1.2.0/go.mod h1:yo6s7OP7yaDglbqo1J04qKzAhqBH6lvTonzMVmEdcZw= +google.golang.org/grpc/cmd/protoc-gen-go-grpc v1.2.0 h1:TLkBREm4nIsEcexnCjgQd5GQWaHcqMzwQV0TX9pq8S0= +google.golang.org/grpc/cmd/protoc-gen-go-grpc v1.2.0/go.mod h1:DNq5QpG7LJqD2AamLZ7zvKE0DEpVl2BSEVjFycAAjRY= +google.golang.org/protobuf v1.27.1 h1:SnqbnDw1V7RiZcXPx5MEeqPv2s79L9i7BJUlG/+RurQ= +google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc= +""" + + +@rule +async def setup_go_protoc_plugin(platform: Platform) -> SetupGoProtocPlugin: + go_mod_digest = await Get( + Digest, + CreateDigest( + [ + FileContent("go.mod", GO_PROTOBUF_GO_MOD.encode()), + FileContent("go.sum", GO_PROTOBUF_GO_SUM.encode()), + ] + ), + ) + + download_sources_result = await Get( + ProcessResult, + GoSdkProcess( + ["mod", "download", "all"], + input_digest=go_mod_digest, + output_directories=("gopath",), + description="Download Go `protoc` plugin sources.", + allow_downloads=True, + ), + ) + + go_plugin_build_result, go_grpc_plugin_build_result = await MultiGet( + Get( + ProcessResult, + GoSdkProcess( + ["install", "google.golang.org/protobuf/cmd/protoc-gen-go@v1.27.1"], + input_digest=download_sources_result.output_digest, + output_files=["gopath/bin/protoc-gen-go"], + description="Build Go protobuf plugin for `protoc`.", + platform=platform, + ), + ), + Get( + ProcessResult, + GoSdkProcess( + [ + "install", + "google.golang.org/grpc/cmd/protoc-gen-go-grpc@v1.2.0", + ], + input_digest=download_sources_result.output_digest, + output_files=["gopath/bin/protoc-gen-go-grpc"], + description="Build Go gRPC protobuf plugin for `protoc`.", + platform=platform, + ), + ), + ) + if go_plugin_build_result.output_digest == EMPTY_DIGEST: + raise AssertionError( + f"Failed to build protoc-gen-go:\n" + f"stdout:\n{go_plugin_build_result.stdout.decode()}\n\n" + f"stderr:\n{go_plugin_build_result.stderr.decode()}" + ) + if go_grpc_plugin_build_result.output_digest == EMPTY_DIGEST: + raise AssertionError( + f"Failed to build protoc-gen-go-grpc:\n" + f"stdout:\n{go_grpc_plugin_build_result.stdout.decode()}\n\n" + f"stderr:\n{go_grpc_plugin_build_result.stderr.decode()}" + ) + + merged_output_digests = await Get( + Digest, + MergeDigests( + [go_plugin_build_result.output_digest, go_grpc_plugin_build_result.output_digest] + ), + ) + plugin_digest = await Get(Digest, RemovePrefix(merged_output_digests, "gopath/bin")) + return SetupGoProtocPlugin(plugin_digest) + + +def rules(): + return ( + *collect_rules(), + UnionRule(GenerateSourcesRequest, GenerateGoFromProtobufRequest), + ) diff --git a/src/python/pants/backend/codegen/protobuf/go/rules_integration_test.py b/src/python/pants/backend/codegen/protobuf/go/rules_integration_test.py new file mode 100644 index 00000000000..6228a9a43b0 --- /dev/null +++ b/src/python/pants/backend/codegen/protobuf/go/rules_integration_test.py @@ -0,0 +1,201 @@ +# Copyright 2021 Pants project contributors (see CONTRIBUTORS.md). +# Licensed under the Apache License, Version 2.0 (see LICENSE). +from __future__ import annotations + +from textwrap import dedent +from typing import Iterable + +import pytest + +from pants.backend.codegen.protobuf.go.rules import GenerateGoFromProtobufRequest +from pants.backend.codegen.protobuf.go.rules import rules as go_protobuf_rules +from pants.backend.codegen.protobuf.target_types import ( + ProtobufSourceField, + ProtobufSourcesGeneratorTarget, +) +from pants.backend.codegen.protobuf.target_types import rules as protobuf_target_types_rules +from pants.backend.go.target_types import GoPackageTarget +from pants.backend.go.util_rules import sdk +from pants.build_graph.address import Address +from pants.core.util_rules import config_files, source_files, stripped_source_files +from pants.core.util_rules.external_tool import rules as external_tool_rules +from pants.engine.fs import Digest, DigestContents +from pants.engine.rules import QueryRule +from pants.engine.target import GeneratedSources, HydratedSources, HydrateSourcesRequest +from pants.jvm.jdk_rules import rules as jdk_rules +from pants.testutil.rule_runner import PYTHON_BOOTSTRAP_ENV, RuleRunner + + +@pytest.fixture +def rule_runner() -> RuleRunner: + rule_runner = RuleRunner( + rules=[ + *config_files.rules(), + *external_tool_rules(), + *source_files.rules(), + *jdk_rules(), + *protobuf_target_types_rules(), + *stripped_source_files.rules(), + *go_protobuf_rules(), + *sdk.rules(), + QueryRule(HydratedSources, [HydrateSourcesRequest]), + QueryRule(GeneratedSources, [GenerateGoFromProtobufRequest]), + QueryRule(DigestContents, (Digest,)), + ], + target_types=[ + GoPackageTarget, + ProtobufSourcesGeneratorTarget, + ], + ) + rule_runner.set_options( + [], + env_inherit=PYTHON_BOOTSTRAP_ENV, + ) + return rule_runner + + +def assert_files_generated( + rule_runner: RuleRunner, + address: Address, + *, + expected_files: list[str], + source_roots: list[str], + extra_args: Iterable[str] = (), +) -> None: + args = [f"--source-root-patterns={repr(source_roots)}", *extra_args] + rule_runner.set_options(args, env_inherit=PYTHON_BOOTSTRAP_ENV) + tgt = rule_runner.get_target(address) + protocol_sources = rule_runner.request( + HydratedSources, [HydrateSourcesRequest(tgt[ProtobufSourceField])] + ) + generated_sources = rule_runner.request( + GeneratedSources, + [GenerateGoFromProtobufRequest(protocol_sources.snapshot, tgt)], + ) + assert set(generated_sources.snapshot.files) == set(expected_files) + + +def test_generates_go(rule_runner: RuleRunner) -> None: + # This tests a few things: + # * We generate the correct file names. + # * Protobuf files can import other protobuf files, and those can import others + # (transitive dependencies). We'll only generate the requested target, though. + # * We can handle multiple source roots, which need to be preserved in the final output. + rule_runner.write_files( + { + "src/protobuf/dir1/f.proto": dedent( + """\ + syntax = "proto3"; + + option go_package = "example.com/dir1"; + + package dir1; + + message Person { + string name = 1; + int32 id = 2; + string email = 3; + } + """ + ), + "src/protobuf/dir1/f2.proto": dedent( + """\ + syntax = "proto3"; + + option go_package = "example.com/dir1"; + + package dir1; + """ + ), + "src/protobuf/dir1/BUILD": "protobuf_sources()", + "src/protobuf/dir2/f.proto": dedent( + """\ + syntax = "proto3"; + + option go_package = "example.com/dir2"; + + package dir2; + + import "dir1/f.proto"; + """ + ), + "src/protobuf/dir2/BUILD": ("protobuf_sources(dependencies=['src/protobuf/dir1'])"), + # Test another source root. + "tests/protobuf/test_protos/f.proto": dedent( + """\ + syntax = "proto3"; + + option go_package = "example.com/test_protos"; + + package test_protos; + + import "dir2/f.proto"; + """ + ), + "tests/protobuf/test_protos/BUILD": ( + "protobuf_sources(dependencies=['src/protobuf/dir2'])" + ), + } + ) + + def assert_gen(addr: Address, expected: Iterable[str]) -> None: + assert_files_generated( + rule_runner, + addr, + source_roots=["src/python", "/src/protobuf", "/tests/protobuf"], + expected_files=list(expected), + ) + + assert_gen( + Address("src/protobuf/dir1", relative_file_path="f.proto"), + ("src/protobuf/dir1/f.pb.go",), + ) + assert_gen( + Address("src/protobuf/dir1", relative_file_path="f2.proto"), + ("src/protobuf/dir1/f2.pb.go",), + ) + assert_gen( + Address("src/protobuf/dir2", relative_file_path="f.proto"), + ("src/protobuf/dir2/f.pb.go",), + ) + assert_gen( + Address("tests/protobuf/test_protos", relative_file_path="f.proto"), + ("tests/protobuf/test_protos/f.pb.go",), + ) + + +def test_generates_go_grpc(rule_runner: RuleRunner) -> None: + rule_runner.write_files( + { + "protos/BUILD": "protobuf_sources(grpc=True)", + "protos/service.proto": dedent( + """\ + syntax = "proto3"; + + option go_package = "example.com/protos"; + + package service; + + message TestMessage { + string foo = 1; + } + + service TestService { + rpc noStreaming (TestMessage) returns (TestMessage); + rpc clientStreaming (stream TestMessage) returns (TestMessage); + rpc serverStreaming (TestMessage) returns (stream TestMessage); + rpc bothStreaming (stream TestMessage) returns (stream TestMessage); + } + """ + ), + } + ) + assert_files_generated( + rule_runner, + Address("protos", relative_file_path="service.proto"), + source_roots=["/"], + expected_files=[ + "protos/service.pb.go", + "protos/service_grpc.pb.go", + ], + ) diff --git a/src/python/pants/backend/go/util_rules/sdk.py b/src/python/pants/backend/go/util_rules/sdk.py index ca891520d00..ca4ccca6357 100644 --- a/src/python/pants/backend/go/util_rules/sdk.py +++ b/src/python/pants/backend/go/util_rules/sdk.py @@ -13,6 +13,7 @@ from pants.engine.environment import Environment, EnvironmentRequest from pants.engine.fs import EMPTY_DIGEST, CreateDigest, Digest, FileContent, MergeDigests from pants.engine.internals.selectors import Get, MultiGet +from pants.engine.platform import Platform from pants.engine.process import Process, ProcessResult from pants.engine.rules import collect_rules, rule from pants.util.frozendict import FrozenDict @@ -26,10 +27,11 @@ class GoSdkProcess: command: tuple[str, ...] description: str env: FrozenDict[str, str] - input_digest: Digest = EMPTY_DIGEST - working_dir: str | None = None - output_files: tuple[str, ...] = () - output_directories: tuple[str, ...] = () + input_digest: Digest + working_dir: str | None + output_files: tuple[str, ...] + output_directories: tuple[str, ...] + platform: Platform | None def __init__( self, @@ -42,6 +44,7 @@ def __init__( output_files: Iterable[str] = (), output_directories: Iterable[str] = (), allow_downloads: bool = False, + platform: Platform | None = None, ) -> None: self.command = tuple(command) self.description = description @@ -54,6 +57,7 @@ def __init__( self.working_dir = working_dir self.output_files = tuple(output_files) self.output_directories = tuple(output_directories) + self.platform = platform @dataclass(frozen=True) @@ -112,6 +116,7 @@ async def setup_go_sdk_process( output_files=request.output_files, output_directories=request.output_directories, level=LogLevel.DEBUG, + platform=request.platform, )