feat(load-task): add a --user flag to specify which user to run as

Author: ahal

src/taskgraph/docker.py

@@ -251,7 +251,8 @@ def _index(l: List, s: str) -> Optional[int]:
         pass
 
 
-def load_task(task_id, remove=True):
+def load_task(task_id, remove=True, user=None):
+    user = user or "worker"
     task_def = get_task_definition(task_id)
 
     if (
@@ -316,7 +317,7 @@ def load_task(task_id, remove=True):
             image_tag,
             "bash",
             "-c",
-            f"{shlex.join(command)} && cd $TASK_WORKDIR && su -p worker",
+            f"{shlex.join(command)} && cd $TASK_WORKDIR && su -p {user}",
         ]
 
         if remove:

src/taskgraph/main.py

@@ -667,11 +667,12 @@ def image_digest(args):
     default=True,
     help="Keep the docker container after exiting.",
 )
+@argument("--user", default=None, help="Container user to start shell with.")
 def load_task(args):
     from taskgraph.docker import load_task
 
     validate_docker()
-    return load_task(args["task_id"], remove=args["remove"])
+    return load_task(args["task_id"], remove=args["remove"], user=args["user"])
 
 
 @command("decision", help="Run the decision task")