uncloud/ucloud/common/storage_handlers.py

import shutil
import subprocess as sp
import os
import stat

from abc import ABC
from . import logger
from os.path import join as join_path

from ucloud.settings import settings as config


class ImageStorageHandler(ABC):
    handler_name = "base"

    def __init__(self, image_base, vm_base):
        self.image_base = image_base
        self.vm_base = vm_base

    def import_image(self, image_src, image_dest, protect=False):
        """Put an image at the destination
        :param image_src: An Image file
        :param image_dest: A path where :param src: is to be put.
        :param protect: If protect is true then the dest is protect (readonly etc)
        The obj must exist on filesystem.
        """

        raise NotImplementedError()

    def make_vm_image(self, image_path, path):
        """Copy image from src to dest

        :param image_path: A path
        :param path: A path

        src and destination must be on same storage system i.e both on file system or both on CEPH etc.
        """
        raise NotImplementedError()

    def resize_vm_image(self, path, size):
        """Resize image located at :param path:
        :param path: The file which is to be resized
        :param size: Size must be in Megabytes
        """
        raise NotImplementedError()

    def delete_vm_image(self, path):
        raise NotImplementedError()

    def execute_command(self, command, report=True, error_origin=None):
        if not error_origin:
            error_origin = self.handler_name

        command = list(map(str, command))
        try:
            sp.check_output(command, stderr=sp.PIPE)
        except sp.CalledProcessError as e:
            _stderr = e.stderr.decode("utf-8").strip()
            if report:
                logger.exception("%s:- %s", error_origin, _stderr)
            return False
        return True

    def vm_path_string(self, path):
        raise NotImplementedError()

    def qemu_path_string(self, path):
        raise NotImplementedError()

    def is_vm_image_exists(self, path):
        raise NotImplementedError()


class FileSystemBasedImageStorageHandler(ImageStorageHandler):
    handler_name = "Filesystem"

    def import_image(self, src, dest, protect=False):
        dest = join_path(self.image_base, dest)
        try:
            shutil.copy(src, dest)
            if protect:
                os.chmod(
                    dest, stat.S_IRUSR | stat.S_IRGRP | stat.S_IROTH
                )
        except Exception as e:
            logger.exception(e)
            return False
        return True

    def make_vm_image(self, src, dest):
        src = join_path(self.image_base, src)
        dest = join_path(self.vm_base, dest)
        try:
            shutil.copyfile(src, dest)
        except Exception as e:
            logger.exception(e)
            return False
        return True

    def resize_vm_image(self, path, size):
        path = join_path(self.vm_base, path)
        command = [
            "qemu-img",
            "resize",
            "-f",
            "raw",
            path,
            "{}M".format(size),
        ]
        if self.execute_command(command):
            return True
        else:
            self.delete_vm_image(path)
            return False

    def delete_vm_image(self, path):
        path = join_path(self.vm_base, path)
        try:
            os.remove(path)
        except Exception as e:
            logger.exception(e)
            return False
        return True

    def vm_path_string(self, path):
        return join_path(self.vm_base, path)

    def qemu_path_string(self, path):
        return self.vm_path_string(path)

    def is_vm_image_exists(self, path):
        path = join_path(self.vm_base, path)
        command = ["ls", path]
        return self.execute_command(command, report=False)


class CEPHBasedImageStorageHandler(ImageStorageHandler):
    handler_name = "Ceph"

    def import_image(self, src, dest, protect=False):
        dest = join_path(self.image_base, dest)
        import_command = ["rbd", "import", src, dest]
        commands = [import_command]
        if protect:
            snap_create_command = [
                "rbd",
                "snap",
                "create",
                "{}@protected".format(dest),
            ]
            snap_protect_command = [
                "rbd",
                "snap",
                "protect",
                "{}@protected".format(dest),
            ]
            commands.append(snap_create_command)
            commands.append(snap_protect_command)

        result = True
        for command in commands:
            result = result and self.execute_command(command)

        return result

    def make_vm_image(self, src, dest):
        src = join_path(self.image_base, src)
        dest = join_path(self.vm_base, dest)

        command = ["rbd", "clone", "{}@protected".format(src), dest]
        return self.execute_command(command)

    def resize_vm_image(self, path, size):
        path = join_path(self.vm_base, path)
        command = ["rbd", "resize", path, "--size", size]
        return self.execute_command(command)

    def delete_vm_image(self, path):
        path = join_path(self.vm_base, path)
        command = ["rbd", "rm", path]
        return self.execute_command(command)

    def vm_path_string(self, path):
        return join_path(self.vm_base, path)

    def qemu_path_string(self, path):
        return "rbd:{}".format(self.vm_path_string(path))

    def is_vm_image_exists(self, path):
        path = join_path(self.vm_base, path)
        command = ["rbd", "info", path]
        return self.execute_command(command, report=False)


def get_storage_handler():
    __storage_backend = config["storage"]["storage_backend"]
    if __storage_backend == "filesystem":
        return FileSystemBasedImageStorageHandler(
            vm_base=config["storage"]["vm_dir"],
            image_base=config["storage"]["image_dir"],
        )
    elif __storage_backend == "ceph":
        return CEPHBasedImageStorageHandler(
            vm_base=config["storage"]["ceph_vm_pool"],
            image_base=config["storage"]["ceph_image_pool"],
        )
    else:
        raise Exception("Unknown Image Storage Handler")