# -*- 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