poyee-data-warehouse/dw_base/datax/plugins/writer/elasticsearch_writer.py

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

from configparser import ConfigParser

from dw_base.datax.plugins.writer.writer import Writer

# elasticsearch writer
ES_WRITER_NAME = 'elasticsearch-writer'
ES_WRITER_PARAMETER_NODES = 'nodes'
ES_WRITER_PARAMETER_BATCH_SIZE = 'batchSize'
ES_WRITER_PARAMETER_INDEX = 'index'
ES_WRITER_PARAMETER_SPEED_PER_SECOND = 'speedPerSecond'
ES_WRITER_PARAMETER_WRITE_MODE = 'writeMode'


class ElasticsearchWriter(Writer):
    def __init__(self, base_dir: str, config_parser: ConfigParser, start_date: str = None, stop_date: str = None):
        super(ElasticsearchWriter, self).__init__(base_dir, config_parser, start_date, stop_date)
        self.plugin_name = ES_WRITER_NAME

    def load_others(self):
        self.parameter[ES_WRITER_PARAMETER_BATCH_SIZE] = self.config_parser.get(self.plugin_type,
                                                                                ES_WRITER_PARAMETER_BATCH_SIZE)
        self.parameter[ES_WRITER_PARAMETER_INDEX] = self.config_parser.get(self.plugin_type, ES_WRITER_PARAMETER_INDEX)
        self.parameter[ES_WRITER_PARAMETER_SPEED_PER_SECOND] = self.config_parser.get(self.plugin_type,
                                                                                      ES_WRITER_PARAMETER_SPEED_PER_SECOND)
        self.parameter[ES_WRITER_PARAMETER_WRITE_MODE] = self.config_parser.get(self.plugin_type,
                                                                                ES_WRITER_PARAMETER_WRITE_MODE)