diff --git a/src/python/pants/bin/local_pants_runner.py b/src/python/pants/bin/local_pants_runner.py index 72d53cc48ea..13b7d132dca 100644 --- a/src/python/pants/bin/local_pants_runner.py +++ b/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( diff --git a/src/python/pants/build_graph/build_configuration.py b/src/python/pants/build_graph/build_configuration.py index b876e61362c..ff577091cd8 100644 --- a/src/python/pants/build_graph/build_configuration.py +++ b/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, ) diff --git a/src/python/pants/init/extension_loader.py b/src/python/pants/init/extension_loader.py index 6e659f22282..0365eeecb3c 100644 --- a/src/python/pants/init/extension_loader.py +++ b/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 diff --git a/src/python/pants/option/global_options.py b/src/python/pants/option/global_options.py index b5de9a21362..df844699a52 100644 --- a/src/python/pants/option/global_options.py +++ b/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,12 +380,9 @@ 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, @@ -371,7 +390,11 @@ def _use_auth_plugin( 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( diff --git a/src/python/pants/option/global_options_test.py b/src/python/pants/option/global_options_test.py index fa70258a6c1..a43f6e954ff 100644 --- a/src/python/pants/option/global_options_test.py +++ b/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: diff --git a/src/python/pants/pantsd/pants_daemon_core.py b/src/python/pants/pantsd/pants_daemon_core.py index 429be6766a0..330d499c007 100644 --- a/src/python/pants/pantsd/pants_daemon_core.py +++ b/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