Processes partitions from a partition file.

One main use of this feature would be to process partitions from a previous run that had errors

Author: pravinbhat

pom.xml

@@ -3,7 +3,7 @@
 
   <groupId>datastax.astra.migrate</groupId>
   <artifactId>cassandra-data-migrator</artifactId>
-  <version>1.7</version>
+  <version>1.8</version>
   <packaging>jar</packaging>
 
   <properties>

src/main/java/datastax/astra/migrate/AbstractJobSession.java

@@ -35,14 +35,14 @@ protected AbstractJobSession(CqlSession sourceSession, CqlSession astraSession,
         sourceKeyspaceTable = sparkConf.get("spark.source.keyspaceTable");
         astraKeyspaceTable = sparkConf.get("spark.destination.keyspaceTable");
 
-        String ttlColsStr = sparkConf.get("spark.query.ttl.cols");
+        String ttlColsStr = sparkConf.get("spark.query.ttl.cols", "");
         if (null != ttlColsStr && ttlColsStr.trim().length() > 0) {
             for (String ttlCol : ttlColsStr.split(",")) {
                 ttlCols.add(Integer.parseInt(ttlCol));
             }
         }
 
-        String writeTimestampColsStr = sparkConf.get("spark.query.writetime.cols");
+        String writeTimestampColsStr = sparkConf.get("spark.query.writetime.cols", "");
         if (null != writeTimestampColsStr && writeTimestampColsStr.trim().length() > 0) {
             for (String writeTimeStampCol : writeTimestampColsStr.split(",")) {
                 writeTimeStampCols.add(Integer.parseInt(writeTimeStampCol));

src/main/java/datastax/astra/migrate/SplitPartitions.java

@@ -3,6 +3,9 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.io.BufferedReader;
+import java.io.FileReader;
+import java.io.IOException;
 import java.io.Serializable;
 import java.math.BigInteger;
 import java.util.ArrayList;
@@ -16,15 +19,16 @@ public class SplitPartitions {
     public final static Long MAX_PARTITION = Long.MAX_VALUE;
     public static Logger logger = LoggerFactory.getLogger(SplitPartitions.class.getName());
 
-    public static void main(String[] args) {
-        Collection<Partition> partitions = getSubPartitions(new BigInteger("20"), BigInteger.valueOf(MIN_PARTITION),
-                BigInteger.valueOf(MAX_PARTITION), 20);
+    public static void main(String[] args) throws IOException {
+        Collection<Partition> partitions = getSubPartitions(2, BigInteger.valueOf(1),
+                BigInteger.valueOf(1000), 100);
+//        Collection<Partition> partitions = getSubPartitionsFromFile(3);
         for (Partition partition : partitions) {
             System.out.println(partition);
         }
     }
 
-    public static Collection<Partition> getRandomSubPartitions(BigInteger splitSize, BigInteger min, BigInteger max, int coveragePercent) {
+    public static Collection<Partition> getRandomSubPartitions(int splitSize, BigInteger min, BigInteger max, int coveragePercent) {
         logger.info("TreadID: " + Thread.currentThread().getId() + " Splitting min: " + min + " max:" + max);
         List<Partition> partitions = getSubPartitions(splitSize, min, max, coveragePercent);
         Collections.shuffle(partitions);
@@ -34,12 +38,30 @@ public static Collection<Partition> getRandomSubPartitions(BigInteger splitSize,
         return partitions;
     }
 
-    private static List<Partition> getSubPartitions(BigInteger splitSize, BigInteger min, BigInteger max, int coveragePercent) {
+    public static List<Partition> getSubPartitionsFromFile(int splitSize) throws IOException {
+        logger.info("TreadID: " + Thread.currentThread().getId() +
+                " Splitting partitions in file: ./partitions.csv using a split-size of " + splitSize);
+        List<Partition> partitions = new ArrayList<Partition>();
+        BufferedReader reader = new BufferedReader(new FileReader("./partitions.csv"));
+        String line = null;
+        while ((line = reader.readLine()) != null) {
+            String[] minMax = line.split(",");
+            try {
+                partitions.addAll(getSubPartitions(splitSize, new BigInteger(minMax[0]), new BigInteger(minMax[1]), 100));
+            } catch (Exception e) {
+                logger.error("Skipping partition: " + line, e);
+            }
+        }
+
+        return partitions;
+    }
+
+    private static List<Partition> getSubPartitions(int splitSize, BigInteger min, BigInteger max, int coveragePercent) {
         if (coveragePercent < 1 || coveragePercent > 100) {
             coveragePercent = 100;
         }
         BigInteger curMax = new BigInteger(min.toString());
-        BigInteger partitionSize = max.subtract(min).divide(splitSize);
+        BigInteger partitionSize = max.subtract(min).divide(BigInteger.valueOf(splitSize));
         List<Partition> partitions = new ArrayList<Partition>();
         if (partitionSize.compareTo(new BigInteger("0")) == 0) {
             partitionSize = new BigInteger("100000");
@@ -64,6 +86,7 @@ private static List<Partition> getSubPartitions(BigInteger splitSize, BigInteger
             if (exausted) {
                 break;
             }
+            curMax = curMax.add(BigInteger.ONE);
         }
 
         return partitions;

src/main/scala/datastax/astra/migrate/AbstractJob.scala

@@ -6,9 +6,8 @@ class AbstractJob extends BaseJob {
 
   abstractLogger.info("PARAM -- Min Partition: " + minPartition)
   abstractLogger.info("PARAM -- Max Partition: " + maxPartition)
-  abstractLogger.info("PARAM -- Split Size: " + coveragePercent)
+  abstractLogger.info("PARAM -- Split Size: " + splitSize)
   abstractLogger.info("PARAM -- Coverage Percent: " + coveragePercent)
-  abstractLogger.info("PARAM Calculated -- Total Partitions: " + partitions.size())
 
   var sourceConnection = getConnection(true, sourceIsAstra, sourceScbPath, sourceHost, sourceUsername, sourcePassword, sourceReadConsistencyLevel,
     sourceTrustStorePath, sourceTrustStorePassword, sourceTrustStoreType, sourceKeyStorePath, sourceKeyStorePassword, sourceEnabledAlgorithms);

src/main/scala/datastax/astra/migrate/BaseJob.scala

@@ -10,8 +10,11 @@ class BaseJob extends App {
 
   val abstractLogger = LoggerFactory.getLogger(this.getClass.getName)
   val spark = SparkSession.builder
-    .appName("Datastax Data Validation")
+    .appName("Cassandra Data Migrator")
     .getOrCreate()
+  abstractLogger.info("################################################################################################")
+  abstractLogger.info("############################## Cassandra Data Migrator - Starting ##############################")
+  abstractLogger.info("################################################################################################")
 
   val sc = spark.sparkContext
 
@@ -44,11 +47,13 @@ class BaseJob extends App {
   val minPartition = new BigInteger(sc.getConf.get("spark.source.minPartition", "-9223372036854775808"))
   val maxPartition = new BigInteger(sc.getConf.get("spark.source.maxPartition", "9223372036854775807"))
   val coveragePercent = sc.getConf.get("spark.coveragePercent", "100")
-  val splitSize = sc.getConf.get("spark.splitSize", "10000")
-  val partitions = SplitPartitions.getRandomSubPartitions(BigInteger.valueOf(Long.parseLong(splitSize)), minPartition, maxPartition, Integer.parseInt(coveragePercent))
+  val splitSize = Integer.parseInt(sc.getConf.get("spark.splitSize", "10000"))
 
   protected def exitSpark() = {
     spark.stop()
+    abstractLogger.info("################################################################################################")
+    abstractLogger.info("############################## Cassandra Data Migrator - Stopped ###############################")
+    abstractLogger.info("################################################################################################")
     sys.exit(0)
   }
 

src/main/scala/datastax/astra/migrate/DiffData.scala

@@ -15,6 +15,8 @@ object DiffData extends AbstractJob {
   exitSpark
 
   private def diffTable(sourceConnection: CassandraConnector, destinationConnection: CassandraConnector) = {
+    val partitions = SplitPartitions.getRandomSubPartitions(splitSize, minPartition, maxPartition, Integer.parseInt(coveragePercent))
+    logger.info("PARAM Calculated -- Total Partitions: " + partitions.size())
     val parts = sc.parallelize(partitions.toSeq, partitions.size);
     logger.info("Spark parallelize created : " + parts.count() + " parts!");
 

src/main/scala/datastax/astra/migrate/Migrate.scala

@@ -17,6 +17,8 @@ object Migrate extends AbstractJob {
   exitSpark
 
   private def migrateTable(sourceConnection: CassandraConnector, destinationConnection: CassandraConnector) = {
+    val partitions = SplitPartitions.getRandomSubPartitions(splitSize, minPartition, maxPartition, Integer.parseInt(coveragePercent))
+    logger.info("PARAM Calculated -- Total Partitions: " + partitions.size())
     val parts = sc.parallelize(partitions.toSeq, partitions.size);
     logger.info("Spark parallelize created : " + parts.count() + " parts!");
 

src/main/scala/datastax/astra/migrate/MigratePartitionsFromFile.scala

@@ -0,0 +1,37 @@
+package datastax.astra.migrate
+
+import com.datastax.spark.connector.cql.CassandraConnector
+import org.slf4j.LoggerFactory
+
+import java.math.BigInteger
+import scala.collection.JavaConversions._
+import java.lang.Long
+
+object MigratePartitionsFromFile extends AbstractJob {
+
+  val logger = LoggerFactory.getLogger(this.getClass.getName)
+  logger.info("Started MigratePartitionsFromFile App")
+
+  migrateTable(sourceConnection, destinationConnection)
+
+  exitSpark
+
+  private def migrateTable(sourceConnection: CassandraConnector, destinationConnection: CassandraConnector) = {
+    val partitions = SplitPartitions.getSubPartitionsFromFile(splitSize)
+    logger.info("PARAM Calculated -- Total Partitions: " + partitions.size())
+    val parts = sc.parallelize(partitions.toSeq, partitions.size);
+    logger.info("Spark parallelize created : " + parts.count() + " parts!");
+
+    parts.foreach(part => {
+      sourceConnection.withSessionDo(sourceSession =>
+        destinationConnection.withSessionDo(destinationSession =>
+          CopyJobSession.getInstance(sourceSession, destinationSession, sc.getConf)
+            .getDataAndInsert(part.getMin, part.getMax)))
+    })
+
+  }
+
+}
+
+
+