Superset/dict_import_export_util.py

# pylint: disable=C,R,W
import logging

from superset.connectors.druid.models import DruidCluster
from superset.models.core import Database


DATABASES_KEY = 'databases'
DRUID_CLUSTERS_KEY = 'druid_clusters'


def export_schema_to_dict(back_references):
    """Exports the supported import/export schema to a dictionary"""
    databases = [Database.export_schema(recursive=True,
                 include_parent_ref=back_references)]
    clusters = [DruidCluster.export_schema(recursive=True,
                include_parent_ref=back_references)]
    data = dict()
    if databases:
        data[DATABASES_KEY] = databases
    if clusters:
        data[DRUID_CLUSTERS_KEY] = clusters
    return data


def export_to_dict(session,
                   recursive,
                   back_references,
                   include_defaults):
    """Exports databases and druid clusters to a dictionary"""
    logging.info('Starting export')
    dbs = session.query(Database)
    databases = [database.export_to_dict(recursive=recursive,
                 include_parent_ref=back_references,
                 include_defaults=include_defaults) for database in dbs]
    logging.info('Exported %d %s', len(databases), DATABASES_KEY)
    cls = session.query(DruidCluster)
    clusters = [cluster.export_to_dict(recursive=recursive,
                include_parent_ref=back_references,
                include_defaults=include_defaults) for cluster in cls]
    logging.info('Exported %d %s', len(clusters), DRUID_CLUSTERS_KEY)
    data = dict()
    if databases:
        data[DATABASES_KEY] = databases
    if clusters:
        data[DRUID_CLUSTERS_KEY] = clusters
    return data


def import_from_dict(session, data, sync=[]):
    """Imports databases and druid clusters from dictionary"""
    if isinstance(data, dict):
        logging.info('Importing %d %s',
                     len(data.get(DATABASES_KEY, [])),
                     DATABASES_KEY)
        for database in data.get(DATABASES_KEY, []):
            Database.import_from_dict(session, database, sync=sync)

        logging.info('Importing %d %s',
                     len(data.get(DRUID_CLUSTERS_KEY, [])),
                     DRUID_CLUSTERS_KEY)
        for datasource in data.get(DRUID_CLUSTERS_KEY, []):
            DruidCluster.import_from_dict(session, datasource, sync=sync)
        session.commit()
    else:
        logging.info('Supplied object is not a dictionary.')