class LocalDbProvider(KeyValueProvider):

    def __init__(self, config, blocking_auth=True):
        self.config = config
        self.bucket = config.get('bucket', 'studioml-meta')

        factory: StorageHandlerFactory = StorageHandlerFactory.get_factory()
        self.meta_store = factory.get_handler(StorageType.storageLocal, config)

        self.endpoint = self.meta_store.get_endpoint()
        self.db_root = os.path.join(self.endpoint, self.bucket)
        self._ensure_path_dirs_exist(self.db_root)

        super().__init__(config, self.meta_store, blocking_auth)

    def _ensure_path_dirs_exist(self, path):
        dirs = os.path.dirname(path)
        os.makedirs(dirs, mode = 0o777, exist_ok = True)

    def _get(self, key, shallow=False):
        file_name = os.path.join(self.db_root, key)
        if not os.path.exists(file_name):
            return None
        try:
            with open(file_name) as infile:
                result = json.load(infile)
        except BaseException as exc:
            self.logger.error("FAILED to load file %s - %s", file_name, exc)
            result = None
        return result

    def _delete(self, key, shallow=True):
        file_name = os.path.join(self.db_root, key)
        if os.path.exists(file_name):
            self.logger.debug("Deleting local database file %s.", file_name)
            util.delete_local_path(file_name, self.db_root, shallow)

    def _set(self, key, value):
        file_name = os.path.join(self.db_root, key)
        self._ensure_path_dirs_exist(file_name)
        with open(file_name, 'w') as outfile:
            json.dump(value, outfile) class S3Provider(KeyValueProvider):

    def __init__(self, config, blocking_auth=True):
        self.config = config
        self.bucket = config.get('bucket', 'studioml-meta')

        factory: StorageHandlerFactory = StorageHandlerFactory.get_factory()
        self.meta_store = factory.get_handler(StorageType.storageS3, config)

        super().__init__(
            config,
            self.meta_store,
            blocking_auth)

    def _get(self, key, shallow=False):
        try:
            response = self.meta_store.client.list_objects(
                Bucket=self.bucket,
                Prefix=key,
                Delimiter='/',
                MaxKeys=1024*16
            )
        except Exception as exc:
            msg: str = "FAILED to list objects in bucket {0}: {1}"\
                .format(self.bucket, exc)
            self._report_fatal(msg)
            return None

        if response is None:
            return None

        if 'Contents' not in response.keys():
            return None

        key_count = len(response['Contents'])

        if key_count == 0:
            return None

        for key_item in response['Contents']:
            if 'Key' in key_item.keys() and key_item['Key'] == key:
                response = self.meta_store.client.get_object(
                    Bucket=self.bucket,
                    Key=key)
                return json.loads(response['Body'].read().decode("utf-8"))

        return None

    def _delete(self, key, shallow=True):
        self.logger.info("S3 deleting object: %s/%s", self.bucket, key)

        try:
            response = self.meta_store.client.delete_object(
                Bucket=self.bucket,
                Key=key)
        except Exception as exc:
            msg: str = "FAILED to delete object {0} in bucket {1}: {2}"\
                .format(key, self.bucket, exc)
            self.logger.info(msg)
            return

        reason = response['ResponseMetadata'] if response else "None"
        if response is None or\
            response['ResponseMetadata']['HTTPStatusCode'] != 204:
            msg: str = ('attempt to delete key {0} in bucket {1}' +
                       ' returned response {2}')\
                           .format(key, self.bucket, reason)
            self.logger.info(msg)

    def _set(self, key, value):
        try:
            response = self.meta_store.client.put_object(
                Bucket=self.bucket,
                Key=key,
                Body=json.dumps(value))
        except Exception as exc:
            msg: str = "FAILED to write object {0} in bucket {1}: {2}"\
                .format(key, self.bucket, exc)
            self._report_fatal(msg)

        reason = response['ResponseMetadata'] if response else "None"
        if response is None or \
                response['ResponseMetadata']['HTTPStatusCode'] != 200:
            msg: str = ('attempt to write key {0} in bucket {1}' +
                       ' returned response {2}')\
                .format(key, self.bucket, reason)
            self._report_fatal(msg) class StorageHandler:
    def __init__(self, storage_type: StorageType,
                 logger,
                 measure_timestamp_diff=False,
                 compression=None):
        self.type = storage_type
        self.logger = logger
        self.compression = compression
        self._timestamp_shift = 0
        if measure_timestamp_diff:
            try:
                self._timestamp_shift = self._measure_timestamp_diff()
            except BaseException:
                check_for_kb_interrupt()
                self._timestamp_shift = 0

    @classmethod
    def get_id(cls, config: Dict) -> str:
        raise NotImplementedError("Not implemented: upload_file")

    def upload_file(self, key, local_path):
        raise NotImplementedError("Not implemented: upload_file")

    def download_file(self, key, local_path):
        raise NotImplementedError("Not implemented: download_file")

    def download_remote_path(self, remote_path, local_path):
        raise NotImplementedError("Not implemented: download_remote_path")

    def delete_file(self, key, shallow=True):
        raise NotImplementedError("Not implemented: delete_file")

    def get_file_url(self, key, method='GET'):
        raise NotImplementedError("Not implemented: get_file_url")

    def get_file_timestamp(self, key):
        raise NotImplementedError("Not implemented: get_file_timestamp")

    def get_qualified_location(self, key):
        raise NotImplementedError("Not implemented: get_qualified_location")

    def get_local_destination(self, remote_path: str):
        raise NotImplementedError("Not implemented: get_local_destination")

    def get_timestamp_shift(self):
        return self._timestamp_shift

    def get_compression(self):
        return self.compression

    def cleanup(self):
        pass

    def _measure_timestamp_diff(self):
        max_diff = 60
        tmpfile = get_temp_filename() + '.txt'
        with open(tmpfile, 'w') as f:
            f.write('timestamp_diff_test')
        key = 'tests/' + str(uuid.uuid4())
        self.upload_file(key, tmpfile)
        remote_timestamp = self.get_file_timestamp(key)

        if remote_timestamp is not None:
            now_remote_diff = time.time() - remote_timestamp
            self.storage_handler.delete_file(key)
            os.remove(tmpfile)

            assert -max_diff < now_remote_diff and \
                now_remote_diff < max_diff, \
                "Timestamp difference is more than 60 seconds. " + \
                "You'll need to adjust local clock for caching " + \
                "to work correctly"

            return -now_remote_diff if now_remote_diff < 0 else 0 class StorageHandlerFactory:
    def __init__(self):
        self.logger = logs.get_logger(self.__class__.__name__)
        self.logger.setLevel(get_storage_verbose_level())
        self.handlers_cache = dict()
        self.cleanup_at_exit: bool = True

    @classmethod
    def get_factory(cls):
        global _storage_factory
        if _storage_factory is None:
            _storage_factory = StorageHandlerFactory()
        return _storage_factory

    def set_cleanup_at_exit(self, value: bool):
        self.cleanup_at_exit = value

    def cleanup(self):
        if not self.cleanup_at_exit:
            return

        for _, handler in self.handlers_cache.items():
            if handler is not None:
                handler.cleanup()

    def get_handler(self, handler_type: StorageType,
                    config: Dict) -> StorageHandler:
        if handler_type == StorageType.storageS3:
            handler_id: str = S3StorageHandler.get_id(config)
            handler = self.handlers_cache.get(handler_id, None)
            if handler is None:
                handler = S3StorageHandler(config)
                self.handlers_cache[handler_id] = handler
            return handler
        if handler_type == StorageType.storageHTTP:
            handler_id: str = HTTPStorageHandler.get_id(config)
            handler = self.handlers_cache.get(handler_id, None)
            if handler is None:
                handler = HTTPStorageHandler(
                    config.get('endpoint', None),
                    config.get('credentials', None))
                self.handlers_cache[handler_id] = handler
            return handler
        if handler_type == StorageType.storageLocal:
            handler_id: str = LocalStorageHandler.get_id(config)
            handler = self.handlers_cache.get(handler_id, None)
            if handler is None:
                handler = LocalStorageHandler(config)
                self.handlers_cache[handler_id] = handler
            return handler
        self.logger("FAILED to get storage handler: unsupported type %s",
                    repr(handler_type))
        return None def setup_storage(db_provider, artifact_store):
    global _storage_setup
    _storage_setup = { DB_KEY: db_provider, STORE_KEY: artifact_store } def get_storage_db_provider():
    global _storage_setup
    if _storage_setup is None:
        return None
    return _storage_setup.get(DB_KEY, None) def set_storage_verbose_level(level: int):
    global _storage_verbose_level
    _storage_verbose_level = level class StorageType(Enum):
    storageHTTP = 1
    storageS3 = 2
    storageLocal = 3
    storageFirebase = 4
    storageDockerHub = 5
    storageSHub = 6

    storageInvalid = 99 DEBUG = logging.DEBUG
INFO = logging.INFO
ERROR = logging.ERROR
def get_logger(name):
def debug(line):
def error(line):
def info(line): def parse_verbosity(verbosity=None):
    if verbosity is None:
        return parse_verbosity('info')

    if verbosity == 'True':
        return parse_verbosity('info')

    logger_levels = {
        'debug': 10,
        'info': 20,
        'warn': 30,
        'error': 40,
        'crit': 50
    }

    if isinstance(verbosity, str) and \
            verbosity in logger_levels.keys():
        return logger_levels[verbosity]
    return int(verbosity)
import os
import yaml
import pyhocon
from studio.db_providers.local_db_provider import LocalDbProvider
from studio.db_providers.s3_provider import S3Provider
from studio.storage.storage_handler import StorageHandler
from studio.storage.storage_handler_factory import StorageHandlerFactory
from studio.storage.storage_setup import setup_storage, get_storage_db_provider,\
    set_storage_verbose_level
from studio.storage.storage_type import StorageType
from studio.util import logs
from studio.util.util import parse_verbosity
"""Data providers."""


def get_config(config_file=None):

    config_paths = []
    if config_file:
        if not os.path.exists(config_file):
            raise ValueError('User config file {} not found'
                             .format(config_file))
        config_paths.append(os.path.expanduser(config_file))

    config_paths.append(os.path.expanduser('~/.studioml/config.yaml'))
    config_paths.append(
        os.path.join(
            os.path.dirname(os.path.realpath(__file__)),
            "default_config.yaml"))

    for path in config_paths:
        if not os.path.exists(path):
            continue

        with(open(path)) as f_in:
            if path.endswith('.hocon'):
                config = pyhocon.ConfigFactory.parse_string(f_in.read())
            else:
                config = yaml.load(f_in.read(), Loader=yaml.FullLoader)

                def replace_with_env(config):
                    for key, value in config.items():
                        if isinstance(value, str):
                            config[key] = os.path.expandvars(value)

                        elif isinstance(value, dict):
                            replace_with_env(value)

                replace_with_env(config)

            return config

    raise ValueError('None of the config paths {0} exists!'
                     .format(config_paths))

def get_artifact_store(config) -> StorageHandler:
    storage_type: str = config['type'].lower()


    factory: StorageHandlerFactory = StorageHandlerFactory.get_factory()