Add initial savepoint infos and partition counts (#48)

Author: tarzanek

src/main/scala/com/scylladb/migrator/Migrator.scala

@@ -3,8 +3,8 @@ package com.scylladb.migrator
 import java.nio.charset.StandardCharsets
 import java.nio.file.{ Files, Paths }
 import java.util.concurrent.{ ScheduledThreadPoolExecutor, TimeUnit }
-
 import com.amazonaws.services.dynamodbv2.streamsadapter.model.RecordAdapter
+import com.datastax.spark.connector.rdd.partitioner.{ CassandraPartition, CqlTokenRange }
 import com.datastax.spark.connector.rdd.partitioner.dht.Token
 import com.datastax.spark.connector.writer._
 import com.scylladb.migrator.config._
@@ -75,6 +75,35 @@ object Migrator {
         Some(tokenRangeAccumulator)
       }
 
+    log.info(
+      "We need to transfer: " + sourceDF.dataFrame.rdd.getNumPartitions + " partitions in total")
+
+    val partitions = sourceDF.dataFrame.rdd.partitions
+    val cassandraPartitions = partitions.map(p => { p.asInstanceOf[CassandraPartition[_, _]] })
+    var allTokenRanges = Set[(Token[_], Token[_])]()
+    cassandraPartitions.foreach(p => {
+      p.tokenRanges
+        .asInstanceOf[Vector[CqlTokenRange[_, _]]]
+        .foreach(tr => {
+          val range =
+            Set((tr.range.start.asInstanceOf[Token[_]], tr.range.end.asInstanceOf[Token[_]]))
+          allTokenRanges = allTokenRanges ++ range
+        })
+
+    })
+
+    log.info("All token ranges extracted from partitions size:" + allTokenRanges.size)
+
+    if (migratorConfig.skipTokenRanges != None) {
+      log.info(
+        "Savepoints array defined, size of the array: " + migratorConfig.skipTokenRanges.size)
+
+      val diff = allTokenRanges.diff(migratorConfig.skipTokenRanges)
+      log.info("Diff ... total diff of full ranges to savepoints is: " + diff.size)
+      log.debug("Dump of the missing tokens: ")
+      log.debug(diff)
+    }
+
     log.info("Starting write...")
 
     try {