poyee-data-warehouse/dw_base/datax/datasources/mongo_data_source.py

# -*- coding:utf-8 -*-


from dw_base.datax.datasources.data_source import DataSource

# Mongo DataSource
DS_TYPE_MONGO = 'mongo'
DS_MONGO_KEYS = ['address']


class MongoDataSource(DataSource):
    def __init__(self, ds_file: str):
        super(MongoDataSource, self).__init__(ds_file)
        self.source_type = DS_TYPE_MONGO
        self.keys = DS_MONGO_KEYS

    def get_datasource_dict(self):
        for key in self.keys:
            try:
                value = self.config_parser.get('base', key)
            except KeyError:
                raise KeyError('%s must specified at %s data source config.' % (key, self.source_type))
            if value:
                if key == 'address':
                    self.ds_dict[key] = [value]
                else:
                    self.ds_dict[key] = value
            else:
                raise KeyError('%s must specified at %s data source config.' % (key, self.source_type))
        return self.ds_dict