CDM-23 allowing bindInsert to return a collection

mieslep · mieslep · commit 5f8aafb9eeb0 · 2023-03-30T18:06:35.000+01:00
diff --git a/src/main/java/datastax/astra/migrate/CopyJobSession.java b/src/main/java/datastax/astra/migrate/CopyJobSession.java
@@ -9,6 +9,7 @@
 
 import java.math.BigInteger;
 import java.util.ArrayList;
+import java.util.List;
 import java.util.Collection;
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.atomic.AtomicLong;
@@ -95,15 +96,17 @@ public void getDataAndInsert(BigInteger min, BigInteger max) {
                             targetRow = targetResultSet.one();
                         }
 
-                        BoundStatement bInsert = cqlHelper.bindInsertOneRow(cqlHelper.getPreparedStatement(CqlHelper.CQL.TARGET_INSERT), originRow, targetRow);
-                        if (null == bInsert) {
+                        List<BoundStatement> bInsertList = cqlHelper.bindInsert(cqlHelper.getPreparedStatement(CqlHelper.CQL.TARGET_INSERT), originRow, targetRow);
+                        if (null == bInsertList || bInsertList.isEmpty()) {
                             skipCnt++;
                             continue;
                         }
-                        CompletionStage<AsyncResultSet> targetWriteResultSet = cqlHelper.getTargetSession().executeAsync(bInsert);
-                        writeResults.add(targetWriteResultSet);
-                        if (writeResults.size() > cqlHelper.getFetchSizeInRows()) {
-                            writeCnt += iterateAndClearWriteResults(writeResults, 1);
+                        for (BoundStatement bInsert : bInsertList) {
+                            CompletionStage<AsyncResultSet> targetWriteResultSet = cqlHelper.getTargetSession().executeAsync(bInsert);
+                            writeResults.add(targetWriteResultSet);
+                            if (writeResults.size() > cqlHelper.getFetchSizeInRows()) {
+                                writeCnt += iterateAndClearWriteResults(writeResults, 1);
+                            }
                         }
                     }
 
@@ -128,22 +131,24 @@ public void getDataAndInsert(BigInteger min, BigInteger max) {
                         }
 
                         writeLimiter.acquire(1);
-                        BoundStatement bInsert = cqlHelper.bindInsertOneRow(cqlHelper.getPreparedStatement(CqlHelper.CQL.TARGET_INSERT), originRow, null);
-                        if (null == bInsert) {
+                        List<BoundStatement> bInsertList = cqlHelper.bindInsert(cqlHelper.getPreparedStatement(CqlHelper.CQL.TARGET_INSERT), originRow, null);
+                        if (null == bInsertList || bInsertList.isEmpty()) {
                             skipCnt++;
                             continue;
                         }
-                        batchStatement = batchStatement.add(bInsert);
-
-                        // if batch threshold is met, send the writes and clear the batch
-                        if (batchStatement.size() >= cqlHelper.getBatchSize()) {
-                            CompletionStage<AsyncResultSet> writeResultSet = cqlHelper.getTargetSession().executeAsync(batchStatement);
-                            writeResults.add(writeResultSet);
-                            batchStatement = BatchStatement.newInstance(BatchType.UNLOGGED);
-                        }
+                        for (BoundStatement bInsert : bInsertList) {
+                            batchStatement = batchStatement.add(bInsert);
+
+                            // if batch threshold is met, send the writes and clear the batch
+                            if (batchStatement.size() >= cqlHelper.getBatchSize()) {
+                                CompletionStage<AsyncResultSet> writeResultSet = cqlHelper.getTargetSession().executeAsync(batchStatement);
+                                writeResults.add(writeResultSet);
+                                batchStatement = BatchStatement.newInstance(BatchType.UNLOGGED);
+                            }
 
-                        if (writeResults.size() * cqlHelper.getBatchSize() > cqlHelper.getFetchSizeInRows()) {
-                            writeCnt += iterateAndClearWriteResults(writeResults, cqlHelper.getBatchSize());
+                            if (writeResults.size() * cqlHelper.getBatchSize() > cqlHelper.getFetchSizeInRows()) {
+                                writeCnt += iterateAndClearWriteResults(writeResults, cqlHelper.getBatchSize());
+                            }
                         }
                     }
 
diff --git a/src/main/java/datastax/astra/migrate/CopyPKJobSession.java b/src/main/java/datastax/astra/migrate/CopyPKJobSession.java
@@ -55,11 +55,15 @@ public void getRowAndInsert(List<SplitPartitions.PKRows> rowsList) {
                     logger.error("Could not find row with primary-key: {}", row);
                     return;
                 }
-                ResultSet targetWriteResultSet = cqlHelper.getTargetSession()
-                        .execute(cqlHelper.bindInsertOneRow(cqlHelper.getPreparedStatement(CqlHelper.CQL.TARGET_INSERT), pkRow, null));
-                writeCounter.incrementAndGet();
-                if (readCounter.get() % printStatsAfter == 0) {
-                    printCounts(false);
+                List<BoundStatement> boundInserts = cqlHelper.bindInsert(cqlHelper.getPreparedStatement(CqlHelper.CQL.TARGET_INSERT), pkRow, null);
+                if (null != boundInserts) {
+                    for (BoundStatement bs : boundInserts) {
+                        ResultSet targetWriteResultSet = cqlHelper.getTargetSession().execute(bs);
+                        writeCounter.incrementAndGet();
+                        if (readCounter.get() % printStatsAfter == 0) {
+                            printCounts(false);
+                        }
+                    }
                 }
             });
         }
diff --git a/src/main/java/datastax/astra/migrate/cql/CqlHelper.java b/src/main/java/datastax/astra/migrate/cql/CqlHelper.java
@@ -96,6 +96,22 @@ public boolean initialize() {
         return validInit;
     }
 
+    public List<BoundStatement> bindInsert(PreparedStatement insertStatement, Row originRow, Row targetRow) {
+        List<BoundStatement> rtnList = new ArrayList<>();
+        if (isFeatureEnabled(Featureset.EXPLODE_MAP)) {
+            Feature explodeMapFeature = getFeature(Featureset.EXPLODE_MAP);
+            Map maptoExplode = (Map) getData(explodeMapFeature.getMigrateDataType(ExplodeMap.Property.MAP_COLUMN_TYPE),
+                    explodeMapFeature.getInteger(ExplodeMap.Property.MAP_COLUMN_INDEX),
+                    originRow);
+            for (Object key : maptoExplode.keySet()) {
+                rtnList.add(bindInsertOneRow(insertStatement, originRow, targetRow, key, maptoExplode.get(key)));
+            }
+        } else {
+            rtnList.add(bindInsertOneRow(insertStatement, originRow, targetRow, null, null));
+        }
+        return rtnList;
+    }
+
     public BoundStatement bindInsertOneRow(PreparedStatement insertStatement, Row originRow, Row targetRow) {
         return bindInsertOneRow(insertStatement, originRow, targetRow, null, null);
     }
@@ -119,7 +135,9 @@ public BoundStatement bindInsertOneRow(PreparedStatement insertStatement, Row or
             // This loops over the selected columns and binds each type to the boundInsertStatement
             Feature explodeMapFeature = getFeature(Featureset.EXPLODE_MAP);
             for (index = 0; index < originColTypesSize; index++) {
-                if (mapKey != null && explodeMapFeature.isEnabled() && index == explodeMapFeature.getInteger(ExplodeMap.Property.MAP_COLUMN_INDEX)) {
+                if (mapKey != null &&
+                        isFeatureEnabled(Featureset.EXPLODE_MAP) &&
+                        index == explodeMapFeature.getInteger(ExplodeMap.Property.MAP_COLUMN_INDEX)) {
                     // This substitutes the map column with the key and value types of the map
                     boundInsertStatement = boundInsertStatement.set(index, mapKey, explodeMapFeature.getMigrateDataType(ExplodeMap.Property.KEY_COLUMN_TYPE).getType());
                     // Add an 'extra' column to the statement, which will also increase the loop limit
diff --git a/src/main/java/datastax/astra/migrate/cql/features/ExplodeMap.java b/src/main/java/datastax/astra/migrate/cql/features/ExplodeMap.java
@@ -16,6 +16,7 @@ public class ExplodeMap extends AbstractFeature {
     public enum Property {
         MAP_COLUMN_NAME,
         MAP_COLUMN_INDEX,
+        MAP_COLUMN_TYPE,
         KEY_COLUMN_NAME,
         KEY_COLUMN_TYPE,
         VALUE_COLUMN_NAME,
@@ -39,6 +40,7 @@ public boolean initialize(PropertyHelper helper) {
             putNumber(Property.MAP_COLUMN_INDEX, helper.getStringList(KnownProperties.ORIGIN_COLUMN_NAMES).indexOf(mapColumnName));
 
             MigrateDataType columnMapDataType = getColumnMapDataType(helper);
+            putMigrateDataType(Property.MAP_COLUMN_TYPE, columnMapDataType);
             putMigrateDataType(Property.KEY_COLUMN_TYPE, columnMapDataType.getSubTypeTypes().get(0));
             putMigrateDataType(Property.VALUE_COLUMN_TYPE, columnMapDataType.getSubTypeTypes().get(1));
         }
diff --git a/src/test/java/datastax/astra/migrate/cql/features/ExplodeMapTest.java b/src/test/java/datastax/astra/migrate/cql/features/ExplodeMapTest.java
@@ -30,6 +30,7 @@ public void setup() {
     public void tearDown() {
         PropertyHelper.destroyInstance();
         validSparkConf = null;
+        feature = null;
     }
 
     private void setValidSparkConf() {
@@ -54,6 +55,7 @@ public void smokeTest_initialize() {
                 () -> assertEquals("map_key", feature.getAsString(ExplodeMap.Property.KEY_COLUMN_NAME), "KEY_COLUMN_NAME"),
                 () -> assertEquals("map_val", feature.getAsString(ExplodeMap.Property.VALUE_COLUMN_NAME), "VALUE_COLUMN_NAME"),
                 () -> assertEquals(2, feature.getNumber(ExplodeMap.Property.MAP_COLUMN_INDEX), "MAP_COLUMN_INDEX"),
+                () -> assertEquals(new MigrateDataType("5%0%3"), feature.getMigrateDataType(ExplodeMap.Property.MAP_COLUMN_TYPE), "MAP_COLUMN_TYPE"),
                 () -> assertEquals(new MigrateDataType("0"), feature.getMigrateDataType(ExplodeMap.Property.KEY_COLUMN_TYPE), "KEY_COLUMN_TYPE"),
                 () -> assertEquals(new MigrateDataType("3"), feature.getMigrateDataType(ExplodeMap.Property.VALUE_COLUMN_TYPE), "VALUE_COLUMN_TYPE")
         );
@@ -65,8 +67,9 @@ public void smokeTest_alterProperties() {
         helper.initializeSparkConf(validSparkConf);
         feature.initialize(helper);
         feature.alterProperties(helper);
+
         assertAll(
-                () -> assertTrue(feature.isEnabled()),
+                () -> assertTrue(feature.isEnabled(), "isEnabled"),
                 () -> assertEquals(Arrays.asList("key","map_key"), helper.getStringList(KnownProperties.TARGET_PRIMARY_KEY), "TARGET_PRIMARY_KEY"),
                 () -> assertEquals(Arrays.asList(new MigrateDataType("4"),new MigrateDataType("0")), helper.getMigrationTypeList(KnownProperties.TARGET_PRIMARY_KEY_TYPES), "TARGET_PRIMARY_KEY_TYPES"),
                 () -> assertEquals(helper.getStringList(KnownProperties.TARGET_PRIMARY_KEY).size(), helper.getMigrationTypeList(KnownProperties.TARGET_PRIMARY_KEY_TYPES).size(), "sizes match"),
