_stream/cli: Add pull option to shell

This will fetch artifacts if they're not local or are incomplete.

Part of #1044

2 files changed, 30 insertions, 15 deletions

diff --git a/src/buildstream/_frontend/cli.py b/src/buildstream/_frontend/cli.py
index b353ced93..53adb188c 100644
--- a/src/buildstream/_frontend/cli.py
+++ b/src/buildstream/_frontend/cli.py
@@ -501,11 +501,13 @@ def show(app, elements, deps, except_, order, format_):
 @click.option('--use-buildtree', '-t', 'cli_buildtree', type=click.Choice(['ask', 'try', 'always', 'never']),
               default='ask',
               help='Defaults to ask but if set to always the function will fail if a build tree is not available')
+@click.option('--pull', 'pull_', is_flag=True, default=False,
+              help='Attempt to pull missing or incomplete artifacts')
 @click.argument('element', required=False,
                 type=click.Path(readable=False))
 @click.argument('command', type=click.STRING, nargs=-1)
 @click.pass_obj
-def shell(app, element, sysroot, mount, isolate, build_, cli_buildtree, command):
+def shell(app, element, sysroot, mount, isolate, build_, cli_buildtree, pull_, command):
     """Run a command in the target element's sandbox environment
 
     When this command is executed from a workspace directory, the default
@@ -535,11 +537,10 @@ def shell(app, element, sysroot, mount, isolate, build_, cli_buildtree, command)
     from .._project import HostMount
     from .._pipeline import PipelineSelection
 
-    if build_:
-        scope = Scope.BUILD
-    else:
-        scope = Scope.RUN
+    scope = Scope.BUILD if build_ else Scope.RUN
 
+    # We may need to fetch dependency artifacts if we're pulling the artifact
+    selection = PipelineSelection.ALL if pull_ else PipelineSelection.NONE
     use_buildtree = None
 
     with app.initialized():
@@ -548,9 +549,14 @@ def shell(app, element, sysroot, mount, isolate, build_, cli_buildtree, command)
             if not element:
                 raise AppError('Missing argument "ELEMENT".')
 
-        dependencies = app.stream.load_selection((element,), selection=PipelineSelection.NONE,
-                                                 use_artifact_config=True)
-        element = dependencies[0]
+        elements = app.stream.load_selection((element,), selection=selection,
+                                             use_artifact_config=True)
+
+        # last one will be the element we want to stage, previous ones are
+        # elements to try and pull
+        element = elements[-1]
+        pull_dependencies = elements[:-1] if pull_ else None
+
         prompt = app.shell_prompt(element)
         mounts = [
             HostMount(path, host_path)
@@ -561,7 +567,7 @@ def shell(app, element, sysroot, mount, isolate, build_, cli_buildtree, command)
         buildtree_exists = element._buildtree_exists()
 
         if cli_buildtree in ("always", "try"):
-            if buildtree_exists:
+            if buildtree_exists or pull_:
                 use_buildtree = cli_buildtree
                 if not cached and use_buildtree == "always":
                     click.echo("WARNING: buildtree is not cached locally, will attempt to pull from available remotes",
@@ -601,7 +607,8 @@ def shell(app, element, sysroot, mount, isolate, build_, cli_buildtree, command)
                                         mounts=mounts,
                                         isolate=isolate,
                                         command=command,
-                                        usebuildtree=use_buildtree)
+                                        usebuildtree=use_buildtree,
+                                        pull_dependencies=pull_dependencies)
         except BstError as e:
             raise AppError("Error launching shell: {}".format(e), detail=e.detail) from e
 
diff --git a/src/buildstream/_stream.py b/src/buildstream/_stream.py
index 6f6d30f17..8097f451d 100644
--- a/src/buildstream/_stream.py
+++ b/src/buildstream/_stream.py
@@ -140,6 +140,7 @@ class Stream():
     #    isolate (bool): Whether to isolate the environment like we do in builds
     #    command (list): An argv to launch in the sandbox, or None
     #    usebuildtree (str): Whether to use a buildtree as the source, given cli option
+    #    pull_dependencies ([Element]|None): Elements to attempt to pull
     #
     # Returns:
     #    (int): The exit code of the launched shell
@@ -149,20 +150,27 @@ class Stream():
               mounts=None,
               isolate=False,
               command=None,
-              usebuildtree=None):
+              usebuildtree=None,
+              pull_dependencies=None):
 
         # Assert we have everything we need built, unless the directory is specified
         # in which case we just blindly trust the directory, using the element
         # definitions to control the execution environment only.
         if directory is None:
             missing_deps = [
-                dep._get_full_name()
-                for dep in self._pipeline.dependencies([element], scope)
+                dep for dep in self._pipeline.dependencies([element], scope)
                 if not dep._cached()
             ]
             if missing_deps:
-                raise StreamError("Elements need to be built or downloaded before staging a shell environment",
-                                  detail="\n".join(missing_deps))
+                if not pull_dependencies:
+                    raise StreamError(
+                        "Elements need to be built or downloaded before staging a shell environment",
+                        detail="\n"
+                        .join(list(map(lambda x: x._get_full_name(), missing_deps))))
+                self._message(MessageType.INFO, "Attempting to fetch missing or incomplete artifacts")
+                self._add_queue(PullQueue(self._scheduler))
+                self._enqueue_plan([element] + missing_deps)
+                self._run()
 
         buildtree = False
         # Check if we require a pull queue attempt, with given artifact state and context