diff --git a/src/main/java/com/contrastsecurity/cassandra/migration/CassandraMigration.java b/src/main/java/com/contrastsecurity/cassandra/migration/CassandraMigration.java index 9e2adac..e7c1043 100644 --- a/src/main/java/com/contrastsecurity/cassandra/migration/CassandraMigration.java +++ b/src/main/java/com/contrastsecurity/cassandra/migration/CassandraMigration.java @@ -1,5 +1,6 @@ package com.contrastsecurity.cassandra.migration; +import com.contrastsecurity.cassandra.migration.action.Clear; import com.contrastsecurity.cassandra.migration.action.Initialize; import com.contrastsecurity.cassandra.migration.action.Migrate; import com.contrastsecurity.cassandra.migration.action.Validate; @@ -108,6 +109,15 @@ public String execute(Session session) { throw new CassandraMigrationException("Validation failed. " + validationError); } } + + public Boolean clear() { + return execute(new Action() { + public Boolean execute(Session session) { + Clear clear = new Clear(session, keyspace); + return clear.run(); + } + }); + } public void baseline() { //TODO diff --git a/src/main/java/com/contrastsecurity/cassandra/migration/CommandLine.java b/src/main/java/com/contrastsecurity/cassandra/migration/CommandLine.java index 782b526..4ff4523 100644 --- a/src/main/java/com/contrastsecurity/cassandra/migration/CommandLine.java +++ b/src/main/java/com/contrastsecurity/cassandra/migration/CommandLine.java @@ -21,6 +21,11 @@ public class CommandLine { */ public static final String VALIDATE = "validate"; + /** + * command to trigger clear action + */ + public static final String CLEAR = "clear"; + /** * logging support */ @@ -45,10 +50,16 @@ public static void main(String[] args) { CassandraMigration cm = new CassandraMigration(); Keyspace ks = new Keyspace(); cm.setKeyspace(ks); - if (MIGRATE.equalsIgnoreCase(operation)) { - cm.migrate(); - } else if (VALIDATE.equalsIgnoreCase(operation)) { - cm.validate(); + switch (operation.toLowerCase()) { + case MIGRATE: + cm.migrate(); + break; + case VALIDATE: + cm.validate(); + break; + case CLEAR: + cm.clear(); + break; } } @@ -92,6 +103,7 @@ private static void printUsage() { LOG.info("========"); LOG.info("migrate : Migrates the database"); LOG.info("validate : Validates the applied migrations against the available ones"); + LOG.info("clear : Clears the whole database"); LOG.info(""); LOG.info("Add -X to print debug output"); LOG.info("Add -q to suppress all output, except for errors and warnings"); diff --git a/src/main/java/com/contrastsecurity/cassandra/migration/action/Clear.java b/src/main/java/com/contrastsecurity/cassandra/migration/action/Clear.java new file mode 100644 index 0000000..568fd69 --- /dev/null +++ b/src/main/java/com/contrastsecurity/cassandra/migration/action/Clear.java @@ -0,0 +1,77 @@ +package com.contrastsecurity.cassandra.migration.action; + +import com.contrastsecurity.cassandra.migration.config.Keyspace; +import com.contrastsecurity.cassandra.migration.logging.Log; +import com.contrastsecurity.cassandra.migration.logging.LogFactory; +import com.contrastsecurity.cassandra.migration.utils.StopWatch; +import com.datastax.driver.core.ResultSet; +import com.datastax.driver.core.Row; +import com.datastax.driver.core.Session; +import com.datastax.driver.core.querybuilder.QueryBuilder; +import com.datastax.driver.core.querybuilder.Select; + +import static com.datastax.driver.core.querybuilder.QueryBuilder.eq; + +public class Clear { + private static final Log LOG = LogFactory.getLog(Clear.class); + + private final Session session; + private final Keyspace keyspace; + + public Clear(Session session, Keyspace keyspace) { + this.session = session; + this.keyspace = keyspace; + } + + public boolean run() { + StopWatch stopWatch = new StopWatch(); + stopWatch.start(); + + for (ObjectType objectType : ObjectType.values()) { + clearObjects(objectType); + } + + stopWatch.stop(); + LOG.info("CLEARED ALL OBJECTS"); + LOG.info(String.format("CLEARING TOOK %d ms", stopWatch.getTotalTimeMillis())); + + return true; + } + + public void clearObjects(ObjectType objectType) { + Select.Where objectsQuery = QueryBuilder.select(objectType.getSchemaColumnName()).from("system_schema", objectType.getSchemaTable()).where(eq("keyspace_name", keyspace.getName())); + ResultSet objects = session.execute(objectsQuery); + for (Row object : objects) { + LOG.info(String.format("Clearing %s of type %s", object.getString(objectType.getSchemaColumnName()), objectType.queryFormat())); + session.execute(String.format("DROP %s IF EXISTS %s", + objectType.queryFormat(), + object.getString(objectType.getSchemaColumnName()))); + } + + } + + public enum ObjectType { + MATERIALIZED_VIEW("views", "view_name"), + TABLE("tables", "table_name"); + + private final String schemaTable; + private final String schemaColumnName; + + ObjectType(String schemaTable, String schemaColumnName) { + this.schemaTable = schemaTable; + this.schemaColumnName = schemaColumnName; + } + + public String queryFormat() { + return name().replace("_", " ").toUpperCase(); + } + + public String getSchemaTable() { + return schemaTable; + } + + public String getSchemaColumnName() { + return schemaColumnName; + } + } +}