package com.lingbobu.flashdb.transfer.impl;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Types;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import javax.sql.DataSource;
import org.springframework.dao.DataAccessException;
import org.springframework.jdbc.core.BatchPreparedStatementSetter;
import org.springframework.jdbc.core.ConnectionCallback;
import org.springframework.jdbc.core.JdbcTemplate;
import com.lingbobu.flashdb.common.ColumnInfo;
import com.lingbobu.flashdb.common.DataTypes;
import com.lingbobu.flashdb.common.ResultSetMeta;
/**
* 基于JDBC的MySQL迁移数据输出器
*/
public class OutputMySQL extends OutputSQL {
public OutputMySQL(DataSource dataSource, String tableName, int mode) {
super(tableName, mode);
this.jdbcTemplate = new JdbcTemplate(dataSource);
}
public OutputMySQL(DataSource dataSource, String tableName, String[] updateColumns) {
super(tableName, updateColumns);
this.jdbcTemplate = new JdbcTemplate(dataSource);
}
private JdbcTemplate jdbcTemplate;
@Override
protected ResultSetMeta getTableMetaInfos(final String tableName) {
return jdbcTemplate.execute(new ConnectionCallback<ResultSetMeta>(){
@Override
public ResultSetMeta doInConnection(Connection conn) throws SQLException, DataAccessException {
return getResultSetMeta(conn, tableName);
}
});
}
private static ResultSetMeta getResultSetMeta(Connection conn, String tableName) throws SQLException, DataAccessException {
DatabaseMetaData dbMetaData = conn.getMetaData();
ResultSet rsColumns = dbMetaData.getColumns(null, null, tableName, "%");
List<ColumnInfo> lcolumnInfos = new ArrayList<ColumnInfo>();
Map<String, Integer> columnIndexs = new HashMap<String, Integer>();
int nAutoIncrement = -1;
while (rsColumns.next()) {
String columnName = rsColumns.getString("COLUMN_NAME");
int columnSqlType = rsColumns.getInt("DATA_TYPE");
String sisAutoIncrement = rsColumns.getString("IS_AUTOINCREMENT");
String sisNullable = rsColumns.getString("IS_NULLABLE");
boolean isAutoIncrement = ((sisAutoIncrement != null) && "YES".equals(sisAutoIncrement));
boolean isNullable = ((sisNullable != null) && (! "NO".equals(sisAutoIncrement)));
if (isAutoIncrement)
nAutoIncrement = lcolumnInfos.size();
int dataType = sqlType2DataType(columnSqlType);
if (dataType == DataTypes.TYPE_OBJECT)
throw new RuntimeException("Unsupport sql type "+columnSqlType+" for column '"+columnName+"'");
columnIndexs.put(columnName, lcolumnInfos.size());
lcolumnInfos.add(new ColumnInfo(columnName, dataType, isNullable));
}
ResultSet rstPrimaryKeys = dbMetaData.getPrimaryKeys(null, null, tableName);
List<String> keyColumns = new ArrayList<String>();
while (rstPrimaryKeys.next()) {
String columnName = rstPrimaryKeys.getString("COLUMN_NAME");
keyColumns.add(columnName);
}
if (nAutoIncrement >= 0) {
if (keyColumns.size() != 1)
throw new RuntimeException("AutoIncrement but no PrimaryKey");
if (! keyColumns.get(0).equals(lcolumnInfos.get(nAutoIncrement).getName()))
throw new RuntimeException("AutoIncrement only support on column of PrimaryKey");
}
// 将主键字段移到最前面
ColumnInfo[] columnInfos = lcolumnInfos.toArray(new ColumnInfo[lcolumnInfos.size()]);
for (int i=0; i < keyColumns.size(); i++) {
String columnName = keyColumns.get(i);
Integer columnIndex = columnIndexs.get(columnName);
if (columnIndex == null)
throw new RuntimeException("Unkown column '"+columnName+"' of PrimaryKey");
int n = columnIndex.intValue();
if (n == i) continue;
else if (n < i)
throw new RuntimeException("Error order of column '"+columnName+"' for PrimaryKey");
ColumnInfo tmp = columnInfos[i];
columnInfos[i] = columnInfos[n];
columnInfos[n] = tmp;
columnIndexs.put(columnInfos[i].getName(), i);
columnIndexs.put(columnInfos[n].getName(), n);
}
return new ResultSetMeta(columnInfos, keyColumns.size(), tableName, (nAutoIncrement >= 0));
}
public static int sqlType2DataType(int sqlType) {
if ((sqlType == Types.VARCHAR)||(sqlType == Types.CHAR)||(sqlType == Types.CLOB)||(sqlType == Types.LONGVARCHAR))
return DataTypes.TYPE_String;
else if ((sqlType == Types.INTEGER)||(sqlType == Types.SMALLINT)||(sqlType == Types.TINYINT))
return DataTypes.TYPE_Integer;
else if (sqlType == Types.BIGINT)
return DataTypes.TYPE_Long;
else if (sqlType == Types.FLOAT)
return DataTypes.TYPE_Float;
else if (sqlType == Types.DOUBLE)
return DataTypes.TYPE_Double;
else if ((sqlType == Types.DECIMAL)||(sqlType == Types.NUMERIC))
return DataTypes.TYPE_Currency;
else if (sqlType == Types.DATE)
return DataTypes.TYPE_Date;
else if (sqlType == Types.TIMESTAMP)
return DataTypes.TYPE_DateTime;
else if (sqlType == Types.BIT)
return DataTypes.TYPE_Boolean;
else if (sqlType == Types.VARBINARY)
return DataTypes.TYPE_Binary;
return DataTypes.TYPE_OBJECT;
}
@Override
protected void flushTable(String tableName) {
jdbcTemplate.execute("optimize table " + tableName);
}
@Override
protected void batchOutput(String sql, final Object[][] batchRows) {
final ColumnInfo[] columnInfos = super.columnInfos;
jdbcTemplate.batchUpdate(sql, new BatchPreparedStatementSetter(){
@Override
public void setValues(PreparedStatement ps, int i) throws SQLException {
Object[] row = batchRows[i];
for (int k=0; k < columnInfos.length; k++) {
Object value = row[k];
ps.setObject(1+k, value);
}
}
@Override
public int getBatchSize() {
return batchRows.length;
}
});
}
}