Enable row count validation for SqlInserts and YCQL workloads (#94)

* Perform row count validation for SqlInserts and YCQL workloads extending from CassandraKeyValue.

Author: ashetkar

src/main/java/com/yugabyte/sample/Main.java

@@ -186,6 +186,7 @@ public void run() throws Exception {
                                        IOType.Read, app.appConfig.printAllExceptions));
       }
 
+      app.recordExistingRowCount();
       // Start the reader and writer threads.
       for (IOPSThread iopsThread : iopsThreads) {
         iopsThread.start();
@@ -201,6 +202,7 @@ public void run() throws Exception {
           LOG.error("Error waiting for thread join()", e);
         }
       }
+      app.verifyTotalRowsWritten();
     } finally {
       terminate();
     }

src/main/java/com/yugabyte/sample/apps/AppBase.java

@@ -20,9 +20,7 @@
 import java.security.KeyStore;
 import java.security.cert.CertificateFactory;
 import java.security.cert.X509Certificate;
-import java.sql.Connection;
-import java.sql.DriverManager;
-import java.sql.PreparedStatement;
+import java.sql.*;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -812,6 +810,18 @@ public void performRead() {
     }
   }
 
+  public void verifyTotalRowsWritten() throws Exception {
+    throw new UnsupportedOperationException("Row count check is not supported for this workload");
+  }
+
+  public void recordExistingRowCount() throws Exception {
+    throw new UnsupportedOperationException("Row count check is not supported for this workload");
+  }
+
+  public String getTableName() {
+    return "AppBaseTable";
+  }
+
   @Override
   public String appenderName() {
     return this.getClass().getSimpleName();

src/main/java/com/yugabyte/sample/apps/CassandraKeyValue.java

@@ -17,13 +17,12 @@
 import java.util.Arrays;
 import java.util.List;
 
-import org.apache.log4j.Logger;
-
 import com.datastax.driver.core.BoundStatement;
-import com.datastax.driver.core.ConsistencyLevel;
 import com.datastax.driver.core.PreparedStatement;
 import com.datastax.driver.core.ResultSet;
 import com.datastax.driver.core.Row;
+import org.apache.log4j.Logger;
+
 import com.yugabyte.sample.common.SimpleLoadGenerator.Key;
 
 /**
@@ -32,6 +31,7 @@
  */
 public class CassandraKeyValue extends CassandraKeyValueBase {
 
+  private static final Logger LOG = Logger.getLogger(CassandraKeyValueBase.class);
   // The default table name to create and use for CRUD ops.
   private static final String DEFAULT_TABLE_NAME = CassandraKeyValue.class.getSimpleName();
   static {
@@ -43,6 +43,8 @@ public class CassandraKeyValue extends CassandraKeyValueBase {
     // updates).
     appConfig.numUniqueKeysToWrite = NUM_UNIQUE_KEYS_FOR_YSQL_AND_YCQL;
   }
+  private long initialRowCount = 0;
+
   @Override
   public List<String> getCreateTableStatements() {
     String create_stmt = String.format(
@@ -74,7 +76,26 @@ protected BoundStatement bindSelect(String key)  {
   @Override
   protected BoundStatement bindInsert(String key, ByteBuffer value)  {
     return getPreparedInsert().bind(key, value);
-   }
+  }
+
+  @Override
+  public void verifyTotalRowsWritten() throws Exception {
+    if (appConfig.numKeysToWrite >= 0) {
+      long actual = getRowCount();
+      LOG.info("Total rows count in table " + getTableName() + ": " + actual);
+      long expected = numKeysWritten.get();
+      if (actual != (expected + initialRowCount)) {
+        LOG.fatal("New rows count does not match! Expected: " + (expected + initialRowCount) + ", actual: " + actual);
+      } else {
+        LOG.info("Table row count verified successfully");
+      }
+    }
+  }
+
+  @Override
+  public void recordExistingRowCount() throws Exception {
+    initialRowCount = getRowCount();
+  }
 
   @Override
   public List<String> getWorkloadDescription() {

src/main/java/com/yugabyte/sample/apps/CassandraKeyValueBase.java

@@ -110,7 +110,7 @@ public long doRead() {
     ResultSet rs = getCassandraClient().execute(select);
     List<Row> rows = rs.all();
     if (rows.size() != 1) {
-      // If TTL is enabled, turn off correctness validation.
+      // If TTL is disabled, turn on correctness validation.
       if (appConfig.tableTTLSeconds <= 0) {
         LOG.fatal("Read key: " + key.asString() + " expected 1 row in result, got " + rows.size());
       }
@@ -173,6 +173,14 @@ public long doWrite(int threadIdx) {
     }
   }
 
+  protected long getRowCount() {
+    ResultSet rs = getCassandraClient().execute("SELECT COUNT(*) FROM " + getTableName());
+    List<Row> rows = rs.all();
+    long actual = rows.get(0).getLong(0);
+    LOG.info("Found " + actual + " rows in " + getTableName());
+    return actual;
+  }
+
   @Override
   public void appendMessage(StringBuilder sb) {
     super.appendMessage(sb);

src/main/java/com/yugabyte/sample/apps/SQLAppBase.java

@@ -0,0 +1,59 @@
+package com.yugabyte.sample.apps;
+
+import org.apache.log4j.Logger;
+
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.Statement;
+
+public class SQLAppBase extends AppBase {
+
+    private Logger LOG = Logger.getLogger(SQLAppBase.class);
+
+    private long initialRowCount = 0;
+
+    protected long getRowCount() throws Exception {
+        long count = 0;
+        String table = getTableName();
+        Connection connection = getPostgresConnection();
+        try {
+            Statement statement = connection.createStatement();
+            String query = "SELECT COUNT(*) FROM " + table;
+            ResultSet rs = statement.executeQuery(query);
+            if (rs.next()) {
+                count = rs.getLong(1);
+                LOG.info("Row count in table " + table + ": " + count);
+            } else {
+                LOG.error("No result received!");
+            }
+        } finally {
+            if (connection != null) {
+                connection.close();
+            }
+        }
+        return count;
+    }
+
+    @Override
+    public void verifyTotalRowsWritten() throws Exception {
+        if (appConfig.numKeysToWrite >= 0) {
+            String table = getTableName();
+            LOG.info("Verifying the inserts on table " + table + " (" + initialRowCount + " rows initially) ...");
+            LOG.info("appConfig.numKeysToWrite: " + appConfig.numKeysToWrite + ", numKeysWritten: " + numKeysWritten);
+            long actual = getRowCount();
+            long expected = numKeysWritten.get();
+            if (actual != (expected + initialRowCount)) {
+                LOG.fatal("New rows count does not match! Expected: " + (expected + initialRowCount) + ", actual: " + actual);
+            } else {
+                LOG.info("Table row count verified successfully");
+            }
+        }
+    }
+
+    @Override
+    public void recordExistingRowCount() throws Exception {
+        LOG.info("Recording the row count in table " + getTableName() + " if it exists ...");
+        initialRowCount = getRowCount();
+    }
+
+}

src/main/java/com/yugabyte/sample/apps/SqlInserts.java

@@ -12,21 +12,18 @@
 //
 package com.yugabyte.sample.apps;
 
-import java.sql.Connection;
-import java.sql.PreparedStatement;
-import java.sql.ResultSet;
+import java.sql.*;
 import java.util.Arrays;
 import java.util.List;
 
 import org.apache.log4j.Logger;
 
-import com.yugabyte.sample.apps.AppBase.TableOp;
 import com.yugabyte.sample.common.SimpleLoadGenerator.Key;
 
 /**
  * This workload writes and reads some random string keys from a postgresql table.
  */
-public class SqlInserts extends AppBase {
+public class SqlInserts extends SQLAppBase {
   private static final Logger LOG = Logger.getLogger(SqlInserts.class);
 
   // Static initialization of this workload's config. These are good defaults for getting a decent