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二次开发
