Merge ~ogayot/curtin:retry-download into curtin:master

[Olivier Gayot]

system-install: make it possible to download packages only and specify retries

Allow curtin system-install to specify the following options:

--download-only: the package(s) specified are downloaded but not installed/unpacked
--assume-downloaded: make curtin assume that the packages have already been downloaded. This is similar in essence to apt-get's --no-download option but only works on some of the package managers.

[Server Team CI bot (server-team-bot)]

FAILED: Continuous integration, rev:d267ca634452869b410b00f4585f8dbe3ba050c4
https://jenkins.canonical.com/server-team/job/curtin-ci/83/
Executed test runs:
    FAILURE: https://jenkins.canonical.com/server-team/job/curtin-ci/nodes=metal-amd64/83/
    FAILURE: https://jenkins.canonical.com/server-team/job/curtin-ci/nodes=metal-arm64/83/
    FAILURE: https://jenkins.canonical.com/server-team/job/curtin-ci/nodes=metal-ppc64el/83/
    FAILURE: https://jenkins.canonical.com/server-team/job/curtin-ci/nodes=metal-s390x/83/

Click here to trigger a rebuild:
https://jenkins.canonical.com/server-team/job/curtin-ci/83//rebuild

review: Needs Fixing (continuous-integration)

[Server Team CI bot (server-team-bot)]

PASSED: Continuous integration, rev:c75bf1d34e8f980a8ca94495bbc61524c4580cdc
https://jenkins.canonical.com/server-team/job/curtin-ci/84/
Executed test runs:
    SUCCESS: https://jenkins.canonical.com/server-team/job/curtin-ci/nodes=metal-amd64/84/
    SUCCESS: https://jenkins.canonical.com/server-team/job/curtin-ci/nodes=metal-arm64/84/
    SUCCESS: https://jenkins.canonical.com/server-team/job/curtin-ci/nodes=metal-ppc64el/84/
    SUCCESS: https://jenkins.canonical.com/server-team/job/curtin-ci/nodes=metal-s390x/84/

Click here to trigger a rebuild:
https://jenkins.canonical.com/server-team/job/curtin-ci/84//rebuild

review: Approve (continuous-integration)

[Michael Hudson-Doyle (mwhudson)]

Seems broadly fine. Are there any docs for these commands to update?

[Olivier Gayot (ogayot)]

> Are there any docs for these commands to update?

I could not find any but I might have overlooked something.

[Michael Hudson-Doyle (mwhudson)]

Oh well then :-)

review: Approve

diff --git a/curtin/commands/__init__.py b/curtin/commands/__init__.py
index 089a166..51b91c6 100644
--- a/curtin/commands/__init__.py
+++ b/curtin/commands/__init__.py
@@ -1,12 +1,27 @@
 # This file is part of curtin. See LICENSE file for copyright and license info.
 
+class MutuallyExclusiveGroup:
+    def __init__(self, entries) -> None:
+        self.entries = entries
+
 
 def populate_one_subcmd(parser, options_dict, handler):
-    for ent in options_dict:
-        args = ent[0]
-        if not isinstance(args, (list, tuple)):
-            args = (args,)
-        parser.add_argument(*args, **ent[1])
+    for entry in options_dict:
+        def add_entry_to_parser(parser, entry):
+            args = entry[0]
+            if not isinstance(args, (list, tuple)):
+                args = (args,)
+            parser.add_argument(*args, **entry[1])
+
+        if isinstance(entry, MutuallyExclusiveGroup):
+            group_parser = parser.add_mutually_exclusive_group()
+            subentries = entry.entries
+        else:
+            group_parser = parser
+            subentries = [entry]
+
+        for subentry in subentries:
+            add_entry_to_parser(group_parser, subentry)
     parser.set_defaults(func=handler)
 
 # vi: ts=4 expandtab syntax=python
diff --git a/curtin/commands/system_install.py b/curtin/commands/system_install.py
index 6d7b736..55bae07 100644
--- a/curtin/commands/system_install.py
+++ b/curtin/commands/system_install.py
@@ -5,7 +5,7 @@ import sys
 
 import curtin.util as util
 
-from . import populate_one_subcmd
+from . import populate_one_subcmd, MutuallyExclusiveGroup
 from curtin.log import LOG
 from curtin import distro
 
@@ -19,7 +19,10 @@ def system_install_pkgs_main(args):
     try:
         distro.install_packages(
             pkglist=args.packages, target=args.target,
-            allow_daemons=args.allow_daemons)
+            allow_daemons=args.allow_daemons,
+            download_retries=args.download_retry_after,
+            download_only=args.download_only,
+            assume_downloaded=args.assume_downloaded)
     except util.ProcessExecutionError as e:
         LOG.warn("system install failed for %s: %s" % (args.packages, e))
         exit_code = e.exit_code
@@ -27,6 +30,19 @@ def system_install_pkgs_main(args):
     sys.exit(exit_code)
 
 
+MUTUALLY_EXCLUSIVE_DOWNLOAD_OPTIONS = (
+    ((('--assume-downloaded',),
+      {'help': ('assume packages to install have already been downloaded.'
+                ' not supported on SUSE distro family.'),
+       'action': 'store_true'}),
+     (('--download-only',),
+      {'help': ('do not install/upgrade packages, only perform download.'
+                ' not supported on SUSE distro family.'),
+       'action': 'store_true'}),
+     )
+)
+
+
 CMD_ARGUMENTS = (
     ((('--allow-daemons',),
       {'help': ('do not disable running of daemons during upgrade.'),
@@ -36,6 +52,12 @@ CMD_ARGUMENTS = (
                 'default is env[TARGET_MOUNT_POINT]'),
        'action': 'store', 'metavar': 'TARGET',
        'default': os.environ.get('TARGET_MOUNT_POINT')}),
+     (('--download-retry-after',),
+      {'help': ('when a download fails, wait N seconds and try again.'
+                ' can be specified multiple times.'
+                ' not supported on SUSE distro family.'),
+       'action': 'append', 'nargs': '*'}),
+     MutuallyExclusiveGroup(MUTUALLY_EXCLUSIVE_DOWNLOAD_OPTIONS),
      ('packages',
       {'help': 'the list of packages to install',
        'metavar': 'PACKAGES', 'action': 'store', 'nargs': '+'}),
diff --git a/curtin/commands/system_upgrade.py b/curtin/commands/system_upgrade.py
index d4f6735..ac9aef6 100644
--- a/curtin/commands/system_upgrade.py
+++ b/curtin/commands/system_upgrade.py
@@ -18,7 +18,8 @@ def system_upgrade_main(args):
     exit_code = 0
     try:
         distro.system_upgrade(target=args.target,
-                              allow_daemons=args.allow_daemons)
+                              allow_daemons=args.allow_daemons,
+                              download_retries=args.download_retry_after)
     except util.ProcessExecutionError as e:
         LOG.warn("system upgrade failed: %s" % e)
         exit_code = e.exit_code
@@ -35,6 +36,11 @@ CMD_ARGUMENTS = (
                 'default is env[TARGET_MOUNT_POINT]'),
        'action': 'store', 'metavar': 'TARGET',
        'default': os.environ.get('TARGET_MOUNT_POINT')}),
+     (('--download-retry-after',),
+      {'help': ('when a download fails, wait N seconds and try again.'
+                ' can be specified multiple times.'
+                ' not supported on SUSE distro family.'),
+       'action': 'append', 'nargs': '*'}),
      )
 )
 
diff --git a/curtin/distro.py b/curtin/distro.py
index 4266278..433004c 100644
--- a/curtin/distro.py
+++ b/curtin/distro.py
@@ -6,6 +6,7 @@ import re
 import shutil
 import tempfile
 import textwrap
+from typing import Optional, Sequence
 
 from .paths import target_path
 from .util import (
@@ -276,7 +277,9 @@ def apt_update(target=None, env=None, force=False, comment=None,
 
 
 def run_apt_command(mode, args=None, opts=None, env=None, target=None,
-                    execute=True, allow_daemons=False, clean=True):
+                    execute=True, allow_daemons=False, clean=True,
+                    download_retries: Optional[Sequence[int]] = None,
+                    download_only=False, assume_downloaded=False):
     defopts = ['--quiet', '--assume-yes',
                '--option=Dpkg::options::=--force-unsafe-io',
                '--option=Dpkg::Options::=--force-confold']
@@ -299,17 +302,69 @@ def run_apt_command(mode, args=None, opts=None, env=None, target=None,
     if not execute:
         return env, cmd
 
-    apt_update(target, env=env, comment=' '.join(cmd))
+    if not assume_downloaded:
+        apt_update(target, env=env, comment=' '.join(cmd))
+    if mode in ['dist-upgrade', 'install', 'upgrade']:
+        cmd_rv = apt_install(mode, args, opts=opts, env=env, target=target,
+                             allow_daemons=allow_daemons,
+                             download_retries=download_retries,
+                             download_only=download_only,
+                             assume_downloaded=assume_downloaded)
+        if clean and not download_only:
+            with ChrootableTarget(
+                    target, allow_daemons=allow_daemons) as inchroot:
+                inchroot.subp(['apt-get', 'clean'])
+        return cmd_rv
+
+    with ChrootableTarget(target, allow_daemons=allow_daemons) as inchroot:
+        return inchroot.subp(cmd, env=env)
+
+
+def apt_install(mode, packages=None, opts=None, env=None, target=None,
+                allow_daemons=False,
+                download_retries: Optional[Sequence[int]] = None,
+                download_only=False, assume_downloaded=False):
+    """ Install or upgrade a set or all the packages using apt-get. """
+    defopts = ['--quiet', '--assume-yes',
+               '--option=Dpkg::options::=--force-unsafe-io',
+               '--option=Dpkg::Options::=--force-confold']
+    if packages is None:
+        packages = []
+
+    if opts is None:
+        opts = []
+
+    if mode not in ['install', 'upgrade', 'dist-upgrade']:
+        raise ValueError(
+            'Unsupported mode "%s" for apt package install/upgrade' % mode)
+
+    if download_only and assume_downloaded:
+        raise ValueError(
+            'download-only and assume-downloaded options are incompatible')
+
+    # download first, then install/upgrade from cache
+    cmd = ['apt-get'] + defopts + opts + [mode]
+    dl_opts = ['--download-only']
+    # NOTE: it would feel natural to use --no-download here but sadly apt-get
+    # can fail with "Internal Error, Pathname to install is not absolute [...]"
+    # when packages were retrieved from the cdrom and not downloaded from the
+    # Internet.
+    inst_opts = []
+
     with ChrootableTarget(target, allow_daemons=allow_daemons) as inchroot:
-        cmd_rv = inchroot.subp(cmd, env=env)
-        if clean and mode in ['dist-upgrade', 'install', 'upgrade']:
-            inchroot.subp(['apt-get', 'clean'])
+        if not assume_downloaded:
+            cmd_rv = inchroot.subp(cmd + dl_opts + packages, env=env,
+                                   retries=download_retries)
+        if not download_only:
+            cmd_rv = inchroot.subp(cmd + inst_opts + packages, env=env)
 
     return cmd_rv
 
 
 def run_yum_command(mode, args=None, opts=None, env=None, target=None,
-                    execute=True, allow_daemons=False):
+                    execute=True, allow_daemons=False,
+                    download_retries: Optional[Sequence[int]] = None,
+                    download_only=False, assume_downloaded=False):
     defopts = ['--assumeyes', '--quiet']
 
     if args is None:
@@ -330,14 +385,19 @@ def run_yum_command(mode, args=None, opts=None, env=None, target=None,
 
     if mode in ["install", "update", "upgrade"]:
         return yum_install(mode, args, opts=opts, env=env, target=target,
-                           allow_daemons=allow_daemons)
+                           allow_daemons=allow_daemons,
+                           download_retries=download_retries,
+                           download_only=download_only,
+                           assume_downloaded=assume_downloaded)
 
     with ChrootableTarget(target, allow_daemons=allow_daemons) as inchroot:
         return inchroot.subp(cmd, env=env)
 
 
 def yum_install(mode, packages=None, opts=None, env=None, target=None,
-                allow_daemons=False):
+                allow_daemons=False,
+                download_retries: Optional[Sequence[int]] = None,
+                download_only=False, assume_downloaded=False):
 
     defopts = ['--assumeyes', '--quiet']
 
@@ -347,10 +407,17 @@ def yum_install(mode, packages=None, opts=None, env=None, target=None,
     if opts is None:
         opts = []
 
+    if download_retries is None:
+        download_retries = [1] * 10
+
     if mode not in ['install', 'update', 'upgrade']:
         raise ValueError(
             'Unsupported mode "%s" for yum package install/upgrade' % mode)
 
+    if download_only and assume_downloaded:
+        raise ValueError(
+            'download-only and assume-downloaded options are incompatible')
+
     # dnf is a drop in replacement for yum. On newer RH based systems yum
     # is just a sym link to dnf.
     if which('dnf', target=target):
@@ -364,9 +431,13 @@ def yum_install(mode, packages=None, opts=None, env=None, target=None,
 
     # rpm requires /dev /sys and /proc be mounted, use ChrootableTarget
     with ChrootableTarget(target, allow_daemons=allow_daemons) as inchroot:
-        inchroot.subp(cmd + dl_opts + packages,
-                      env=env, retries=[1] * 10)
-        return inchroot.subp(cmd + inst_opts + packages, env=env)
+        if not assume_downloaded:
+            cmd_rv = inchroot.subp(cmd + dl_opts + packages,
+                                   env=env, retries=download_retries)
+        if not download_only:
+            cmd_rv = inchroot.subp(cmd + inst_opts + packages, env=env)
+
+    return cmd_rv
 
 
 def rpm_get_dist_id(target=None):
@@ -380,7 +451,9 @@ def rpm_get_dist_id(target=None):
 
 
 def run_zypper_command(mode, args=None, opts=None, env=None, target=None,
-                       execute=True, allow_daemons=False):
+                       execute=True, allow_daemons=False,
+                       download_retries: Optional[Sequence[int]] = None,
+                       download_only=False, assume_downloaded=False):
     defopts = ['--non-interactive', '--non-interactive-include-reboot-patches',
                '--quiet']
 
@@ -395,12 +468,15 @@ def run_zypper_command(mode, args=None, opts=None, env=None, target=None,
     if not execute:
         return env, cmd
 
+    # TODO add support for retried downloads, download-only and no-download.
+
     with ChrootableTarget(target, allow_daemons=allow_daemons) as inchroot:
         return inchroot.subp(cmd, env=env)
 
 
 def system_upgrade(opts=None, target=None, env=None, allow_daemons=False,
-                   osfamily=None):
+                   osfamily=None,
+                   download_retries: Optional[Sequence[int]] = None):
     LOG.debug("Upgrading system in %s", target)
 
     if not osfamily:
@@ -421,12 +497,15 @@ def system_upgrade(opts=None, target=None, env=None, allow_daemons=False,
     for mode in distro_cfg[osfamily]['subcommands']:
         ret = distro_cfg[osfamily]['function'](
                 mode, opts=opts, target=target,
-                env=env, allow_daemons=allow_daemons)
+                env=env, allow_daemons=allow_daemons,
+                download_retries=download_retries)
     return ret
 
 
 def install_packages(pkglist, osfamily=None, opts=None, target=None, env=None,
-                     allow_daemons=False):
+                     allow_daemons=False,
+                     download_retries: Optional[Sequence[int]] = None,
+                     download_only=False, assume_downloaded=False):
     if isinstance(pkglist, str):
         pkglist = [pkglist]
 
@@ -445,7 +524,10 @@ def install_packages(pkglist, osfamily=None, opts=None, target=None, env=None,
                          osfamily)
 
     return install_cmd('install', args=pkglist, opts=opts, target=target,
-                       env=env, allow_daemons=allow_daemons)
+                       env=env, allow_daemons=allow_daemons,
+                       download_retries=download_retries,
+                       download_only=download_only,
+                       assume_downloaded=assume_downloaded)
 
 
 def has_pkg_available(pkg, target=None, osfamily=None):
diff --git a/tests/unittests/test_distro.py b/tests/unittests/test_distro.py
index 9851650..5743475 100644
--- a/tests/unittests/test_distro.py
+++ b/tests/unittests/test_distro.py
@@ -2,6 +2,7 @@
 
 from unittest import skipIf
 import mock
+import os
 import sys
 
 from curtin import distro
@@ -293,6 +294,107 @@ class TestDistroIdentity(CiTestCase):
             self.mock_os_path.assert_called_with('/etc/redhat-release')
 
 
+class TestAptInstall(CiTestCase):
+    @mock.patch.object(util.ChrootableTarget, "__enter__", new=lambda a: a)
+    @mock.patch.dict(os.environ, clear=True)
+    @mock.patch.object(distro, 'apt_install')
+    @mock.patch.object(distro, 'apt_update')
+    @mock.patch('curtin.util.subp')
+    def test_run_apt_command(self, m_subp, m_apt_update, m_apt_install):
+        # install with defaults
+        expected_env = {'DEBIAN_FRONTEND': 'noninteractive'}
+        expected_calls = [
+            mock.call('install', ['foobar', 'wark'],
+                      opts=[], env=expected_env, target=None,
+                      allow_daemons=False, download_retries=None,
+                      download_only=False, assume_downloaded=False)
+        ]
+
+        distro.run_apt_command('install', ['foobar', 'wark'])
+        m_apt_update.assert_called_once()
+        m_apt_install.assert_has_calls(expected_calls)
+        m_subp.assert_called_once_with(['apt-get', 'clean'], target='/')
+
+        m_subp.reset_mock()
+        m_apt_install.reset_mock()
+        m_apt_update.reset_mock()
+
+        # no clean option
+        distro.run_apt_command('install', ['foobar', 'wark'], clean=False)
+        m_apt_update.assert_called_once()
+        m_subp.assert_has_calls(expected_calls[:-1])
+
+    @mock.patch.object(util.ChrootableTarget, "__enter__", new=lambda a: a)
+    @mock.patch('curtin.util.subp')
+    def test_apt_install(self, m_subp):
+        cmd_prefix = [
+            'apt-get', '--quiet', '--assume-yes',
+            '--option=Dpkg::options::=--force-unsafe-io',
+            '--option=Dpkg::Options::=--force-confold',
+        ]
+
+        expected_calls = [
+            mock.call(cmd_prefix + ['install', '--download-only']
+                                 + ['foobar', 'wark'],
+                      env=None, target='/', retries=None),
+            mock.call(cmd_prefix + ['install']
+                                 + ['foobar', 'wark'],
+                      env=None, target='/'),
+        ]
+
+        distro.apt_install('install', packages=['foobar', 'wark'])
+        m_subp.assert_has_calls(expected_calls)
+
+        expected_calls = [
+            mock.call(cmd_prefix + ['upgrade', '--download-only'],
+                      env=None, target='/', retries=None),
+            mock.call(cmd_prefix + ['upgrade'],
+                      env=None, target='/'),
+        ]
+
+        m_subp.reset_mock()
+        distro.apt_install('upgrade')
+        m_subp.assert_has_calls(expected_calls)
+
+        expected_calls = [
+            mock.call(cmd_prefix + ['dist-upgrade', '--download-only'],
+                      env=None, target='/', retries=None),
+            mock.call(cmd_prefix + ['dist-upgrade'],
+                      env=None, target='/'),
+        ]
+
+        m_subp.reset_mock()
+        distro.apt_install('dist-upgrade')
+        m_subp.assert_has_calls(expected_calls)
+
+        expected_dl_cmd = cmd_prefix + ['install', '--download-only', 'git']
+        expected_inst_cmd = cmd_prefix + ['install', 'git']
+
+        m_subp.reset_mock()
+        distro.apt_install('install', ['git'], download_only=True)
+        m_subp.assert_called_once_with(expected_dl_cmd, env=None, target='/',
+                                       retries=None)
+
+        m_subp.reset_mock()
+        distro.apt_install('install', ['git'], assume_downloaded=True)
+        m_subp.assert_called_once_with(expected_inst_cmd, env=None, target='/')
+
+    @mock.patch.object(util.ChrootableTarget, "__enter__", new=lambda a: a)
+    @mock.patch('curtin.util.subp')
+    def test_apt_install_invalid_mode(self, m_subp):
+        with self.assertRaisesRegex(ValueError, 'Unsupported mode.*'):
+            distro.apt_install('update')
+        m_subp.assert_not_called()
+
+    @mock.patch.object(util.ChrootableTarget, "__enter__", new=lambda a: a)
+    @mock.patch('curtin.util.subp')
+    def test_apt_install_conflict(self, m_subp):
+        with self.assertRaisesRegex(ValueError, '.*incompatible.*'):
+            distro.apt_install('install', ['git'],
+                               download_only=True, assume_downloaded=True)
+        m_subp.assert_not_called()
+
+
 class TestYumInstall(CiTestCase):
 
     @mock.patch.object(util.ChrootableTarget, "__enter__", new=lambda a: a)
@@ -482,9 +584,13 @@ class TestSystemUpgrade(CiTestCase):
             '--option=Dpkg::options::=--force-unsafe-io',
             '--option=Dpkg::Options::=--force-confold']
         apt_cmd = apt_base + ['dist-upgrade'] + pkglist
+        dl_apt_cmd = apt_base + ['dist-upgrade', '--download-only'] + pkglist
+        inst_apt_cmd = apt_base + ['dist-upgrade'] + pkglist
         auto_remove = apt_base + ['autoremove']
         expected_calls = [
-            mock.call(apt_cmd, env=env, target=paths.target_path(target)),
+            mock.call(dl_apt_cmd, env=env, retries=None,
+                      target=paths.target_path(target)),
+            mock.call(inst_apt_cmd, env=env, target=paths.target_path(target)),
             mock.call(['apt-get', 'clean'], target=paths.target_path(target)),
             mock.call(auto_remove, env=env, target=paths.target_path(target)),
         ]