Add pants-plugins/stevedore_extensions to add teach pants' dependency inference about our runtime-loaded plugins

CHANGELOG.rst

@@ -14,7 +14,7 @@ Added
   working on StackStorm, improve our security posture, and improve CI reliability thanks in part
   to pants' use of PEX lockfiles. This is not a user-facing addition.
   #5778 #5789 #5817 #5795 #5830 #5833 #5834 #5841 #5840 #5838 #5842 #5837 #5849 #5850
-  #5846 #5853 #5848 #5847 #5858 #5857 #5860 #5868 #5871
+  #5846 #5853 #5848 #5847 #5858 #5857 #5860 #5868 #5871 #5869
   Contributed by @cognifloyd
 
 * Added a joint index to solve the problem of slow mongo queries for scheduled executions. #5805

contrib/runners/noop_runner/BUILD

@@ -0,0 +1,7 @@
+# stevedore_extension(
+#    name="runner",
+#    namespace="st2common.runners.runner",
+#    entry_points={
+#        "noop": "noop_runner.noop_runner",
+#    },
+# )

pants-plugins/README.md

@@ -8,6 +8,9 @@ The plugins here add custom goals or other logic into pants.
 
 To see available goals, do "./pants help goals" and "./pants help $goal".
 
+These plugins might be useful outside of the StackStorm project:
+- `stevedore_extensions`
+
 These StackStorm-specific plugins might be useful in other StackStorm-related repos.
 - `pack_metadata`
 
@@ -66,3 +69,12 @@ the `fmt` goal (eg `./pants fmt contrib/schemas::`), the schemas will
 be regenerated if any of the files used to generate them have changed.
 Also, running the `lint` goal will fail if the schemas need to be
 regenerated.
+
+### `stevedore_extensions` plugin
+
+This plugin teaches pants how to infer dependencies on stevedore
+extensions (python plugins loaded at runtime via entry points).
+This includes the `stevedore_extensions` target and the
+`stevedore_namespaces` field. When necessary, it generates an
+`entry_points.txt` file so that stevedore can work correctly
+within pants sandboxes.

pants-plugins/stevedore_extensions/BUILD

@@ -0,0 +1,5 @@
+python_sources()
+
+python_tests(
+    name="tests",
+)

pants-plugins/stevedore_extensions/python_target_dependencies.py

@@ -0,0 +1,123 @@
+# Copyright 2023 The StackStorm Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import List, Mapping, Tuple
+
+from pants.backend.python.target_types import (
+    PythonTestTarget,
+    PythonTestsGeneratorTarget,
+    PythonTestsDependenciesField,
+)
+from pants.engine.addresses import Address
+from pants.engine.rules import collect_rules, rule, UnionRule
+from pants.engine.target import (
+    AllTargets,
+    FieldSet,
+    InferDependenciesRequest,
+    InferredDependencies,
+)
+from pants.util.frozendict import FrozenDict
+from pants.util.logging import LogLevel
+from pants.util.ordered_set import OrderedSet
+
+from stevedore_extensions.target_types import (
+    AllStevedoreExtensionTargets,
+    StevedoreEntryPointsField,
+    StevedoreExtension,
+    StevedoreNamespaceField,
+    StevedoreNamespacesField,
+)
+
+
+# -----------------------------------------------------------------------------------------------
+# Utility rules to analyze all `StevedoreExtension` targets
+# -----------------------------------------------------------------------------------------------
+
+
+@rule(desc="Find all StevedoreExtension targets in project", level=LogLevel.DEBUG)
+def find_all_stevedore_extension_targets(
+    targets: AllTargets,
+) -> AllStevedoreExtensionTargets:
+    return AllStevedoreExtensionTargets(
+        tgt for tgt in targets if tgt.has_field(StevedoreEntryPointsField)
+    )
+
+
+@dataclass(frozen=True)
+class StevedoreExtensions:
+    """A mapping of stevedore namespaces to a list of StevedoreExtension targets that provide them"""
+
+    mapping: FrozenDict[str, Tuple[StevedoreExtension]]
+
+
+@rule(
+    desc="Creating map of stevedore_extension namespaces to StevedoreExtension targets",
+    level=LogLevel.DEBUG,
+)
+async def map_stevedore_extensions(
+    stevedore_extensions: AllStevedoreExtensionTargets,
+) -> StevedoreExtensions:
+    mapping: Mapping[str, List[StevedoreExtension]] = defaultdict(list)
+    for extension in stevedore_extensions:
+        mapping[extension[StevedoreNamespaceField].value].append(extension)
+    return StevedoreExtensions(
+        FrozenDict((k, tuple(v)) for k, v in sorted(mapping.items()))
+    )
+
+
+# -----------------------------------------------------------------------------------------------
+# Dependencies for `python_test` and `python_tests` targets
+# -----------------------------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class PythonTestsStevedoreNamespaceInferenceFieldSet(FieldSet):
+    required_fields = (PythonTestsDependenciesField, StevedoreNamespacesField)
+
+    stevedore_namespaces: StevedoreNamespacesField
+
+
+class InferStevedoreNamespaceDependencies(InferDependenciesRequest):
+    infer_from = PythonTestsStevedoreNamespaceInferenceFieldSet
+
+
+@rule(
+    desc="Infer stevedore_extension target dependencies based on namespace list.",
+    level=LogLevel.DEBUG,
+)
+async def infer_stevedore_namespace_dependencies(
+    request: InferStevedoreNamespaceDependencies,
+    stevedore_extensions: StevedoreExtensions,
+) -> InferredDependencies:
+    namespaces: StevedoreNamespacesField = request.field_set.stevedore_namespaces
+    if namespaces.value is None:
+        return InferredDependencies(())
+
+    addresses = []
+    for namespace in namespaces.value:
+        extensions = stevedore_extensions.mapping.get(namespace, ())
+        addresses.extend(extension.address for extension in extensions)
+
+    result: OrderedSet[Address] = OrderedSet(addresses)
+    return InferredDependencies(sorted(result))
+
+
+def rules():
+    return [
+        *collect_rules(),
+        PythonTestsGeneratorTarget.register_plugin_field(StevedoreNamespacesField),
+        PythonTestTarget.register_plugin_field(StevedoreNamespacesField),
+        UnionRule(InferDependenciesRequest, InferStevedoreNamespaceDependencies),
+    ]

pants-plugins/stevedore_extensions/python_target_dependencies_test.py

@@ -0,0 +1,182 @@
+# Copyright 2023 The StackStorm Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from textwrap import dedent
+
+import pytest
+
+from pants.backend.python.target_types import (
+    PythonSourceTarget,
+    PythonSourcesGeneratorTarget,
+    PythonTestTarget,
+    PythonTestsGeneratorTarget,
+)
+from pants.backend.python.target_types_rules import rules as python_target_types_rules
+from pants.engine.addresses import Address
+from pants.engine.target import InferredDependencies
+from pants.testutil.rule_runner import QueryRule, RuleRunner
+from pants.util.frozendict import FrozenDict
+
+from .python_target_dependencies import (
+    InferStevedoreNamespaceDependencies,
+    PythonTestsStevedoreNamespaceInferenceFieldSet,
+    StevedoreExtensions,
+    rules as stevedore_dep_rules,
+)
+from .target_types import (
+    AllStevedoreExtensionTargets,
+    StevedoreExtension,
+)
+
+
+# random set of runner names to use in tests
+st2_runners = ["noop", "python", "foobar"]
+
+
+@pytest.fixture
+def rule_runner() -> RuleRunner:
+    rule_runner = RuleRunner(
+        rules=[
+            *python_target_types_rules(),
+            *stevedore_dep_rules(),
+            QueryRule(AllStevedoreExtensionTargets, ()),
+            QueryRule(StevedoreExtensions, ()),
+            QueryRule(InferredDependencies, (InferStevedoreNamespaceDependencies,)),
+        ],
+        target_types=[
+            PythonSourceTarget,
+            PythonSourcesGeneratorTarget,
+            PythonTestTarget,
+            PythonTestsGeneratorTarget,
+            StevedoreExtension,
+        ],
+    )
+    for runner in st2_runners:
+        rule_runner.write_files(
+            {
+                f"runners/{runner}_runner/BUILD": dedent(
+                    f"""\
+                    stevedore_extension(
+                        name="runner",
+                        namespace="st2common.runners.runner",
+                        entry_points={{
+                            "{runner}": "{runner}_runner.{runner}_runner",
+                        }},
+                    )
+                    stevedore_extension(
+                        name="thing",
+                        namespace="some.thing.else",
+                        entry_points={{
+                            "{runner}": "{runner}_runner.thing",
+                        }},
+                    )
+                    """
+                ),
+                f"runners/{runner}_runner/{runner}_runner/BUILD": "python_sources()",
+                f"runners/{runner}_runner/{runner}_runner/__init__.py": "",
+                f"runners/{runner}_runner/{runner}_runner/{runner}_runner.py": "",
+                f"runners/{runner}_runner/{runner}_runner/thing.py": "",
+            }
+        )
+    args = [
+        "--source-root-patterns=runners/*_runner",
+    ]
+    rule_runner.set_options(args, env_inherit={"PATH", "PYENV_ROOT", "HOME"})
+    return rule_runner
+
+
+# -----------------------------------------------------------------------------------------------
+# Tests for utility rules
+# -----------------------------------------------------------------------------------------------
+
+
+def test_find_all_stevedore_extension_targets(rule_runner: RuleRunner) -> None:
+    assert rule_runner.request(
+        AllStevedoreExtensionTargets, []
+    ) == AllStevedoreExtensionTargets(
+        rule_runner.get_target(
+            Address(f"runners/{runner}_runner", target_name=target_name)
+        )
+        for runner in sorted(st2_runners)
+        for target_name in ["runner", "thing"]
+    )
+
+
+def test_map_stevedore_extensions(rule_runner: RuleRunner) -> None:
+    assert rule_runner.request(StevedoreExtensions, []) == StevedoreExtensions(
+        FrozenDict(
+            {
+                "some.thing.else": tuple(
+                    rule_runner.get_target(
+                        Address(f"runners/{runner}_runner", target_name="thing")
+                    )
+                    for runner in sorted(st2_runners)
+                ),
+                "st2common.runners.runner": tuple(
+                    rule_runner.get_target(
+                        Address(f"runners/{runner}_runner", target_name="runner")
+                    )
+                    for runner in sorted(st2_runners)
+                ),
+            }
+        )
+    )
+
+
+# -----------------------------------------------------------------------------------------------
+# Tests for dependency inference of python targets (python_tests, etc)
+# -----------------------------------------------------------------------------------------------
+
+
+def test_infer_stevedore_namespace_dependencies(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            "src/foobar/BUILD": dedent(
+                """\
+                python_tests(
+                    name="tests",
+                    stevedore_namespaces=["st2common.runners.runner"],
+                )
+                """
+            ),
+            "src/foobar/test_something.py": "",
+        }
+    )
+
+    def run_dep_inference(address: Address) -> InferredDependencies:
+        target = rule_runner.get_target(address)
+        return rule_runner.request(
+            InferredDependencies,
+            [
+                InferStevedoreNamespaceDependencies(
+                    PythonTestsStevedoreNamespaceInferenceFieldSet.create(target)
+                )
+            ],
+        )
+
+    # this asserts that only the st2common.runners.runner namespace gets selected.
+    assert run_dep_inference(
+        Address(
+            "src/foobar", target_name="tests", relative_file_path="test_something.py"
+        ),
+    ) == InferredDependencies(
+        [
+            Address(
+                f"runners/{runner}_runner",
+                target_name="runner",
+            )
+            for runner in st2_runners
+        ],
+    )

pants-plugins/stevedore_extensions/register.py

@@ -0,0 +1,39 @@
+# Copyright 2023 The StackStorm Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pants.backend.codegen import export_codegen_goal
+
+from stevedore_extensions import (
+    target_types_rules,
+    rules as stevedore_rules,
+    python_target_dependencies,
+)
+from stevedore_extensions.target_types import StevedoreExtension
+
+
+# TODO: add the entry_points automatically to setup_py
+# TODO: add stevedore_namespaces field to python_sources?
+
+
+def rules():
+    return [
+        *target_types_rules.rules(),
+        *stevedore_rules.rules(),
+        *python_target_dependencies.rules(),
+        *export_codegen_goal.rules(),
+    ]
+
+
+def target_types():
+    return [StevedoreExtension]