many: support for stage-snaps

Add support for stage-snaps, with support for advanced grammar. Support for a snap
source was added to be able to extract the downloaded snaps.

LP: #1805214

Signed-off-by: Sergio Schvezov <sergio.schvezov@canonical.com>

schema/snapcraft.json

@@ -651,6 +651,11 @@
                             },
                             "default": []
                         },
+                        "stage-snaps": {
+                            "$comment": "For some reason 'default' doesn't work if in the ref",
+                            "$ref": "#/definitions/grammar-array",
+                            "default": []
+                        },
                         "stage-packages": {
                             "$comment": "For some reason 'default' doesn't work if in the ref",
                             "$ref": "#/definitions/grammar-array",

snapcraft/_baseplugin.py

@@ -72,6 +72,7 @@ def stage_packages(self, value):
     def __init__(self, name, options, project=None):
         self.name = name
         self.build_snaps = []
+        self.stage_snaps = []
         self.build_packages = []
         self._stage_packages = []
 
@@ -81,6 +82,8 @@ def __init__(self, name, options, project=None):
             self.build_packages = options.build_packages.copy()
         with contextlib.suppress(AttributeError):
             self.build_snaps = options.build_snaps.copy()
+        with contextlib.suppress(AttributeError):
+            self.stage_snaps = options.stage_snaps.copy()
 
         self.project = project
         self.options = options
@@ -94,6 +97,7 @@ def __init__(self, name, options, project=None):
         self.installdir = os.path.join(self.partdir, "install")
         self.statedir = os.path.join(self.partdir, "state")
         self.osrepodir = os.path.join(self.partdir, "ubuntu")
+        self.snapsdir = os.path.join(self.partdir, "snaps")
 
         self.build_basedir = os.path.join(self.partdir, "build")
         source_subdir = getattr(self.options, "source_subdir", None)

snapcraft/internal/pluginhandler/__init__.py

@@ -382,6 +382,28 @@ def mark_cleaned(self, step):
         if os.path.isdir(self.plugin.statedir) and not os.listdir(self.plugin.statedir):
             os.rmdir(self.plugin.statedir)
 
+    def _fetch_stage_snaps(self):
+        stage_snaps = self._grammar_processor.get_stage_snaps()
+        if stage_snaps:
+            repo.snaps.download_snaps(
+                snaps_list=stage_snaps, directory=self.plugin.snapsdir
+            )
+
+    def _unpack_stage_snaps(self):
+        stage_snaps = self._grammar_processor.get_stage_snaps()
+        if not stage_snaps:
+            return
+
+        logger.debug("Unpacking stage-snaps to {!r}".format(self.plugin.stage_snaps))
+        snap_files = iglob(os.path.join(self.plugin.snapsdir, "*.snap"))
+        snap_sources = (
+            sources.Snap(source=s, source_dir=self.plugin.snapsdir) for s in snap_files
+        )
+        for snap_source in snap_sources:
+            snap_source.provision(
+                self.plugin.installdir, clean_target=False, keep_snap=True
+            )
+
     def _fetch_stage_packages(self):
         stage_packages = self._grammar_processor.get_stage_packages()
         if stage_packages:
@@ -402,7 +424,9 @@ def _unpack_stage_packages(self):
     def prepare_pull(self, force=False):
         self.makedirs()
         self._fetch_stage_packages()
+        self._fetch_stage_snaps()
         self._unpack_stage_packages()
+        self._unpack_stage_snaps()
 
     def pull(self, force=False):
         # Ensure any previously-failed pull is cleared out before we try again
@@ -479,6 +503,10 @@ def clean_pull(self):
         if os.path.exists(self.plugin.osrepodir):
             shutil.rmtree(self.plugin.osrepodir)
 
+        # Remove snaps dir (where stage snaps are fetched)
+        if os.path.exists(self.plugin.snapsdir):
+            shutil.rmtree(self.plugin.snapsdir)
+
         if os.path.exists(self.plugin.sourcedir):
             if os.path.islink(self.plugin.sourcedir):
                 os.remove(self.plugin.sourcedir)

snapcraft/internal/project_loader/grammar_processing/_part_grammar_processor.py

@@ -85,6 +85,7 @@ def __init__(
         self._repo = repo
 
         self.__build_snaps = set()  # type: Set[str]
+        self.__stage_snaps = set()  # type: Set[str]
         self.__build_packages = set()  # type: Set[str]
         self.__stage_packages = set()  # type: Set[str]
 
@@ -119,6 +120,17 @@ def get_build_snaps(self) -> Set[str]:
 
         return self.__build_snaps
 
+    def get_stage_snaps(self) -> Set[str]:
+        if not self.__stage_snaps:
+            processor = grammar.GrammarProcessor(
+                getattr(self._plugin, "stage_snaps", []),
+                self._project,
+                repo.snaps.SnapPackage.is_valid_snap,
+            )
+            self.__stage_snaps = processor.process()
+
+        return self.__stage_snaps
+
     def get_build_packages(self) -> Set[str]:
         if not self.__build_packages:
             processor = grammar.GrammarProcessor(

snapcraft/internal/repo/errors.py

@@ -141,6 +141,14 @@ def __init__(self, *, snap_name, snap_channel):
         super().__init__(snap_name=snap_name, snap_channel=snap_channel)
 
 
+class SnapDownloadError(RepoError):
+
+    fmt = "Error while downloading snap {snap_name!r} from channel {snap_channel!r}"
+
+    def __init__(self, *, snap_name, snap_channel):
+        super().__init__(snap_name=snap_name, snap_channel=snap_channel)
+
+
 class SnapGetAssertionError(RepoError):
 
     fmt = (

snapcraft/internal/repo/snaps.py

@@ -1,6 +1,6 @@
 # -*- Mode:Python; indent-tabs-mode:nil; tab-width:4 -*-
 #
-# Copyright (C) 2017-2018 Canonical Ltd
+# Copyright (C) 2017-2019 Canonical Ltd
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License version 3 as
@@ -13,8 +13,10 @@
 #
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
 import contextlib
 import logging
+import os
 import sys
 from subprocess import check_call, check_output, CalledProcessError
 from typing import Sequence
@@ -158,6 +160,20 @@ def is_valid(self):
         store_channels = self._get_store_channels()
         return self.channel in store_channels.keys()
 
+    def download(self, *, directory: str = None):
+        """Downloads a given snap."""
+        # We use the `snap download` command here on recommendation
+        # of the snapd team.
+        snap_download_cmd = ["snap", "download", self.name]
+        if self._original_channel:
+            snap_download_cmd.extend(["--channel", self._original_channel])
+        try:
+            check_output(snap_download_cmd, cwd=directory)
+        except CalledProcessError:
+            raise errors.SnapDownloadError(
+                snap_name=self.name, snap_channel=self.channel
+            )
+
     def install(self):
         """Installs the snap onto the system."""
         snap_install_cmd = []
@@ -195,6 +211,27 @@ def refresh(self):
             )
 
 
+def download_snaps(*, snaps_list: Sequence[str], directory: str) -> None:
+    """
+    Download snaps of the format <snap-name>/<channel> into directory.
+
+    The target directory is created if it does not exist.
+    """
+    # TODO manifest.yaml with snap revision from future machine output
+    # for `snap download`.
+    os.makedirs(directory, exist_ok=True)
+    for snap in snaps_list:
+        snap_pkg = SnapPackage(snap)
+        if not snap_pkg.is_valid():
+            raise errors.SnapUnavailableError(
+                snap_name=snap_pkg.name, snap_channel=snap_pkg.channel
+            )
+
+        # TODO: use dependency injected echoer
+        logger.info("Downloading snap {!r}".format(snap_pkg.name))
+        snap_pkg.download(directory=directory)
+
+
 def install_snaps(snaps_list):
     """Install snaps of the format <snap-name>/<channel>.
 

snapcraft/internal/sources/__init__.py

@@ -93,6 +93,7 @@
     from ._7z import SevenZip  # noqa
     from ._deb import Deb  # noqa
     from ._rpm import Rpm  # noqa
+    from ._snap import Snap  # noqa: F401
 
     _source_handler = {
         "bzr": Bazaar,
@@ -107,6 +108,7 @@
         "local": Local,
         "deb": Deb,
         "rpm": Rpm,
+        "snap": Snap,
         "": Local,
     }
 
@@ -167,7 +169,7 @@ def get_source_handler(source, *, source_type=""):
 
 
 def _get_source_type_from_uri(source, ignore_errors=False):  # noqa: C901
-    for extension in ["zip", "deb", "rpm", "7z"]:
+    for extension in ["zip", "deb", "rpm", "7z", "snap"]:
         if source.endswith(".{}".format(extension)):
             return extension
     source_type = ""

snapcraft/internal/sources/_snap.py

@@ -0,0 +1,116 @@
+# -*- Mode:Python; indent-tabs-mode:nil; tab-width:4 -*-
+#
+# Copyright (C) 2019 Canonical
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 3 as
+# published by the Free Software Foundation.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import os
+import shutil
+import tempfile
+
+from . import errors
+from ._base import FileBase
+from snapcraft import file_utils, yaml_utils
+
+
+class Snap(FileBase):
+    """Handles downloading and extractions for a snap source.
+
+    On provision, the meta directory is renamed to meta.<snap-name> and, if present,
+    the same applies for the snap directory which shall be renamed to
+    snap.<snap-name>.
+    """
+
+    def __init__(
+        self,
+        source: str,
+        source_dir: str,
+        source_tag: str = None,
+        source_commit: str = None,
+        source_branch: str = None,
+        source_depth: str = None,
+        source_checksum: str = None,
+    ) -> None:
+        super().__init__(
+            source,
+            source_dir,
+            source_tag,
+            source_commit,
+            source_branch,
+            source_depth,
+            source_checksum,
+            "unsquashfs",
+        )
+        if source_tag:
+            raise errors.SnapcraftSourceInvalidOptionError("snap", "source-tag")
+        elif source_commit:
+            raise errors.SnapcraftSourceInvalidOptionError("snap", "source-commit")
+        elif source_branch:
+            raise errors.SnapcraftSourceInvalidOptionError("snap", "source-branch")
+
+    def provision(
+        self,
+        dst: str,
+        clean_target: bool = True,
+        keep_snap: bool = False,
+        src: str = None,
+    ) -> None:
+        """
+        Provision the snap source to dst.
+
+        :param str dst: the destination directory to provision to.
+        :param bool clean_target: clean dst before provisioning if True.
+        :param bool keep_snap: keep the snap after provisioning is done if
+                               True.
+        :param str src: force a new source to use for extraction.
+        raises errors.InvalidSnapError: when trying to provision an invalid
+                                        snap.
+        """
+        if src:
+            snap_file = src
+        else:
+            snap_file = os.path.join(self.source_dir, os.path.basename(self.source))
+        snap_file = os.path.realpath(snap_file)
+
+        if clean_target:
+            tmp_snap = tempfile.NamedTemporaryFile().name
+            shutil.move(snap_file, tmp_snap)
+            shutil.rmtree(dst)
+            os.makedirs(dst)
+            shutil.move(tmp_snap, snap_file)
+
+        # unsquashfs [options] filesystem [directories or files to extract]
+        # options:
+        # -force: if file already exists then overwrite
+        # -dest <pathname>: unsquash to <pathname>
+        with tempfile.TemporaryDirectory(prefix=os.path.dirname(snap_file)) as temp_dir:
+            extract_command = ["unsquashfs", "-force", "-dest", temp_dir, snap_file]
+            self._run_output(extract_command)
+            snap_name = _get_snap_name(temp_dir)
+            # Rename meta and snap dirs from the snap
+            rename_paths = (os.path.join(temp_dir, d) for d in ["meta", "snap"])
+            rename_paths = (d for d in rename_paths if os.path.exists(d))
+            for rename in rename_paths:
+                shutil.move(rename, "{}.{}".format(rename, snap_name))
+            file_utils.link_or_copy_tree(source_tree=temp_dir, destination_tree=dst)
+
+        if not keep_snap:
+            os.remove(snap_file)
+
+
+def _get_snap_name(snap_dir: str) -> str:
+    try:
+        with open(os.path.join(snap_dir, "meta", "snap.yaml")) as snap_yaml:
+            return yaml_utils.load(snap_yaml)["name"]
+    except (FileNotFoundError, KeyError) as snap_error:
+        raise errors.InvalidSnapError() from snap_error

snapcraft/internal/sources/errors.py

@@ -87,6 +87,15 @@ class InvalidDebError(SnapcraftSourceError):
     )
 
 
+class InvalidSnapError(SnapcraftSourceError):
+
+    fmt = (
+        "The snap file used does not contain valid data. "
+        "Ensure a proper snap file is passed for .snap files "
+        "as sources."
+    )
+
+
 class SourceUpdateUnsupportedError(SnapcraftSourceError):
 
     fmt = "Failed to update source: {source!s} sources don't support updating."

tests/spread/general/stage_snaps/task.yaml

@@ -0,0 +1,34 @@
+summary: stage-snaps
+
+# We currently only have core18 on a stable channel
+systems: [ubuntu-18*]
+
+environment:
+  STAGE_SNAPS/default: "hello"
+  STAGE_SNAPS/channel: "hello/latest/stable"
+  STAGE_SNAPS/n_snaps: "hello, black/latest/edge"
+
+prepare: |
+  mkdir test-snap
+  cd test-snap
+  snapcraft init
+  
+  echo "    stage-snaps: [$STAGE_SNAPS]" >> snap/snapcraft.yaml
+
+restore: |
+  rm -rf test-snap
+
+execute: |
+  cd test-snap
+
+  snapcraft stage
+
+  if [ ! -f stage/meta.hello ]; then
+      echo "Missing expected stage-snaps payload from hello"
+      exit 1
+  fi
+
+  if grep black snap/snapcraft.yaml && [ ! -f stage/meta.black ]; then
+      echo "Missing expected stage-snaps payload from black"
+      exit 1
+  fi