diff --git a/paasta_tools/cli/cmds/spark_run.py b/paasta_tools/cli/cmds/spark_run.py index 88bb90d4bf..6da09265d9 100644 --- a/paasta_tools/cli/cmds/spark_run.py +++ b/paasta_tools/cli/cmds/spark_run.py @@ -26,6 +26,7 @@ from paasta_tools.cli.cmds.check import makefile_responds_to from paasta_tools.cli.cmds.cook_image import paasta_cook_image from paasta_tools.cli.utils import get_instance_config +from paasta_tools.cli.utils import get_service_auth_token from paasta_tools.cli.utils import lazy_choices_completer from paasta_tools.cli.utils import list_instances from paasta_tools.clusterman import get_clusterman_metrics @@ -337,6 +338,18 @@ def add_subparser(subparsers): default=None, ) + list_parser.add_argument( + "--use-service-auth-token", + help=( + "Acquire service authentication token for the underlying instance," + " and set it in the container environment" + ), + action="store_true", + dest="use_service_auth_token", + required=False, + default=False, + ) + aws_group = list_parser.add_argument_group( title="AWS credentials options", description="If --aws-credentials-yaml is specified, it overrides all " @@ -795,6 +808,9 @@ def configure_and_run_docker_container( ) # type:ignore environment.update(extra_driver_envs) + if args.use_service_auth_token: + environment["YELP_SVC_AUTHZ_TOKEN"] = get_service_auth_token() + webui_url = get_webui_url(spark_conf["spark.ui.port"]) webui_url_msg = PaastaColors.green(f"\nSpark monitoring URL: ") + f"{webui_url}\n" diff --git a/tests/cli/test_cmds_spark_run.py b/tests/cli/test_cmds_spark_run.py index 97f81dd432..d969e5664e 100644 --- a/tests/cli/test_cmds_spark_run.py +++ b/tests/cli/test_cmds_spark_run.py @@ -493,6 +493,7 @@ def test_configure_and_run_docker_container( args.docker_shm_size = False args.tronfig = None args.job_id = None + args.use_service_auth_token = False with mock.patch.object( self.instance_config, "get_env_dictionary", return_value={"env1": "val1"} ): @@ -607,6 +608,7 @@ def test_configure_and_run_docker_driver_resource_limits_config( args.docker_cpu_limit = 3 args.docker_memory_limit = "4g" args.docker_shm_size = "1g" + args.use_service_auth_token = False with mock.patch.object( self.instance_config, "get_env_dictionary", return_value={"env1": "val1"} ): @@ -721,6 +723,7 @@ def test_configure_and_run_docker_driver_resource_limits( args.docker_cpu_limit = False args.docker_memory_limit = False args.docker_shm_size = False + args.use_service_auth_token = False with mock.patch.object( self.instance_config, "get_env_dictionary", return_value={"env1": "val1"} ):