1/ Improved readability added a timeout parameter

2/ Added a comment for min pages per second (11 per second)
3/ processedPage++ moved to whenCompeted
4/ Added JSON ser for persisting a state of the exporting process

Author: Nikolai Kolesnikov

.gitignore

@@ -2,6 +2,7 @@
 .idea/
 *.iml
 *.iws
+*.conf
 
 # Maven
 target/

src/main/java/com/amazon/aws/keyspaces/PersistCassandraRowToParquet.java

@@ -173,8 +173,12 @@ void processRowsAsync(AsyncResultSet rs, Throwable error) {
 
             if (rs.hasMorePages()) {
                 pageLimiter.acquire();
-                rs.fetchNextPage().whenComplete(this::processRowsAsync);
-                processedPage++;
+                //rs.fetchNextPage().whenComplete(this::processRowsAsync);
+                //processedPage++;
+                rs.fetchNextPage().whenComplete((thisRs, thisError) ->{
+                    processRowsAsync(thisRs, thisError);
+                    processedPage++;
+                });
             } else {
                 LOG.info("Completed to read pages");
                 LOG.info("Starting page was " + startingPage);

src/main/java/com/amazon/aws/keyspaces/Runner.java

@@ -21,16 +21,22 @@
 public class Runner implements Callable<Integer> {
 
     private static final Logger LOG = Logger.getLogger(Runner.class.getName());
+    // Default timeout is 86,400 seconds
+    private static final long ONE_DAY = 1*24*60*60;
+
+
     @CommandLine.Option(names = "--recover", description = "Recovering mode")
     boolean recoveringMode = false;
     @CommandLine.Parameters(index = "0", description = "Destination of parquet files")
-    private final String dataLocation = System.getProperty("user.dir") + "/" + "output";
+    private String dataLocation = System.getProperty("user.dir") + "/" + "output";
     @CommandLine.Parameters(index = "1", description = "Query")
     private String query;
     @CommandLine.Option(names = {"-p", "--page"}, description = "Starting page")
-    private final int startingPage = 0;
+    private int startingPage = 0;
     @CommandLine.Option(names = {"-r", "--rateLimiter"}, description = "A rate limiter")
-    private final int rateLimiter = 3000;
+    private int rateLimiter = 3000;
+    @CommandLine.Option(names = {"-t", "--timeout"}, description = "Override default timeout 86400 seconds")
+    private long timeout = ONE_DAY;
     private final File configFile = new File(System.getProperty("user.dir") + "/application.conf");
 
     public static void main(String[] args) throws IOException, NoSuchAlgorithmException, InterruptedException {
@@ -53,7 +59,7 @@ public Integer call() throws Exception {
             LOG.info("Export in progress...");
             LOG.info("Dest:" + dataLocation);
             persistCassandraRowToParquet.start();
-            countDownLatch.await(604800, TimeUnit.SECONDS);
+            countDownLatch.await(timeout, TimeUnit.SECONDS);
             persistCassandraRowToParquet.close();
             LOG.info("Amazon Keyspaces Session is closed");
             long elapsedTime = System.nanoTime() - startTime;
@@ -76,13 +82,14 @@ public Integer call() throws Exception {
                     countDownLatch,
                     configFile,
                     state.getProcessedPages(),
+                    // 11 pages * 3500 rows per page * 100 bytes per row = 3,759 kilobytes per partition
                     Math.min(11, (int) Math.round(state.getPageRate()))
             );
 
             LOG.info("Export in progress...");
             LOG.info("Dest:" + dataLocation);
             persistCassandraRowToParquet.start();
-            countDownLatch.await(604800, TimeUnit.SECONDS);
+            countDownLatch.await(timeout, TimeUnit.SECONDS);
             persistCassandraRowToParquet.close();
             LOG.info("Amazon Keyspaces Session is closed");
             long elapsedTime = System.nanoTime() - startTime;

src/main/java/com/amazon/aws/keyspaces/Utils.java

@@ -3,21 +3,20 @@
 
 package com.amazon.aws.keyspaces;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
+
 import java.io.FileInputStream;
 import java.io.FileOutputStream;
-import java.io.ObjectInputStream;
-import java.io.ObjectOutputStream;
 import java.util.logging.Logger;
 
 public class Utils {
     private static final Logger LOG = Logger.getLogger(Utils.class.getName());
 
     public static void writeState(State state) {
         try {
+            ObjectMapper mapper = new ObjectMapper();
             FileOutputStream myFileOutputStream = new FileOutputStream(System.getProperty("user.dir") + "/state.ser");
-            ObjectOutputStream myObjectOutputStream = new ObjectOutputStream(myFileOutputStream);
-            myObjectOutputStream.writeObject(state);
-            myObjectOutputStream.close();
+            mapper.writeValue(myFileOutputStream, state);
         } catch (Exception e) {
             LOG.severe("Error when saving to file." + e.getMessage());
         }
@@ -26,13 +25,13 @@ public static void writeState(State state) {
     public static State readState() {
         State state = new State();
         try {
+            ObjectMapper mapper = new ObjectMapper();
             FileInputStream myFileInputStream = new FileInputStream(System.getProperty("user.dir") + "/state.ser");
-            ObjectInputStream myObjectInputStream = new ObjectInputStream(myFileInputStream);
-            state = (State) myObjectInputStream.readObject();
-            myObjectInputStream.close();
+            state = mapper.readValue(myFileInputStream, State.class);
         } catch (Exception e) {
             LOG.severe("Error when loading from file." + e.getMessage());
         }
         return state;
     }
+
 }