From 2092e746ebd85e5647044e82d7cde164721444ce Mon Sep 17 00:00:00 2001
From: Sean Ryan <sean.ryan@uptake.io>
Date: Mon, 1 Feb 2016 13:42:58 -0600
Subject: [PATCH] Update to support Cassandra 3.x -
 Contrast-Security-OSS/cassandra-migration#12

---
 pom.xml                                       |  2 +-
 .../migration/dao/SchemaVersionDAO.java       | 41 ++++++++++++-------
 2 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/pom.xml b/pom.xml
index 8d18157..f8d590e 100644
--- a/pom.xml
+++ b/pom.xml
@@ -81,7 +81,7 @@
         <dependency>
             <groupId>com.datastax.cassandra</groupId>
             <artifactId>cassandra-driver-core</artifactId>
-            <version>2.1.9</version>
+            <version>3.0.0</version>
         </dependency>
         <dependency>
             <groupId>junit</groupId>
diff --git a/src/main/java/com/contrastsecurity/cassandra/migration/dao/SchemaVersionDAO.java b/src/main/java/com/contrastsecurity/cassandra/migration/dao/SchemaVersionDAO.java
index f25e796..921831e 100644
--- a/src/main/java/com/contrastsecurity/cassandra/migration/dao/SchemaVersionDAO.java
+++ b/src/main/java/com/contrastsecurity/cassandra/migration/dao/SchemaVersionDAO.java
@@ -8,6 +8,7 @@
 import com.contrastsecurity.cassandra.migration.logging.LogFactory;
 import com.contrastsecurity.cassandra.migration.utils.CachePrepareStatement;
 import com.datastax.driver.core.*;
+import com.datastax.driver.core.exceptions.InvalidQueryException;
 import com.datastax.driver.core.querybuilder.QueryBuilder;
 import com.datastax.driver.core.querybuilder.Select;
 
@@ -15,10 +16,8 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 
 import static com.datastax.driver.core.querybuilder.QueryBuilder.eq;
-import static com.datastax.driver.core.querybuilder.QueryBuilder.in;
 
 public class SchemaVersionDAO {
 
@@ -78,23 +77,37 @@ public boolean tablesExist() {
         boolean schemaVersionTableExists = false;
         boolean schemaVersionCountsTableExists = false;
 
-        Statement statement = QueryBuilder
+        Statement schemaVersionStatement = QueryBuilder
                 .select()
-                .column("columnfamily_name")
-                .from("System", "schema_columnfamilies")
-                .where(eq("keyspace_name", keyspace.getName()))
-                .and(in("columnfamily_name", tableName, tableName + COUNTS_TABLE_NAME_SUFFIX));
-        statement.setConsistencyLevel(ConsistencyLevel.ALL);
-        ResultSet results = session.execute(statement);
-        for (Row row : results) {
-            String table = row.getString("columnfamily_name");
-            if (null != table && table.equalsIgnoreCase(tableName)) {
+                .countAll()
+                .from(keyspace.getName(), tableName);
+
+        Statement schemaVersionCountsStatement = QueryBuilder
+                .select()
+                .countAll()
+                .from(keyspace.getName(), tableName + COUNTS_TABLE_NAME_SUFFIX);
+
+        schemaVersionStatement.setConsistencyLevel(ConsistencyLevel.ALL);
+        schemaVersionCountsStatement.setConsistencyLevel(ConsistencyLevel.ALL);
+
+        try {
+            ResultSet resultsSchemaVersion = session.execute(schemaVersionStatement);
+            if (resultsSchemaVersion.one() != null) {
                 schemaVersionTableExists = true;
             }
-            if (null != table && table.equalsIgnoreCase(tableName + COUNTS_TABLE_NAME_SUFFIX)) {
+        } catch (InvalidQueryException e) {
+            LOG.debug("No schema version table found with a name of " + tableName);
+        }
+
+        try {
+            ResultSet resultsSchemaVersionCounts = session.execute(schemaVersionCountsStatement);
+            if (resultsSchemaVersionCounts.one() != null) {
                 schemaVersionCountsTableExists = true;
             }
+        } catch (InvalidQueryException e) {
+            LOG.debug("No schema version counts table found with a name of " + tableName + COUNTS_TABLE_NAME_SUFFIX);
         }
+
         return schemaVersionTableExists && schemaVersionCountsTableExists;
     }
 
@@ -165,7 +178,7 @@ public List<AppliedMigration> findAppliedMigrations() {
                     MigrationType.valueOf(row.getString("type")),
                     row.getString("script"),
                     row.isNull("checksum") ? null : row.getInt("checksum"),
-                    row.getDate("installed_on"),
+                    row.getTimestamp("installed_on"),
                     row.getString("installed_by"),
                     row.getInt("execution_time"),
                     row.getBool("success")