Update to support Cassandra 3.x - Contrast-Security-OSS#12

seanryan41 · Feb 1, 2016 · 2092e74 · 2092e74
1 parent 1ec13c3
commit 2092e74
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 15 deletions.
diff --git 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
@@ -8,17 +8,16 @@
 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;
 
 import java.util.ArrayList;
 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")