vstf/vstf/agent/env/basic/image_manager.py

   1 """
   2 Created on 2015-7-28
   3
   4 @author: y00228926
   5 """
   6 from vstf.common.utils import check_call
   7 import os
   8 import logging
   9
  10 LOG = logging.getLogger(__name__)
  11
  12
  13 class _ImageManager(object):
  14     """
  15     A qemu-img wrapper to create qcow2 child image from a parent image.
  16
  17     """
  18     def __init__(self, parent_image_path, child_image_dir):
  19         """
  20         :param parent_image_path    str: the parent image path.
  21         :param child_image_dir      str: the destination path to put child images.
  22         """
  23         self._create_child_str = 'qemu-img create -f %(image_type)s %(child_path)s -o backing_file=%(parent_path)s'
  24         self._convert_str = "qemu-img convert -O %(image_type)s %(parent_path)s %(child_path)s"
  25         self.child_image_dir = child_image_dir
  26         self.parent_image_path = parent_image_path
  27         assert os.path.isfile(self.parent_image_path)
  28         assert os.path.isdir(self.child_image_dir)
  29
  30     def create_child_image(self, child_name, full_clone=False, image_type='qcow2'):
  31         """
  32         create a child image and put it in self.child_image_dir.
  33
  34         :param child_name:  the image name to be created..
  35         :return: return the path of child image.
  36         """
  37
  38         image_path = os.path.join(self.child_image_dir, child_name) + '.' + image_type
  39         if full_clone:
  40             cmd = self._convert_str % {'image_type': image_type, 'child_path': image_path, 'parent_path': self.parent_image_path}
  41         else:
  42             cmd = self._create_child_str % {'child_path': image_path, 'parent_path': self.parent_image_path, 'image_type':image_type}
  43         check_call(cmd.split())
  44         return image_path
  45
  46
  47 class ImageManager(object):
  48     def __init__(self, cfg):
  49         """
  50         ImageManager creates images from configuration context.
  51
  52         :param cfg: dict, example:
  53         {
  54             'parent_image': "/mnt/sdb/ubuntu_salt_master.img",
  55             'dst_location': "/mnt/sdb",
  56             'full_clone':False,
  57             'type': "qcow2",
  58             'names': ['vm1','vm2','vm3','vm4']
  59         }
  60         :return:
  61         """
  62         super(ImageManager, self).__init__()
  63         cfg = self._check_cfg(cfg)
  64         self.parent_image = cfg['parent_image']
  65         self.image_dir = cfg['dst_location']
  66         self.full_clone = cfg['full_clone']
  67         self.image_type = cfg['type']
  68         self.names = cfg['names']
  69         self.mgr = _ImageManager(self.parent_image, self.image_dir)
  70
  71     @staticmethod
  72     def _check_cfg(cfg):
  73         for key in ('parent_image', 'dst_location', 'full_clone', 'type', 'names'):
  74             if key not in cfg:
  75                 raise Exception("does't find %s config" % key)
  76         if cfg['type'] not in ('raw', 'qcow2'):
  77             raise Exception("type:%s not supported, only support 'raw' and 'qcow2'" % cfg['type'])
  78         if not cfg['full_clone'] and cfg['type'] == 'raw':
  79             raise Exception("only support 'qcow2' for not full_clone image creation" % cfg['type'])
  80         return cfg
  81
  82     def create_all(self):
  83         """
  84         create images by configuration context.
  85
  86         :return: True for success, False for failure.
  87         """
  88         for name in self.names:
  89             image = self.mgr.create_child_image(name, self.full_clone, self.image_type)
  90             LOG.info("image: %s created", image)
  91         return True
  92
  93     def clean_all(self):
  94         """
  95         remove all the images created in one go.
  96
  97         :return: True for success. Raise exception otherwise.
  98         """
  99         for name in self.names:
 100             image_path = os.path.join(self.image_dir, name + '.' + self.image_type)
 101             try:
 102                 os.unlink(image_path)
 103                 LOG.info("remove:%s successfully", image_path)
 104             except Exception:
 105                 LOG.info("cann't find path:%s", image_path)
 106         return True
 107
 108
 109 if __name__ == '__main__':
 110     import argparse
 111     import json
 112     parser = argparse.ArgumentParser()
 113     parser.add_argument('action', choices = ('create','clean'), help='action:create|clean')
 114     parser.add_argument('--config', help='config file to parse')
 115     args = parser.parse_args()
 116     logging.basicConfig(level=logging.INFO)
 117     image_cfg = json.load(open(args.config))
 118     mgr = ImageManager(image_cfg)
 119     if args.action == 'create':
 120         mgr.create_all()
 121     if args.action == 'clean':
 122         mgr.clean_all()
 123
 124