add --use-service-auth-token option to spark-run

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"
 

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"}
         ):