Simplified TTL & Writetime logic - No internal commands (ttl or writetimes) on the select-query.

TTL & Writetime properties are optional

Author: pravinbhat

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

@@ -26,7 +26,6 @@ public abstract class AbstractJobSession {
     protected final RateLimiter writeLimiter;
     public Logger logger = LoggerFactory.getLogger(this.getClass().getName());
     protected PreparedStatement sourceSelectStatement;
-    protected String sourceSelectCondition;
     protected PreparedStatement astraSelectStatement;
     protected PreparedStatement astraInsertStatement;
     protected Integer maxRetries = 10;
@@ -73,25 +72,26 @@ protected AbstractJobSession(CqlSession sourceSession, CqlSession astraSession,
 
         isPreserveTTLWritetime = Boolean.parseBoolean(sparkConf.get("spark.preserveTTLWriteTime", "false"));
         if (isPreserveTTLWritetime) {
-            String ttlColsStr = sparkConf.get("spark.source.ttl.cols");
+            String ttlColsStr = sparkConf.get("spark.preserveTTLWriteTime.ttl.cols");
             if (null != ttlColsStr && ttlColsStr.trim().length() > 0) {
                 for (String ttlCol : ttlColsStr.split(",")) {
                     ttlCols.add(Integer.parseInt(ttlCol));
                 }
             }
+
+            String writeTimestampColsStr = sparkConf.get("spark.preserveTTLWriteTime.writetime.cols");
+            if (null != writeTimestampColsStr && writeTimestampColsStr.trim().length() > 0) {
+                for (String writeTimeStampCol : writeTimestampColsStr.split(",")) {
+                    writeTimeStampCols.add(Integer.parseInt(writeTimeStampCol));
+                }
+            }
         }
 
         writeTimeStampFilter = Boolean
                 .parseBoolean(sparkConf.get("spark.source.writeTimeStampFilter", "false"));
         // batchsize set to 1 if there is a writeFilter
         if (writeTimeStampFilter) {
             batchSize = 1;
-            String writeTimestampColsStr = sparkConf.get("spark.source.writeTimeStampFilter.cols");
-            if (null != writeTimestampColsStr && writeTimestampColsStr.trim().length() > 0) {
-                for (String writeTimeStampCol : writeTimestampColsStr.split(",")) {
-                    writeTimeStampCols.add(Integer.parseInt(writeTimeStampCol));
-                }
-            }
         }
 
         String minWriteTimeStampFilterStr =
@@ -118,13 +118,27 @@ protected AbstractJobSession(CqlSession sourceSession, CqlSession astraSession,
 
         String selectCols = sparkConf.get("spark.query.source");
         String partionKey = sparkConf.get("spark.query.source.partitionKey");
+        String sourceSelectCondition = sparkConf.get("spark.query.condition", "");
+
+
+        final StringBuilder selectTTLWriteTimeCols = new StringBuilder();
+        if (isPreserveTTLWritetime) {
+            String[] allCols = selectCols.split(",");
+            ttlCols.forEach(col -> {
+                selectTTLWriteTimeCols.append(",ttl(" + allCols[col] + ")");
+            });
+            writeTimeStampCols.forEach(col -> {
+                selectTTLWriteTimeCols.append(",writetime(" + allCols[col] + ")");
+            });
+        }
+        String fullSelectQuery = "select " + selectCols + selectTTLWriteTimeCols.toString() + " from " + sourceKeyspaceTable + " where token(" + partionKey.trim()
+                + ") >= ? and token(" + partionKey.trim() + ") <= ?  " + sourceSelectCondition + " ALLOW FILTERING";
+        sourceSelectStatement = sourceSession.prepare(fullSelectQuery);
+        logger.info("PARAM -- Query used: " + fullSelectQuery);
+
         selectColTypes = getTypes(sparkConf.get("spark.query.types"));
         String idCols = sparkConf.get("spark.query.destination.id", "");
         idColTypes = selectColTypes.subList(0, idCols.split(",").length);
-        sourceSelectCondition = sparkConf.get("spark.query.condition", "");
-        sourceSelectStatement = sourceSession.prepare(
-                "select " + selectCols + " from " + sourceKeyspaceTable + " where token(" + partionKey.trim()
-                        + ") >= ? and token(" + partionKey.trim() + ") <= ?  " + sourceSelectCondition + " ALLOW FILTERING");
 
         String insertCols = sparkConf.get("spark.query.destination", "");
         if (null == insertCols || insertCols.trim().isEmpty()) {
@@ -182,15 +196,15 @@ public List<MigrateDataType> getTypes(String types) {
     public int getLargestTTL(Row sourceRow) {
         int ttl = 0;
         for (Integer ttlCol : ttlCols) {
-            ttl = Math.max(ttl, sourceRow.getInt(ttlCol));
+            ttl = Math.max(ttl, sourceRow.getInt(selectColTypes.size() + ttlCol - 1));
         }
         return ttl;
     }
 
     public long getLargestWriteTimeStamp(Row sourceRow) {
         long writeTimestamp = 0;
         for (Integer writeTimeStampCol : writeTimeStampCols) {
-            writeTimestamp = Math.max(writeTimestamp, sourceRow.getLong(writeTimeStampCol));
+            writeTimestamp = Math.max(writeTimestamp, sourceRow.getLong(selectColTypes.size() + ttlCols.size() + writeTimeStampCol - 1));
         }
         return writeTimestamp;
     }

src/resources/sparkConf.properties

@@ -22,7 +22,7 @@ spark.batchSize                                         5
 spark.coveragePercent                                   100
 spark.printStatsAfter                                   100000
 
-spark.query.source                                      partition-key,clustering-key,order-date,amount,writetime(order-date),writetime(amount),ttl(order-date),ttl(amount)
+spark.query.source                                      partition-key,clustering-key,order-date,amount
 spark.query.source.partitionKey                         partition-key
 spark.query.destination                                 partition-key,clustering-key,order-date,amount
 spark.query.destination.id                              partition-key,clustering-key
@@ -33,10 +33,10 @@ spark.counterTable.cql
 spark.counterTable.cql.index                            0
 
 spark.preserveTTLWriteTime                              true
-spark.source.ttl.cols                                   6,7
+spark.preserveTTLWriteTime.ttl.cols                     2,3
+spark.preserveTTLWriteTime.writetime.cols               2,3
 
 spark.source.writeTimeStampFilter                       false
-spark.source.writeTimeStampFilter.cols                  4,5
 spark.source.minWriteTimeStampFilter                    0
 spark.source.maxWriteTimeStampFilter                    9223372036854775807