Milan Verespej | 6e372ca | 2019-07-08 12:42:30 +0200 | [diff] [blame] | 1 | #! /usr/bin/env python3 |
Milan Verespej | da9e6b9 | 2019-06-18 15:09:41 +0200 | [diff] [blame] | 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) |
Samuli Silvius | 3c6f676 | 2019-11-25 09:13:41 +0200 | [diff] [blame^] | 43 | self._docker_client = docker.from_env(timeout=300) |
Milan Verespej | da9e6b9 | 2019-06-18 15:09:41 +0200 | [diff] [blame] | 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() |