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

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

import re
from configparser import ConfigParser

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

# postgresql writer
POSTGRE_SQL_WRITER_NAME = 'postgresqlwriter'
POSTGRE_SQL_WRITER_PARAMETER_BATCH_SIZE = 'batchSize'
POSTGRE_SQL_WRITER_PARAMETER_CONNECTION = 'connection'
POSTGRE_SQL_WRITER_PARAMETER_COLUMN = 'column'
POSTGRE_SQL_WRITER_PARAMETER_DATABASE = 'database'
POSTGRE_SQL_WRITER_PARAMETER_POST_SQL = 'postSql'
POSTGRE_SQL_WRITER_PARAMETER_PRE_SQL = 'preSql'
POSTGRE_SQL_WRITER_PARAMETER_TABLE = 'table'
POSTGRE_SQL_WRITER_PARAMETER_WRITE_MODE = 'writeMode'


class PostgreSQLWriter(Writer):
    def __init__(self, base_dir: str, config_parser: ConfigParser, start_time: str = None, stop_time: str = None):
        super(PostgreSQLWriter, self).__init__(base_dir, config_parser, start_time, stop_time)
        self.plugin_name = POSTGRE_SQL_WRITER_NAME

    def load_others(self):
        database = self.config_parser.get(self.plugin_type, POSTGRE_SQL_WRITER_PARAMETER_DATABASE)
        self.check_config(POSTGRE_SQL_WRITER_PARAMETER_DATABASE, database)
        table = self.config_parser.get(self.plugin_type, POSTGRE_SQL_WRITER_PARAMETER_TABLE)
        self.check_config(POSTGRE_SQL_WRITER_PARAMETER_TABLE, table)
        jdbc_url: str = self.parameter[DS_POSTGRE_SQL_JDBC_URL]
        matcher = re.search('jdbc:postgresql://(.+?)/(.+)', jdbc_url)
        if matcher:
            if database:
                jdbc_url = jdbc_url.replace(matcher.group(2), database)
        elif jdbc_url.endswith('/'):
            jdbc_url = f'{jdbc_url}{database}'
        else:
            jdbc_url = f'{jdbc_url}/{database}'
        connection = {
            DS_POSTGRE_SQL_JDBC_URL: jdbc_url,
            POSTGRE_SQL_WRITER_PARAMETER_TABLE: table.split(',')
        }
        self.parameter[POSTGRE_SQL_WRITER_PARAMETER_CONNECTION] = [connection]
        del self.parameter[DS_POSTGRE_SQL_JDBC_URL]
        self.parameter[POSTGRE_SQL_WRITER_PARAMETER_POST_SQL] = \
            self.config_parser.get(self.plugin_type, POSTGRE_SQL_WRITER_PARAMETER_POST_SQL).split(';') or []
        self.parameter[POSTGRE_SQL_WRITER_PARAMETER_PRE_SQL] = \
            self.config_parser.get(self.plugin_type, POSTGRE_SQL_WRITER_PARAMETER_PRE_SQL).split(';') or []
        self.parameter[POSTGRE_SQL_WRITER_PARAMETER_WRITE_MODE] = \
            self.config_parser.get(self.plugin_type, POSTGRE_SQL_WRITER_PARAMETER_WRITE_MODE) or 'insert'
        self.parameter[POSTGRE_SQL_WRITER_PARAMETER_BATCH_SIZE] = \
            self.config_parser.get(self.plugin_type, POSTGRE_SQL_WRITER_PARAMETER_BATCH_SIZE) or '1024'

    def load_column(self):
        columns = self.config_parser.get(self.plugin_type, POSTGRE_SQL_WRITER_PARAMETER_COLUMN).split(',')
        self.parameter[POSTGRE_SQL_WRITER_PARAMETER_COLUMN] = columns