_stream.py: Ability to pull missing buildtrees outside of pull/buildtpollard/774

Adds helper function _buildtree_pull_required() to determine if a
pullqueue should be constructed, for commands outside of bst pull
and build where it is determined that an element's buildtree
artifact is to be required given the respective semantics and
config. Utilised in push() to attempt to mitigate skipping the push
of partial elements without the user having to have preceded it with
an explicit pull.

cli.py: Add new behaviour to push command description

element.py: Move _cached_buildtree() to be non local private method,
use _KeyStrength types to reduce duplication.

tests/integration/pullbuildtrees.py also updated to cover this
use-case.

4 files changed, 97 insertions, 26 deletions

diff --git a/buildstream/_frontend/cli.py b/buildstream/_frontend/cli.py
index b1b4e03b0..c7d4c1eed 100644
--- a/buildstream/_frontend/cli.py
+++ b/buildstream/_frontend/cli.py
@@ -469,6 +469,10 @@ def push(app, elements, deps, remote):
     The default destination is the highest priority configured cache. You can
     override this by passing a different cache URL with the `--remote` flag.
 
+    If bst has been configured to include build trees on artifact pulls,
+    an attempt will be made to pull any required build trees to avoid the
+    skipping of partial artifacts being pushed.
+
     Specify `--deps` to control which artifacts to push:
 
     \b
diff --git a/buildstream/_stream.py b/buildstream/_stream.py
index 76f1d67aa..2f9799178 100644
--- a/buildstream/_stream.py
+++ b/buildstream/_stream.py
@@ -327,6 +327,10 @@ class Stream():
     # If `remote` specified as None, then regular configuration will be used
     # to determine where to push artifacts to.
     #
+    # If any of the given targets are missing their expected buildtree artifact,
+    # a pull queue will be created if user context and available remotes allow for
+    # attempting to fetch them.
+    #
     def push(self, targets, *,
              selection=PipelineSelection.NONE,
              remote=None):
@@ -345,8 +349,17 @@ class Stream():
             raise StreamError("No artifact caches available for pushing artifacts")
 
         self._pipeline.assert_consistent(elements)
-        self._add_queue(PushQueue(self._scheduler))
-        self._enqueue_plan(elements)
+
+        # Check if we require a pull queue, with given artifact state and context
+        require_buildtrees = self._buildtree_pull_required(elements)
+        if require_buildtrees:
+            self._message(MessageType.INFO, "Attempting to fetch missing artifact buildtrees")
+            self._add_queue(PullQueue(self._scheduler))
+            self._enqueue_plan(require_buildtrees)
+
+        push_queue = PushQueue(self._scheduler)
+        self._add_queue(push_queue)
+        self._enqueue_plan(elements, queue=push_queue)
         self._run()
 
     # checkout()
@@ -1237,3 +1250,28 @@ class Stream():
             parts.append(element.normal_name)
 
         return os.path.join(directory, *reversed(parts))
+
+    # _buildtree_pull_required()
+    #
+    # Check if current task, given config, requires element buildtree artifact
+    #
+    # Args:
+    #    elements (list): elements to check if buildtrees are required
+    #
+    # Returns:
+    #    (list): elements requiring buildtrees
+    #
+    def _buildtree_pull_required(self, elements):
+        required_list = []
+
+        # If context is set to not pull buildtrees, or no fetch remotes, return empty list
+        if not (self._context.pull_buildtrees or self._artifacts.has_fetch_remotes()):
+            return required_list
+
+        for element in elements:
+            # Check if element is partially cached without its buildtree, as the element
+            # artifact may not be cached at all
+            if element._cached() and not element._cached_buildtree():
+                required_list.append(element)
+
+        return required_list
diff --git a/buildstream/element.py b/buildstream/element.py
index 7c647b323..d7072dc8c 100644
--- a/buildstream/element.py
+++ b/buildstream/element.py
@@ -1422,7 +1422,7 @@ class Element(Plugin):
                                                  .format(workspace.get_absolute_path())):
                             workspace.stage(temp_staging_directory)
                 # Check if we have a cached buildtree to use
-                elif self.__cached_buildtree():
+                elif self._cached_buildtree():
                     artifact_base, _ = self.__extract()
                     import_dir = os.path.join(artifact_base, 'buildtree')
                 else:
@@ -1808,7 +1808,7 @@ class Element(Plugin):
 
         # Do not push elements that aren't cached, or that are cached with a dangling buildtree
         # artifact unless element type is expected to have an an empty buildtree directory
-        if not self.__cached_buildtree():
+        if not self._cached_buildtree():
             return True
 
         # Do not push tainted artifact
@@ -1998,6 +1998,29 @@ class Element(Plugin):
     def _get_source_element(self):
         return self
 
+    # _cached_buildtree()
+    #
+    # Check if element artifact contains expected buildtree. An
+    # element's buildtree artifact will not be present if the rest
+    # of the partial artifact is not cached.
+    #
+    # Returns:
+    #     (bool): True if artifact cached with buildtree, False if
+    #             element not cached or missing expected buildtree.
+    #
+    def _cached_buildtree(self):
+        context = self._get_context()
+
+        if not self._cached():
+            return False
+
+        key_strength = _KeyStrength.STRONG if context.get_strict() else _KeyStrength.WEAK
+        if not self.__artifacts.contains_subdir_artifact(self, self._get_cache_key(strength=key_strength),
+                                                         'buildtree'):
+            return False
+
+        return True
+
     #############################################################
     #                   Private Local Methods                   #
     #############################################################
@@ -2764,27 +2787,6 @@ class Element(Plugin):
 
         return True
 
-    # __cached_buildtree():
-    #
-    # Check if cached element artifact contains expected buildtree
-    #
-    # Returns:
-    #     (bool): True if artifact cached with buildtree, False if
-    #             element not cached or missing expected buildtree
-    #
-    def __cached_buildtree(self):
-        context = self._get_context()
-
-        if not self._cached():
-            return False
-        elif context.get_strict():
-            if not self.__artifacts.contains_subdir_artifact(self, self.__strict_cache_key, 'buildtree'):
-                return False
-        elif not self.__artifacts.contains_subdir_artifact(self, self.__weak_cache_key, 'buildtree'):
-            return False
-
-        return True
-
     # __pull_directories():
     #
     # Which directories to include or exclude given the current
diff --git a/tests/integration/pullbuildtrees.py b/tests/integration/pullbuildtrees.py
index 0f9397251..f6fc71226 100644
--- a/tests/integration/pullbuildtrees.py
+++ b/tests/integration/pullbuildtrees.py
@@ -38,7 +38,8 @@ def test_pullbuildtrees(cli, tmpdir, datafiles, integration_cache):
 
     # Create artifact shares for pull & push testing
     with create_artifact_share(os.path.join(str(tmpdir), 'share1')) as share1,\
-        create_artifact_share(os.path.join(str(tmpdir), 'share2')) as share2:
+        create_artifact_share(os.path.join(str(tmpdir), 'share2')) as share2,\
+        create_artifact_share(os.path.join(str(tmpdir), 'share3')) as share3:
         cli.configure({
             'artifacts': {'url': share1.repo, 'push': True},
             'artifactdir': os.path.join(str(tmpdir), 'artifacts')
@@ -123,6 +124,32 @@ def test_pullbuildtrees(cli, tmpdir, datafiles, integration_cache):
         assert share2.has_artifact('test', element_name, cli.get_element_key(project, element_name))
         default_state(cli, tmpdir, share1)
 
+        # Assert that bst push will automatically attempt to pull a missing buildtree
+        # if pull-buildtrees is set, however as share3 is the only defined remote and is empty,
+        # assert that no element artifact buildtrees are pulled (no available remote buildtree) and thus the
+        # artifact cannot be pushed.
+        result = cli.run(project=project, args=['pull', element_name])
+        assert element_name in result.get_pulled_elements()
+        cli.configure({'artifacts': {'url': share3.repo, 'push': True}})
+        result = cli.run(project=project, args=['--pull-buildtrees', 'push', element_name])
+        assert "Attempting to fetch missing artifact buildtrees" in result.stderr
+        assert element_name not in result.get_pulled_elements()
+        assert not os.path.isdir(buildtreedir)
+        assert element_name not in result.get_pushed_elements()
+        assert not share3.has_artifact('test', element_name, cli.get_element_key(project, element_name))
+
+        # Assert that if we add an extra remote that has the buildtree artfact cached, bst push will
+        # automatically attempt to pull it and will be successful, leading to the full artifact being pushed
+        # to the empty share3. This gives the ability to attempt push currently partial artifacts to a remote,
+        # without exlipictly requiring a bst pull.
+        cli.configure({'artifacts': [{'url': share1.repo, 'push': False}, {'url': share3.repo, 'push': True}]})
+        result = cli.run(project=project, args=['--pull-buildtrees', 'push', element_name])
+        assert "Attempting to fetch missing artifact buildtrees" in result.stderr
+        assert element_name in result.get_pulled_elements()
+        assert os.path.isdir(buildtreedir)
+        assert element_name in result.get_pushed_elements()
+        assert share3.has_artifact('test', element_name, cli.get_element_key(project, element_name))
+
 
 # Ensure that only valid pull-buildtrees boolean options make it through the loading
 # process.