DataX二次开发-支持writeMode配置update

背景#

目前很多主流数据库支持on duplicate key update（当主键冲突update数据）模式，DataX也支持通过配置writeMode来配置写入模式。但是目前仅仅只支持MySQL的实现，这里以支持PostgreSQL介绍如何适配数据库on duplicate key update模式。

环境准备#

从github上拉取最新的DataX源码#

git clone https://github.com/alibaba/DataX.git 

需求#

• 当PostgreSQL 同步数据时出现主键冲突时，对字段进行update。
• 如果没有这样一条记录，则新增。

SQL语法#

 INSERT INTO table_name(column_list) VALUES(value_list)
 ON CONFLICT target action;

举例

INSERT INTO %s (id,name,data_time,remark) VALUES ( ?,?,?,? ) 
ON CONFLICT (id,name) 
DO  UPDATE SET id=excluded.id,name=excluded.name,data_time=excluded.data_time,remark=excluded.remark

代码#

通过分析DataX源码可知，update模式主要是在com.alibaba.datax.plugin.rdbms.writer.util.WriterUtil中通过生成update sql语句来实现存在则更新。

MySQL语句不需要填充主键或唯一键字段，但PostgreSQL需要显式指定主键或唯一键。

以此可修改代码

    public static String getWriteTemplate(List<String> columnHolders, List<String> valueHolders, String writeMode, DataBaseType dataBaseType, boolean forceUseUpdate) {
        boolean isWriteModeLegal = writeMode.trim().toLowerCase().startsWith("insert")
                || writeMode.trim().toLowerCase().startsWith("replace")
                || writeMode.trim().toLowerCase().startsWith("update");
        String mode = writeMode.trim().toLowerCase();
        String columns = StringUtils.join(columnHolders, ",");
        String placeHolders = StringUtils.join(valueHolders, ",");
        if (!isWriteModeLegal) {
            throw DataXException.asDataXException(DBUtilErrorCode.ILLEGAL_VALUE,
                    String.format("您所配置的 writeMode:%s 错误. 因为DataX 目前仅支持replace,update 或 insert 方式. 请检查您的配置并作出修改.", writeMode));
        }
        // && writeMode.trim().toLowerCase().startsWith("replace")
        //适配pgsql
        String writeDataSqlTemplate;
        if (forceUseUpdate || mode.startsWith("update")) {
            if (dataBaseType == DataBaseType.MySql || dataBaseType == DataBaseType.Tddl) {
                writeDataSqlTemplate = new StringBuilder()
                        .append("INSERT INTO %s (").append(StringUtils.join(columnHolders, ","))
                        .append(") VALUES(").append(StringUtils.join(valueHolders, ","))
                        .append(")")
                        .append(onDuplicateKeyUpdateString(columnHolders))
                        .toString();
            } else if (dataBaseType == DataBaseType.PostgreSQL) {
                writeDataSqlTemplate = "INSERT INTO %s (" + columns + ") VALUES ( " + placeHolders + " )" +
                        doPostgresqlUpdate(writeMode, columnHolders);
            } else if (dataBaseType == DataBaseType.Oracle) {
                writeDataSqlTemplate = doOracleUpdate(writeMode, columnHolders) + "INSERT (" +
                        StringUtils.join(columnHolders, ",") +
                        ") VALUES(" + StringUtils.join(valueHolders, ",") +
                        ")";
            } else {
                throw DataXException.asDataXException(DBUtilErrorCode.ILLEGAL_VALUE,
                        String.format("当前数据库不支持 writeMode:%s 模式.", writeMode));
            }

        } else {

            //这里是保护,如果其他错误的使用了update,需要更换为replace
            if (writeMode.trim().toLowerCase().startsWith("update")) {
                writeMode = "replace";
            }
            writeDataSqlTemplate = new StringBuilder().append(writeMode)
                    .append(" INTO %s (").append(StringUtils.join(columnHolders, ","))
                    .append(") VALUES(").append(StringUtils.join(valueHolders, ","))
                    .append(")").toString();
        }

        return writeDataSqlTemplate;
    }

    /**
     *  INSERT INTO table_name(column_list) VALUES(value_list)
     *  ON CONFLICT target action;
     *
     * Postgresql 存在则用主键或唯一键更新
     *
     * @param writeMode
     * @param columnHolders
     * @return
     */
    private static String doPostgresqlUpdate(String writeMode, List<String> columnHolders) {

        String conflict = writeMode.replace("update", "");
        StringBuilder sb = new StringBuilder();
        sb.append(" ON CONFLICT ");
        sb.append(conflict);
        sb.append(" DO ");
        if (columnHolders == null || columnHolders.size() < 1) {
            sb.append("NOTHING");
            return sb.toString();
        }
        sb.append(" UPDATE SET ");
        boolean first = true;
        for (String column : columnHolders) {
            if (!first) {
                sb.append(",");
            } else {
                first = false;
            }
            sb.append(column);
            sb.append("=excluded.");
            sb.append(column);
        }
        return sb.toString();
    }

在PostgresqlWriter中，解除对writeMode的判断

public void init() {
	this.originalConfig = super.getPluginJobConf();
	String writeMode = this.originalConfig.getString(Key.WRITE_MODE);
	if (null != writeMode) {
		if (!"insert".equalsIgnoreCase(writeMode)
				&& !writeMode.startsWith("update")) {
			throw DataXException.asDataXException(
					DBUtilErrorCode.CONF_ERROR,
					String.format("写入模式(writeMode)配置错误. PostgreSQL 仅支持insert, update两种模式." +
									" %s 不支持",
							writeMode));
		}
	}
	this.commonRdbmsWriterMaster = new CommonRdbmsWriter.Job(DATABASE_TYPE);
	this.commonRdbmsWriterMaster.init(this.originalConfig);
}

至此，对于PostgreSQL update模式已经支持，重写打包即可使用。对于其他数据源实现步骤基本相同，原理就是生成对应的SQL即可。

使用#

在job.json配置writeMode为update(key1,key2,key3)即可使用，例如：

{
    "job": {
        "setting": {
            "speed": {
                "channel": 1
            }
        },
        "content": [
            {
                 "reader": {
                    "name": "streamreader",
                    "parameter": {
                        "column" : [
                            {
                                "value": "DataX",
                                "type": "string"
                            },
                            {
                                "value": 19880808,
                                "type": "long"
                            },
                            {
                                "value": "1988-08-08 08:08:08",
                                "type": "date"
                            },
                            {
                                "value": true,
                                "type": "bool"
                            },
                            {
                                "value": "test",
                                "type": "bytes"
                            }
                        ],
                        "sliceRecordCount": 1000
                    }
                },
                "writer": {
                    "name": "postgresqlwriter",
                    "parameter": {
                        "username": "xx",
                        "password": "xx",
                        "writeMode": "update(id,name)",
                        "column": [
                            "id",
                            "name"
                        ],
                        "preSql": [
                            "delete from test"
                        ],
                        "connection": [
                            {
                                "jdbcUrl": "jdbc:postgresql://127.0.0.1:3002/datax",
                                "table": [
                                    "test"
                                ]
                            }
                        ]
                    }
                }
            }
        ]
    }
}

源码#

github源码：DataX二次开发