Add entry point support for remote auth plugin (#16212)

This will remove the need to specify the remote_auth_plugin in the pants config and will allow the auth plugin to be auto discovered using the entry point mechanism.

Once the updated Toolchain Pants Plugin with the entry point support is released and updated in this repo I will mark the `[GLOBAL].remote_auth_plugin` as deprecated.

src/python/pants/bin/local_pants_runner.py

@@ -69,7 +69,9 @@ def _init_graph_session(
     ) -> GraphSession:
         native_engine.maybe_set_panic_handler()
         if scheduler is None:
-            dynamic_remote_options, _ = DynamicRemoteOptions.from_options(options, env)
+            dynamic_remote_options, _ = DynamicRemoteOptions.from_options(
+                options, env, remote_auth_plugin_func=build_config.remote_auth_plugin_func
+            )
             bootstrap_options = options.bootstrap_option_values()
             assert bootstrap_options is not None
             scheduler = EngineInitializer.setup_graph(

src/python/pants/build_graph/build_configuration.py

@@ -8,7 +8,7 @@
 from collections.abc import Iterable
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, DefaultDict
+from typing import Any, Callable, DefaultDict
 
 from pants.backend.project_info.filter_targets import FilterSubsystem
 from pants.build_graph.build_file_aliases import BuildFileAliases
@@ -46,6 +46,7 @@ class BuildConfiguration:
     rule_to_providers: FrozenDict[Rule, tuple[str, ...]]
     union_rule_to_providers: FrozenDict[UnionRule, tuple[str, ...]]
     allow_unknown_options: bool
+    remote_auth_plugin_func: Callable | None
 
     @property
     def all_subsystems(self) -> tuple[type[Subsystem], ...]:
@@ -121,6 +122,7 @@ class Builder:
             default_factory=lambda: defaultdict(list)
         )
         _allow_unknown_options: bool = False
+        _remote_auth_plugin: Callable | None = None
 
         def registered_aliases(self) -> BuildFileAliases:
             """Return the registered aliases exposed in BUILD files.
@@ -248,6 +250,9 @@ def register_target_types(
                 # walked during union membership setup.
                 _ = target_type._plugin_field_cls
 
+        def register_remote_auth_plugin(self, remote_auth_plugin: Callable) -> None:
+            self._remote_auth_plugin = remote_auth_plugin
+
         def allow_unknown_options(self, allow: bool = True) -> None:
             """Allows overriding whether Options parsing will fail for unrecognized Options.
 
@@ -276,4 +281,5 @@ def create(self) -> BuildConfiguration:
                     (k, tuple(v)) for k, v in self._union_rule_to_providers.items()
                 ),
                 allow_unknown_options=self._allow_unknown_options,
+                remote_auth_plugin_func=self._remote_auth_plugin,
             )

src/python/pants/init/extension_loader.py

@@ -98,6 +98,10 @@ def load_plugins(
         if "rules" in entries:
             rules = entries["rules"].load()()
             build_configuration.register_rules(req.key, rules)
+        if "remote_auth" in entries:
+            remote_auth_func = entries["remote_auth"].load()
+            build_configuration.register_remote_auth_plugin(remote_auth_func)
+
         loaded[dist.as_requirement().key] = dist
 
 

src/python/pants/option/global_options.py

@@ -14,7 +14,7 @@
 from datetime import datetime
 from enum import Enum
 from pathlib import Path, PurePath
-from typing import Any, Type, cast
+from typing import Any, Callable, Type, cast
 
 from pants.base.build_environment import (
     get_buildroot,
@@ -227,6 +227,18 @@ def disabled(cls) -> DynamicRemoteOptions:
             execution_rpc_concurrency=DEFAULT_EXECUTION_OPTIONS.remote_execution_rpc_concurrency,
         )
 
+    @classmethod
+    def _get_auth_plugin_from_option(cls, remote_auth_plugin_option_value: str) -> Callable:
+        if ":" not in remote_auth_plugin_option_value:
+            raise OptionsError(
+                "Invalid value for `--remote-auth-plugin`: "
+                f"{remote_auth_plugin_option_value}. Please use the format "
+                "`path.to.module:my_func`."
+            )
+        auth_plugin_path, auth_plugin_func = remote_auth_plugin_option_value.split(":")
+        auth_plugin_module = importlib.import_module(auth_plugin_path)
+        return cast(Callable, getattr(auth_plugin_module, auth_plugin_func))
+
     @classmethod
     def _use_oauth_token(cls, bootstrap_options: OptionValueContainer) -> DynamicRemoteOptions:
         oauth_token = (
@@ -278,6 +290,7 @@ def from_options(
         full_options: Options,
         env: CompleteEnvironment,
         prior_result: AuthPluginResult | None = None,
+        remote_auth_plugin_func: Callable | None = None,
     ) -> tuple[DynamicRemoteOptions, AuthPluginResult | None]:
         bootstrap_options = full_options.bootstrap_option_values()
         assert bootstrap_options is not None
@@ -295,10 +308,13 @@ def from_options(
             )
         if bootstrap_options.remote_oauth_bearer_token_path:
             return cls._use_oauth_token(bootstrap_options), None
-
-        if bootstrap_options.remote_auth_plugin:
+        if bootstrap_options.remote_auth_plugin or remote_auth_plugin_func is not None:
             return cls._use_auth_plugin(
-                bootstrap_options, full_options=full_options, env=env, prior_result=prior_result
+                bootstrap_options,
+                full_options=full_options,
+                env=env,
+                prior_result=prior_result,
+                remote_auth_plugin_func_from_entry_point=remote_auth_plugin_func,
             )
         return cls._use_no_auth(bootstrap_options), None
 
@@ -338,14 +354,20 @@ def _use_auth_plugin(
         full_options: Options,
         env: CompleteEnvironment,
         prior_result: AuthPluginResult | None,
+        remote_auth_plugin_func_from_entry_point: Callable | None,
     ) -> tuple[DynamicRemoteOptions, AuthPluginResult | None]:
         auth_plugin_result: AuthPluginResult | None = None
-        if ":" not in bootstrap_options.remote_auth_plugin:
-            raise OptionsError(
-                "Invalid value for `[GLOBAL].remote_auth_plugin`: "
-                f"{bootstrap_options.remote_auth_plugin}. Please use the format "
-                "`path.to.module:my_func`."
+        if not remote_auth_plugin_func_from_entry_point:
+            remote_auth_plugin_func = cls._get_auth_plugin_from_option(
+                bootstrap_options.remote_auth_plugin
             )
+        else:
+            remote_auth_plugin_func = remote_auth_plugin_func_from_entry_point
+            if bootstrap_options.remote_auth_plugin:
+                raise OptionsError(
+                    "remote auth plugin already provided via entry point of a plugin. `[GLOBAL].remote_auth_plugin` should not be specified in options."
+                )
+
         execution = cast(bool, bootstrap_options.remote_execution)
         cache_read = cast(bool, bootstrap_options.remote_cache_read)
         cache_write = cast(bool, bootstrap_options.remote_cache_write)
@@ -358,20 +380,21 @@ def _use_auth_plugin(
         store_rpc_concurrency = cast(int, bootstrap_options.remote_store_rpc_concurrency)
         cache_rpc_concurrency = cast(int, bootstrap_options.remote_cache_rpc_concurrency)
         execution_rpc_concurrency = cast(int, bootstrap_options.remote_execution_rpc_concurrency)
-        auth_plugin_path, _, auth_plugin_func = bootstrap_options.remote_auth_plugin.partition(":")
-        auth_plugin_module = importlib.import_module(auth_plugin_path)
-        auth_plugin_func = getattr(auth_plugin_module, auth_plugin_func)
         auth_plugin_result = cast(
             AuthPluginResult,
-            auth_plugin_func(
+            remote_auth_plugin_func(
                 initial_execution_headers=execution_headers,
                 initial_store_headers=store_headers,
                 options=full_options,
                 env=dict(env),
                 prior_result=prior_result,
             ),
         )
-        plugin_name = auth_plugin_result.plugin_name or bootstrap_options.remote_auth_plugin
+        plugin_name = (
+            auth_plugin_result.plugin_name
+            or bootstrap_options.remote_auth_plugin
+            or f"{remote_auth_plugin_func.__module__}.{remote_auth_plugin_func.__name__}"
+        )
         if not auth_plugin_result.is_available:
             # NB: This is debug because we expect plugins to log more informative messages.
             logger.debug(

src/python/pants/option/global_options_test.py

@@ -42,7 +42,9 @@ def create_dynamic_remote_options(
     ob = create_options_bootstrapper(args)
     env = CompleteEnvironment({})
     _build_config, options = OptionsInitializer(ob).build_config_and_options(ob, env, raise_=False)
-    return DynamicRemoteOptions.from_options(options, env)[0]
+    return DynamicRemoteOptions.from_options(
+        options, env, remote_auth_plugin_func=_build_config.remote_auth_plugin_func
+    )[0]
 
 
 def test_dynamic_remote_options_oauth_bearer_token_path(tmp_path: Path) -> None:

src/python/pants/pantsd/pants_daemon_core.py

@@ -136,7 +136,10 @@ def prepare(
         # they need to be re-evaluated every run. We only reinitialize the scheduler if changes
         # were made, though.
         dynamic_remote_options, auth_plugin_result = DynamicRemoteOptions.from_options(
-            options, env, self._prior_auth_plugin_result
+            options,
+            env,
+            self._prior_auth_plugin_result,
+            remote_auth_plugin_func=build_config.remote_auth_plugin_func,
         )
         remote_options_changed = (
             self._prior_dynamic_remote_options is not None