build/download/docker_downloader.py

   1 #! /usr/bin/env python3
   2 # -*- coding: utf-8 -*-
   3
   4 #   COPYRIGHT NOTICE STARTS HERE
   5
   6 #   Copyright 2019 © Samsung Electronics Co., Ltd.
   7 #
   8 #   Licensed under the Apache License, Version 2.0 (the "License");
   9 #   you may not use this file except in compliance with the License.
  10 #   You may obtain a copy of the License at
  11 #
  12 #       http://www.apache.org/licenses/LICENSE-2.0
  13 #
  14 #   Unless required by applicable law or agreed to in writing, software
  15 #   distributed under the License is distributed on an "AS IS" BASIS,
  16 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  17 #   See the License for the specific language governing permissions and
  18 #   limitations under the License.
  19
  20 #   COPYRIGHT NOTICE ENDS HERE
  21
  22 import argparse
  23 import datetime
  24 import itertools
  25 import logging
  26 import os
  27 import sys
  28 import timeit
  29
  30 import docker
  31 from retrying import retry
  32
  33 from concurrent_downloader import ConcurrentDownloader
  34
  35 log = logging.getLogger(__name__)
  36
  37
  38 class DockerDownloader(ConcurrentDownloader):
  39     def __init__(self, save, *list_args, workers=3):
  40         self._save = save
  41         try:
  42             # big timeout in case of massive images like pnda-mirror-container:5.0.0 (11.4GB)
  43             self._docker_client = docker.client.DockerClient(version='auto', timeout=300)
  44         except docker.errors.DockerException as err:
  45             log.exception(
  46                 'Error creating docker client. Check if is docker installed and running'
  47                 ' or if you have right permissions.')
  48             raise err
  49         self._pulled_images = set(itertools.chain.from_iterable((image.tags for image
  50                                                                  in self._docker_client.images.list())))
  51         list_args = ([*x, None] if len(x) < 2 else x for x in list_args)
  52         super().__init__('docker images', *list_args, workers=workers)
  53
  54     @staticmethod
  55     def image_registry_name(image_name):
  56         """
  57         Get the name as shown in local registry. Since some strings are not part of name
  58         when using default registry e.g. docker.io
  59         :param image_name: name of the image from the list
  60         :return: name of the image as it is shown by docker
  61         """
  62         name = image_name
  63
  64         if name.startswith('docker.io/'):
  65             name = name.replace('docker.io/', '')
  66
  67         if name.startswith('library/'):
  68             name = name.replace('library/', '')
  69
  70         if ':' not in name.rsplit('/')[-1]:
  71             name = '{}:latest'.format(name)
  72
  73         return name
  74
  75     @property
  76     def check_table(self):
  77         """
  78         Table showing information of which images are pulled/saved
  79         """
  80         self.missing()
  81         return self._table(self._data_list)
  82
  83     @property
  84     def fail_table(self):
  85         """
  86         Table showing information about state of download of images
  87         that encountered problems while downloading
  88         """
  89         return self._table(self.missing())
  90
  91     @staticmethod
  92     def _image_filename(image_name):
  93         """
  94         Get a name of a file where image will be saved.
  95         :param image_name: Name of the image from list
  96         :return: Filename of the image
  97         """
  98         return '{}.tar'.format(image_name.replace(':', '_').replace('/', '_'))
  99
 100     def _table(self, images):
 101         """
 102         Get table in format for images
 103         :param images: images to put into table
 104         :return: check table format with specified images
 105         """
 106         header = ['Name', 'Pulled', 'Saved']
 107         data = []
 108         for item in images:
 109             if item not in self._missing:
 110                 data.append((item, True, True if self._save else 'N/A'))
 111             else:
 112                 data.append((item, self._missing[item]['pulled'], self._missing[item]['saved']))
 113         return self._check_table(header, {'Name': 'l'}, data)
 114
 115     def _is_pulled(self, image):
 116         return self.image_registry_name(image) in self._pulled_images
 117
 118     def _is_saved(self, image):
 119         dst = '{}/{}'.format(self._data_list[image], self._image_filename(image))
 120         return os.path.isfile(dst)
 121
 122     def _is_missing(self, item):
 123         """
 124         Missing docker images are checked slightly differently.
 125         """
 126         pass
 127
 128     def missing(self):
 129         """
 130         Get dictionary of images not present locally.
 131         """
 132         missing = dict()
 133         for image, dst in self._data_list.items():
 134             pulled = self._is_pulled(image)
 135             if self._save:
 136                 # if pulling and save is True. Save every pulled image to assure parity
 137                 saved = False if not pulled else self._is_saved(image)
 138             else:
 139                 saved = 'N/A'
 140             if not pulled or not saved:
 141                 missing[image] = {'dst': dst, 'pulled': pulled, 'saved': saved}
 142         self._missing = missing
 143         return self._missing
 144
 145     @retry(stop_max_attempt_number=5, wait_fixed=5000)
 146     def _pull_image(self, image_name):
 147         """
 148         Pull docker image.
 149         :param image_name: name of the image to be pulled
 150         :return: pulled image (image object)
 151         :raises docker.errors.APIError: after unsuccessful retries
 152         """
 153         if ':' not in image_name.rsplit('/')[-1]:
 154             image_name = '{}:latest'.format(image_name)
 155         try:
 156             image = self._docker_client.images.pull(image_name)
 157             log.info('Image {} pulled'.format(image_name))
 158             return image
 159         except docker.errors.APIError as err:
 160             log.warning('Failed: {}: {}. Retrying...'.format(image_name, err))
 161             raise err
 162
 163     def _save_image(self, image_name, image, output_dir):
 164         """
 165         Save image to tar.
 166         :param output_dir: path to destination directory
 167         :param image: image object from pull_image function
 168         :param image_name: name of the image from list
 169         """
 170         dst = '{}/{}'.format(output_dir, self._image_filename(image_name))
 171         if not os.path.exists(output_dir):
 172             os.makedirs(output_dir)
 173         try:
 174             with open(dst, 'wb') as f:
 175                 for chunk in image.save(named=self.image_registry_name(image_name)):
 176                     f.write(chunk)
 177             log.info('Image {} saved as {}'.format(image_name, dst))
 178         except Exception as err:
 179             if os.path.isfile(dst):
 180                 os.remove(dst)
 181             raise err
 182
 183     def _download_item(self, image):
 184         """ Pull and save docker image from specified docker registry
 185         :param image: image to be downloaded
 186         """
 187         image_name, image_dict = image
 188         log.info('Downloading image: {}'.format(image_name))
 189         try:
 190             if image_dict['pulled']:
 191                 image_to_save = self._docker_client.images.get(image_name)
 192             else:
 193                 image_to_save = self._pull_image(image_name)
 194             if self._save:
 195                 self._save_image(image_name, image_to_save, image_dict['dst'])
 196         except Exception as err:
 197             log.exception('Error downloading {}: {}'.format(image_name, err))
 198             raise err
 199
 200
 201 def run_cli():
 202     parser = argparse.ArgumentParser(description='Download docker images from list')
 203     parser.add_argument('image_list', metavar='image-list',
 204                         help='File with list of images to download.')
 205     parser.add_argument('--save', '-s', action='store_true', default=False,
 206                         help='Save images (without it only pull is executed)')
 207     parser.add_argument('--output-dir', '-o', default=os.getcwd(),
 208                         help='Download destination')
 209     parser.add_argument('--check', '-c', action='store_true', default=False,
 210                         help='Check what is missing. No download.'
 211                              'Use with combination with -s to check saved images as well.')
 212     parser.add_argument('--debug', action='store_true', default=False,
 213                         help='Turn on debug output')
 214     parser.add_argument('--workers', type=int, default=3,
 215                         help='Set maximum workers for parallel download (default: 3)')
 216
 217     args = parser.parse_args()
 218
 219     if args.debug:
 220         logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
 221     else:
 222         logging.basicConfig(stream=sys.stdout, level=logging.INFO, format='%(message)s')
 223
 224     downloader = DockerDownloader(args.save, [args.file_list, args.output_dir], workers=args.workers)
 225
 226     if args.check:
 227         log.info('Check mode. No download will be executed.')
 228         log.info(downloader.check_table)
 229         sys.exit(0)
 230
 231     timer_start = timeit.default_timer()
 232     try:
 233         downloader.download()
 234     except RuntimeError:
 235         sys.exit(1)
 236     finally:
 237         log.info('Downloading finished in {}'.format(
 238             datetime.timedelta(seconds=timeit.default_timer() - timer_start)))
 239
 240
 241 if __name__ == '__main__':
 242     run_cli()