protobuf: infer dependencies on targets with more than one resolve-like field

docs/notes/2.34.x.md

@@ -160,6 +160,8 @@ The plugins a template names can come from:
 
 If `buf.yaml` declares BSR `deps:`, a `buf.lock` is required: create it with `pants generate-lockfiles --resolve=buf`.
 
+Protobuf dependency inference now works on targets with more than one resolve-like field, such as `python_resolve` and `jvm_resolve` when Python and Java or Scala codegen are enabled together. It used to fail with `NotImplementedError: TODO: Multiple resolve-like fields on target`. An import is now inferred to the owner that shares every one of the importing target's resolves, so a `protobuf_source` parametrized over both Python and JVM resolves depends on the matching variant of each proto it imports. Disabling one language's resolves also no longer stops inference from partitioning by the other's.
+
 #### S3
 
 Fixed S3 downloads failing with HTTP 403 `SignatureDoesNotMatch` when the object key contains a character that must be percent-encoded in a URL, such as `+`. The key is now percent-encoded before the request is signed, and SigV4 requests are signed with botocore's S3-specific signer, so the signature covers exactly the path S3 receives. Keys that are already percent-encoded in the requested URL are unaffected.

src/python/pants/backend/codegen/protobuf/protobuf_dependency_inference.py

@@ -245,36 +245,36 @@ class InferProtobufDependencies(InferDependenciesRequest):
     infer_from = ProtobufDependencyInferenceFieldSet
 
 
-async def get_resolve_key_from_target(address: Address) -> ProtobufMappingResolveKey:
+async def get_resolve_keys_from_target(address: Address) -> tuple[ProtobufMappingResolveKey, ...]:
+    """The key of each resolve the target is in: one per resolve-like field on it."""
     wrapped_target = await resolve_target(
         WrappedTargetRequest(address=address, description_of_origin="protobuf"), **implicitly()
     )
-    resolve_field_type: type[Field] | None = None
-    for field_type in wrapped_target.target.field_types:
-        if issubclass(field_type, ResolveLikeField):
-            if resolve_field_type is not None:
-                raise NotImplementedError(
-                    f"TODO: Multiple resolve-like fields on target at address `{address}`."
-                )
-            resolve_field_type = field_type
-    if resolve_field_type is None:
+    target = wrapped_target.target
+    resolve_field_types = [
+        field_type for field_type in target.field_types if issubclass(field_type, ResolveLikeField)
+    ]
+    if not resolve_field_types:
         raise ValueError(f"Failed to find resolve-like field on target at address `{address}.")
 
-    resolve_request_type = typing.cast(
-        ResolveLikeField, wrapped_target.target[resolve_field_type]
-    ).get_resolve_like_field_to_value_request()
-    resolve_request = resolve_request_type(target=wrapped_target.target)
-    resolve_result = await get_resolve_from_resolve_like_field_request(
-        **implicitly({resolve_request: ResolveLikeFieldToValueRequest})
+    resolve_requests = [
+        typing.cast(ResolveLikeField, target[field_type]).get_resolve_like_field_to_value_request()(
+            target=target
+        )
+        for field_type in resolve_field_types
+    ]
+    resolve_results = await concurrently(
+        get_resolve_from_resolve_like_field_request(
+            **implicitly({resolve_request: ResolveLikeFieldToValueRequest})
+        )
+        for resolve_request in resolve_requests
     )
-
-    # When resolves are disabled, return the sentinel key
-    if resolve_result.value is None:
-        return _NO_RESOLVE_LIKE_FIELDS_DEFINED
-
-    return ProtobufMappingResolveKey(
-        field_type=resolve_field_type,
-        resolve=resolve_result.value,
+    return tuple(
+        # When resolves are disabled, return the sentinel key
+        _NO_RESOLVE_LIKE_FIELDS_DEFINED
+        if resolve_result.value is None
+        else ProtobufMappingResolveKey(field_type=field_type, resolve=resolve_result.value)
+        for field_type, resolve_result in zip(resolve_field_types, resolve_results)
     )
 
 
@@ -287,11 +287,11 @@ async def infer_protobuf_dependencies(
 
     address = request.field_set.address
 
-    resolve_key: ProtobufMappingResolveKey
-    if _NO_RESOLVE_LIKE_FIELDS_DEFINED in protobuf_mapping.mapping:
-        resolve_key = _NO_RESOLVE_LIKE_FIELDS_DEFINED
+    resolve_keys: tuple[ProtobufMappingResolveKey, ...]
+    if set(protobuf_mapping.mapping) <= {_NO_RESOLVE_LIKE_FIELDS_DEFINED}:
+        resolve_keys = (_NO_RESOLVE_LIKE_FIELDS_DEFINED,)
     else:
-        resolve_key = await get_resolve_key_from_target(address)
+        resolve_keys = await get_resolve_keys_from_target(address)
 
     explicitly_provided_deps, hydrated_sources = await concurrently(
         determine_explicitly_provided_dependencies(
@@ -305,17 +305,19 @@ async def infer_protobuf_dependencies(
 
     result: OrderedSet[Address] = OrderedSet()
     for import_path in parse_proto_imports(file_content.content.decode()):
-        mapping_in_resolve = protobuf_mapping.mapping.get(resolve_key)
-        unambiguous = mapping_in_resolve.get(import_path) if mapping_in_resolve else None
-
-        ambiguous_modules_in_resolve = protobuf_mapping.ambiguous_modules.get(resolve_key)
-        ambiguous = (
-            ambiguous_modules_in_resolve.get(import_path) if ambiguous_modules_in_resolve else None
-        )
-
-        if unambiguous:
-            result.add(unambiguous)
-        elif ambiguous:
+        # A target in several resolves, one per language, can only depend on an owner that is in
+        # all of them, so intersect the owners each resolve has for the import.
+        owners: set[Address] | None = None
+        for resolve_key in resolve_keys:
+            owners_in_resolve = _owners_in_resolve(protobuf_mapping, resolve_key, import_path)
+            owners = owners_in_resolve if owners is None else owners & owners_in_resolve
+
+        if not owners:
+            continue
+        if len(owners) == 1:
+            result.add(next(iter(owners)))
+        else:
+            ambiguous = tuple(sorted(owners))
             explicitly_provided_deps.maybe_warn_of_ambiguous_dependency_inference(
                 ambiguous,
                 address,
@@ -333,5 +335,16 @@ async def infer_protobuf_dependencies(
     return InferredDependencies(sorted(result))
 
 
+def _owners_in_resolve(
+    protobuf_mapping: ProtobufMapping, resolve_key: ProtobufMappingResolveKey, import_path: str
+) -> set[Address]:
+    unambiguous = protobuf_mapping.mapping.get(resolve_key, FrozenDict()).get(import_path)
+    if unambiguous:
+        return {unambiguous}
+    return set(
+        protobuf_mapping.ambiguous_modules.get(resolve_key, FrozenDict()).get(import_path, ())
+    )
+
+
 def rules():
     return (*collect_rules(), UnionRule(InferDependenciesRequest, InferProtobufDependencies))

src/python/pants/backend/codegen/protobuf/protobuf_dependency_inference_test.py

@@ -13,11 +13,19 @@
     ProtobufMapping,
     parse_proto_imports,
 )
-from pants.backend.codegen.protobuf.target_types import ProtobufSourcesGeneratorTarget
+from pants.backend.codegen.protobuf.python import additional_fields as python_fields
+from pants.backend.codegen.protobuf.target_types import (
+    ProtobufSourcesGeneratorTarget,
+    ProtobufSourceTarget,
+)
 from pants.backend.codegen.protobuf.target_types import rules as target_types_rules
+from pants.backend.python import target_types_rules as python_target_types_rules
 from pants.core.util_rules import stripped_source_files
 from pants.engine.addresses import Address
+from pants.engine.internals.parametrize import Parametrize
 from pants.engine.target import InferredDependencies
+from pants.jvm import target_types as jvm_target_types
+from pants.jvm.target_types import PrefixedJvmResolveField
 from pants.testutil.rule_runner import QueryRule, RuleRunner
 from pants.util.frozendict import FrozenDict
 
@@ -232,3 +240,133 @@ def run_dep_inference(address: Address) -> InferredDependencies:
 
     deps = run_dep_inference(Address("protos", relative_file_path="a.proto"))
     assert deps == InferredDependencies([Address("protos", relative_file_path="b.proto")])
+
+
+@pytest.fixture
+def rule_runner_with_python_and_jvm_resolves() -> RuleRunner:
+    """Both Python's and the JVM's resolve fields on protobuf targets, as when Python and Java or
+    Scala codegen are enabled together."""
+    return RuleRunner(
+        rules=[
+            *stripped_source_files.rules(),
+            *protobuf_dependency_inference.rules(),
+            *target_types_rules(),
+            *python_fields.rules(),
+            *python_target_types_rules.rules(),
+            *jvm_target_types.rules(),
+            ProtobufSourceTarget.register_plugin_field(PrefixedJvmResolveField),
+            ProtobufSourcesGeneratorTarget.register_plugin_field(PrefixedJvmResolveField),
+            QueryRule(ProtobufMapping, []),
+            QueryRule(InferredDependencies, [InferProtobufDependencies]),
+        ],
+        target_types=[ProtobufSourceTarget, ProtobufSourcesGeneratorTarget],
+        objects={"parametrize": Parametrize},
+    )
+
+
+def _infer(rule_runner: RuleRunner, address: Address) -> InferredDependencies:
+    tgt = rule_runner.get_target(address)
+    return rule_runner.request(
+        InferredDependencies,
+        [InferProtobufDependencies(ProtobufDependencyInferenceFieldSet.create(tgt))],
+    )
+
+
+def test_dependency_inference_with_two_resolve_fields(
+    rule_runner_with_python_and_jvm_resolves: RuleRunner,
+) -> None:
+    rule_runner = rule_runner_with_python_and_jvm_resolves
+    rule_runner.set_options(
+        [
+            "--source-root-patterns=['protos']",
+            "--python-enable-resolves",
+            "--python-resolves={'a': ''}",
+            "--jvm-resolves={'j': 'j.lock'}",
+            "--jvm-default-resolve=j",
+        ]
+    )
+    rule_runner.write_files(
+        {
+            "protos/a.proto": "import 'b.proto';",
+            "protos/b.proto": "",
+            "protos/BUILD": "protobuf_sources(python_resolve='a', jvm_resolve='j')",
+        }
+    )
+    assert _infer(rule_runner, Address("protos", relative_file_path="a.proto")) == (
+        InferredDependencies([Address("protos", relative_file_path="b.proto")])
+    )
+
+
+def test_dependency_inference_with_two_parametrized_resolve_fields(
+    rule_runner_with_python_and_jvm_resolves: RuleRunner, caplog
+) -> None:
+    """The import is owned by one variant per pair of resolves, though by two in each resolve."""
+    rule_runner = rule_runner_with_python_and_jvm_resolves
+    rule_runner.set_options(
+        [
+            "--source-root-patterns=['protos']",
+            "--python-enable-resolves",
+            "--python-resolves={'a': '', 'b': ''}",
+            "--jvm-resolves={'j1': 'j1.lock', 'j2': 'j2.lock'}",
+            "--jvm-default-resolve=j1",
+        ]
+    )
+    rule_runner.write_files(
+        {
+            "protos/a.proto": "import 'b.proto';",
+            "protos/b.proto": "",
+            "protos/BUILD": dedent(
+                """\
+                for name in ("a", "b"):
+                    protobuf_source(
+                        name=name,
+                        source=f"{name}.proto",
+                        python_resolve=parametrize("a", "b"),
+                        jvm_resolve=parametrize("j1", "j2"),
+                    )
+                """
+            ),
+        }
+    )
+    for python_resolve in ("a", "b"):
+        for jvm_resolve in ("j1", "j2"):
+            parameters = {"python_resolve": python_resolve, "jvm_resolve": jvm_resolve}
+            assert _infer(
+                rule_runner, Address("protos", target_name="a", parameters=parameters)
+            ) == InferredDependencies([Address("protos", target_name="b", parameters=parameters)])
+    assert not caplog.records
+
+
+def test_dependency_inference_with_one_resolve_field_disabled(
+    rule_runner_with_python_and_jvm_resolves: RuleRunner,
+) -> None:
+    """Python's resolves being disabled leaves inference partitioned by the JVM's."""
+    rule_runner = rule_runner_with_python_and_jvm_resolves
+    rule_runner.set_options(
+        [
+            "--source-root-patterns=['protos']",
+            "--jvm-resolves={'j1': 'j1.lock', 'j2': 'j2.lock'}",
+            "--jvm-default-resolve=j1",
+        ]
+    )
+    rule_runner.write_files(
+        {
+            "protos/a.proto": "import 'b.proto';",
+            "protos/b.proto": "",
+            "protos/BUILD": dedent(
+                """\
+                for name in ("a", "b"):
+                    protobuf_source(
+                        name=name,
+                        source=f"{name}.proto",
+                        jvm_resolve=parametrize("j1", "j2"),
+                    )
+                """
+            ),
+        }
+    )
+    for jvm_resolve in ("j1", "j2"):
+        parameters = {"jvm_resolve": jvm_resolve}
+        assert _infer(
+            rule_runner, Address("protos", target_name="a", parameters=parameters)
+        ) == InferredDependencies([Address("protos", target_name="b", parameters=parameters)])