/cassandrawriter/src/main/java/com/alibaba/datax/plugin/writer/cassandrawriter/CassandraWriterHelper.java

https://github.com/alibaba/DataX · Java · 351 lines · 277 code · 71 blank · 3 comment · 28 complexity · 1d0e0ae1fd38e17a103aacf45e86f18e MD5 · raw file

  1. package com.alibaba.datax.plugin.writer.cassandrawriter;
  2. import java.math.BigDecimal;
  3. import java.math.BigInteger;
  4. import java.net.InetAddress;
  5. import java.nio.ByteBuffer;
  6. import java.text.SimpleDateFormat;
  7. import java.util.ArrayList;
  8. import java.util.Arrays;
  9. import java.util.Date;
  10. import java.util.HashMap;
  11. import java.util.HashSet;
  12. import java.util.Iterator;
  13. import java.util.List;
  14. import java.util.Map;
  15. import java.util.Set;
  16. import java.util.UUID;
  17. import com.alibaba.datax.common.element.Column;
  18. import com.alibaba.datax.common.exception.DataXException;
  19. import com.alibaba.fastjson.JSON;
  20. import com.alibaba.fastjson.JSONArray;
  21. import com.alibaba.fastjson.JSONException;
  22. import com.alibaba.fastjson.JSONObject;
  23. import com.datastax.driver.core.BoundStatement;
  24. import com.datastax.driver.core.CodecRegistry;
  25. import com.datastax.driver.core.DataType;
  26. import com.datastax.driver.core.DataType.Name;
  27. import com.datastax.driver.core.Duration;
  28. import com.datastax.driver.core.LocalDate;
  29. import com.datastax.driver.core.TupleType;
  30. import com.datastax.driver.core.TupleValue;
  31. import com.datastax.driver.core.UDTValue;
  32. import com.datastax.driver.core.UserType;
  33. import com.datastax.driver.core.UserType.Field;
  34. import com.google.common.base.Splitter;
  35. import org.apache.commons.codec.binary.Base64;
  36. /**
  37. * Created by mazhenlin on 2019/8/21.
  38. */
  39. public class CassandraWriterHelper {
  40. static CodecRegistry registry = new CodecRegistry();
  41. public static Object parseFromString(String s, DataType sqlType ) throws Exception {
  42. if (s == null || s.isEmpty()) {
  43. if (sqlType.getName() == Name.ASCII || sqlType.getName() == Name.TEXT ||
  44. sqlType.getName() == Name.VARCHAR) {
  45. return s;
  46. } else {
  47. return null;
  48. }
  49. }
  50. switch (sqlType.getName()) {
  51. case ASCII:
  52. case TEXT:
  53. case VARCHAR:
  54. return s;
  55. case BLOB:
  56. if (s.length() == 0) {
  57. return new byte[0];
  58. }
  59. byte[] byteArray = new byte[s.length() / 2];
  60. for (int i = 0; i < byteArray.length; i++) {
  61. String subStr = s.substring(2 * i, 2 * i + 2);
  62. byteArray[i] = ((byte) Integer.parseInt(subStr, 16));
  63. }
  64. return ByteBuffer.wrap(byteArray);
  65. case BOOLEAN:
  66. return Boolean.valueOf(s);
  67. case TINYINT:
  68. return Byte.valueOf(s);
  69. case SMALLINT:
  70. return Short.valueOf(s);
  71. case INT:
  72. return Integer.valueOf(s);
  73. case BIGINT:
  74. return Long.valueOf(s);
  75. case VARINT:
  76. return new BigInteger(s, 10);
  77. case FLOAT:
  78. return Float.valueOf(s);
  79. case DOUBLE:
  80. return Double.valueOf(s);
  81. case DECIMAL:
  82. return new BigDecimal(s);
  83. case DATE: {
  84. String[] a = s.split("-");
  85. if (a.length != 3) {
  86. throw new Exception(String.format("DATE类型数据 '%s' 格式不正确,必须为yyyy-mm-dd格式", s));
  87. }
  88. return LocalDate.fromYearMonthDay(Integer.valueOf(a[0]), Integer.valueOf(a[1]),
  89. Integer.valueOf(a[2]));
  90. }
  91. case TIME:
  92. return Long.valueOf(s);
  93. case TIMESTAMP:
  94. return new Date(Long.valueOf(s));
  95. case UUID:
  96. case TIMEUUID:
  97. return UUID.fromString(s);
  98. case INET:
  99. String[] b = s.split("/");
  100. if (b.length < 2) {
  101. return InetAddress.getByName(s);
  102. }
  103. byte[] addr = InetAddress.getByName(b[1]).getAddress();
  104. return InetAddress.getByAddress(b[0], addr);
  105. case DURATION:
  106. return Duration.from(s);
  107. case LIST:
  108. case MAP:
  109. case SET:
  110. case TUPLE:
  111. case UDT:
  112. Object jsonObject = JSON.parse(s);
  113. return parseFromJson(jsonObject,sqlType);
  114. default:
  115. throw DataXException.asDataXException(CassandraWriterErrorCode.CONF_ERROR,
  116. "不支持您配置的列类型:" + sqlType + ", 请检查您的配置 或者 联系 管理员.");
  117. } // end switch
  118. }
  119. public static Object parseFromJson(Object jsonObject,DataType type) throws Exception {
  120. if( jsonObject == null ) return null;
  121. switch (type.getName()) {
  122. case ASCII:
  123. case TEXT:
  124. case VARCHAR:
  125. case BOOLEAN:
  126. case TIME:
  127. return jsonObject;
  128. case TINYINT:
  129. return ((Number)jsonObject).byteValue();
  130. case SMALLINT:
  131. return ((Number)jsonObject).shortValue();
  132. case INT:
  133. return ((Number)jsonObject).intValue();
  134. case BIGINT:
  135. return ((Number)jsonObject).longValue();
  136. case VARINT:
  137. return new BigInteger(jsonObject.toString());
  138. case FLOAT:
  139. return ((Number)jsonObject).floatValue();
  140. case DOUBLE:
  141. return ((Number)jsonObject).doubleValue();
  142. case DECIMAL:
  143. return new BigDecimal(jsonObject.toString());
  144. case BLOB:
  145. return ByteBuffer.wrap(Base64.decodeBase64((String)jsonObject));
  146. case DATE:
  147. return LocalDate.fromMillisSinceEpoch(((Number)jsonObject).longValue());
  148. case TIMESTAMP:
  149. return new Date(((Number)jsonObject).longValue());
  150. case DURATION:
  151. return Duration.from(jsonObject.toString());
  152. case UUID:
  153. case TIMEUUID:
  154. return UUID.fromString(jsonObject.toString());
  155. case INET:
  156. return InetAddress.getByName((String)jsonObject);
  157. case LIST:
  158. List l = new ArrayList();
  159. for( Object o : (JSONArray)jsonObject ) {
  160. l.add(parseFromJson(o,type.getTypeArguments().get(0)));
  161. }
  162. return l;
  163. case MAP: {
  164. Map m = new HashMap();
  165. for (JSONObject.Entry e : ((JSONObject)jsonObject).entrySet()) {
  166. Object k = parseFromString((String) e.getKey(), type.getTypeArguments().get(0));
  167. Object v = parseFromJson(e.getValue(), type.getTypeArguments().get(1));
  168. m.put(k,v);
  169. }
  170. return m;
  171. }
  172. case SET:
  173. Set s = new HashSet();
  174. for( Object o : (JSONArray)jsonObject ) {
  175. s.add(parseFromJson(o,type.getTypeArguments().get(0)));
  176. }
  177. return s;
  178. case TUPLE: {
  179. TupleValue t = ((TupleType) type).newValue();
  180. int j = 0;
  181. for (Object e : (JSONArray)jsonObject) {
  182. DataType eleType = ((TupleType) type).getComponentTypes().get(j);
  183. t.set(j, parseFromJson(e, eleType), registry.codecFor(eleType).getJavaType());
  184. j++;
  185. }
  186. return t;
  187. }
  188. case UDT: {
  189. UDTValue t = ((UserType) type).newValue();
  190. UserType userType = t.getType();
  191. for (JSONObject.Entry e : ((JSONObject)jsonObject).entrySet()) {
  192. DataType eleType = userType.getFieldType((String)e.getKey());
  193. t.set((String)e.getKey(), parseFromJson(e.getValue(), eleType), registry.codecFor(eleType).getJavaType());
  194. }
  195. return t;
  196. }
  197. }
  198. return null;
  199. }
  200. public static void setupColumn(BoundStatement ps, int pos, DataType sqlType, Column col) throws Exception {
  201. if (col.getRawData() != null) {
  202. switch (sqlType.getName()) {
  203. case ASCII:
  204. case TEXT:
  205. case VARCHAR:
  206. ps.setString(pos, col.asString());
  207. break;
  208. case BLOB:
  209. ps.setBytes(pos, ByteBuffer.wrap(col.asBytes()));
  210. break;
  211. case BOOLEAN:
  212. ps.setBool(pos, col.asBoolean());
  213. break;
  214. case TINYINT:
  215. ps.setByte(pos, col.asLong().byteValue());
  216. break;
  217. case SMALLINT:
  218. ps.setShort(pos, col.asLong().shortValue());
  219. break;
  220. case INT:
  221. ps.setInt(pos, col.asLong().intValue());
  222. break;
  223. case BIGINT:
  224. ps.setLong(pos, col.asLong());
  225. break;
  226. case VARINT:
  227. ps.setVarint(pos, col.asBigInteger());
  228. break;
  229. case FLOAT:
  230. ps.setFloat(pos, col.asDouble().floatValue());
  231. break;
  232. case DOUBLE:
  233. ps.setDouble(pos, col.asDouble());
  234. break;
  235. case DECIMAL:
  236. ps.setDecimal(pos, col.asBigDecimal());
  237. break;
  238. case DATE:
  239. ps.setDate(pos, LocalDate.fromMillisSinceEpoch(col.asDate().getTime()));
  240. break;
  241. case TIME:
  242. ps.setTime(pos, col.asLong());
  243. break;
  244. case TIMESTAMP:
  245. ps.setTimestamp(pos, col.asDate());
  246. break;
  247. case UUID:
  248. case TIMEUUID:
  249. ps.setUUID(pos, UUID.fromString(col.asString()));
  250. break;
  251. case INET:
  252. ps.setInet(pos, InetAddress.getByName(col.asString()));
  253. break;
  254. case DURATION:
  255. ps.set(pos, Duration.from(col.asString()), Duration.class);
  256. break;
  257. case LIST:
  258. ps.setList(pos, (List<?>) parseFromString(col.asString(), sqlType));
  259. break;
  260. case MAP:
  261. ps.setMap(pos, (Map) parseFromString(col.asString(), sqlType));
  262. break;
  263. case SET:
  264. ps.setSet(pos, (Set) parseFromString(col.asString(), sqlType));
  265. break;
  266. case TUPLE:
  267. ps.setTupleValue(pos, (TupleValue) parseFromString(col.asString(), sqlType));
  268. break;
  269. case UDT:
  270. ps.setUDTValue(pos, (UDTValue) parseFromString(col.asString(), sqlType));
  271. break;
  272. default:
  273. throw DataXException.asDataXException(CassandraWriterErrorCode.CONF_ERROR,
  274. "不支持您配置的列类型:" + sqlType + ", 请检查您的配置 或者 联系 管理员.");
  275. } // end switch
  276. } else {
  277. ps.setToNull(pos);
  278. }
  279. }
  280. }