dvc/neatlynx/cmd_data_import.py

import sys
import os
from shutil import copyfile
import re
import fasteners
import requests

from neatlynx.cmd_base import CmdBase
from neatlynx.cmd_data_sync import sizeof_fmt
from neatlynx.logger import Logger
from neatlynx.data_file_obj import DataFileObj
from neatlynx.exceptions import NeatLynxException
from neatlynx.state_file import StateFile


class DataImportError(NeatLynxException):
    def __init__(self, msg):
        NeatLynxException.__init__(self, 'Import error: {}'.format(msg))


class CmdDataImport(CmdBase):
    def __init__(self):
        CmdBase.__init__(self)
        pass

    def define_args(self, parser):
        self.set_skip_git_actions(parser)

        parser.add_argument('input',
                            metavar='',
                            help='Input file',
                            nargs='*')

        self.add_string_arg(parser, 'output', 'Output file')
        pass

    def run(self):
        lock = fasteners.InterProcessLock(self.git.lock_file)
        gotten = lock.acquire(timeout=5)
        if not gotten:
            Logger.printing('Cannot perform the command since NLX is busy and locked. Please retry the command later.')
            return 1

        try:
            if not self.skip_git_actions and not self.git.is_ready_to_go():
                return 1

            output = self.args.output
            for file in self.args.input:
                self.import_file(file, output)

            if self.skip_git_actions:
                self.not_committed_changes_warning()
                return 0

            message = 'NLX data import: {} {}'.format(' '.join(self.args.input), self.args.output)
            self.git.commit_all_changes_and_log_status(message)
        finally:
            lock.release()
        pass

    def import_file(self, input, output):
        if not CmdDataImport.is_url(input):
            if not os.path.exists(input):
                raise DataImportError('Input file "{}" does not exist'.format(input))
            if not os.path.isfile(input):
                raise DataImportError('Input file "{}" has to be a regular file'.format(input))

        if os.path.isdir(output):
            output = os.path.join(output, os.path.basename(input))

        dobj = DataFileObj(output, self.git, self.config)

        if os.path.exists(dobj.data_file_relative):
            raise DataImportError('Output file "{}" already exists'.format(dobj.data_file_relative))
        if not os.path.isdir(os.path.dirname(dobj.data_file_abs)):
            raise DataImportError('Output file directory "{}" does not exists'.format(
                os.path.dirname(dobj.data_file_relative)))

        os.makedirs(os.path.dirname(dobj.cache_file_relative), exist_ok=True)
        if CmdDataImport.is_url(input):
            Logger.debug('Downloading file {} ...'.format(input))
            self.download_file(input, dobj.cache_file_relative)
            Logger.debug('Input file "{}" was downloaded to cache "{}"'.format(
                input, dobj.cache_file_relative))
        else:
            copyfile(input, dobj.cache_file_relative)
            Logger.debug('Input file "{}" was copied to cache "{}"'.format(
                input, dobj.cache_file_relative))

        dobj.create_symlink()
        Logger.debug('Symlink from data file "{}" to the cache file "{}" was created'.
                     format(dobj.data_file_relative, dobj.cache_file_relative))

        state_file = StateFile(dobj.state_file_relative, self.git)
        state_file.save()
        Logger.debug('State file "{}" was created'.format(dobj.state_file_relative))
        pass

    URL_REGEX = re.compile(
        r'^(?:http|ftp)s?://'  # http:// or https://
        r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+(?:[A-Z]{2,6}\.?|[A-Z0-9-]{2,}\.?)|'  # domain...
        r'localhost|'  # localhost...
        r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})'  # ...or ip
        r'(?::\d+)?'  # optional port
        r'(?:/?|[/?]\S+)$', re.IGNORECASE)

    @staticmethod
    def is_url(url):
        return CmdDataImport.URL_REGEX.match(url) is not None

    @staticmethod
    def download_file(from_url, to_file):
        r = requests.get(from_url, stream=True)

        chunk_size = 1024 * 100
        downloaded = 0
        last_reported = 0
        report_bucket = 100*1024*1024
        with open(to_file, 'wb') as f:
            for chunk in r.iter_content(chunk_size=chunk_size):
                if chunk:  # filter out keep-alive new chunks
                    downloaded += chunk_size
                    last_reported += chunk_size
                    if last_reported >= report_bucket:
                        last_reported = 0
                        Logger.debug('Downloaded {}'.format(sizeof_fmt(downloaded)))
                    f.write(chunk)
        return


if __name__ == '__main__':
    try:
        sys.exit(CmdDataImport().run())
    except NeatLynxException as e:
        Logger.error(e)
        sys.exit(1)