CDM-37 consolidating registry to cqlHelper

Author: mieslep

src/main/java/datastax/cdm/cql/CqlHelper.java

@@ -117,8 +117,7 @@ public Boolean isFeatureEnabled(Featureset featureEnum) {
     private void registerTargetCodecs() {
         List<String> codecList = propertyHelper.getStringList(KnownProperties.TRANSFORM_CODECS);
         if (null!=codecList && !codecList.isEmpty()) {
-            MutableCodecRegistry registry =
-                    (MutableCodecRegistry) targetSession.getContext().getCodecRegistry();
+            MutableCodecRegistry registry = getCodecRegistry();
 
             StringBuilder sb = new StringBuilder("PARAM -- Codecs Enabled: ");
             for (String codecString : codecList) {
@@ -136,6 +135,10 @@ public boolean isCodecRegistered(Codecset codecEnum) {
         return codecMap.containsKey(codecEnum);
     }
 
+    public MutableCodecRegistry getCodecRegistry() {
+        return (MutableCodecRegistry) targetSession.getContext().getCodecRegistry();
+    }
+
     // --------------- Session and Performance -------------------------
     public void setOriginSession(CqlSession originSession) {
         this.originSession = originSession;

src/main/java/datastax/cdm/job/DiffJobSession.java

@@ -51,8 +51,7 @@ public class DiffJobSession extends CopyJobSession {
     private final int explodeMapKeyIndex;
     private final int explodeMapValueIndex;
 
-    private final CodecRegistry originCodecRegistry;
-    private final CodecRegistry targetCodecRegistry;
+    private final CodecRegistry codecRegistry;
 
     private DiffJobSession(CqlSession originSession, CqlSession targetSession, SparkConf sc) {
         super(originSession, targetSession, sc);
@@ -81,8 +80,7 @@ private DiffJobSession(CqlSession originSession, CqlSession targetSession, Spark
             this.explodeMapValueIndex = -1;
         }
 
-        this.originCodecRegistry = originSession.getContext().getCodecRegistry();
-        this.targetCodecRegistry = targetSession.getContext().getCodecRegistry();
+        this.codecRegistry = cqlHelper.getCodecRegistry();
     }
 
     public static DiffJobSession getInstance(CqlSession originSession, CqlSession targetSession, SparkConf sparkConf) {
@@ -232,7 +230,7 @@ private String isDifferent(EnhancedPK pk, Row originRow, Row targetRow) {
                     MigrateDataType originDataTypeObj = originColumnTypes.get(originIndex);
                     origin = cqlHelper.getData(originDataTypeObj, originIndex, originRow);
                     if (!originDataTypeObj.equals(targetDataTypeObj)) {
-                        origin = MigrateDataType.convert(origin, originDataTypeObj, originCodecRegistry, targetDataTypeObj, targetCodecRegistry);
+                        origin = MigrateDataType.convert(origin, originDataTypeObj, targetDataTypeObj, codecRegistry);
                     }
                 }
             }

src/main/java/datastax/cdm/job/MigrateDataType.java

@@ -89,18 +89,18 @@ public boolean diff(Object obj1, Object obj2) {
     }
 
     @SuppressWarnings("unchecked")
-    public static Object convert(Object value, MigrateDataType fromDataType, CodecRegistry fromRegistry, MigrateDataType toDataType, CodecRegistry toRegistry) {
+    public static Object convert(Object value, MigrateDataType fromDataType, MigrateDataType toDataType, CodecRegistry codecRegistry) {
         Class<?> fromClass = fromDataType.getTypeClass();
         Class<?> toClass = toDataType.getTypeClass();
         DataType cqlType = toDataType.getCqlDataType();
 
-        TypeCodec<Object> fromCodec = (TypeCodec<Object>) fromRegistry.codecFor(cqlType, fromClass);
+        TypeCodec<Object> fromCodec = (TypeCodec<Object>) codecRegistry.codecFor(cqlType, fromClass);
         if (fromCodec == null) {
-            throw new IllegalArgumentException("No codec found in fromRegistry for Java type " + fromClass.getName() + " to CQL type " + toDataType);
+            throw new IllegalArgumentException("No codec found in codecRegistry for Java type " + fromClass.getName() + " to CQL type " + toDataType);
         }
-        TypeCodec<Object> toCodec = (TypeCodec<Object>) toRegistry.codecFor(cqlType, toClass);
+        TypeCodec<Object> toCodec = (TypeCodec<Object>) codecRegistry.codecFor(cqlType, toClass);
         if (toCodec == null) {
-            throw new IllegalArgumentException("No codec found in toRegistry for Java type " + toClass.getName() + " to CQL type " + toDataType);
+            throw new IllegalArgumentException("No codec found in codecRegistry for Java type " + toClass.getName() + " to CQL type " + toDataType);
         }
         ByteBuffer encoded = fromCodec.encode(value, ProtocolVersion.DEFAULT);
         return toCodec.decode(encoded, ProtocolVersion.DEFAULT);